Update app.py

app.py

@@ -18,7 +18,40 @@ class CustomLCMScheduler(LCMScheduler):
         return {k: v for k, v in super().config.items() if k != "skip_prk_steps"}
 
 def get_image(image_data):
-    # ... (keep the get_image function as is)
+    if isinstance(image_data, str):
+        return image_data
+    
+    if isinstance(image_data, dict):
+        local_path = image_data.get('local_path')
+        hf_url = image_data.get('hf_url')
+    else:
+        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:
+                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"Error loading image from URL {hf_url}: {e}")
+
+    print(f"Failed to load image for {image_data}")
+    return None
 
 with open("sdxl_lora.json", "r") as file:
     data = json.load(file)