Update app.py

app.py

@@ -83,6 +83,7 @@ def preprocess_images(examples):
             return None
         image = image.convert("RGB").resize((512, 512))
         image = np.array(image).astype(np.float32) / 255.0
+        image = image.transpose(2, 0, 1)  # Change to channel-first format
         return image
 
     processed = [process_image(img) for img in examples["image"]]
@@ -121,7 +122,6 @@ if len(sample_batch['pixel_values']) > 0:
 def train_step(state, batch, rng):
     def compute_loss(params, pixel_values, rng):
         pixel_values = jnp.array(pixel_values, dtype=jnp.float32)
-        pixel_values = jnp.expand_dims(pixel_values, axis=0)  # Add batch dimension
         print(f"pixel_values shape in compute_loss: {pixel_values.shape}")
         
         latents = pipeline.vae.apply(