update app [adapter : hf download]

app.py

@@ -1,10 +1,5 @@
 import os
 import gc
-
-# 1. FIX: Set memory allocation configuration BEFORE importing torch
-# 'expandable_segments:True' prevents the specific CUDACachingAllocator assertion failure
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-
 import gradio as gr
 import numpy as np
 import spaces
@@ -15,7 +10,6 @@ from typing import Iterable
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 
-# Define Theme
 colors.orange_red = colors.Color(
     name="orange_red",
     c50="#FFF0E5",
@@ -104,51 +98,65 @@ pipe = QwenImageEditPlusPipeline.from_pretrained(
     torch_dtype=dtype
 ).to(device)
 
-# 2. FIX: Enable VAE Tiling. This is crucial for decoding large images without OOM.
 try:
     pipe.enable_vae_tiling()
     print("VAE Tiling enabled.")
 except Exception as e:
     print(f"Warning: Could not enable VAE tiling: {e}")
 
-print("Loading and Fusing Lightning LoRA...")
+print("Loading and Fusing Lightning LoRA (Base Optimization)...")
 pipe.load_lora_weights("lightx2v/Qwen-Image-Lightning",
                        weight_name="Qwen-Image-Lightning-4steps-V2.0-bf16.safetensors",
                        adapter_name="lightning")
 pipe.fuse_lora(adapter_names=["lightning"], lora_scale=1.0)
 
-print("Loading Task Adapters...")
-
-pipe.load_lora_weights("tarn59/apply_texture_qwen_image_edit_2509", 
-                       weight_name="apply_texture_v2_qwen_image_edit_2509.safetensors", 
-                       adapter_name="texture")
-
-pipe.load_lora_weights("ostris/qwen_image_edit_inpainting", 
-                       weight_name="qwen_image_edit_inpainting.safetensors", 
-                       adapter_name="fusion")
-
-pipe.load_lora_weights("ostris/qwen_image_edit_2509_shirt_design", 
-                       weight_name="qwen_image_edit_2509_shirt_design.safetensors", 
-                       adapter_name="shirt_design")
-
-pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Fusion", 
-                       weight_name="溶图.safetensors", 
-                       adapter_name="fusion-x")
-
-pipe.load_lora_weights("oumoumad/Qwen-Edit-2509-Material-transfer", 
-                       weight_name="material-transfer_000004769.safetensors", 
-                       adapter_name="material-transfer")
-
-pipe.load_lora_weights("dx8152/Qwen-Edit-2509-Light-Migration", 
-                       weight_name="参考色调.safetensors", 
-                       adapter_name="light-migration")
-
 try:
     pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
     print("Flash Attention 3 Processor set successfully.")
 except Exception as e:
-    print(f"Could not set FA3 processor (likely hardware mismatch): {e}. using default attention.")
+    print(f"Could not set FA3 processor: {e}. using default attention.")
+
+
+ADAPTER_SPECS = {
+    "Texture Edit": {
+        "repo": "tarn59/apply_texture_qwen_image_edit_2509",
+        "weights": "apply_texture_v2_qwen_image_edit_2509.safetensors",
+        "adapter_name": "texture",
+        "default_prompt": "Apply texture to object."
+    },
+    "Fuse-Objects": {
+        "repo": "ostris/qwen_image_edit_inpainting",
+        "weights": "qwen_image_edit_inpainting.safetensors",
+        "adapter_name": "fusion",
+        "default_prompt": "Fuse object into background."
+    },
+    "Cloth-Design-Fuse": {
+        "repo": "ostris/qwen_image_edit_2509_shirt_design",
+        "weights": "qwen_image_edit_2509_shirt_design.safetensors",
+        "adapter_name": "shirt_design",
+        "default_prompt": "Put this design on their shirt."
+    },
+    "Super-Fusion": {
+        "repo": "dx8152/Qwen-Image-Edit-2509-Fusion",
+        "weights": "溶图.safetensors",
+        "adapter_name": "fusion-x",
+        "default_prompt": "Blend the product into the background, correct its perspective and lighting."
+    },
+    "Material-Transfer": {
+        "repo": "oumoumad/Qwen-Edit-2509-Material-transfer",
+        "weights": "material-transfer_000004769.safetensors",
+        "adapter_name": "material-transfer",
+        "default_prompt": "Change materials of image1 to match the reference in image2."
+    },
+    "Light-Migration": {
+        "repo": "dx8152/Qwen-Edit-2509-Light-Migration",
+        "weights": "参考色调.safetensors",
+        "adapter_name": "light-migration",
+        "default_prompt": "Relight Image 1 based on the lighting and color tone of Image 2."
+    }
+}
 
