Skip to content

Commit 8c938fb

Browse files
a-r-r-o-wstevhliu
andauthored
[docs] Add a note of _keep_in_fp32_modules (#11851)
* update * Update docs/source/en/using-diffusers/schedulers.md Co-authored-by: Steven Liu <[email protected]> * Update schedulers.md --------- Co-authored-by: Steven Liu <[email protected]>
1 parent f864a9a commit 8c938fb

File tree

1 file changed

+12
-0
lines changed

1 file changed

+12
-0
lines changed

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -242,3 +242,15 @@ unet = UNet2DConditionModel.from_pretrained(
242242
)
243243
unet.save_pretrained("./local-unet", variant="non_ema")
244244
```
245+
246+
Use the `torch_dtype` argument in [`~ModelMixin.from_pretrained`] to specify the dtype to load a model in.
247+
248+
```py
249+
from diffusers import AutoModel
250+
251+
unet = AutoModel.from_pretrained(
252+
"stabilityai/stable-diffusion-xl-base-1.0", subfolder="unet", torch_dtype=torch.float16
253+
)
254+
```
255+
256+
You can also use the [torch.Tensor.to](https://docs.pytorch.org/docs/stable/generated/torch.Tensor.to.html) method to convert to the specified dtype on the fly. It converts *all* weights unlike the `torch_dtype` argument that respects the `_keep_in_fp32_modules`. This is important for models whose layers must remain in fp32 for numerical stability and best generation quality (see example [here](https://github.com/huggingface/diffusers/blob/f864a9a352fa4a220d860bfdd1782e3e5af96382/src/diffusers/models/transformers/transformer_wan.py#L374)).

0 commit comments

Comments
 (0)