Update app.py

app.py

@@ -155,23 +155,27 @@ state = train_state.TrainState.create(
 )
 
 # Modify the train_step function
-def train_step(state, batch, rng):
+ef 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
         timesteps = jax.random.randint(
             rng, (latents.shape[0],), 0, pipeline.scheduler.config.num_train_timesteps
         )
+        timesteps = jnp.array(timesteps, dtype=jnp.float32)
         
         # Add noise to latents
         noisy_latents = pipeline.scheduler.add_noise(
@@ -184,7 +188,8 @@ 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
@@ -199,12 +204,16 @@ 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
 
+
+
 # Training loop (remains the same)
 num_epochs = 3
 batch_size = 1