+LOADED_ADAPTERS = set()
 MAX_SEED = np.iinfo(np.int32).max
 
 def update_dimensions_on_upload(image):
@@ -172,7 +180,7 @@ def update_dimensions_on_upload(image):
     
     return new_width, new_height
 
-@spaces.GPU
+@spaces.GPU(duration=30)
 def infer(
     image_1,
     image_2,
@@ -184,42 +192,42 @@ def infer(
     steps,
     progress=gr.Progress(track_tqdm=True)
 ):
-    # 3. FIX: Aggressive Garbage Collection before run
     gc.collect()
     torch.cuda.empty_cache()
 
     if image_1 is None or image_2 is None:
         raise gr.Error("Please upload both images for Fusion/Texture/FaceSwap tasks.")
     
-    if not prompt:
-        if lora_adapter == "Cloth-Design-Fuse":
-            prompt = "Put this design on their shirt."
-        elif lora_adapter == "Texture Edit":
-            prompt = "Apply texture to object."
-        elif lora_adapter == "Fuse-Objects":
-            prompt = "Fuse object into background."
-        elif lora_adapter == "Super-Fusion":
-            prompt = "Blend the product into the background, correct its perspective and lighting, and make it naturally integrated with the scene."
-        elif lora_adapter == "Material-Transfer":
-            prompt = "change materials of image1 to match the reference in image2"
-        elif lora_adapter == "Light-Migration":
-            prompt = "Refer to the color tone, remove the original lighting from Image 1, and relight Image 1 based on the lighting and color tone of Image 2."
-            
-    adapters_map = {
-        "Texture Edit": "texture",
-        "Fuse-Objects": "fusion",
-        "Cloth-Design-Fuse": "shirt_design",
-        "Super-Fusion": "fusion-x",
-        "Material-Transfer": "material-transfer",
-        "Light-Migration": "light-migration",
-    }
+    # 1. Get Adapter Spec
+    spec = ADAPTER_SPECS.get(lora_adapter)
+    if not spec:
+        raise gr.Error(f"Invalid Adapter Selection: {lora_adapter}")
     
-    active_adapter = adapters_map.get(lora_adapter)
-    
-    if active_adapter:
-        pipe.set_adapters([active_adapter], adapter_weights=[1.0])
+    adapter_name = spec["adapter_name"]
+
+    # 2. Dynamic Loading Logic
+    if adapter_name not in LOADED_ADAPTERS:
+        print(f"--- Downloading and Loading Adapter: {lora_adapter} ---")
+        try:
+            pipe.load_lora_weights(
+                spec["repo"], 
+                weight_name=spec["weights"], 
+                adapter_name=adapter_name
+            )
+            LOADED_ADAPTERS.add(adapter_name)
+        except Exception as e:
+            raise gr.Error(f"Failed to load adapter {lora_adapter}: {e}")
     else:
-        pipe.set_adapters([], adapter_weights=[])
+        print(f"--- Adapter {lora_adapter} already loaded. Activating. ---")
+
+    # 3. Handle Default Prompts
+    if not prompt:
+        prompt = spec["default_prompt"]
+            
+    # 4. Activate specific adapter
+    # Note: We do not fuse these task adapters, we just activate them.
+    # Lightning is already fused.
+    pipe.set_adapters([adapter_name], adapter_weights=[1.0])
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -233,7 +241,6 @@ def infer(
     width, height = update_dimensions_on_upload(img1_pil)
 
     try:
-        # 3. FIX: Use inference_mode for better memory efficiency
         with torch.inference_mode():
             result = pipe(
                 image=[img1_pil, img2_pil],
@@ -249,11 +256,9 @@ def infer(
         return result, seed
         
     except Exception as e:
-        # Rethrow so Gradio sees the error, but allow finally block to run
         raise e
         
     finally:
-        # 3. FIX: Cleanup after run regardless of success or failure
         gc.collect()
         torch.cuda.empty_cache()
 
@@ -314,7 +319,7 @@ with gr.Blocks() as demo:
                 with gr.Row():
                     lora_adapter = gr.Dropdown(
                         label="Choose Editing Style",
-                        choices=["Texture Edit", "Cloth-Design-Fuse", "Fuse-Objects", "Super-Fusion", "Light-Migration", "Material-Transfer"],
+                        choices=list(ADAPTER_SPECS.keys()),
                         value="Texture Edit",
                     )