update app

app.py

@@ -14,12 +14,12 @@ try:
 except ImportError:
     class spaces:
         @staticmethod
-        def GPU(duration=30):
+        def GPU(duration=60):
             def decorator(func):
                 return func
             return decorator
 
-# --- Custom Theme Setup (Steel Blue) ---
+# --- Custom Theme Setup ---
 colors.steel_blue = colors.Color(
     name="steel_blue",
     c50="#EBF3F8",
@@ -90,15 +90,7 @@ steel_blue_theme = SteelBlueTheme()
 
 # --- Hardware Setup ---
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
-print("torch.__version__ =", torch.__version__)
-print("cuda available:", torch.cuda.is_available())
-if torch.cuda.is_available():
-    print("current device:", torch.cuda.current_device())
-    print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))
-
-print("Using device:", device)
+dtype = torch.bfloat16
 
 # --- Imports for Custom Pipeline ---
 from diffusers import FlowMatchEulerDiscreteScheduler
@@ -106,9 +98,8 @@ from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
 from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 
-dtype = torch.bfloat16
-
-# Load Pipeline with Rapid-AIO Transformer (Fast Version)
+# --- Model Initialization ---
+print("Loading Qwen Image Edit Pipeline...")
 pipe = QwenImageEditPlusPipeline.from_pretrained(
     "Qwen/Qwen-Image-Edit-2509",
     transformer=QwenImageTransformer2DModel.from_pretrained(
@@ -120,31 +111,37 @@ pipe = QwenImageEditPlusPipeline.from_pretrained(
     torch_dtype=dtype
 ).to(device)
 
-# --- Load Fusion/Texture/Face-Swap LoRAs ---
-print("Loading LoRA adapters...")
+# 1. Load and Fuse Lightning (for speed)
+print("Loading and Fusing Lightning LoRA...")
+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)
 
-# 1. Texture Edit
-pipe.load_lora_weights("tarn59/apply_texture_qwen_image_edit_2509",
-                       weight_name="apply_texture_v2_qwen_image_edit_2509.safetensors",
-                       adapter_name="texture-edit")
+# 2. Load Task Specific LoRAs
+print("Loading Task Adapters...")
 
-# 2. Fuse Objects
-pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Fusion",
-                       weight_name="溶图.safetensors",
-                       adapter_name="fuse-objects")
+# Texture
+pipe.load_lora_weights("tarn59/apply_texture_qwen_image_edit_2509", 
+                       weight_name="apply_texture_v2_qwen_image_edit_2509.safetensors", 
+                       adapter_name="texture")
 
-# 3. Face Swap
-pipe.load_lora_weights("Alissonerdx/BFS-Best-Face-Swap",
-                       weight_name="bfs_face_v1_qwen_image_edit_2509.safetensors",
-                       adapter_name="face-swap")
+# Fusion (Fuse-Objects)
+pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Fusion", 
+                       weight_name="溶图.safetensors", 
+                       adapter_name="fusion")
 
+# Face Swap
+pipe.load_lora_weights("Alissonerdx/BFS-Best-Face-Swap", 
+                       weight_name="bfs_head_v3_qwen_image_edit_2509.safetensors", 
+                       adapter_name="faceswap")
 
-# Attempt to set Flash Attention 3
+# Attempt Flash Attention 3
 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 (likely hardware mismatch): {e}. using default attention.")
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -163,15 +160,16 @@ def update_dimensions_on_upload(image):
         aspect_ratio = original_width / original_height
         new_width = int(new_height * aspect_ratio)
         
