Update app.py

app.py

@@ -113,7 +113,7 @@ def train_step(state, batch, rng):
         batch_size = pixel_values.shape[0]
         
         # Generate random noise
-        noise_rng, timestep_rng = jax.random.split(rng)
+        noise_rng, timestep_rng, latents_rng = jax.random.split(rng, 3)
         noise = jax.random.normal(noise_rng, pixel_values.shape)
         
         # Sample random timesteps
@@ -132,11 +132,15 @@ def train_step(state, batch, rng):
             timesteps=timesteps
         )
         
+        # Generate random latents for text encoder
+        latents = jax.random.normal(latents_rng, (batch_size, pipeline.text_encoder.config.hidden_size))
+        
         # Predict noise
         model_output = state.apply_fn.apply(
             {'params': params},
             jnp.array(noisy_images),
             jnp.array(timesteps),
+            encoder_hidden_states=latents,
             train=True,
         )
         
@@ -173,7 +177,6 @@ for epoch in range(num_epochs):
     avg_loss = epoch_loss / num_batches
     print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss}")
 
-
 # Save the fine-tuned model
 output_dir = "/tmp/montevideo_fine_tuned_model"
 os.makedirs(output_dir, exist_ok=True)