AutoModel

AutoModel 클래스는 config.json 파일에서 올바른 모델 클래스(UNet, transformer, VAE)를 자동으로 감지하고 로드합니다. 특정 모델 클래스 이름을 미리 알 필요가 없습니다. 데이터 타입과 디바이스 배치를 지원하며, 다양한 모델 유형과 라이브러리에서 작동합니다.

아래 예시에서는 Diffusers에서 transformer를, Transformers에서 text encoder를 로드합니다. config.json 파일을 로드할 위치를 지정하려면 subfolder 매개변수를 사용하세요.

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

커스텀 모델

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.

커스텀 모델 저장소에는 모델 클래스가 포함된 Python 모듈과 "AutoModel"을 "module_file.ClassName"으로 매핑하는 auto_map 항목이 있는 config.json이 필요합니다.

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

config.json에는 커스텀 클래스를 가리키는 auto_map 필드가 포함되어 있습니다.

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

그 다음 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"
)

실제 예시로, Overworld/Waypoint-1-Small은 transformer 하위 폴더의 여러 모듈에 걸쳐 사용자 정의 WorldModel 클래스를 호스팅합니다.

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

커스텀 모델이 ModelMixin 클래스를 상속하면 regional compilation 및 group offloading과 같은 Diffusers 모델 클래스와 동일한 기능에 접근할 수 있습니다.

trust_remote_code=True를 사용할 때, 모델 소유자를 완전히 신뢰하지 않는 한 코드가 새로운 악성 코드로 업데이트되지 않았는지 확인하기 위해 AutoModel.from_pretrained()의 revision 인수에 커밋 해시를 전달하는 것이 좋습니다.
transformer = AutoModel.from_pretrained(
    "Overworld/Waypoint-1-Small", subfolder="transformer", trust_remote_code=True, revision="a3d8cb2"
)

커스텀 모델 저장

register_for_auto_class()를 사용하면 저장 시 config.json에 auto_map 항목이 자동으로 추가됩니다. 이를 통해 설정 파일을 수동으로 편집할 필요가 없습니다.

# 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에는 auto_map 필드가 포함됩니다.

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

다음 단계

커뮤니티 컴포넌트 가이드에서 커스텀 모델 구현에 대해 자세히 알아보세요.

Update on GitHub