-    # Ensure dimensions are multiples of 16 (safer for transformers)
+    # Ensure dimensions are multiples of 16
     new_width = (new_width // 16) * 16
     new_height = (new_height // 16) * 16
     
     return new_width, new_height
 
-@spaces.GPU(duration=30)
+@spaces.GPU(duration=60)
 def infer(
-    input_gallery_items,
+    image_1,
+    image_2,
     prompt,
     lora_adapter,
     seed,
@@ -180,34 +178,28 @@ def infer(
     steps,
     progress=gr.Progress(track_tqdm=True)
 ):
-    """
-    Input:
-    input_gallery_items: Since type="pil", this is a List[Tuple[PIL.Image, str]] or List[PIL.Image]
-    """
-    if not input_gallery_items:
-        raise gr.Error("Please upload an image to edit.")
-
-    # Extract the image from the Gallery input
-    # When type='pil', Gradio Gallery returns a list of tuples (image, caption) or just images
-    first_item = input_gallery_items[0]
+    if image_1 is None or image_2 is None:
+        raise gr.Error("Please upload both images for Fusion/Texture/FaceSwap tasks.")
     
-    if isinstance(first_item, tuple):
-        # Format is (PIL.Image, Caption)
-        input_pil = first_item[0]
-    else:
-        # Format is PIL.Image directly
-        input_pil = first_item
-
-    # Map Dropdown choices to internal Adapter names
+    if not prompt:
+        # Add default prompts based on mode if user leaves empty (optional helper)
+        if lora_adapter == "Face-Swap":
+            prompt = "Swap the face."
+        elif lora_adapter == "Texture Edit":
+            prompt = "Apply texture to object."
+        elif lora_adapter == "Fuse-Objects":
+            prompt = "Fuse object into background."
+
+    # Switch Adapters
+    # Note: Lightning is already fused, so we just enable the style adapter
     adapters_map = {
-        "Texture Edit": "texture-edit",
-        "Fuse-Objects": "fuse-objects",
-        "Face-Swap": "face-swap",
+        "Texture Edit": "texture",
+        "Fuse-Objects": "fusion",
+        "Face-Swap": "faceswap",
     }
     
     active_adapter = adapters_map.get(lora_adapter)
     
-    # Reset adapters first, then activate selected
     if active_adapter:
         pipe.set_adapters([active_adapter], adapter_weights=[1.0])
     else:
@@ -219,11 +211,15 @@ def infer(
     generator = torch.Generator(device=device).manual_seed(seed)
     negative_prompt = "worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry"
 
-    original_image = input_pil.convert("RGB")
-    width, height = update_dimensions_on_upload(original_image)
+    # Prepare Images
+    img1_pil = image_1.convert("RGB")
+    img2_pil = image_2.convert("RGB")
+    
+    # Calculate dimensions based on the primary image (Image 1)
+    width, height = update_dimensions_on_upload(img1_pil)
 
     result = pipe(
-        image=original_image,
+        image=[img1_pil, img2_pil], # Pass both images
         prompt=prompt,
         negative_prompt=negative_prompt,
         height=height,
@@ -235,32 +231,27 @@ def infer(
 
     return result, seed
 
-@spaces.GPU(duration=30)
-def infer_example(input_gallery_items, prompt, lora_adapter):
-    # input_gallery_items will be the list structure from gr.Examples
-    if not input_gallery_items:
+@spaces.GPU(duration=60)
+def infer_example(image_1, image_2, prompt, lora_adapter):
+    # Wrapper for examples that sets defaults
+    if image_1 is None or image_2 is None:
         return None, 0
-        
-    # When passed from gr.Examples with type="pil" and a Gallery component, 
-    # we might need to handle file paths if cache_examples=False or PIL if processed.
-    # However, since we use infer_example as the fn, we mimic the infer logic.
-    
-    # For examples with type="pil", gradio usually converts paths to PIL.
-    return infer(
-        input_gallery_items, 
+    result, seed = infer(
+        image_1.convert("RGB"), 
+        image_2.convert("RGB"), 
         prompt, 
         lora_adapter, 
-        seed=0, 
-        randomize_seed=True, 
-        guidance_scale=1.0, 
-        steps=4
+        0, # seed
+        True, # randomize
+        1.0, # guidance
+        4 # steps (Lightning optimized)
     )
-
+    return result, seed
 
 css="""
 #col-container {
     margin: 0 auto;
-    max-width: 960px;
+    max-width: 1100px;
 }
 #main-title h1 {font-size: 2.1em !important;}
 """
@@ -268,72 +259,57 @@ css="""
 with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown("# **Qwen-Image-Edit-2509-LoRAs-Fast-Fusion**", elem_id="main-title")
-        gr.Markdown("Perform advanced image manipulation including Texture editing, Object Fusion, and Face Swapping using specialized [LoRA](https://huggingface.co/models?other=base_model:adapter:Qwen/Qwen-Image-Edit-2509) adapters.")
+        gr.Markdown("Advanced dual-image editing: **Texture Application**, **Object Fusion**, and **Face Swapping** using Qwen-Image-Edit-2509 + Lightning ⚡ (4 Steps).")
 
         with gr.Row(equal_height=True):
-            with gr.Column():
-                # Changed to Gallery to support potential multi-image flows (conceptually) and match user request
-                input_image = gr.Gallery(
-                    label="Input Images", 
-                    show_label=False, 
-                    type="pil", 
-                    interactive=True,
-                    height=290,
-                    columns=1
-                )
+            # Left Column: Inputs
+            with gr.Column(scale=1):
+                with gr.Row():
+                    image_1 = gr.Image(label="Base / Background / Body", type="pil", height=250)
+                    image_2 = gr.Image(label="Reference / Texture / Face", type="pil", height=250)
                 
+                lora_adapter = gr.Dropdown(
+                    label="Choose Editing Style",
+                    choices=["Texture Edit", "Fuse-Objects", "Face-Swap"],
+                    value="Texture Edit",
+                    info="Select the operation to perform."
+                )
+
                 prompt = gr.Text(
                     label="Edit Prompt",
                     show_label=True,
-                    placeholder="e.g., Change the material to wooden texture...",
+                    placeholder="e.g., Apply wood texture to the mug...",
                 )
-
-                run_button = gr.Button("Edit Image", variant="primary")
-
-            with gr.Column():
-                output_image = gr.Image(label="Output Image", interactive=False, format="png", height=350)
                 
-                with gr.Row():
-                    lora_adapter = gr.Dropdown(
-                        label="Choose Editing Style",
-                        choices=["Texture Edit", "Fuse-Objects", "Face-Swap"],
-                        value="Texture Edit"
-                    )
-                with gr.Accordion("Advanced Settings", open=False, visible=False):
+                run_button = gr.Button("Generate Fusion", variant="primary")
+
+                with gr.Accordion("Advanced Settings", open=False):
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                     randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
-                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0)
+                    guidance_scale = gr.Slider(label="True Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0)
                     steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=4)
+
+            # Right Column: Output
+            with gr.Column(scale=1):
+                output_image = gr.Image(label="Output Image", interactive=False, format="png", height=550)
         
+        # Examples
         gr.Examples(
             examples=[
-                # Format: [ [Image_List], Prompt, Adapter ]
-                [
-                    ["examples/texture_sample.jpg", "examples/texture_sample.2jpg"], 
-                    "Change the material of the object to rusted metal texture.", 
-                    "Texture Edit"
-                ],
-                [
-                    ["examples/fusion_sample.jpg"], 
-                    "Fuse the product naturally into the background.", 
-                    "Fuse-Objects"
-                ],
-                [
-                    ["examples/face_sample.jpg"], 
-                    "Swap the face with a cyberpunk robot face.", 
-                    "Face-Swap"
-                ],
+                ["examples/mug.jpg", "examples/wood.jpg", "Apply wood texture to the mug.", "Texture Edit"],
+                ["examples/room.jpg", "examples/chair.jpg", "Put the chair in the room naturally.", "Fuse-Objects"],
+                ["examples/body.jpg", "examples/face.jpg", "Swap the face.", "Face-Swap"],
             ],
-            inputs=[input_image, prompt, lora_adapter],
+            inputs=[image_1, image_2, prompt, lora_adapter],
             outputs=[output_image, seed],
             fn=infer_example,
             cache_examples=False,
-            label="Examples (Ensure images exist in 'examples/' folder)"
+            label="Examples (Ensure files exist in 'examples/' folder)"
         )
 
     run_button.click(
         fn=infer,
-        inputs=[input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
+        inputs=[image_1, image_2, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
         outputs=[output_image, seed]
     )