Update app.py

app.py

@@ -149,14 +149,65 @@ def train_step(state, batch, rng):
 learning_rate = 1e-5
 optimizer = optax.adam(learning_rate)
 state = train_state.TrainState.create(
-    apply_fn=unet.__call__,
-    params={"unet": params["unet"], "vae": params["vae"]},
+    apply_fn=unet.__call__,  # Use __call__ directly
+    params=params,  # Pass all params
     tx=optimizer,
 )
 
-# Training loop
+# Modify the train_step function
+def train_step(state, batch, rng):
+    def compute_loss(params, pixel_values, rng):
+        # Encode images to latent space
+        latents = pipeline.vae.apply(
+            {"params": params["vae"]},
+            pixel_values,
+            method=pipeline.vae.encode
+        ).latent_dist.sample(rng)
+        latents = latents * 0.18215
+
+        # Generate random noise
+        noise = jax.random.normal(rng, latents.shape)
+        
+        # Sample random timesteps
+        timesteps = jax.random.randint(
+            rng, (latents.shape[0],), 0, pipeline.scheduler.config.num_train_timesteps
+        )
+        
+        # Add noise to latents
+        noisy_latents = pipeline.scheduler.add_noise(
+            pipeline.scheduler.create_state(),
+            original_samples=latents,
+            noise=noise,
+            timesteps=timesteps
+        )
+        
+        # Generate random encoder hidden states (simulating text embeddings)
+        encoder_hidden_states = jax.random.normal(
+            rng, 
+            (latents.shape[0], pipeline.text_encoder.config.hidden_size)
+        )
+        
+        # Predict noise
+        model_output = state.apply_fn(
+            {'params': params["unet"]},
+            noisy_latents,
+            timesteps,
+            encoder_hidden_states=encoder_hidden_states,
+            train=True,
+        )
+        
+        # Compute loss
+        return jnp.mean((model_output - noise) ** 2)
+
+    grad_fn = jax.value_and_grad(compute_loss)
+    rng, step_rng = jax.random.split(rng)
+    loss, grads = grad_fn(state.params, batch["pixel_values"], step_rng)
+    state = state.apply_gradients(grads=grads)
+    return state, loss
+
+# Training loop (remains the same)
 num_epochs = 3
-batch_size = 1  # Reduced batch size due to memory constraints
+batch_size = 1
 rng = jax.random.PRNGKey(0)
 
 for epoch in range(num_epochs):