Update app.py

app.py

@@ -56,7 +56,7 @@ def preprocess_images(examples):
         image = image.convert("RGB").resize((512, 512))
         # Convert to numpy array and normalize
         image = np.array(image).astype(np.float32) / 127.5 - 1.0
-        # Ensure the image has the shape (height, width, 3)
+        # Ensure the image has the shape (3, height, width)
         return image.transpose(2, 0, 1)  # Change to channel-first format
 
     return {"pixel_values": [process_image(img) for img in examples["image"]]}
@@ -178,12 +178,22 @@ print(processed_dataset)
 print("First batch:")
 first_batch = next(iter(processed_dataset.batch(batch_size)))
 print(f"Batch keys: {first_batch.keys()}")
-print(f"Pixel values shape: {first_batch['pixel_values'].shape}")
+print(f"Type of pixel_values: {type(first_batch['pixel_values'])}")
+if isinstance(first_batch['pixel_values'], list):
+    print(f"Length of pixel_values list: {len(first_batch['pixel_values'])}")
+    if len(first_batch['pixel_values']) > 0:
+        print(f"Shape of first item in pixel_values: {np.array(first_batch['pixel_values'][0]).shape}")
+
+# Convert the list of pixel values to a numpy array
+first_batch['pixel_values'] = np.array(first_batch['pixel_values'])
+print(f"Pixel values shape after conversion: {first_batch['pixel_values'].shape}")
 
 for epoch in range(num_epochs):
     epoch_loss = 0
     num_batches = 0
     for batch in tqdm(processed_dataset.batch(batch_size)):
+        # Convert the list of pixel values to a numpy array for each batch
+        batch['pixel_values'] = np.array(batch['pixel_values'])
         rng, step_rng = jax.random.split(rng)
         state, loss = train_step(state, batch, step_rng)
         epoch_loss += loss
@@ -192,6 +202,7 @@ for epoch in range(num_epochs):
     print(f"Epoch {epoch+1}/{num_epochs}, Average Loss: {avg_loss}")
 
 
+
 # Save the fine-tuned model
 output_dir = "/tmp/montevideo_fine_tuned_model"
 os.makedirs(output_dir, exist_ok=True)