Update app.py

app.py

@@ -185,7 +185,9 @@ optimizer = optax.adam(learning_rate)
 float32_params = jax.tree_util.tree_map(lambda x: x.astype(jnp.float32) if x.dtype != jnp.int32 else x, params)
 
 # Create a new UNet with the adjusted parameters
-adjusted_unet = FlaxUNet2DConditionModel(**unet.config)
+unet_config = unet.config.to_dict()
+unet_config.pop('_use_default_values', None)  # Remove the unexpected argument
+adjusted_unet = FlaxUNet2DConditionModel(**unet_config)
 adjusted_params = adjusted_unet.init(jax.random.PRNGKey(0), jnp.ones((1, 4, 64, 64)), jnp.ones((1,)), jnp.ones((1, 77, 768)))
 adjusted_params['params'] = float32_params['unet']