AutoModel

Класс AutoModel автоматически определяет и загружает правильный класс модели (UNet, transformer, VAE) из файла config.json. Вам не нужно заранее знать конкретное имя класса модели. Он поддерживает типы данных и размещение на устройствах, а также работает с различными типами моделей и библиотеками.

Пример ниже загружает трансформер из Diffusers и текстовый энкодер из 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 содержит пользовательский класс WorldModel, распределённый по нескольким модулям в подпапке transformer.

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, например региональная компиляция и групповая разгрузка.

В качестве меры предосторожности при использовании trust_remote_code=True передайте хеш коммита в аргумент revision в AutoModel.from_pretrained(), чтобы убедиться, что код не был обновлён с добавлением нового вредоносного кода (если вы полностью доверяете владельцам модели).

Скопировано

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"
  }
}

Дальнейшие шаги

Подробнее о реализации пользовательских моделей см. в руководстве Компоненты сообщества.