AutoModel

La classe AutoModel détecte et charge automatiquement la classe de modèle appropriée (UNet, transformer, VAE) à partir d’un fichier config.json. Il n’est pas nécessaire de connaître à l’avance le nom spécifique de la classe de modèle. Elle prend en charge les types de données et le placement sur les appareils, et fonctionne entre les différents types de modèles et les bibliothèques.

L’exemple ci-dessous charge un transformer depuis Diffusers et un encodeur de texte depuis Transformers. Utilisez le paramètre subfolder pour spécifier l’emplacement du fichier config.json à charger.

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

Modèles personnalisés

AutoModel charge également des modèles du Hub qui ne sont pas inclus dans Diffusers. Définissez trust_remote_code=True dans AutoModel.from_pretrained() pour charger des modèles personnalisés.

Un dépôt de modèle personnalisé nécessite un module Python contenant la classe du modèle, ainsi qu'un config.json avec une entrée auto_map qui associe "AutoModel" à "module_file.ClassName".

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

Le config.json contient le champ auto_map pointant vers la classe personnalisée.

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

Puis chargez-le avec 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"
)

Pour un exemple concret, Overworld/Waypoint-1-Small héberge une classe WorldModel personnalisée répartie sur plusieurs modules dans son sous-dossier 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 le modèle personnalisé hérite de la classe ModelMixin, il accède aux mêmes fonctionnalités que les classes de modèles Diffusers, comme la compilation régionale et le déchargement par groupe.

Par précaution avec trust_remote_code=True, passez un hachage de commit à l’argument revision dans AutoModel.from_pretrained() pour vous assurer que le code n’a pas été mis à jour avec du nouveau code malveillant (sauf si vous faites entièrement confiance aux propriétaires du modèle).

Copié

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

Enregistrer des modèles personnalisés

Utilisez register_for_auto_class() pour ajouter automatiquement l’entrée auto_map à config.json lors de l’enregistrement. Cela évite d’avoir à modifier manuellement le fichier de configuration.

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

Le fichier config.json enregistré inclura le champ auto_map.

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

Étapes suivantes

En savoir plus sur l'implémentation de modèles personnalisés dans le guide Composants de la communauté.