Update app.py

app.py

@@ -121,8 +121,12 @@ def train_step(state, batch, rng):
             timestep_rng, (batch_size,), 0, pipeline.scheduler.config.num_train_timesteps
         )
         
+        # Create scheduler state
+        scheduler_state = pipeline.scheduler.create_state()
+        
         # Add noise to images using the scheduler
         noisy_images = pipeline.scheduler.add_noise(
+            scheduler_state,
             original_samples=pixel_values,
             noise=noise,
             timesteps=timesteps
@@ -169,6 +173,7 @@ 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)