AutoModel

Lớp AutoModel tự động phát hiện và nạp đúng lớp mô hình (UNet, transformer, VAE) từ tệp config.json. Bạn không cần biết trước tên lớp mô hình cụ thể. Nó hỗ trợ kiểu dữ liệu và vị trí thiết bị, đồng thời hoạt động xuyên suốt các loại mô hình và thư viện.

Ví dụ dưới đây tải một transformer từ Diffusers và một text encoder từ Transformers. Sử dụng tham số subfolder để chỉ định vị trí tải tệp 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"
)

Mô hình tùy chỉnh

AutoModel cũng tải các mô hình từ Hub không có sẵn trong Diffusers. Đặt trust_remote_code=True trong AutoModel.from_pretrained() để tải các mô hình tùy chỉnh.

Một kho lưu trữ model tùy chỉnh cần một module Python chứa class của model, và một config.json có mục auto_map ánh xạ "AutoModel" tới "module_file.ClassName".

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

config.json bao gồm trường auto_map trỏ đến lớp tùy chỉnh.

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

Sau đó, tải nó với 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"
)

Một ví dụ thực tế là Overworld/Waypoint-1-Small chứa một lớp WorldModel tùy chỉnh trải dài qua một số module trong thư mục con transformer của nó.

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

Nếu mô hình tùy chỉnh kế thừa từ lớp ModelMixin, nó sẽ được truy cập các tính năng giống như các lớp mô hình của Diffusers, chẳng hạn như biên dịch khu vực và chuyển nhóm.

Như một biện pháp phòng ngừa khi sử dụng trust_remote_code=True, hãy truyền một commit hash vào tham số revision trong AutoModel.from_pretrained() để đảm bảo rằng mã nguồn không bị cập nhật với mã độc mới (trừ khi bạn hoàn toàn tin tưởng các chủ sở hữu mô hình).

Đã sao chép

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

Lưu các mô hình tùy chỉnh

Sử dụng register_for_auto_class() để tự động thêm mục auto_map vào config.json khi lưu. Điều này giúp tránh phải chỉnh sửa thủ công tệp cấu hình.

# 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 được lưu sẽ bao gồm trường auto_map.

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

Các bước tiếp theo

Tìm hiểu thêm về việc triển khai các mô hình tùy chỉnh trong hướng dẫn Thành phần cộng đồng.