Update app.py

app.py

@@ -186,7 +186,11 @@ float32_params = jax.tree_util.tree_map(lambda x: x.astype(jnp.float32) if x.dty
 
 # Create a new UNet with the adjusted parameters
 unet_config = dict(unet.config)
-unet_config.pop('_use_default_values', None)  # Remove the unexpected argument
+# Remove unexpected keys
+unexpected_keys = ['_class_name', '_diffusers_version', '_use_default_values', '_name_or_path']
+for key in unexpected_keys:
+    unet_config.pop(key, None)
+
 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']