Skip to content

Commit 5ffc621

Browse files
authored
Merge branch 'main' into modular-wan-more-pipelines
2 parents 7f97055 + 10dc589 commit 5ffc621

File tree

67 files changed

+10439
-1264
lines changed

Some content is hidden

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

67 files changed

+10439
-1264
lines changed

docs/source/en/_toctree.yml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,8 @@
114114
title: Guiders
115115
- local: modular_diffusers/custom_blocks
116116
title: Building Custom Blocks
117+
- local: modular_diffusers/mellon
118+
title: Using Custom Blocks with Mellon
117119
title: Modular Diffusers
118120
- isExpanded: false
119121
sections:

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

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,41 @@ image = pipe(
5353
image.save("zimage_img2img.png")
5454
```
5555

56+
## Inpainting
57+
58+
Use [`ZImageInpaintPipeline`] to inpaint specific regions of an image based on a text prompt and mask.
59+
60+
```python
61+
import torch
62+
import numpy as np
63+
from PIL import Image
64+
from diffusers import ZImageInpaintPipeline
65+
from diffusers.utils import load_image
66+
67+
pipe = ZImageInpaintPipeline.from_pretrained("Tongyi-MAI/Z-Image-Turbo", torch_dtype=torch.bfloat16)
68+
pipe.to("cuda")
69+
70+
url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
71+
init_image = load_image(url).resize((1024, 1024))
72+
73+
# Create a mask (white = inpaint, black = preserve)
74+
mask = np.zeros((1024, 1024), dtype=np.uint8)
75+
mask[256:768, 256:768] = 255 # Inpaint center region
76+
mask_image = Image.fromarray(mask)
77+
78+
prompt = "A beautiful lake with mountains in the background"
79+
image = pipe(
80+
prompt,
81+
image=init_image,
82+
mask_image=mask_image,
83+
strength=1.0,
84+
num_inference_steps=9,
85+
guidance_scale=0.0,
86+
generator=torch.Generator("cuda").manual_seed(42),
87+
).images[0]
88+
image.save("zimage_inpaint.png")
89+
```
90+
5691
## ZImagePipeline
5792

5893
[[autodoc]] ZImagePipeline
@@ -64,3 +99,9 @@ image.save("zimage_img2img.png")
6499
[[autodoc]] ZImageImg2ImgPipeline
65100
- all
66101
- __call__
102+
103+
## ZImageInpaintPipeline
104+
105+
[[autodoc]] ZImageInpaintPipeline
106+
- all
107+
- __call__

docs/source/en/modular_diffusers/components_manager.md

Lines changed: 42 additions & 136 deletions
Original file line numberDiff line numberDiff line change
@@ -12,179 +12,85 @@ specific language governing permissions and limitations under the License.
1212

1313
# ComponentsManager
1414

15-
The [`ComponentsManager`] is a model registry and management system for Modular Diffusers. It adds and tracks models, stores useful metadata (model size, device placement, adapters), prevents duplicate model instances, and supports offloading.
15+
The [`ComponentsManager`] is a model registry and management system for Modular Diffusers. It adds and tracks models, stores useful metadata (model size, device placement, adapters), and supports offloading.
1616

1717
This guide will show you how to use [`ComponentsManager`] to manage components and device memory.
1818

19-
## Add a component
19+
## Connect to a pipeline
2020

21-
The [`ComponentsManager`] should be created alongside a [`ModularPipeline`] in either [`~ModularPipeline.from_pretrained`] or [`~ModularPipelineBlocks.init_pipeline`].
21+
Create a [`ComponentsManager`] and pass it to a [`ModularPipeline`] with either [`~ModularPipeline.from_pretrained`] or [`~ModularPipelineBlocks.init_pipeline`].
2222

23-
> [!TIP]
24-
> The `collection` parameter is optional but makes it easier to organize and manage components.
2523

2624
<hfoptions id="create">
2725
<hfoption id="from_pretrained">
2826

2927
```py
3028
from diffusers import ModularPipeline, ComponentsManager
29+
import torch
3130

32-
comp = ComponentsManager()
33-
pipe = ModularPipeline.from_pretrained("YiYiXu/modular-demo-auto", components_manager=comp, collection="test1")
31+
manager = ComponentsManager()
32+
pipe = ModularPipeline.from_pretrained("Tongyi-MAI/Z-Image-Turbo", components_manager=manager)
33+
pipe.load_components(torch_dtype=torch.bfloat16)
3434
```
3535

3636
</hfoption>
3737
<hfoption id="init_pipeline">
3838

3939
```py
40-
from diffusers import ComponentsManager
41-
from diffusers.modular_pipelines import SequentialPipelineBlocks
42-
from diffusers.modular_pipelines.stable_diffusion_xl import TEXT2IMAGE_BLOCKS
43-
44-
t2i_blocks = SequentialPipelineBlocks.from_blocks_dict(TEXT2IMAGE_BLOCKS)
45-
46-
modular_repo_id = "YiYiXu/modular-loader-t2i-0704"
47-
components = ComponentsManager()
48-
t2i_pipeline = t2i_blocks.init_pipeline(modular_repo_id, components_manager=components)
40+
from diffusers import ModularPipelineBlocks, ComponentsManager
41+
import torch
42+
manager = ComponentsManager()
43+
blocks = ModularPipelineBlocks.from_pretrained("diffusers/Florence2-image-Annotator", trust_remote_code=True)
44+
pipe= blocks.init_pipeline(components_manager=manager)
45+
pipe.load_components(torch_dtype=torch.bfloat16)
4946
```
5047

5148
</hfoption>
5249
</hfoptions>
5350

54-
Components are only loaded and registered when using [`~ModularPipeline.load_components`] or [`~ModularPipeline.load_components`]. The example below uses [`~ModularPipeline.load_components`] to create a second pipeline that reuses all the components from the first one, and assigns it to a different collection
55-
56-
```py
57-
pipe.load_components()
58-
pipe2 = ModularPipeline.from_pretrained("YiYiXu/modular-demo-auto", components_manager=comp, collection="test2")
59-
```
60-
61-
Use the [`~ModularPipeline.null_component_names`] property to identify any components that need to be loaded, retrieve them with [`~ComponentsManager.get_components_by_names`], and then call [`~ModularPipeline.update_components`] to add the missing components.
62-
63-
```py
64-
pipe2.null_component_names
65-
['text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'image_encoder', 'unet', 'vae', 'scheduler', 'controlnet']
66-
67-
comp_dict = comp.get_components_by_names(names=pipe2.null_component_names)
68-
pipe2.update_components(**comp_dict)
69-
```
70-
71-
To add individual components, use the [`~ComponentsManager.add`] method. This registers a component with a unique id.
72-
73-
```py
74-
from diffusers import AutoModel
75-
76-
text_encoder = AutoModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder")
77-
component_id = comp.add("text_encoder", text_encoder)
78-
comp
79-
```
80-
81-
Use [`~ComponentsManager.remove`] to remove a component using their id.
82-
83-
```py
84-
comp.remove("text_encoder_139917733042864")
85-
```
86-
87-
## Retrieve a component
88-
89-
The [`ComponentsManager`] provides several methods to retrieve registered components.
90-
91-
### get_one
92-
93-
The [`~ComponentsManager.get_one`] method returns a single component and supports pattern matching for the `name` parameter. If multiple components match, [`~ComponentsManager.get_one`] returns an error.
94-
95-
| Pattern | Example | Description |
96-
|-------------|----------------------------------|-------------------------------------------|
97-
| exact | `comp.get_one(name="unet")` | exact name match |
98-
| wildcard | `comp.get_one(name="unet*")` | names starting with "unet" |
99-
| exclusion | `comp.get_one(name="!unet")` | exclude components named "unet" |
100-
| or | `comp.get_one(name="unet&#124;vae")` | name is "unet" or "vae" |
101-
102-
[`~ComponentsManager.get_one`] also filters components by the `collection` argument or `load_id` argument.
103-
104-
```py
105-
comp.get_one(name="unet", collection="sdxl")
106-
```
107-
108-
### get_components_by_names
109-
110-
The [`~ComponentsManager.get_components_by_names`] method accepts a list of names and returns a dictionary mapping names to components. This is especially useful with [`ModularPipeline`] since they provide lists of required component names and the returned dictionary can be passed directly to [`~ModularPipeline.update_components`].
111-
112-
```py
113-
component_dict = comp.get_components_by_names(names=["text_encoder", "unet", "vae"])
114-
{"text_encoder": component1, "unet": component2, "vae": component3}
115-
```
116-
117-
## Duplicate detection
118-
119-
It is recommended to load model components with [`ComponentSpec`] to assign components with a unique id that encodes their loading parameters. This allows [`ComponentsManager`] to automatically detect and prevent duplicate model instances even when different objects represent the same underlying checkpoint.
120-
121-
```py
122-
from diffusers import ComponentSpec, ComponentsManager
123-
from transformers import CLIPTextModel
124-
125-
comp = ComponentsManager()
126-
127-
# Create ComponentSpec for the first text encoder
128-
spec = ComponentSpec(name="text_encoder", repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder", type_hint=AutoModel)
129-
# Create ComponentSpec for a duplicate text encoder (it is same checkpoint, from the same repo/subfolder)
130-
spec_duplicated = ComponentSpec(name="text_encoder_duplicated", repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder", type_hint=CLIPTextModel)
131-
132-
# Load and add both components - the manager will detect they're the same model
133-
comp.add("text_encoder", spec.load())
134-
comp.add("text_encoder_duplicated", spec_duplicated.load())
135-
```
136-
137-
This returns a warning with instructions for removing the duplicate.
51+
Components loaded by the pipeline are automatically registered in the manager. You can inspect them right away.
13852

139-
```py
140-
ComponentsManager: adding component 'text_encoder_duplicated_139917580682672', but it has duplicate load_id 'stabilityai/stable-diffusion-xl-base-1.0|text_encoder|null|null' with existing components: text_encoder_139918506246832. To remove a duplicate, call `components_manager.remove('<component_id>')`.
141-
'text_encoder_duplicated_139917580682672'
142-
```
143-
144-
You could also add a component without using [`ComponentSpec`] and duplicate detection still works in most cases even if you're adding the same component under a different name.
145-
146-
However, [`ComponentManager`] can't detect duplicates when you load the same component into different objects. In this case, you should load a model with [`ComponentSpec`].
147-
148-
```py
149-
text_encoder_2 = AutoModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder")
150-
comp.add("text_encoder", text_encoder_2)
151-
'text_encoder_139917732983664'
152-
```
53+
## Inspect components
15354

154-
## Collections
55+
Print the [`ComponentsManager`] to see all registered components, including their class, device placement, dtype, memory size, and load ID.
15556

156-
Collections are labels assigned to components for better organization and management. Add a component to a collection with the `collection` argument in [`~ComponentsManager.add`].
157-
158-
Only one component per name is allowed in each collection. Adding a second component with the same name automatically removes the first component.
57+
The output below corresponds to the `from_pretrained` example above.
15958

16059
```py
161-
from diffusers import ComponentSpec, ComponentsManager
162-
163-
comp = ComponentsManager()
164-
# Create ComponentSpec for the first UNet
165-
spec = ComponentSpec(name="unet", repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", type_hint=AutoModel)
166-
# Create ComponentSpec for a different UNet
167-
spec2 = ComponentSpec(name="unet", repo="RunDiffusion/Juggernaut-XL-v9", subfolder="unet", type_hint=AutoModel, variant="fp16")
168-
169-
# Add both UNets to the same collection - the second one will replace the first
170-
comp.add("unet", spec.load(), collection="sdxl")
171-
comp.add("unet", spec2.load(), collection="sdxl")
60+
Components:
61+
=============================================================================================================================
62+
Models:
63+
-----------------------------------------------------------------------------------------------------------------------------
64+
Name_ID | Class | Device: act(exec) | Dtype | Size (GB) | Load ID
65+
-----------------------------------------------------------------------------------------------------------------------------
66+
text_encoder_140458257514752 | Qwen3Model | cpu | torch.bfloat16 | 7.49 | Tongyi-MAI/Z-Image-Turbo|text_encoder|null|null
67+
vae_140458257515376 | AutoencoderKL | cpu | torch.bfloat16 | 0.16 | Tongyi-MAI/Z-Image-Turbo|vae|null|null
68+
transformer_140458257515616 | ZImageTransformer2DModel | cpu | torch.bfloat16 | 11.46 | Tongyi-MAI/Z-Image-Turbo|transformer|null|null
69+
-----------------------------------------------------------------------------------------------------------------------------
70+
71+
Other Components:
72+
-----------------------------------------------------------------------------------------------------------------------------
73+
ID | Class | Collection
74+
-----------------------------------------------------------------------------------------------------------------------------
75+
scheduler_140461023555264 | FlowMatchEulerDiscreteScheduler | N/A
76+
tokenizer_140458256346432 | Qwen2Tokenizer | N/A
77+
-----------------------------------------------------------------------------------------------------------------------------
17278
```
17379

174-
This makes it convenient to work with node-based systems because you can:
175-
176-
- Mark all models as loaded from one node with the `collection` label.
177-
- Automatically replace models when new checkpoints are loaded under the same name.
178-
- Batch delete all models in a collection when a node is removed.
80+
The table shows models (with device, dtype, and memory info) separately from other components like schedulers and tokenizers. If any models have LoRA adapters, IP-Adapters, or quantization applied, that information is displayed in an additional section at the bottom.
17981

18082
## Offloading
18183

18284
The [`~ComponentsManager.enable_auto_cpu_offload`] method is a global offloading strategy that works across all models regardless of which pipeline is using them. Once enabled, you don't need to worry about device placement if you add or remove components.
18385

18486
```py
185-
comp.enable_auto_cpu_offload(device="cuda")
87+
manager.enable_auto_cpu_offload(device="cuda")
18688
```
18789

18890
All models begin on the CPU and [`ComponentsManager`] moves them to the appropriate device right before they're needed, and moves other models back to the CPU when GPU memory is low.
18991

190-
You can set your own rules for which models to offload first.
92+
Call [`~ComponentsManager.disable_auto_cpu_offload`] to disable offloading.
93+
94+
```py
95+
manager.disable_auto_cpu_offload()
96+
```

0 commit comments

Comments
 (0)