FlaxUNet2DConditionModel is not initialized with correct dtypes

[wittenator]

Describe the bug

The FlaxUNet2DConditionModel allows specifying the dtype of the weights. Supplying a dtype different from float32 does not seem to be propagated to the actual model. This is imo different from #2068 since the afaik the code has correct dtype initialization. but the result is still incorrect. So this is not connected to loading FP32 weights or something similar.

Reproduction

import diffusers
from jax import random, numpy as jnp

dummy_input = jnp.zeros((2, 4, 32, 32), dtype=jnp.bfloat16)
dummy_t = jnp.zeros(2, dtype=jnp.bfloat16)
model = diffusers.FlaxUNet2DConditionModel(dtype=jnp.bfloat16)
key1, key2 = random.split(random.key(0))
params = model.init(key1, dummy_input, dummy_t, None)
print(jax.tree_util.tree_map(jnp.dtype, params))

Logs

{'params': {'conv_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_norm_out': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'conv_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'down_blocks_0': {'attentions_0': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'attentions_1': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'downsamplers_0': {'conv': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_0': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_1': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}}, 'down_blocks_1': {'attentions_0': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'attentions_1': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'downsamplers_0': {'conv': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_0': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_shortcut': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_1': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}}, 'down_blocks_2': {'attentions_0': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'attentions_1': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'downsamplers_0': {'conv': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_0': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_shortcut': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_1': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}}, 'down_blocks_3': {'resnets_0': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_1': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}}, 'mid_block': {'attentions_0': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'resnets_0': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_1': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}}, 'time_embedding': {'linear_1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'linear_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'up_blocks_0': {'resnets_0': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_shortcut': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_1': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_shortcut': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_2': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_shortcut': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'upsamplers_0': {'conv': {'bias': dtype('float32'), 'kernel': dtype('float32')}}}, 'up_blocks_1': {'attentions_0': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'attentions_1': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'attentions_2': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'resnets_0': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_shortcut': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_1': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_shortcut': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_2': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_shortcut': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'upsamplers_0': {'conv': {'bias': dtype('float32'), 'kernel': dtype('float32')}}}, 'up_blocks_2': {'attentions_0': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'attentions_1': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'attentions_2': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'resnets_0': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_shortcut': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_1': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_shortcut': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_2': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_shortcut': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'upsamplers_0': {'conv': {'bias': dtype('float32'), 'kernel': dtype('float32')}}}, 'up_blocks_3': {'attentions_0': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'attentions_1': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'attentions_2': {'norm': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'proj_in': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'proj_out': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'transformer_blocks_0': {'attn1': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'attn2': {'to_k': {'kernel': dtype('float32')}, 'to_out_0': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'to_q': {'kernel': dtype('float32')}, 'to_v': {'kernel': dtype('float32')}}, 'ff': {'net_0': {'proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'net_2': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm3': {'bias': dtype('float32'), 'scale': dtype('float32')}}}, 'resnets_0': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_shortcut': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_1': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_shortcut': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}, 'resnets_2': {'conv1': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv2': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'conv_shortcut': {'bias': dtype('float32'), 'kernel': dtype('float32')}, 'norm1': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'norm2': {'bias': dtype('float32'), 'scale': dtype('float32')}, 'time_emb_proj': {'bias': dtype('float32'), 'kernel': dtype('float32')}}}}}

System Info

• 🤗 Diffusers version: 0.32.2
• Platform: Linux-6.1.85+-x86_64-with-glibc2.35
• Running on Google Colab?: Yes
• Python version: 3.11.11
• PyTorch version (GPU?): 2.6.0+cu124 (True)
• Flax version (CPU?/GPU?/TPU?): 0.10.4 (gpu)
• Jax version: 0.5.2
• JaxLib version: 0.5.1
• Huggingface_hub version: 0.29.3
• Transformers version: 4.49.0
• Accelerate version: 1.5.2
• PEFT version: 0.14.0
• Bitsandbytes version: not installed
• Safetensors version: 0.5.3
• xFormers version: not installed
• Accelerator: Tesla T4, 15360 MiB
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

Who can help?

No response

[ishan-modi]

I think following code is why you are seeing the above behaviour

Since we have used JAX to create the FlaxUnet2D model it needs to return casted value since the params are immutable please follow this to cast model weights code

[wittenator]

Ah this looks actually like this is the culprit. Do you know why FlaxUNet2DConditionModel even has the dtype argument and the code that propagates it to the submodules? I'll have a look at the function, but I think that it is not compatible with the new Flax NNX API.

[ishan-modi]

You are correct, I dont think the dtype argument is of much use in FlaxUNet2DConditionModel and also its Pipeline, it is difficult to say why it was added because the commit is quite old.

for making it compatible with Flax NNX you will have to make quite a lot of changes to the existing code, instead of returning params during type_casting you will have to do it inplace. This is because Flax NNX supports mutable variables

[ishan-modi]

@yiyixuxu, feel free to close this one as well if you would like

[wittenator]

I also think this can be closed, but maybe it would be a good idea to add to the documentation of the parameter that it has no effect and deprecate it.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[wittenator]

@ishan-modi I just saw that Huggingface models actually has a little disclaimer about this at e.g. https://huggingface.co/docs/transformers/v4.24.0/model_doc/beit#transformers.FlaxBeitModel.dtype :

dtype (jax.numpy.dtype, optional, defaults to jax.numpy.float32) — The data type of the computation. Can be one of jax.numpy.float32, jax.numpy.float16 (on GPUs) and jax.numpy.bfloat16 (on TPUs).

This can be used to enable mixed-precision training or half-precision inference on GPUs or TPUs. If specified all the computation will be performed with the given dtype.

Note that this only specifies the dtype of the computation and does not influence the dtype of model parameters.

If you wish to change the dtype of the model parameters, see [to_fp16()](https://huggingface.co/docs/transformers/v4.24.0/en/main_classes/model#transformers.FlaxPreTrainedModel.to_fp16) and [to_bf16()](https://huggingface.co/docs/transformers/v4.24.0/en/main_classes/model#transformers.FlaxPreTrainedModel.to_bf16).

Maybe this can be included with diffusers as well? I assume that the difference is clear to people that use a lot of Jax, but to people not as fluent in it as me that was pretty non-intuitive at first.

[ishan-modi]

Yes, I think we should add this to docs for FlaxUNet2DConditionModel, would you like to create a PR for it ?

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.