AutoModel

A classe AutoModel detecta e carrega automaticamente a classe de modelo correta (UNet, transformer, VAE) a partir de um arquivo config.json. Você não precisa conhecer o nome específico da classe do modelo antecipadamente. Ela suporta tipos de dados e posicionamento de dispositivo, e funciona entre diferentes tipos de modelos e bibliotecas.

O exemplo abaixo carrega um transformer do Diffusers e um codificador de texto do Transformers. Use o parâmetro subfolder para especificar de onde carregar o arquivo 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"
)

Modelos personalizados

AutoModel também carrega modelos do Hub que não estão incluídos no Diffusers. Defina trust_remote_code=True em AutoModel.from_pretrained() para carregar modelos personalizados.

Um repositório de modelo personalizado precisa de um módulo Python com a classe do modelo e de um config.json com uma entrada auto_map que mapeia "AutoModel" para "module_file.ClassName".

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

O config.json inclui o campo auto_map apontando para a classe personalizada.

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

Em seguida, carregue-o com 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"
)

Como exemplo do mundo real, Overworld/Waypoint-1-Small hospeda uma classe WorldModel personalizada em vários módulos em sua subpasta 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"
)

Se o modelo personalizado herda da classe ModelMixin, ele obtém acesso aos mesmos recursos das classes de modelos do Diffusers, como compilação regional e group offloading.

Como medida de precaução ao usar trust_remote_code=True, passe um hash de commit para o argumento revision em AutoModel.from_pretrained() para garantir que o código não tenha sido atualizado com novo código malicioso (a menos que você confie totalmente nos proprietários do modelo).

Copiado

transformer = AutoModel.from_pretrained(
    "Overworld/Waypoint-1-Small", subfolder="transformer", trust_remote_code=True, revision="a3d8cb2"
)

Salvando modelos personalizados

Use register_for_auto_class() para adicionar a entrada auto_map ao config.json automaticamente ao salvar. Isso evita a necessidade de editar manualmente o arquivo de configuração.

# 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")

O config.json salvo incluirá o campo auto_map.

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

Próximas etapas

Saiba mais sobre a implementação de modelos personalizados no guia Componentes da comunidade.