Skip to content

Commit be3c2a0

Browse files
authored
[WIP] Add Flux2 modular (#12763)
* update * update * update * update * update * update * update * update * update * update
1 parent 8b4722d commit be3c2a0

File tree

15 files changed

+1882
-2
lines changed

15 files changed

+1882
-2
lines changed

src/diffusers/__init__.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -404,6 +404,8 @@
404404
else:
405405
_import_structure["modular_pipelines"].extend(
406406
[
407+
"Flux2AutoBlocks",
408+
"Flux2ModularPipeline",
407409
"FluxAutoBlocks",
408410
"FluxKontextAutoBlocks",
409411
"FluxKontextModularPipeline",
@@ -1111,6 +1113,8 @@
11111113
from .utils.dummy_torch_and_transformers_objects import * # noqa F403
11121114
else:
11131115
from .modular_pipelines import (
1116+
Flux2AutoBlocks,
1117+
Flux2ModularPipeline,
11141118
FluxAutoBlocks,
11151119
FluxKontextAutoBlocks,
11161120
FluxKontextModularPipeline,

src/diffusers/modular_pipelines/__init__.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,10 @@
5252
"FluxKontextAutoBlocks",
5353
"FluxKontextModularPipeline",
5454
]
55+
_import_structure["flux2"] = [
56+
"Flux2AutoBlocks",
57+
"Flux2ModularPipeline",
58+
]
5559
_import_structure["qwenimage"] = [
5660
"QwenImageAutoBlocks",
5761
"QwenImageModularPipeline",
@@ -75,6 +79,7 @@
7579
else:
7680
from .components_manager import ComponentsManager
7781
from .flux import FluxAutoBlocks, FluxKontextAutoBlocks, FluxKontextModularPipeline, FluxModularPipeline
82+
from .flux2 import Flux2AutoBlocks, Flux2ModularPipeline
7883
from .modular_pipeline import (
7984
AutoPipelineBlocks,
8085
BlockState,
Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
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["encoders"] = [
25+
"Flux2TextEncoderStep",
26+
"Flux2RemoteTextEncoderStep",
27+
"Flux2VaeEncoderStep",
28+
]
29+
_import_structure["before_denoise"] = [
30+
"Flux2SetTimestepsStep",
31+
"Flux2PrepareLatentsStep",
32+
"Flux2RoPEInputsStep",
33+
"Flux2PrepareImageLatentsStep",
34+
]
35+
_import_structure["denoise"] = [
36+
"Flux2LoopDenoiser",
37+
"Flux2LoopAfterDenoiser",
38+
"Flux2DenoiseLoopWrapper",
39+
"Flux2DenoiseStep",
40+
]
41+
_import_structure["decoders"] = ["Flux2DecodeStep"]
42+
_import_structure["inputs"] = [
43+
"Flux2ProcessImagesInputStep",
44+
"Flux2TextInputStep",
45+
]
46+
_import_structure["modular_blocks"] = [
47+
"ALL_BLOCKS",
48+
"AUTO_BLOCKS",
49+
"REMOTE_AUTO_BLOCKS",
50+
"TEXT2IMAGE_BLOCKS",
51+
"IMAGE_CONDITIONED_BLOCKS",
52+
"Flux2AutoBlocks",
53+
"Flux2AutoVaeEncoderStep",
54+
"Flux2BeforeDenoiseStep",
55+
"Flux2VaeEncoderSequentialStep",
56+
]
57+
_import_structure["modular_pipeline"] = ["Flux2ModularPipeline"]
58+
59+
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
60+
try:
61+
if not (is_transformers_available() and is_torch_available()):
62+
raise OptionalDependencyNotAvailable()
63+
except OptionalDependencyNotAvailable:
64+
from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
65+
else:
66+
from .before_denoise import (
67+
Flux2PrepareImageLatentsStep,
68+
Flux2PrepareLatentsStep,
69+
Flux2RoPEInputsStep,
70+
Flux2SetTimestepsStep,
71+
)
72+
from .decoders import Flux2DecodeStep
73+
from .denoise import (
74+
Flux2DenoiseLoopWrapper,
75+
Flux2DenoiseStep,
76+
Flux2LoopAfterDenoiser,
77+
Flux2LoopDenoiser,
78+
)
79+
from .encoders import (
80+
Flux2RemoteTextEncoderStep,
81+
Flux2TextEncoderStep,
82+
Flux2VaeEncoderStep,
83+
)
84+
from .inputs import (
85+
Flux2ProcessImagesInputStep,
86+
Flux2TextInputStep,
87+
)
88+
from .modular_blocks import (
89+
ALL_BLOCKS,
90+
AUTO_BLOCKS,
91+
IMAGE_CONDITIONED_BLOCKS,
92+
REMOTE_AUTO_BLOCKS,
93+
TEXT2IMAGE_BLOCKS,
94+
Flux2AutoBlocks,
95+
Flux2AutoVaeEncoderStep,
96+
Flux2BeforeDenoiseStep,
97+
Flux2VaeEncoderSequentialStep,
98+
)
99+
from .modular_pipeline import Flux2ModularPipeline
100+
else:
101+
import sys
102+
103+
sys.modules[__name__] = _LazyModule(
104+
__name__,
105+
globals()["__file__"],
106+
_import_structure,
107+
module_spec=__spec__,
108+
)
109+
110+
for name, value in _dummy_objects.items():
111+
setattr(sys.modules[__name__], name, value)

0 commit comments

Comments
 (0)