Update app.py

app.py

@@ -1,26 +1,54 @@
 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 PIL import Image
-import os
+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):
+    # ... (keep the get_image function as is)
 
-# Load the JSON data
 with open("sdxl_lora.json", "r") as file:
     data = json.load(file)
-    sdxl_loras_raw = sorted(data, key=lambda x: x["likes"], reverse=True)
+    sdxl_loras_raw = [
+        {
+            "image": get_image(item["image"]),
+            "title": item["title"],
+            "repo": item["repo"],
+            "trigger_word": item["trigger_word"],
+            "weights": item["weights"],
+            "is_pivotal": item.get("is_pivotal", False),
+            "text_embedding_weights": item.get("text_embedding_weights", None),
+            "likes": item.get("likes", 0),
+        }
+        for item in data
+    ]
+
+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 = LCMScheduler.from_config(pipe.scheduler.config)
+pipe.scheduler = CustomLCMScheduler.from_config(pipe.scheduler.config)
 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
 
@@ -29,46 +57,50 @@ def update_selection(selected_state: gr.SelectData, gr_sdxl_loras):
     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("jasperai/flash-sdxl", 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)}")
+@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),
+):
+    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)
+
+        # Set adapter weights
+        pipe.set_adapters(["flash_lora", new_adapter_id], adapter_weights=[1.0, user_lora_weight])
+        gr.Info("LoRA setup complete")
+
+        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
 
-    # 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(
     pre_prompt,