Update app.py

app.py

@@ -97,18 +97,21 @@ def clear_jit_cache():
 # Training function
 def train_step(state, batch, rng):
     def compute_loss(params, pixel_values, rng):
+        # Ensure pixel_values are float32
+        pixel_values = jnp.array(pixel_values, dtype=jnp.float32)
+        
         # 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
+        latents = latents * jnp.float32(0.18215)
 
         # Generate random noise
-        noise = jax.random.normal(rng, latents.shape)
+        noise = jax.random.normal(rng, latents.shape, dtype=jnp.float32)
         
-        # Sample random timesteps
+        # Sample random timesteps (keep as integers)
         timesteps = jax.random.randint(
             rng, (latents.shape[0],), 0, pipeline.scheduler.config.num_train_timesteps
         )
@@ -124,11 +127,12 @@ def train_step(state, batch, rng):
         # Generate random encoder hidden states (simulating text embeddings)
         encoder_hidden_states = jax.random.normal(
             rng, 
-            (latents.shape[0], pipeline.text_encoder.config.hidden_size)
+            (latents.shape[0], pipeline.text_encoder.config.hidden_size),
+            dtype=jnp.float32
         )
         
         # Predict noise
-        model_output = state.apply_fn.apply(
+        model_output = state.apply_fn(
             {'params': params["unet"]},
             noisy_latents,
             timesteps,
@@ -139,9 +143,11 @@ def train_step(state, batch, rng):
         # Compute loss
         return jnp.mean((model_output - noise) ** 2)
 
-    grad_fn = jax.value_and_grad(compute_loss)
+    grad_fn = jax.grad(compute_loss, argnums=0, allow_int=True)
     rng, step_rng = jax.random.split(rng)
-    loss, grads = grad_fn(state.params, batch["pixel_values"], step_rng)
+    
+    grads = grad_fn(state.params, batch["pixel_values"], step_rng)
+    loss = compute_loss(state.params, batch["pixel_values"], step_rng)
     state = state.apply_gradients(grads=grads)
     return state, loss
 
@@ -236,7 +242,6 @@ for epoch in range(num_epochs):
     print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss}")
     clear_jit_cache()
 
-
     
 # Save the fine-tuned model
 output_dir = "/tmp/montevideo_fine_tuned_model"