Update app.py

app.py

@@ -129,13 +129,14 @@ def train_step(state, batch, rng):
             dtype=jnp.float32
         )
         
+        # Use the correct method to call the UNet
         model_output = unet.apply(
             {'params': params["unet"]},
-            noisy_latents,
-            timesteps,
+            jnp.array(noisy_latents),
+            jnp.array(timesteps, dtype=jnp.int32),
             encoder_hidden_states,
             train=True,
-        )
+        ).sample
         
         return jnp.mean((model_output - noise) ** 2)
 
@@ -160,28 +161,4 @@ state = train_state.TrainState.create(
 # Training loop
 num_epochs = 3
 batch_size = 1
-rng = jax.random.PRNGKey(0)
-
-for epoch in range(num_epochs):
-    epoch_loss = 0
-    num_batches = 0
-    for batch in tqdm(processed_dataset.batch(batch_size)):
-        batch['pixel_values'] = jnp.array(batch['pixel_values'], dtype=jnp.float32)
-        rng, step_rng = jax.random.split(rng)
-        state, loss = train_step(state, batch, step_rng)
-        epoch_loss += loss
-        num_batches += 1
-        
-        if num_batches % 10 == 0:
-            jax.clear_caches()
-    
-    avg_loss = epoch_loss / num_batches
-    print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss}")
-    jax.clear_caches()
-
-# Save the fine-tuned model
-output_dir = "/tmp/montevideo_fine_tuned_model"
-os.makedirs(output_dir, exist_ok=True)
-unet.save_pretrained(output_dir, params=state.params["unet"])
-
-print(f"Model saved to {output_dir}")
+rng = jax.random.PRNGKey(