Update app.py

app.py

@@ -116,6 +116,9 @@ def train_step(state, batch, rng):
             rng, (latents.shape[0],), 0, pipeline.scheduler.config.num_train_timesteps
         )
         
+        # Explicitly cast timesteps to int32
+        timesteps = timesteps.astype(jnp.int32)
+        
         # Add noise to latents
         noisy_latents = pipeline.scheduler.add_noise(
             pipeline.scheduler.create_state(),