AutoModel

Kelas AutoModel secara otomatis mendeteksi dan memuat kelas model yang benar (UNet, transformer, VAE) dari file config.json. Anda tidak perlu mengetahui nama kelas model spesifik sebelumnya. Kelas ini mendukung tipe data dan penempatan perangkat, serta berfungsi di berbagai jenis model dan pustaka.

Contoh di bawah ini memuat transformer dari Diffusers dan text encoder dari Transformers. Gunakan parameter subfolder untuk menentukan dari mana file config.json dimuat.

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

Model kustom

AutoModel juga memuat model dari Hub yang tidak disertakan dalam Diffusers. Atur trust_remote_code=True di AutoModel.from_pretrained() untuk memuat model kustom.

Repositori model kustom memerlukan modul Python dengan kelas model, serta config.json dengan entri auto_map yang memetakan "AutoModel" ke "module_file.ClassName".

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

config.json mencakup bidang auto_map yang menunjuk ke kelas kustom.

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

Lalu muat dengan 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"
)

Sebagai contoh nyata, Overworld/Waypoint-1-Small menyediakan kelas WorldModel kustom di beberapa modul dalam subfolder transformer-nya.

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

Jika model kustom mewarisi dari kelas ModelMixin, model tersebut mendapatkan akses ke fitur yang sama seperti kelas model Diffusers, seperti regional compilation dan group offloading.

Sebagai tindakan pencegahan dengan trust_remote_code=True, berikan commit hash ke argumen revision di AutoModel.from_pretrained() untuk memastikan kode tidak diperbarui dengan kode berbahaya baru (kecuali Anda sepenuhnya mempercayai pemilik model).

Disalin

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

Menyimpan model kustom

Gunakan register_for_auto_class() untuk menambahkan entri auto_map ke config.json secara otomatis saat menyimpan. Ini menghindari kebutuhan untuk mengedit file konfigurasi secara manual.

# 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 yang disimpan akan mencakup bidang auto_map.

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

Langkah selanjutnya

Pelajari lebih lanjut tentang implementasi model kustom di panduan Komponen komunitas.