Update app.py

app.py

@@ -138,7 +138,7 @@ def train_step(state, batch, rng):
         ).latent_dist.sample(rng)
         latents = latents * jnp.float32(0.18215)
         print(f"latents shape: {latents.shape}")
-    
+
         noise = jax.random.normal(rng, latents.shape, dtype=jnp.float32)
         
         timesteps = jax.random.randint(
@@ -172,7 +172,7 @@ def train_step(state, batch, rng):
         ).sample
         
         return jnp.mean((model_output - noise) ** 2)
-    
+
     grad_fn = jax.grad(compute_loss, argnums=0, allow_int=True)
     rng, step_rng = jax.random.split(rng)
     
@@ -181,11 +181,6 @@ 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)
@@ -199,7 +194,7 @@ print("Filtered UNet config keys:", filtered_unet_config.keys())
 
 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']
+adjusted_params = float32_params['unet']  # Use only UNet params
 
 state = train_state.TrainState.create(
     apply_fn=adjusted_unet.apply,