Krea2 Modular Pipeline Support, Documentation, Test Cases

docs/source/en/api/pipelines/krea2.md

@@ -35,7 +35,7 @@ import torch
 from diffusers import Krea2Pipeline
 
 # Load from a local directory produced by the Krea 2 conversion (no hub repo yet).
-pipe = Krea2Pipeline.from_pretrained("path/to/krea2-diffusers", torch_dtype=torch.bfloat16)
+pipe = Krea2Pipeline.from_pretrained("krea/Krea-2-Raw", torch_dtype=torch.bfloat16)
 pipe.to("cuda")
 
 prompt = "a fox in the snow"
@@ -50,6 +50,27 @@ image = pipe(
 image.save("krea2.png")
 ```
 
+We additionally provide an example for using Krea2 Turbo :
+
+```python
+import torch
+from diffusers import Krea2Pipeline
+
+pipe = Krea2Pipeline.from_pretrained("krea/Krea-2-Turbo", torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+
+image = pipe(
+    "a fox in the snow",
+    height=1024,
+    width=1024,
+    num_inference_steps=8,
+    guidance_scale=0.0,
+    generator=torch.Generator("cuda").manual_seed(0),
+).images[0]
+image.save("krea2_turbo.png")
+```
+
+
 ## Krea2Pipeline
 
 [[autodoc]] Krea2Pipeline
@@ -59,3 +80,61 @@ image.save("krea2.png")
 ## Krea2PipelineOutput
 
 [[autodoc]] pipelines.krea2.pipeline_output.Krea2PipelineOutput
+
+## Modular
+
+Krea 2 is also available as a [modular pipeline](../../modular_diffusers/overview). Classifier-free guidance is
+configured through the `guider` component rather than a `guidance_scale` call argument. Krea 2 uses cond-anchored CFG,
+which is [`ClassifierFreeGuidance`] with `use_original_formulation=True`; a scale of `0.0` disables guidance (the TDM
+checkpoint).
+
+```python
+import torch
+from diffusers import ClassifierFreeGuidance, ModularPipeline
+
+pipe = ModularPipeline.from_pretrained("krea/Krea-2-Raw")
+pipe.load_components(torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+
+pipe.update_components(guider=ClassifierFreeGuidance(guidance_scale=4.5, use_original_formulation=True))
+
+image = pipe(
+    prompt="a fox in the snow",
+    height=1024,
+    width=1024,
+    num_inference_steps=28,
+    generator=torch.Generator("cuda").manual_seed(0),
+).images[0]
+image.save("krea2.png")
+```
+
+We additionally provide an example for using Krea2 Turbo. The distilled checkpoint disables guidance by
+setting the `guider` scale to `0.0` and samples in a few steps:
+
+```python
+import torch
+from diffusers import ClassifierFreeGuidance, ModularPipeline
+
+pipe = ModularPipeline.from_pretrained("krea/Krea-2-Turbo")
+pipe.load_components(torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+
+pipe.update_components(guider=ClassifierFreeGuidance(guidance_scale=0.0, use_original_formulation=True))
+
+image = pipe(
+    prompt="a fox in the snow",
+    height=1024,
+    width=1024,
+    num_inference_steps=8,
+    generator=torch.Generator("cuda").manual_seed(0),
+).images[0]
+image.save("krea2_turbo.png")
+```
+
+## Krea2ModularPipeline
+
+[[autodoc]] Krea2ModularPipeline
+
+## Krea2AutoBlocks
+
+[[autodoc]] Krea2AutoBlocks

src/diffusers/__init__.py

@@ -501,6 +501,8 @@
             "HunyuanVideo15ModularPipeline",
             "Ideogram4AutoBlocks",
             "Ideogram4ModularPipeline",
+            "Krea2AutoBlocks",
+            "Krea2ModularPipeline",
             "LTXAutoBlocks",
             "LTXModularPipeline",
             "QwenImageAutoBlocks",
@@ -1353,6 +1355,8 @@
             HunyuanVideo15ModularPipeline,
             Ideogram4AutoBlocks,
             Ideogram4ModularPipeline,
+            Krea2AutoBlocks,
+            Krea2ModularPipeline,
             LTXAutoBlocks,
             LTXModularPipeline,
             QwenImageAutoBlocks,

src/diffusers/modular_pipelines/__init__.py

@@ -83,6 +83,10 @@
         "Ideogram4AutoBlocks",
         "Ideogram4ModularPipeline",
     ]
+    _import_structure["krea2"] = [
+        "Krea2AutoBlocks",
+        "Krea2ModularPipeline",
+    ]
     _import_structure["qwenimage"] = [
         "QwenImageAutoBlocks",
         "QwenImageModularPipeline",
@@ -155,6 +159,10 @@
             Ideogram4AutoBlocks,
             Ideogram4ModularPipeline,
         )
+        from .krea2 import (
+            Krea2AutoBlocks,
+            Krea2ModularPipeline,
+        )
         from .ltx import LTXAutoBlocks, LTXModularPipeline
         from .modular_pipeline import (
             AutoPipelineBlocks,

src/diffusers/modular_pipelines/krea2/__init__.py

@@ -0,0 +1,47 @@
+from typing import TYPE_CHECKING
+
+from ...utils import (
+    DIFFUSERS_SLOW_IMPORT,
+    OptionalDependencyNotAvailable,
+    _LazyModule,
+    get_objects_from_module,
+    is_torch_available,
+    is_transformers_available,
+)
+
+
+_dummy_objects = {}
+_import_structure = {}
+
+try:
+    if not (is_transformers_available() and is_torch_available()):
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    from ...utils import dummy_torch_and_transformers_objects  # noqa F403
+
+    _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
+else:
+    _import_structure["modular_blocks_krea2"] = ["Krea2AutoBlocks"]
+    _import_structure["modular_pipeline"] = ["Krea2ModularPipeline"]
+
+if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
+    try:
+        if not (is_transformers_available() and is_torch_available()):
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
+    else:
+        from .modular_blocks_krea2 import Krea2AutoBlocks
+        from .modular_pipeline import Krea2ModularPipeline
+else:
+    import sys
+
+    sys.modules[__name__] = _LazyModule(
+        __name__,
+        globals()["__file__"],
+        _import_structure,
+        module_spec=__spec__,
+    )
+
+    for name, value in _dummy_objects.items():
+        setattr(sys.modules[__name__], name, value)