Buckets:

mishig's picture
|
download
raw
5.11 kB

ComponentsManager

ComponentsManager 仍处于积极开发阶段,其 API 可能会发生变化。

ComponentsManager 是 Modular Diffusers 的模型注册与管理系统。它负责添加和跟踪模型,存储有用的元数据(模型大小、设备放置、适配器),并支持卸载。

本指南将向您展示如何使用 ComponentsManager 来管理组件和设备内存。

连接到管道

Create a ComponentsManager and pass it to a ModularPipeline with either from_pretrained() or init_pipeline().

from diffusers import ModularPipeline, ComponentsManager
import torch

manager = ComponentsManager()
pipe = ModularPipeline.from_pretrained("Tongyi-MAI/Z-Image-Turbo", components_manager=manager)
pipe.load_components(dtype=torch.bfloat16)
from diffusers import ModularPipelineBlocks, ComponentsManager
import torch
manager = ComponentsManager()
blocks = ModularPipelineBlocks.from_pretrained("diffusers/Florence2-image-Annotator", trust_remote_code=True)
pipe= blocks.init_pipeline(components_manager=manager)
pipe.load_components(dtype=torch.bfloat16)

由 pipeline 加载的组件会自动注册到管理器中。您可以立即查看它们。

检查组件

打印 ComponentsManager 以查看所有已注册的组件,包括其类、设备放置、dtype、内存大小和加载 ID。

以下输出对应于上述 from_pretrained 示例。

Components:
=============================================================================================================================
Models:
-----------------------------------------------------------------------------------------------------------------------------
Name_ID                      | Class                    | Device: act(exec) | Dtype          | Size (GB) | Load ID
-----------------------------------------------------------------------------------------------------------------------------
text_encoder_140458257514752 | Qwen3Model               | cpu               | torch.bfloat16 | 7.49      | Tongyi-MAI/Z-Image-Turbo|text_encoder|null|null
vae_140458257515376          | AutoencoderKL            | cpu               | torch.bfloat16 | 0.16      | Tongyi-MAI/Z-Image-Turbo|vae|null|null
transformer_140458257515616  | ZImageTransformer2DModel | cpu               | torch.bfloat16 | 11.46     | Tongyi-MAI/Z-Image-Turbo|transformer|null|null
-----------------------------------------------------------------------------------------------------------------------------

Other Components:
-----------------------------------------------------------------------------------------------------------------------------
ID                           | Class                           | Collection
-----------------------------------------------------------------------------------------------------------------------------
scheduler_140461023555264    | FlowMatchEulerDiscreteScheduler | N/A
tokenizer_140458256346432    | Qwen2Tokenizer                  | N/A
-----------------------------------------------------------------------------------------------------------------------------

该表格将模型(包含设备、dtype 和内存信息)与其他组件(如调度器和分词器)分开显示。如果任何模型应用了 LoRA 适配器、IP-Adapter 或量化,相关信息将在底部的附加部分中显示。

卸载

enable_auto_cpu_offload() 方法是一种全局卸载策略,可跨所有模型生效,无论它们被哪个 pipeline 使用。启用后,即使添加或移除组件,您也无需担心设备放置问题。

manager.enable_auto_cpu_offload(device="cuda")  # or "mps", "xpu", "cpu"

所有模型最初都位于 CPU 上,ComponentsManager 会在模型即将被使用之前将其移动到合适的设备上,并在 GPU 内存不足时将其他模型移回 CPU。

调用 disable_auto_cpu_offload() 以禁用卸载。

manager.disable_auto_cpu_offload()

Xet Storage Details

Size:
5.11 kB
·
Xet hash:
7d148f6c5e0bdbfa9c76ad65f4f605e165ede12d925dc9d1ebbf651ed541e65f

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.