AutoModel

La clase AutoModel detecta y carga automáticamente la clase de modelo correcta (UNet, transformer, VAE) a partir de un archivo config.json. No es necesario conocer el nombre específico de la clase de modelo de antemano. Admite tipos de datos y colocación en dispositivos, y funciona entre diferentes tipos de modelos y bibliotecas.

El siguiente ejemplo carga un transformer de Diffusers y un codificador de texto de Transformers. Use el parámetro subfolder para especificar desde dónde cargar el archivo 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 también carga modelos del Hub que no están incluidos en Diffusers. Establece trust_remote_code=True en AutoModel.from_pretrained() para cargar modelos personalizados.

Un repositorio de modelos personalizado necesita un módulo de Python con la clase del modelo y un config.json con una entrada auto_map que asocia "AutoModel" con "module_file.ClassName".

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

El config.json incluye el campo auto_map que apunta a la clase personalizada.

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

Luego, cárgalo con 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 ejemplo real, Overworld/Waypoint-1-Small aloja una clase personalizada WorldModel en varios módulos de su subcarpeta 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"
)

Si el modelo personalizado hereda de la clase ModelMixin, obtiene acceso a las mismas características que las clases de modelos de Diffusers, como la compilación regional y el descarga por grupos.

Como precaución con trust_remote_code=True, pase un hash de commit al argumento revision en AutoModel.from_pretrained() para asegurarse de que el código no se haya actualizado con nuevo código malicioso (a menos que confíe por completo en los propietarios del modelo).

Copiado

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

Guardar modelos personalizados

Use register_for_auto_class() para añadir la entrada auto_map a config.json automáticamente al guardar. Esto evita tener que editar el archivo de configuración manualmente.

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

El config.json guardado incluirá el campo auto_map.

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

Próximos pasos

Obtenga más información sobre la implementación de modelos personalizados en la guía de Componentes de la comunidad.