Update app.py

app.py

@@ -124,21 +124,21 @@ if len(sample_batch['pixel_values']) > 0:
 
 # Training function
 def train_step(state, batch, rng):
-    def compute_loss(params, pixel_values, rng):
+    def compute_loss(unet_params, pixel_values, rng):
         pixel_values = jnp.array(pixel_values, dtype=jnp.float32)
         if pixel_values.ndim == 3:
             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
+        # Use the VAE from params
         latents = pipeline.vae.apply(
-            {"params": pipeline.params["vae"]},
+            {"params": params["vae"]},
             pixel_values,
             method=pipeline.vae.encode
         ).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(
@@ -164,7 +164,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': params},
+            {'params': unet_params},
             noisy_latents,
             jnp.array(timesteps, dtype=jnp.int32),
             encoder_hidden_states,
@@ -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)