AutoModel

AutoModel 类会自动检测并从 config.json 文件中加载正确的模型类（UNet、transformer、VAE）。您无需提前知道具体的模型类名称。它支持数据类型和设备放置，并适用于不同的模型类型和库。

以下示例从 Diffusers 加载一个 transformer，并从 Transformers 加载一个文本编码器。使用 subfolder 参数来指定从何处加载 config.json 文件。

import torch
from diffusers import AutoModel, DiffusionPipeline

transformer = AutoModel.from_pretrained(
    "Qwen/Qwen-Image", subfolder="transformer", dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)

text_encoder = AutoModel.from_pretrained(
    "Qwen/Qwen-Image", subfolder="text_encoder", dtype=torch.bfloat16, device_map="cuda"
)

自定义模型

AutoModel also loads models from the Hub that aren’t included in Diffusers. Set trust_remote_code=True in AutoModel.from_pretrained() to load custom models.

自定义模型仓库需要一个包含模型类的 Python 模块，以及一个 config.json，其中包含一个 auto_map 条目，将 "AutoModel" 映射到 "module_file.ClassName"。

custom/custom-transformer-model/
├── config.json
├── my_model.py
└── diffusion_pytorch_model.safetensors

config.json 包含指向自定义类的 auto_map 字段。

{
  "auto_map": {
    "AutoModel": "my_model.MyCustomModel"
  }
}

然后使用 trust_remote_code=True 加载它。

import torch
from diffusers import AutoModel

transformer = AutoModel.from_pretrained(
    "custom/custom-transformer-model", trust_remote_code=True, dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)

以一个实际案例为例，Overworld/Waypoint-1-Small 在其 transformer 子文件夹的多个模块中托管了一个自定义的 WorldModel 类。

transformer/
├── config.json          # auto_map: "model.WorldModel"
├── model.py
├── attn.py
├── nn.py
├── cache.py
├── quantize.py
├── __init__.py
└── diffusion_pytorch_model.safetensors

import torch
from diffusers import AutoModel

transformer = AutoModel.from_pretrained(
    "Overworld/Waypoint-1-Small", subfolder="transformer", trust_remote_code=True, dtype=torch.bfloat16, device_map="cuda"  # or "mps", "xpu", "cpu"
)

如果自定义模型继承自 ModelMixin 类，它就能获得与 Diffusers 模型类相同的功能，例如 regional compilation 和 group offloading。

作为使用 trust_remote_code=True 时的预防措施，请在 AutoModel.from_pretrained() 中向 revision 参数传入一个 commit hash，以确保代码未被更新为包含新的恶意代码（除非你完全信任模型所有者）。
transformer = AutoModel.from_pretrained(
    "Overworld/Waypoint-1-Small", subfolder="transformer", trust_remote_code=True, revision="a3d8cb2"
)

保存自定义模型

使用 register_for_auto_class() 可以在保存时自动将 auto_map 条目添加到 config.json 中，从而避免手动编辑配置文件。

# my_model.py
from diffusers import ModelMixin, ConfigMixin

class MyCustomModel(ModelMixin, ConfigMixin):
    ...

MyCustomModel.register_for_auto_class("AutoModel")

model = MyCustomModel(...)
model.save_pretrained("./my_model")

保存的 config.json 将包含 auto_map 字段。

{
  "auto_map": {
    "AutoModel": "my_model.MyCustomModel"
  }
}

后续步骤

有关在社区组件指南中实现自定义模型的更多信息。

Update on GitHub