Update app.py

app.py

@@ -60,7 +60,7 @@ unet = pipeline.unet
 def adjust_unet_input_layer(params):
     conv_in_weight = params['unet']['conv_in']['kernel']
     new_conv_in_weight = jnp.zeros((3, 3, 4, 320), dtype=jnp.float32)
-    new_conv_in_weight = new_conv_in_weight.at[:, :, :4, :].set(conv_in_weight[:, :, :4, :])
+    new_conv_in_weight = new_conv_in_weight.at[:, :, :3, :].set(conv_in_weight[:, :, :3, :])
     params['unet']['conv_in']['kernel'] = new_conv_in_weight
     return params
 
@@ -70,14 +70,17 @@ params = adjust_unet_input_layer(params)
 def preprocess_images(examples):
     def process_image(image):
         if isinstance(image, str):
+            if not image.lower().endswith('.jpg') and not image.lower().endswith('.jpeg'):
+                return None
             image = Image.open(image)
         if not isinstance(image, Image.Image):
-            raise ValueError(f"Unexpected image type: {type(image)}")
+            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 {"pixel_values": [process_image(img) for img in examples["image"]]}
+    processed = [process_image(img) for img in examples["image"]]
+    return {"pixel_values": [img for img in processed if img is not None]}
 
 # Load dataset from Hugging Face
 dataset_name = "uruguayai/montevideo"
@@ -94,6 +97,7 @@ else:
     print("Processing dataset...")
     dataset = load_dataset(dataset_name)
     processed_dataset = dataset["train"].map(preprocess_images, batched=True, remove_columns=dataset["train"].column_names)
+    processed_dataset = processed_dataset.filter(lambda example: len(example['pixel_values']) > 0)
     with open(dataset_cache_file, 'wb') as f:
         pickle.dump(processed_dataset, f)