AutoModel

AutoModel クラスは、config.json ファイルから正しいモデルクラス（UNet、transformer、VAE）を自動的に検出し、読み込みます。事前に具体的なモデルクラス名を知る必要はありません。データ型とデバイスの配置をサポートしており、モデルタイプやライブラリをまたいで動作します。

以下の例では、Diffusersからtransformerを読み込み、Transformersからテキストエンコーダーを読み込みます。subfolderパラメータを使用して、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"
)

カスタムモデル

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