Update app.py

app.py

@@ -188,7 +188,7 @@ def train_step(state, batch, rng):
         
         # Use the state's apply_fn (which should be the adjusted UNet)
         model_output = state.apply_fn(
-            {'params': unet_params},
+            {"params": unet_params},
             noisy_latents,
             jnp.array(timesteps, dtype=jnp.int32),
             encoder_hidden_states,
@@ -200,9 +200,15 @@ def train_step(state, batch, rng):
     grad_fn = jax.grad(compute_loss, argnums=0, allow_int=True)
     rng, step_rng = jax.random.split(rng)
     
-    grads = grad_fn(state.params, batch["pixel_values"], step_rng)
-    loss = compute_loss(state.params, batch["pixel_values"], step_rng)
-    state = state.apply_gradients(grads=grads)
+    # Ensure we're passing the correct structure to grad_fn and compute_loss
+    unet_params = state.params["params"] if "params" in state.params else state.params
+    grads = grad_fn(unet_params, batch["pixel_values"], step_rng)
+    loss = compute_loss(unet_params, batch["pixel_values"], step_rng)
+    
+    # Update the state with the correct structure
+    new_params = optax.apply_updates(state.params, grads)
+    state = state.replace(params=new_params)
+    
     return state, loss
 
 # Initialize training state
@@ -220,9 +226,10 @@ 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 = adjust_unet_input_layer(adjusted_params)  # Adjust the input layer
 
+# Adjust the state creation
 state = train_state.TrainState.create(
     apply_fn=adjusted_unet.apply,
-    params=adjusted_params,
+    params={"params": adjusted_params},  # Wrap params in a dict with "params" key
     tx=optimizer,
 )