Update app.py

app.py

@@ -1,12 +1,3 @@
-import json
-import random
-
-import gradio as gr
-import numpy as np
-import spaces
-import torch
-from diffusers import DiffusionPipeline, LCMScheduler
-
 with open("sdxl_lora.json", "r") as file:
     data = json.load(file)
     sdxl_loras_raw = [
@@ -38,17 +29,50 @@ pipe.to(device=DEVICE, dtype=torch.float16)
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-
-def update_selection(
-    selected_state: gr.SelectData,
-    gr_sdxl_loras,
-):
-
+def update_selection(selected_state: gr.SelectData, gr_sdxl_loras):
     lora_id = gr_sdxl_loras[selected_state.index]["repo"]
     trigger_word = gr_sdxl_loras[selected_state.index]["trigger_word"]
-
     return lora_id, trigger_word
 
+def load_lora_for_style(style_repo):
+    pipe.unload_lora_weights()
+    pipe.load_lora_weights(style_repo, adapter_name="lora")
+
+def get_image(image_data):
+    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
 
 @spaces.GPU
 def infer(
@@ -63,19 +87,7 @@ def infer(
     user_lora_weight,
     progress=gr.Progress(track_tqdm=True),
 ):
-    flash_sdxl_id = "jasperai/flash-sdxl"
-
-    new_adapter_id = user_lora_selector.replace("/", "_")
-    loaded_adapters = pipe.get_list_adapters()
-
-    if new_adapter_id not in loaded_adapters["unet"]:
-        gr.Info("Swapping LoRA")
-        pipe.unload_lora_weights()
-        pipe.load_lora_weights(flash_sdxl_id, adapter_name="lora")
-        pipe.load_lora_weights(user_lora_selector, adapter_name=new_adapter_id)
-
-    pipe.set_adapters(["lora", new_adapter_id], adapter_weights=[1.0, user_lora_weight])
-    gr.Info("LoRA setup done")
+    load_lora_for_style(user_lora_selector)
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)