Skip to content

Commit 7bb9918

Browse files
committed
[docs] torch_dtype map
1 parent d8c617c commit 7bb9918

File tree

1 file changed

+17
-0
lines changed

1 file changed

+17
-0
lines changed

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,23 @@ Use the Space below to gauge a pipeline's memory requirements before you downloa
9595
></iframe>
9696
</div>
9797

98+
### Specifying Component-Specific Data Types
99+
100+
You can customize the data types for individual sub-models by passing a dictionary to the `torch_dtype` parameter. This allows you to load different components of a pipeline in different floating point precisions. For instance, if you want to load the transformer with `torch.bfloat16` and all other components with `torch.float16`, you can pass a dictionary mapping:
101+
102+
```python
103+
from diffusers import HunyuanVideoPipeline
104+
import torch
105+
106+
pipe = HunyuanVideoPipeline.from_pretrained(
107+
"hunyuanvideo-community/HunyuanVideo",
108+
torch_dtype={'transformer': torch.bfloat16, 'default': torch.float16},
109+
)
110+
print(pipe.transformer.dtype, pipe.vae.dtype) # (torch.bfloat16, torch.float16)
111+
```
112+
113+
If a component is not explicitly specified in the dictionary and no `default` is provided, it will be loaded with `torch.float32`.
114+
98115
### Local pipeline
99116

100117
To load a pipeline locally, use [git-lfs](https://git-lfs.github.com/) to manually download a checkpoint to your local disk.

0 commit comments

Comments
 (0)