Update app.py

app.py

@@ -11,6 +11,7 @@ import os
 import pickle
 from PIL import Image
 import numpy as np
+from inspect import signature
 
 # Custom Scheduler
 class CustomFlaxPNDMScheduler(FlaxPNDMScheduler):
@@ -179,6 +180,11 @@ def train_step(state, batch, rng):
     state = state.apply_gradients(grads=grads)
     return state, loss
 
+def filter_dict(dict_to_filter, target_callable):
+    """Filter a dictionary to only include keys that are valid parameters for the target callable."""
+    valid_params = signature(target_callable).parameters.keys()
+    return {k: v for k, v in dict_to_filter.items() if k in valid_params}
+
 # Initialize training state
 learning_rate = 1e-5
 optimizer = optax.adam(learning_rate)
@@ -186,12 +192,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)
-# 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)
+filtered_unet_config = filter_dict(unet_config, FlaxUNet2DConditionModel.__init__)
+
+print("Filtered UNet config keys:", filtered_unet_config.keys())
 
-adjusted_unet = FlaxUNet2DConditionModel(**unet_config)
+adjusted_unet = FlaxUNet2DConditionModel(**filtered_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']