Update app.py

app.py

@@ -1,63 +1,17 @@
 import json
 import random
-import requests
-import os
-from PIL import Image
 
 import gradio as gr
 import numpy as np
 import spaces
 import torch
 from diffusers import DiffusionPipeline, LCMScheduler
-from peft import PeftModel
-
-# Custom LCMScheduler to ignore unexpected attributes
-class CustomLCMScheduler(LCMScheduler):
-    @property
-    def config(self):
-        return {k: v for k, v in super().config.items() if k != "skip_prk_steps"}
-
-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
 
 with open("sdxl_lora.json", "r") as file:
     data = json.load(file)
     sdxl_loras_raw = [
         {
-            "image": get_image(item["image"]),
+            "image": item["image"],
             "title": item["title"],
             "repo": item["repo"],
             "trigger_word": item["trigger_word"],
@@ -69,27 +23,33 @@ with open("sdxl_lora.json", "r") as file:
         for item in data
     ]
 
+# Sort the loras by likes
 sdxl_loras_raw = sorted(sdxl_loras_raw, key=lambda x: x["likes"], reverse=True)
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 model_id = "stabilityai/stable-diffusion-xl-base-1.0"
 
 pipe = DiffusionPipeline.from_pretrained(model_id, variant="fp16")
-pipe.scheduler = CustomLCMScheduler.from_config(pipe.scheduler.config)
+pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+pipe.load_lora_weights("jasperai/flash-sdxl", adapter_name="lora")
 pipe.to(device=DEVICE, dtype=torch.float16)
 
-# Load Flash SDXL LoRA
-flash_sdxl_id = "jasperai/flash-sdxl"
-pipe.load_lora_weights(flash_sdxl_id, adapter_name="flash_lora")
 
 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
 
+
 @spaces.GPU
 def infer(
     pre_prompt,
@@ -103,51 +63,19 @@ def infer(
     user_lora_weight,
     progress=gr.Progress(track_tqdm=True),
 ):
-    try:
-        # Load the user-selected LoRA
-        new_adapter_id = user_lora_selector.replace("/", "_")
-        pipe.load_lora_weights(user_lora_selector, adapter_name=new_adapter_id)
+    flash_sdxl_id = "jasperai/flash-sdxl"
 
-        # Set adapter weights
-        pipe.set_adapters(["flash_lora", new_adapter_id], adapter_weights=[1.0, user_lora_weight])
-        gr.Info("LoRA setup complete")
+    new_adapter_id = user_lora_selector.replace("/", "_")
+    loaded_adapters = pipe.get_list_adapters()
 
-        if randomize_seed:
-            seed = random.randint(0, MAX_SEED)
-
-        generator = torch.Generator().manual_seed(seed)
-
-        if pre_prompt != "":
-            prompt = f"{pre_prompt} {prompt}"
-
-        # Use Flash Diffusion settings
-        image = pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            guidance_scale=1.0,  # Flash Diffusion typically uses guidance_scale=1
-            num_inference_steps=4,  # Flash Diffusion uses fewer steps
-            generator=generator,
-        ).images[0]
-
-        return image
-    except Exception as e:
-        gr.Error(f"An error occurred: {str(e)}")
-        return None
+    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)
 
-@spaces.GPU
-def infer(
-    pre_prompt,
-    prompt,
-    seed,
-    randomize_seed,
-    num_inference_steps,
-    negative_prompt,
-    guidance_scale,
-    user_lora_selector,
-    user_lora_weight,
-    progress=gr.Progress(track_tqdm=True),
-):
-    load_lora_for_style(user_lora_selector)
+    pipe.set_adapters(["lora", new_adapter_id], adapter_weights=[1.0, user_lora_weight])
+    gr.Info("LoRA setup done")
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)