[Community] Specifying dtype in from_pretrained does not convert the model's datatype

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
and here 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

Comienza en FlaxUNet2DConditionModel.from_pretrained y compara el manejo de dtype con los ejemplos de entrenamiento de Flax train_dreambooth_flax.py y train_text_to_image_flax.py. Reproduce el problema usando el fragmento proporcionado para cargar el modelo y verifica que los parámetros devueltos usen jnp.bfloat16 en lugar de float32.