"vscode:/vscode.git/clone" did not exist on "1216a3b122ba7f13062619c33254d2e737f379dc"
components_manager.md 8.78 KB
Newer Older
YiYi Xu's avatar
YiYi Xu committed
1
2
3
4
5
6
7
8
9
10
11
12
<!--Copyright 2025 The HuggingFace Team. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
specific language governing permissions and limitations under the License.
-->

Steven Liu's avatar
Steven Liu committed
13
# ComponentsManager
YiYi Xu's avatar
YiYi Xu committed
14

Steven Liu's avatar
Steven Liu committed
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.
YiYi Xu's avatar
YiYi Xu committed
16

Steven Liu's avatar
Steven Liu committed
17
This guide will show you how to use [`ComponentsManager`] to manage components and device memory.
YiYi Xu's avatar
YiYi Xu committed
18

Steven Liu's avatar
Steven Liu committed
19
## Add a component
YiYi Xu's avatar
YiYi Xu committed
20

Steven Liu's avatar
Steven Liu committed
21
The [`ComponentsManager`] should be created alongside a [`ModularPipeline`] in either [`~ModularPipeline.from_pretrained`] or [`~ModularPipelineBlocks.init_pipeline`].
YiYi Xu's avatar
YiYi Xu committed
22

Steven Liu's avatar
Steven Liu committed
23
24
> [!TIP]
> The `collection` parameter is optional but makes it easier to organize and manage components.
YiYi Xu's avatar
YiYi Xu committed
25

Steven Liu's avatar
Steven Liu committed
26
27
<hfoptions id="create">
<hfoption id="from_pretrained">
YiYi Xu's avatar
YiYi Xu committed
28
29

```py
Steven Liu's avatar
Steven Liu committed
30
31
from diffusers import ModularPipeline, ComponentsManager

YiYi Xu's avatar
YiYi Xu committed
32
comp = ComponentsManager()
Steven Liu's avatar
Steven Liu committed
33
pipe = ModularPipeline.from_pretrained("YiYiXu/modular-demo-auto", components_manager=comp, collection="test1")
YiYi Xu's avatar
YiYi Xu committed
34
35
```

Steven Liu's avatar
Steven Liu committed
36
37
</hfoption>
<hfoption id="init_pipeline">
YiYi Xu's avatar
YiYi Xu committed
38
39

```py
Steven Liu's avatar
Steven Liu committed
40
41
42
from diffusers import ComponentsManager
from diffusers.modular_pipelines import SequentialPipelineBlocks
from diffusers.modular_pipelines.stable_diffusion_xl import TEXT2IMAGE_BLOCKS
YiYi Xu's avatar
YiYi Xu committed
43

Steven Liu's avatar
Steven Liu committed
44
t2i_blocks = SequentialPipelineBlocks.from_blocks_dict(TEXT2IMAGE_BLOCKS)
YiYi Xu's avatar
YiYi Xu committed
45

Steven Liu's avatar
Steven Liu committed
46
47
48
modular_repo_id = "YiYiXu/modular-loader-t2i-0704"
components = ComponentsManager()
t2i_pipeline = t2i_blocks.init_pipeline(modular_repo_id, components_manager=components)
YiYi Xu's avatar
YiYi Xu committed
49
50
```

Steven Liu's avatar
Steven Liu committed
51
52
53
</hfoption>
</hfoptions>

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
YiYi Xu's avatar
YiYi Xu committed
55
56

```py
57
pipe.load_components()
Steven Liu's avatar
Steven Liu committed
58
pipe2 = ModularPipeline.from_pretrained("YiYiXu/modular-demo-auto", components_manager=comp, collection="test2")
YiYi Xu's avatar
YiYi Xu committed
59
60
```

Steven Liu's avatar
Steven Liu committed
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.
YiYi Xu's avatar
YiYi Xu committed
62

Steven Liu's avatar
Steven Liu committed
63
64
65
```py
pipe2.null_component_names 
['text_encoder', 'text_encoder_2', 'tokenizer', 'tokenizer_2', 'image_encoder', 'unet', 'vae', 'scheduler', 'controlnet']
YiYi Xu's avatar
YiYi Xu committed
66

Steven Liu's avatar
Steven Liu committed
67
68
69
70
71
comp_dict = comp.get_components_by_names(names=pipe2.null_component_names)
pipe2.update_components(**comp_dict)
```

To add individual components, use the [`~ComponentsManager.add`] method. This registers a component with a unique id.
YiYi Xu's avatar
YiYi Xu committed
72
73

```py
Steven Liu's avatar
Steven Liu committed
74
75
76
77
78
from diffusers import AutoModel

text_encoder = AutoModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder")
component_id = comp.add("text_encoder", text_encoder)
comp
YiYi Xu's avatar
YiYi Xu committed
79
80
```

Steven Liu's avatar
Steven Liu committed
81
Use [`~ComponentsManager.remove`] to remove a component using their id.
YiYi Xu's avatar
YiYi Xu committed
82
83

```py
Steven Liu's avatar
Steven Liu committed
84
comp.remove("text_encoder_139917733042864")
YiYi Xu's avatar
YiYi Xu committed
85
86
```

Steven Liu's avatar
Steven Liu committed
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
## Retrieve a component

The [`ComponentsManager`] provides several methods to retrieve registered components.

### get_one

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.

