提前加载lora

app.py

@@ -32,14 +32,12 @@ pipe = FluxPipeline.from_pretrained(
 set_single_lora(pipe.transformer, omni_consistency_path,
                 lora_weights=[1], cond_size=512)
 
+lora_names = [
+    "3D_Chibi", "American_Cartoon", "Macaron",
+    "Pixel", "Poly", "Van_Gogh"
+]
+
 def download_all_loras():
-    lora_names = [
-        "3D_Chibi", "American_Cartoon", "Chinese_Ink", "Clay_Toy",
-        "Fabric", "Ghibli", "Irasutoya", "Jojo", "LEGO", "Line",
-        "Macaron", "Oil_Painting", "Origami", "Paper_Cutting",
-        "Picasso", "Pixel", "Poly", "Pop_Art", "Rick_Morty",
-        "Snoopy", "Van_Gogh", "Vector"
-    ]
     for name in lora_names:
         hf_hub_download(
             repo_id="showlab/OmniConsistency",
@@ -48,15 +46,23 @@ def download_all_loras():
         )
 download_all_loras()
 
+def reload_all_loras():
+    pipe.unload_lora_weights()
+    for name in lora_names:
+        pipe.load_lora_weights(
+            "./LoRAs",
+            weight_name=f"{name}_rank128_bf16.safetensors",
+            adapter_name=name,
+        )
+reload_all_loras()
+        
 def clear_cache(transformer):
     for _, attn_processor in transformer.attn_processors.items():
         attn_processor.bank_kv.clear()
 
 @spaces.GPU()
 def generate_image(
-    lora_name,        
-    custom_repo_id,  
-    custom_weight_name, 
+    lora_name,
     prompt,
     uploaded_image,
     width, height,
@@ -65,31 +71,9 @@ def generate_image(
     seed
 ):
     width, height = int(width), int(height)
-    generator = torch.Generator("cpu").manual_seed(seed)
-
-    if custom_repo_id and custom_repo_id.strip():
-        repo_id = custom_repo_id.strip()
-        try:
-            lora_path = hf_hub_download(
-                repo_id=repo_id,
-                filename=custom_weight_name,
-                local_dir=CUSTOM_LORA_DIR,
-            )
-        except Exception as e:
-            raise gr.Error(f"Load custom LoRA failed: {e}")
-    else:
-        lora_path = os.path.join(
-            f"{LOCAL_LORA_DIR}/LoRAs", f"{lora_name}_rank128_bf16.safetensors"
-        )
+    generator = torch.Generator("cpu").manual_seed(seed) 
 
-    pipe.unload_lora_weights()
-    try:
-        pipe.load_lora_weights(
-            os.path.dirname(lora_path),
-            weight_name=os.path.basename(lora_path)
-        )
-    except Exception as e:
-        raise gr.Error(f"Load LoRA failed: {e}")
+    pipe.set_adapters(lora_name)
 
     spatial_image  = [uploaded_image.convert("RGB")]
     subject_images = []
@@ -113,16 +97,9 @@ def generate_image(
 
 # =============== Gradio UI ===============
 def create_interface():
-    demo_lora_names = [
-        "3D_Chibi", "American_Cartoon", "Chinese_Ink", "Clay_Toy",
-        "Fabric", "Ghibli", "Irasutoya", "Jojo", "LEGO", "Line",
-        "Macaron", "Oil_Painting", "Origami", "Paper_Cutting",
-        "Picasso", "Pixel", "Poly", "Pop_Art", "Rick_Morty",
-        "Snoopy", "Van_Gogh", "Vector"
-    ]
 
     def update_trigger_word(lora_name, prompt):
-      for name in demo_lora_names:
+      for name in lora_names:
         trigger = " ".join(name.split("_")) + " style,"
         prompt = prompt.replace(trigger, "")
       new_trigger = " ".join(lora_name.split("_"))+ " style,"