Update app.py

app.py

@@ -130,8 +130,9 @@ def train_step(state, batch, rng):
             pixel_values = jnp.expand_dims(pixel_values, axis=0)  # Add batch dimension if needed
         print(f"pixel_values shape in compute_loss: {pixel_values.shape}")
         
+        # Use the pipeline's VAE directly
         latents = pipeline.vae.apply(
-            {"params": params["vae"]},
+            {"params": pipeline.params["vae"]},
             pixel_values,
             method=pipeline.vae.encode
         ).latent_dist.sample(rng)
@@ -161,9 +162,9 @@ def train_step(state, batch, rng):
         print(f"timesteps shape: {timesteps.shape}")
         print(f"encoder_hidden_states shape: {encoder_hidden_states.shape}")
         
-        # Use the correct method to call the UNet
+        # Use the state's apply_fn (which should be the adjusted UNet)
         model_output = state.apply_fn(
-            {'params': params["unet"]},
+            {'params': params},
             noisy_latents,
             jnp.array(timesteps, dtype=jnp.int32),
             encoder_hidden_states,