Update app.py

app.py

@@ -57,7 +57,7 @@ def preprocess_images(examples):
             image = Image.open(image)
         if not isinstance(image, Image.Image):
             raise ValueError(f"Unexpected image type: {type(image)}")
-        image = image.convert("RGB").resize((512, 512))  # Keep original size
+        image = image.convert("RGB").resize((512, 512))
         image = np.array(image).astype(np.float32) / 255.0
         return image.transpose(2, 0, 1)
 
@@ -214,7 +214,7 @@ 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'])
+        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