Update app.py

app.py

@@ -7,6 +7,7 @@ import spaces
 import torch
 from diffusers import DiffusionPipeline, LCMScheduler
 from PIL import Image
+import os
 
 # Load the JSON data
 with open("sdxl_lora.json", "r") as file:
@@ -40,21 +41,33 @@ def get_image(image_data):
         local_path = image_data.get('local_path')
         hf_url = image_data.get('hf_url')
     else:
-        return None  # or a default image path
-    
-    try:
-        return local_path  # Return the local path string
-    except:
+        print(f"Unexpected image_data format: {type(image_data)}")
+        return None
+
+    # Try loading from local path first
+    if local_path and os.path.exists(local_path):
+        try:
+            Image.open(local_path).verify()  # Verify that it's a valid image
+            return local_path
+        except Exception as e:
+            print(f"Error loading local image {local_path}: {e}")
+
+    # If local path fails or doesn't exist, try URL
+    if hf_url:
         try:
             response = requests.get(hf_url)
             if response.status_code == 200:
-                with open(local_path, 'wb') as f:
-                    f.write(response.content)
-                return local_path  # Return the local path string
+                img = Image.open(requests.get(hf_url, stream=True).raw)
+                img.verify()  # Verify that it's a valid image
+                img.save(local_path)  # Save for future use
+                return local_path
+            else:
+                print(f"Failed to fetch image from URL {hf_url}. Status code: {response.status_code}")
         except Exception as e:
-            print(f"Failed to load image: {e}")
-    
-    return None  # or a default image path
+            print(f"Error loading image from URL {hf_url}: {e}")
+
+    print(f"Failed to load image for {image_data}")
+    return None
 
 @spaces.GPU
 def infer(
@@ -146,9 +159,7 @@ with gr.Blocks(css=css) as demo:
     with gr.Row():
         with gr.Column(scale=2):
             gallery = gr.Gallery(
-                value=[(img, title) for img, title in 
-                       ((get_image(item["image"]), item["title"]) for item in sdxl_loras_raw) 
-                       if img is not None],
+                value=[(get_image(item["image"]), item["title"]) for item in sdxl_loras_raw if get_image(item["image"]) is not None],
                 label="SDXL LoRA Gallery",
                 show_label=False,
                 elem_id="gallery",