Update app.py

app.py

@@ -83,7 +83,7 @@ def preprocess_images(examples):
             return None
         image = image.convert("RGB").resize((512, 512))
         image = np.array(image).astype(np.float32) / 255.0
-        return image.transpose(2, 0, 1)
+        return image
 
     processed = [process_image(img) for img in examples["image"]]
     return {"pixel_values": [img for img in processed if img is not None]}
@@ -121,6 +121,7 @@ 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(
@@ -176,7 +177,7 @@ def train_step(state, batch, rng):
 # Initialize training state
 learning_rate = 1e-5
 optimizer = optax.adam(learning_rate)
-float32_params = jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype != jnp.int32 else x, params)
+float32_params = jax.tree_util.tree_map(lambda x: x.astype(jnp.float32) if x.dtype != jnp.int32 else x, params)
 state = train_state.TrainState.create(
     apply_fn=unet.apply,
     params=float32_params,