Diffusers documentation
AutoModel
AutoModel
The AutoModel class automatically detects and loads the correct model class (UNet, transformer, VAE) from a config.json file. You don’t need to know the specific model class name ahead of time. It supports data types and device placement, and works across model types and libraries.
The example below loads a transformer from Diffusers and a text encoder from Transformers. Use the subfolder parameter to specify where to load the config.json file from.
import torch
from diffusers import AutoModel, DiffusionPipeline
transformer = AutoModel.from_pretrained(
"Qwen/Qwen-Image", subfolder="transformer", torch_dtype=torch.bfloat16, device_map="cuda"
)
text_encoder = AutoModel.from_pretrained(
"Qwen/Qwen-Image", subfolder="text_encoder", torch_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.
import torch
from diffusers import AutoModel
transformer = AutoModel.from_pretrained(
"custom/custom-transformer-model", trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="cuda"
)If the custom model inherits from the ModelMixin class, it gets access to the same features as Diffusers model classes, like regional compilation and group offloading.
Update on GitHubLearn more about implementing custom models in the Community components guide.