Skip to content

Commit bedc67c

Browse files
authored
[Docs] Add guide for AutoModel with custom code (#13099)
update
1 parent 20efb79 commit bedc67c

File tree

1 file changed

+55
-0
lines changed

1 file changed

+55
-0
lines changed

docs/source/en/using-diffusers/automodel.md

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,31 @@ text_encoder = AutoModel.from_pretrained(
2929
)
3030
```
3131

32+
## Custom models
33+
3234
[`AutoModel`] also loads models from the [Hub](https://huggingface.co/models) that aren't included in Diffusers. Set `trust_remote_code=True` in [`AutoModel.from_pretrained`] to load custom models.
3335

36+
A custom model repository needs a Python module with the model class, and a `config.json` with an `auto_map` entry that maps `"AutoModel"` to `"module_file.ClassName"`.
37+
38+
```
39+
custom/custom-transformer-model/
40+
├── config.json
41+
├── my_model.py
42+
└── diffusion_pytorch_model.safetensors
43+
```
44+
45+
The `config.json` includes the `auto_map` field pointing to the custom class.
46+
47+
```json
48+
{
49+
"auto_map": {
50+
"AutoModel": "my_model.MyCustomModel"
51+
}
52+
}
53+
```
54+
55+
Then load it with `trust_remote_code=True`.
56+
3457
```py
3558
import torch
3659
from diffusers import AutoModel
@@ -40,7 +63,39 @@ transformer = AutoModel.from_pretrained(
4063
)
4164
```
4265

66+
For a real-world example, [Overworld/Waypoint-1-Small](https://huggingface.co/Overworld/Waypoint-1-Small/tree/main/transformer) hosts a custom `WorldModel` class across several modules in its `transformer` subfolder.
67+
68+
```
69+
transformer/
70+
├── config.json # auto_map: "model.WorldModel"
71+
├── model.py
72+
├── attn.py
73+
├── nn.py
74+
├── cache.py
75+
├── quantize.py
76+
├── __init__.py
77+
└── diffusion_pytorch_model.safetensors
78+
```
79+
80+
```py
81+
import torch
82+
from diffusers import AutoModel
83+
84+
transformer = AutoModel.from_pretrained(
85+
"Overworld/Waypoint-1-Small", subfolder="transformer", trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="cuda"
86+
)
87+
```
88+
4389
If the custom model inherits from the [`ModelMixin`] class, it gets access to the same features as Diffusers model classes, like [regional compilation](../optimization/fp16#regional-compilation) and [group offloading](../optimization/memory#group-offloading).
4490

91+
> [!WARNING]
92+
> As a precaution with `trust_remote_code=True`, pass a commit hash to the `revision` argument in [`AutoModel.from_pretrained`] to make sure the code hasn't been updated with new malicious code (unless you fully trust the model owners).
93+
>
94+
> ```py
95+
> transformer = AutoModel.from_pretrained(
96+
> "Overworld/Waypoint-1-Small", subfolder="transformer", trust_remote_code=True, revision="a3d8cb2"
97+
> )
98+
> ```
99+
45100
> [!NOTE]
46101
> Learn more about implementing custom models in the [Community components](../using-diffusers/custom_pipeline_overview#community-components) guide.

0 commit comments

Comments
 (0)