huggingface / huggingface/diffusers
[Community] Specifying dtype in from_pretrained does not convert the model's datatype
- Dominant language
- Python
- Stars
- 34.5k
- Forks
- 7.3k
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 91
Description
### Describe the bug
When I grab a model using "from_pretrained" and specifying the datatype, the returned model's parameters do not have that type.
This causes problems with flax text to image and flax dreambooth training scripts [here](https://github.com/huggingface/diffusers/blob/ce1c27adc85916a802db579df789d990fd14e8bc/examples/dreambooth/train_dreambooth_flax.py)
and [here](https://github.com/huggingface/diffusers/blob/ce1c27adc85916a802db579df789d990fd14e8bc/examples/text_to_image/train_text_to_image_flax.py) since the "mixed_precision" argument does not convert the pretrained model, my TPU was running out of memory when training.
### Reproduction
```
from diffusers import FlaxUNet2DConditionModel
import jax.numpy as jnp
unet, unet_params = FlaxUNet2DConditionModel.from_pretrained(
"duongna/stable-diffusion-v1-4-flax", subfolder="unet", dtype=jnp.bfloat16
)
print(unet_params)
```
shows
`
{conv_in': {'bias': Array([...] dtype=float32) } [...]}`
Expected result: The datatype for the models parameters are jnp.bfloat16
Actual result: The datatype for the models parameters are float32.
Possible solution: Can be fixed by using jax tree utils
`
def convert_dtype(xs, dtype): return jax.tree_util.tree_map(lambda x: jnp.asarray(x, dtype=dtype), xs)
`
### Logs
_No response_
### System Info
- `diffusers` version: 0.12.0.dev0
- Platform: Linux-5.13.0-1027-gcp-x86_64-with-glibc2.31
- Python version: 3.10.8
- PyTorch version (GPU?): 1.13.1+cu117 (False)
- Huggingface_hub version: 0.11.1
- Transformers version: 4.25.1
- Accelerate version: not installed
- xFormers version: not installed
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: N/A
Contributor guide
Assessment
This issue has not been assessed yet.