Skip to content

Commit a8e853b

Browse files
yiyixuxua-r-r-o-wDN6sayakpaul
authored
[modular diffusers] more refactor (#11235)
* add componentspec and configspec * up * up * move methods to blocks * Modular Diffusers Guiders (#11311) * cfg; slg; pag; sdxl without controlnet * support sdxl controlnet * support controlnet union * update * update * cfg zero* * use unwrap_module for torch compiled modules * remove guider kwargs * remove commented code * remove old guider * fix slg bug * remove debug print * autoguidance * smoothed energy guidance * add note about seg * tangential cfg * cfg plus plus * support cfgpp in ddim * apply review suggestions * refactor * rename enable/disable * remove cfg++ for now * rename do_classifier_free_guidance->prepare_unconditional_embeds * remove unused * [modular diffusers] introducing ModularLoader (#11462) * cfg; slg; pag; sdxl without controlnet --------- Co-authored-by: Aryan <[email protected]> * make loader optional * remove lora step and ip-adapter step -> no longer needed * rename pipeline -> components, data -> block_state * seperate controlnet step into input + denoise * refactor controlnet union * reefactor pipeline/block states so that it can dynamically accept kwargs * remove controlnet union denoise step, refactor & reuse controlnet denoisee step to accept aditional contrlnet kwargs * allow input_fields as input & update message * update input formating, consider kwarggs_type inputs with no name, e/g *_controlnet_kwargs * refactor the denoiseestep using LoopSequential! also add a new file for denoise step * change warning to debug * fix get_execusion blocks with loopsequential * fix auto denoise so all tests pass * update imports on guiders * remove modular reelated change from pipelines folder * made a modular_pipelines folder! * update __init__ * add notes * add block state will also make sure modifed intermediates_inputs will be updated * move block mappings to its own file * make inputs truly immutable, remove the output logic in sequential pipeline, and update so that intermediates_outputs are only new variables * decode block, if skip decoding do not need to update latent * fix imports * fix import * fix more * remove the output step * make generator intermediates (it is mutable) * after_denoise -> decoders * add a to-do for guider cconfig mixin * refactor component spec: replace create/create_from_pretrained/create_from_config to just create and load method * refactor modular loader: 1. load only load (pretrained components only if not specific names) 2. update acceept create spec 3. move the updte _componeent_spec logic outside register_components to each method that create/update the component: __init__/update/load * update components manager * up * [WIP] Modular Diffusers support custom code/pipeline blocks (#11539) * update * update * remove the duplicated components_manager file I forgot to deletee * fix import in block mapping * add a to-do for modular loader * prepare_latents_img2img pipeline method -> function, maybe do the same for others? * update input for loop blocks, do not need to include intermediate * solve merge conflict: manually add back the remote code change to modular_pipeline * add node_utils * modular node! * add * refator based on dhruv's feedbacks * update doc format for kwargs_type * up * updatee modular_pipeline.from_pretrained, modular_repo ->pretrained_model_name_or_path * save_pretrained for serializing config. (#11603) * save_pretrained for serializing config. * remove pushtohub * diffusers-cli rough --------- Co-authored-by: YiYi Xu <[email protected]> --------- Co-authored-by: Aryan <[email protected]> Co-authored-by: Dhruv Nair <[email protected]> Co-authored-by: Sayak Paul <[email protected]>
1 parent 6a509ba commit a8e853b

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

42 files changed

+11037
-6489
lines changed

src/diffusers/__init__.py

Lines changed: 72 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -34,10 +34,12 @@
3434

3535
_import_structure = {
3636
"configuration_utils": ["ConfigMixin"],
37+
"guiders": [],
3738
"hooks": [],
3839
"loaders": ["FromOriginalModelMixin"],
3940
"models": [],
4041
"pipelines": [],
42+
"modular_pipelines": [],
4143
"quantizers.quantization_config": [],
4244
"schedulers": [],
4345
"utils": [
@@ -130,12 +132,26 @@
130132
_import_structure["utils.dummy_pt_objects"] = [name for name in dir(dummy_pt_objects) if not name.startswith("_")]
131133

132134
else:
135+
_import_structure["guiders"].extend(
136+
[
137+
"AdaptiveProjectedGuidance",
138+
"AutoGuidance",
139+
"ClassifierFreeGuidance",
140+
"ClassifierFreeZeroStarGuidance",
141+
"SkipLayerGuidance",
142+
"SmoothedEnergyGuidance",
143+
"TangentialClassifierFreeGuidance",
144+
]
145+
)
133146
_import_structure["hooks"].extend(
134147
[
135148
"FasterCacheConfig",
136149
"HookRegistry",
137150
"PyramidAttentionBroadcastConfig",
151+
"LayerSkipConfig",
152+
"SmoothedEnergyGuidanceConfig",
138153
"apply_faster_cache",
154+
"apply_layer_skip",
139155
"apply_pyramid_attention_broadcast",
140156
]
141157
)
@@ -239,13 +255,21 @@
239255
"KarrasVePipeline",
240256
"LDMPipeline",
241257
"LDMSuperResolutionPipeline",
242-
"ModularPipeline",
243258
"PNDMPipeline",
244259
"RePaintPipeline",
245260
"ScoreSdeVePipeline",
246261
"StableDiffusionMixin",
247262
]
248263
)
264+
_import_structure["modular_pipelines"].extend(
265+
[
266+
"ModularLoader",
267+
"ModularPipeline",
268+
"ModularPipelineBlocks",
269+
"ComponentSpec",
270+
"ComponentsManager",
271+
]
272+
)
249273
_import_structure["quantizers"] = ["DiffusersQuantizer"]
250274
_import_structure["schedulers"].extend(
251275
[
@@ -494,12 +518,10 @@
494518
"StableDiffusionXLImg2ImgPipeline",
495519
"StableDiffusionXLInpaintPipeline",
496520
"StableDiffusionXLInstructPix2PixPipeline",
497-
"StableDiffusionXLModularPipeline",
498521
"StableDiffusionXLPAGImg2ImgPipeline",
499522
"StableDiffusionXLPAGInpaintPipeline",
500523
"StableDiffusionXLPAGPipeline",
501524
"StableDiffusionXLPipeline",
502-
"StableDiffusionXLAutoPipeline",
503525
"StableUnCLIPImg2ImgPipeline",
504526
"StableUnCLIPPipeline",
505527
"StableVideoDiffusionPipeline",
@@ -526,6 +548,24 @@
526548
]
527549
)
528550

551+
552+
try:
553+
if not (is_torch_available() and is_transformers_available()):
554+
raise OptionalDependencyNotAvailable()
555+
except OptionalDependencyNotAvailable:
556+
from .utils import dummy_torch_and_transformers_objects # noqa F403
557+
558+
_import_structure["utils.dummy_torch_and_transformers_objects"] = [
559+
name for name in dir(dummy_torch_and_transformers_objects) if not name.startswith("_")
560+
]
561+
562+
else:
563+
_import_structure["modular_pipelines"].extend(
564+
[
565+
"StableDiffusionXLAutoPipeline",
566+
"StableDiffusionXLModularLoader",
567+
]
568+
)
529569
try:
530570
if not (is_torch_available() and is_transformers_available() and is_opencv_available()):
531571
raise OptionalDependencyNotAvailable()
@@ -731,10 +771,22 @@
731771
except OptionalDependencyNotAvailable:
732772
from .utils.dummy_pt_objects import * # noqa F403
733773
else:
774+
from .guiders import (
775+
AdaptiveProjectedGuidance,
776+
AutoGuidance,
777+
ClassifierFreeGuidance,
778+
ClassifierFreeZeroStarGuidance,
779+
SkipLayerGuidance,
780+
SmoothedEnergyGuidance,
781+
TangentialClassifierFreeGuidance,
782+
)
734783
from .hooks import (
735784
FasterCacheConfig,
736785
HookRegistry,
786+
LayerSkipConfig,
737787
PyramidAttentionBroadcastConfig,
788+
SmoothedEnergyGuidanceConfig,
789+
apply_layer_skip,
738790
apply_faster_cache,
739791
apply_pyramid_attention_broadcast,
740792
)
@@ -837,12 +889,18 @@
837889
KarrasVePipeline,
838890
LDMPipeline,
839891
LDMSuperResolutionPipeline,
840-
ModularPipeline,
841892
PNDMPipeline,
842893
RePaintPipeline,
843894
ScoreSdeVePipeline,
844895
StableDiffusionMixin,
845896
)
897+
from .modular_pipelines import (
898+
ModularLoader,
899+
ModularPipeline,
900+
ModularPipelineBlocks,
901+
ComponentSpec,
902+
ComponentsManager,
903+
)
846904
from .quantizers import DiffusersQuantizer
847905
from .schedulers import (
848906
AmusedScheduler,
@@ -1070,12 +1128,10 @@
10701128
StableDiffusionXLImg2ImgPipeline,
10711129
StableDiffusionXLInpaintPipeline,
10721130
StableDiffusionXLInstructPix2PixPipeline,
1073-
StableDiffusionXLModularPipeline,
10741131
StableDiffusionXLPAGImg2ImgPipeline,
10751132
StableDiffusionXLPAGInpaintPipeline,
10761133
StableDiffusionXLPAGPipeline,
10771134
StableDiffusionXLPipeline,
1078-
StableDiffusionXLAutoPipeline,
10791135
StableUnCLIPImg2ImgPipeline,
10801136
StableUnCLIPPipeline,
10811137
StableVideoDiffusionPipeline,
@@ -1100,7 +1156,16 @@
11001156
WuerstchenDecoderPipeline,
11011157
WuerstchenPriorPipeline,
11021158
)
1103-
1159+
try:
1160+
if not (is_torch_available() and is_transformers_available()):
1161+
raise OptionalDependencyNotAvailable()
1162+
except OptionalDependencyNotAvailable:
1163+
from .utils.dummy_torch_and_transformers_objects import * # noqa F403
1164+
else:
1165+
from .modular_pipelines import (
1166+
StableDiffusionXLAutoPipeline,
1167+
StableDiffusionXLModularLoader,
1168+
)
11041169
try:
11051170
if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
11061171
raise OptionalDependencyNotAvailable()
Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
# Copyright 2025 The HuggingFace Team. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""
16+
Usage example:
17+
TODO
18+
"""
19+
20+
import ast
21+
from argparse import ArgumentParser, Namespace
22+
from pathlib import Path
23+
import importlib.util
24+
import os
25+
from ..utils import logging
26+
from . import BaseDiffusersCLICommand
27+
28+
29+
EXPECTED_PARENT_CLASSES = ["PipelineBlock"]
30+
CONFIG = "config.json"
31+
32+
def conversion_command_factory(args: Namespace):
33+
return CustomBlocksCommand(args.block_module_name, args.block_class_name)
34+
35+
36+
class CustomBlocksCommand(BaseDiffusersCLICommand):
37+
@staticmethod
38+
def register_subcommand(parser: ArgumentParser):
39+
conversion_parser = parser.add_parser("custom_blocks")
40+
conversion_parser.add_argument(
41+
"--block_module_name",
42+
type=str,
43+
default="block.py",
44+
help="Module filename in which the custom block will be implemented.",
45+
)
46+
conversion_parser.add_argument(
47+
"--block_class_name", type=str, default=None, help="Name of the custom block. If provided None, we will try to infer it."
48+
)
49+
conversion_parser.set_defaults(func=conversion_command_factory)
50+
51+
def __init__(self, block_module_name: str = "block.py", block_class_name: str = None):
52+
self.logger = logging.get_logger("diffusers-cli/custom_blocks")
53+
self.block_module_name = Path(block_module_name)
54+
self.block_class_name = block_class_name
55+
56+
def run(self):
57+
# determine the block to be saved.
58+
out = self._get_class_names(self.block_module_name)
59+
classes_found = list({cls for cls, _ in out})
60+
61+
if self.block_class_name is not None:
62+
child_class, parent_class = self._choose_block(out, self.block_class_name)
63+
if child_class is None and parent_class is None:
64+
raise ValueError(
65+
"`block_class_name` could not be retrieved. Available classes from "
66+
f"{self.block_module_name}:\n{classes_found}"
67+
)
68+
else:
69+
self.logger.info(
70+
f"Found classes: {classes_found} will be using {classes_found[0]}. "
71+
"If this needs to be changed, re-run the command specifying `block_class_name`."
72+
)
73+
child_class, parent_class = out[0][0], out[0][1]
74+
75+
# dynamically get the custom block and initialize it to call `save_pretrained` in the current directory.
76+
# the user is responsible for running it, so I guess that is safe?
77+
module_name = f"__dynamic__{self.block_module_name.stem}"
78+
spec = importlib.util.spec_from_file_location(module_name, str(self.block_module_name))
79+
module = importlib.util.module_from_spec(spec)
80+
spec.loader.exec_module(module)
81+
getattr(module, child_class)().save_pretrained(os.getcwd())
82+
83+
# or, we could create it manually.
84+
# automap = self._create_automap(parent_class=parent_class, child_class=child_class)
85+
# with open(CONFIG, "w") as f:
86+
# json.dump(automap, f)
87+
with open("requirements.txt", "w") as f:
88+
f.write("")
89+
90+
def _choose_block(self, candidates, chosen=None):
91+
for cls, base in candidates:
92+
if cls == chosen:
93+
return cls, base
94+
return None, None
95+
96+
def _get_class_names(self, file_path):
97+
source = file_path.read_text(encoding="utf-8")
98+
try:
99+
tree = ast.parse(source, filename=file_path)
100+
except SyntaxError as e:
101+
raise ValueError(f"Could not parse {file_path!r}: {e}") from e
102+
103+
results: list[tuple[str, str]] = []
104+
for node in tree.body:
105+
if not isinstance(node, ast.ClassDef):
106+
continue
107+
108+
# extract all base names for this class
109+
base_names = [
110+
bname for b in node.bases
111+
if (bname := self._get_base_name(b)) is not None
112+
]
113+
114+
# for each allowed base that appears in the class's bases, emit a tuple
115+
for allowed in EXPECTED_PARENT_CLASSES:
116+
if allowed in base_names:
117+
results.append((node.name, allowed))
118+
119+
return results
120+
121+
def _get_base_name(self, node: ast.expr):
122+
if isinstance(node, ast.Name):
123+
return node.id
124+
elif isinstance(node, ast.Attribute):
125+
val = self._get_base_name(node.value)
126+
return f"{val}.{node.attr}" if val else node.attr
127+
return None
128+
129+
def _create_automap(self, parent_class, child_class):
130+
module = str(self.block_module_name).replace(".py", "").rsplit(".", 1)[-1]
131+
auto_map = {f"{parent_class}": f"{module}.{child_class}"}
132+
return {"auto_map": auto_map}
133+

src/diffusers/commands/diffusers_cli.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
from .env import EnvironmentCommand
1919
from .fp16_safetensors import FP16SafetensorsCommand
20+
from .custom_blocks import CustomBlocksCommand
2021

2122

2223
def main():
@@ -26,6 +27,7 @@ def main():
2627
# Register commands
2728
EnvironmentCommand.register_subcommand(commands_parser)
2829
FP16SafetensorsCommand.register_subcommand(commands_parser)
30+
CustomBlocksCommand.register_subcommand(commands_parser)
2931

3032
# Let's go
3133
args = parser.parse_args()

0 commit comments

Comments
 (0)