| Pattern     | Example                          | Description                               |
|-------------|----------------------------------|-------------------------------------------|
| exact       | `comp.get_one(name="unet")`      | exact name match                          |
| wildcard    | `comp.get_one(name="unet*")`     | names starting with "unet"                |
| exclusion   | `comp.get_one(name="!unet")`     | exclude components named "unet"           |
| or          | `comp.get_one(name="unet&#124;vae")`  | name is "unet" or "vae"                   |

[`~ComponentsManager.get_one`] also filters components by the `collection` argument or `load_id` argument.
YiYi Xu's avatar
YiYi Xu committed
103
104

```py
Steven Liu's avatar
Steven Liu committed
105
comp.get_one(name="unet", collection="sdxl")
YiYi Xu's avatar
YiYi Xu committed
106
107
```

Steven Liu's avatar
Steven Liu committed
108
109
110
### get_components_by_names

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`].
YiYi Xu's avatar
YiYi Xu committed
111
112

```py
Steven Liu's avatar
Steven Liu committed
113
114
component_dict = comp.get_components_by_names(names=["text_encoder", "unet", "vae"])
{"text_encoder": component1, "unet": component2, "vae": component3}
YiYi Xu's avatar
YiYi Xu committed
115
116
```

Steven Liu's avatar
Steven Liu committed
117
118
119
## Duplicate detection

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.
YiYi Xu's avatar
YiYi Xu committed
120
121
122
123

```py
from diffusers import ComponentSpec, ComponentsManager
from transformers import CLIPTextModel
Steven Liu's avatar
Steven Liu committed
124

YiYi Xu's avatar
YiYi Xu committed
125
126
127
128
comp = ComponentsManager()

# Create ComponentSpec for the first text encoder
spec = ComponentSpec(name="text_encoder", repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder", type_hint=AutoModel)
Steven Liu's avatar
Steven Liu committed
129
# Create ComponentSpec for a duplicate text encoder (it is same checkpoint, from the same repo/subfolder)
YiYi Xu's avatar
YiYi Xu committed
130
131
132
133
134
135
136
spec_duplicated = ComponentSpec(name="text_encoder_duplicated", repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder", type_hint=CLIPTextModel)

# Load and add both components - the manager will detect they're the same model
comp.add("text_encoder", spec.load())
comp.add("text_encoder_duplicated", spec_duplicated.load())
```

Steven Liu's avatar
Steven Liu committed
137
This returns a warning with instructions for removing the duplicate.
YiYi Xu's avatar
YiYi Xu committed
138

Steven Liu's avatar
Steven Liu committed
139
```py
YiYi Xu's avatar
YiYi Xu committed
140
141
142
143
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>')`.
'text_encoder_duplicated_139917580682672'
```

Steven Liu's avatar
Steven Liu committed
144
145
146
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.

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`].
YiYi Xu's avatar
YiYi Xu committed
147
148

```py
Steven Liu's avatar
Steven Liu committed
149
150
151
text_encoder_2 = AutoModel.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="text_encoder")
comp.add("text_encoder", text_encoder_2)
'text_encoder_139917732983664'
YiYi Xu's avatar
YiYi Xu committed
152
153
154
155
```

## Collections

Steven Liu's avatar
Steven Liu committed
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`].
YiYi Xu's avatar
YiYi Xu committed
157

Steven Liu's avatar
Steven Liu committed
158
Only one component per name is allowed in each collection. Adding a second component with the same name automatically removes the first component.
YiYi Xu's avatar
YiYi Xu committed
159
160

```py
Steven Liu's avatar
Steven Liu committed
161
162
from diffusers import ComponentSpec, ComponentsManager

YiYi Xu's avatar
YiYi Xu committed
163
comp = ComponentsManager()
Steven Liu's avatar
Steven Liu committed
164
# Create ComponentSpec for the first UNet
YiYi Xu's avatar
YiYi Xu committed
165
spec = ComponentSpec(name="unet", repo="stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", type_hint=AutoModel)
Steven Liu's avatar
Steven Liu committed
166
# Create ComponentSpec for a different UNet
YiYi Xu's avatar
YiYi Xu committed
167
168
169
170
171
172
173
spec2 = ComponentSpec(name="unet", repo="RunDiffusion/Juggernaut-XL-v9", subfolder="unet", type_hint=AutoModel, variant="fp16")

# Add both UNets to the same collection - the second one will replace the first
comp.add("unet", spec.load(), collection="sdxl")
comp.add("unet", spec2.load(), collection="sdxl")
```

Steven Liu's avatar
Steven Liu committed
174
This makes it convenient to work with node-based systems because you can:
YiYi Xu's avatar
YiYi Xu committed
175

Steven Liu's avatar
Steven Liu committed
176
177
178
- Mark all models as loaded from one node with the `collection` label.
- Automatically replace models when new checkpoints are loaded under the same name.
- Batch delete all models in a collection when a node is removed.
YiYi Xu's avatar
YiYi Xu committed
179

Steven Liu's avatar
Steven Liu committed
180
## Offloading
YiYi Xu's avatar
YiYi Xu committed
181

Steven Liu's avatar
Steven Liu committed
182
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.
YiYi Xu's avatar
YiYi Xu committed
183
184
185
186
187

```py
comp.enable_auto_cpu_offload(device="cuda")
```

Steven Liu's avatar
Steven Liu committed
188
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.
YiYi Xu's avatar
YiYi Xu committed
189

190
You can set your own rules for which models to offload first.