Skip to content

Commit 623213d

Browse files
Cedric-Perauerdg845yiyixuxuclaude
authored
Krea2 Modular Pipeline Support, Documentation, Test Cases (#14083)
* towards krea2 modular * add test cases * self-review fixes: test model id to hf internal, guider change * remove pnd * modular autodoc * modular autodoc * Update src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/modular_pipelines/krea2/before_denoise.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update docs/source/en/api/pipelines/krea2.md Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * input params, dim multiple of scale factor * patch * docs for krea2 turbo modular * docs fix * Update tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Update src/diffusers/modular_pipelines/krea2/before_denoise.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * modular blocks krea2 turbo fix * Update docs/source/en/api/pipelines/krea2.md Co-authored-by: YiYi Xu <yixu310@gmail.com> * Update src/diffusers/modular_pipelines/krea2/before_denoise.py Co-authored-by: YiYi Xu <yixu310@gmail.com> * fix default height/width and inference time steps * Apply suggestions from code review Co-authored-by: YiYi Xu <yixu310@gmail.com> * Regenerate stale krea2 auto docstrings Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> --------- Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: YiYi Xu <yixu310@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
1 parent af95be6 commit 623213d

16 files changed

Lines changed: 2099 additions & 1 deletion

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

Lines changed: 86 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ import torch
3535
from diffusers import Krea2Pipeline
3636

3737
# Load from a local directory produced by the Krea 2 conversion (no hub repo yet).
38-
pipe = Krea2Pipeline.from_pretrained("path/to/krea2-diffusers", torch_dtype=torch.bfloat16)
38+
pipe = Krea2Pipeline.from_pretrained("krea/Krea-2-Raw", torch_dtype=torch.bfloat16)
3939
pipe.to("cuda")
4040

4141
prompt = "a fox in the snow"
@@ -50,6 +50,27 @@ image = pipe(
5050
image.save("krea2.png")
5151
```
5252

53+
We additionally provide an example for using Krea2 Turbo :
54+
55+
```python
56+
import torch
57+
from diffusers import Krea2Pipeline
58+
59+
pipe = Krea2Pipeline.from_pretrained("krea/Krea-2-Turbo", torch_dtype=torch.bfloat16)
60+
pipe.to("cuda")
61+
62+
image = pipe(
63+
"a fox in the snow",
64+
height=1024,
65+
width=1024,
66+
num_inference_steps=8,
67+
guidance_scale=0.0,
68+
generator=torch.Generator("cuda").manual_seed(0),
69+
).images[0]
70+
image.save("krea2_turbo.png")
71+
```
72+
73+
5374
## Krea2Pipeline
5475

5576
[[autodoc]] Krea2Pipeline
@@ -59,3 +80,67 @@ image.save("krea2.png")
5980
## Krea2PipelineOutput
6081

6182
[[autodoc]] pipelines.krea2.pipeline_output.Krea2PipelineOutput
83+
84+
## Modular
85+
86+
Krea 2 is also available as a [modular pipeline](../../modular_diffusers/overview). Classifier-free guidance is
87+
configured through the `guider` component rather than a `guidance_scale` call argument. Krea 2 uses cond-anchored CFG,
88+
which is [`ClassifierFreeGuidance`] with `use_original_formulation=True`.
89+
90+
```python
91+
import torch
92+
from diffusers import ClassifierFreeGuidance, ModularPipeline
93+
94+
pipe = ModularPipeline.from_pretrained("krea/Krea-2-Raw")
95+
pipe.load_components(torch_dtype=torch.bfloat16)
96+
pipe.to("cuda")
97+
98+
99+
image = pipe(
100+
prompt="a fox in the snow",
101+
height=1024,
102+
width=1024,
103+
num_inference_steps=28,
104+
generator=torch.Generator("cuda").manual_seed(0),
105+
).images[0]
106+
image.save("krea2.png")
107+
```
108+
109+
We additionally provide an example for using Krea2 Turbo. The distilled checkpoint maps to its own set of blocks
110+
([`Krea2TurboAutoBlocks`]): it runs guidance-free (no `guider`), takes no negative prompt, and samples in a few steps.
111+
`ModularPipeline.from_pretrained` picks the turbo blocks automatically from the checkpoint's `is_distilled` config, so
112+
no guidance configuration is needed:
113+
114+
```python
115+
import torch
116+
from diffusers import ModularPipeline
117+
118+
pipe = ModularPipeline.from_pretrained("krea/Krea-2-Turbo")
119+
pipe.load_components(torch_dtype=torch.bfloat16)
120+
pipe.to("cuda")
121+
122+
image = pipe(
123+
prompt="a fox in the snow",
124+
height=1024,
125+
width=1024,
126+
num_inference_steps=8,
127+
generator=torch.Generator("cuda").manual_seed(0),
128+
).images[0]
129+
image.save("krea2_turbo.png")
130+
```
131+
132+
## Krea2ModularPipeline
133+
134+
[[autodoc]] Krea2ModularPipeline
135+
136+
## Krea2AutoBlocks
137+
138+
[[autodoc]] Krea2AutoBlocks
139+
140+
## Krea2TurboModularPipeline
141+
142+
[[autodoc]] Krea2TurboModularPipeline
143+
144+
## Krea2TurboAutoBlocks
145+
146+
[[autodoc]] Krea2TurboAutoBlocks

‎src/diffusers/__init__.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -512,6 +512,10 @@
512512
"HunyuanVideo15ModularPipeline",
513513
"Ideogram4AutoBlocks",
514514
"Ideogram4ModularPipeline",
515+
"Krea2AutoBlocks",
516+
"Krea2ModularPipeline",
517+
"Krea2TurboAutoBlocks",
518+
"Krea2TurboModularPipeline",
515519
"LTXAutoBlocks",
516520
"LTXModularPipeline",
517521
"QwenImageAutoBlocks",
@@ -1320,6 +1324,10 @@
13201324
HunyuanVideo15ModularPipeline,
13211325
Ideogram4AutoBlocks,
13221326
Ideogram4ModularPipeline,
1327+
Krea2AutoBlocks,
1328+
Krea2ModularPipeline,
1329+
Krea2TurboAutoBlocks,
1330+
Krea2TurboModularPipeline,
13231331
LTXAutoBlocks,
13241332
LTXModularPipeline,
13251333
QwenImageAutoBlocks,

‎src/diffusers/modular_pipelines/__init__.py‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,12 @@
8383
"Ideogram4AutoBlocks",
8484
"Ideogram4ModularPipeline",
8585
]
86+
_import_structure["krea2"] = [
87+
"Krea2AutoBlocks",
88+
"Krea2ModularPipeline",
89+
"Krea2TurboAutoBlocks",
90+
"Krea2TurboModularPipeline",
91+
]
8692
_import_structure["qwenimage"] = [
8793
"QwenImageAutoBlocks",
8894
"QwenImageModularPipeline",
@@ -162,6 +168,12 @@
162168
Ideogram4AutoBlocks,
163169
Ideogram4ModularPipeline,
164170
)
171+
from .krea2 import (
172+
Krea2AutoBlocks,
173+
Krea2ModularPipeline,
174+
Krea2TurboAutoBlocks,
175+
Krea2TurboModularPipeline,
176+
)
165177
from .ltx import LTXAutoBlocks, LTXModularPipeline
166178
from .modular_pipeline import (
167179
AutoPipelineBlocks,

‎src/diffusers/modular_pipelines/krea2/__init__.py‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
from typing import TYPE_CHECKING
2+
3+
from ...utils import (
4+
DIFFUSERS_SLOW_IMPORT,
5+
OptionalDependencyNotAvailable,
6+
_LazyModule,
7+
get_objects_from_module,
8+
is_torch_available,
9+
is_transformers_available,
10+
)
11+
12+
13+
_dummy_objects = {}
14+
_import_structure = {}
15+
16+
try:
17+
if not (is_transformers_available() and is_torch_available()):
18+
raise OptionalDependencyNotAvailable()
19+
except OptionalDependencyNotAvailable:
20+
from ...utils import dummy_torch_and_transformers_objects # noqa F403
21+
22+
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
23+
else:
24+
_import_structure["modular_blocks_krea2"] = ["Krea2AutoBlocks"]
25+
_import_structure["modular_blocks_krea2_turbo"] = ["Krea2TurboAutoBlocks"]
26+
_import_structure["modular_pipeline"] = ["Krea2ModularPipeline", "Krea2TurboModularPipeline"]
27+
28+
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
29+
try:
30+
if not (is_transformers_available() and is_torch_available()):
31+
raise OptionalDependencyNotAvailable()
32+
except OptionalDependencyNotAvailable:
33+
from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
34+
else:
35+
from .modular_blocks_krea2 import Krea2AutoBlocks
36+
from .modular_blocks_krea2_turbo import Krea2TurboAutoBlocks
37+
from .modular_pipeline import Krea2ModularPipeline, Krea2TurboModularPipeline
38+
else:
39+
import sys
40+
41+
sys.modules[__name__] = _LazyModule(
42+
__name__,
43+
globals()["__file__"],
44+
_import_structure,
45+
module_spec=__spec__,
46+
)
47+
48+
for name, value in _dummy_objects.items():
49+
setattr(sys.modules[__name__], name, value)

0 commit comments

Comments
 (0)