Update app

app.py

@@ -19,7 +19,7 @@ except ImportError:
                 return func
             return decorator
 
-# --- Custom Theme Setup ---
+# --- Custom Theme Setup (Steel Blue) ---
 colors.steel_blue = colors.Color(
     name="steel_blue",
     c50="#EBF3F8",
@@ -92,9 +92,16 @@ steel_blue_theme = SteelBlueTheme()
 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)
 
 # --- Imports for Custom Pipeline ---
+# Note: These require the local 'qwenimage' folder to be present
 from diffusers import FlowMatchEulerDiscreteScheduler
 from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
@@ -102,7 +109,7 @@ from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 
 dtype = torch.bfloat16
 
-# Load Pipeline
+# Load Pipeline with Rapid-AIO Transformer (Fast Version)
 pipe = QwenImageEditPlusPipeline.from_pretrained(
     "Qwen/Qwen-Image-Edit-2509",
     transformer=QwenImageTransformer2DModel.from_pretrained(
@@ -114,35 +121,31 @@ pipe = QwenImageEditPlusPipeline.from_pretrained(
     torch_dtype=dtype
 ).to(device)
 
-# --- Load LoRAs ---
+# --- Load Fusion/Texture/Face-Swap LoRAs ---
 print("Loading LoRA adapters...")
 
-# 1. Lightning (Required for fast 4-step inference)
-pipe.load_lora_weights("lightx2v/Qwen-Image-Lightning",
-                       weight_name="Qwen-Image-Lightning-4steps-V2.0-bf16.safetensors", 
-                       adapter_name="lightning")
-
-# 2. 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")
+# 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")
 
-# 3. Fuse Objects
+# 2. Fuse Objects (Note: Filename contains non-ascii characters, handled as string)
 pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Fusion",
                        weight_name="溶图.safetensors",
-                       adapter_name="fusion")
+                       adapter_name="fuse-objects")
 
-# 4. Face Swap
+# 3. Face Swap
 pipe.load_lora_weights("Alissonerdx/BFS-Best-Face-Swap",
                        weight_name="bfs_head_v3_qwen_image_edit_2509.safetensors",
-                       adapter_name="faceswap")
+                       adapter_name="face-swap")
+
 
-# Attempt to set Flash Attention 3
+# Attempt to set Flash Attention 3 (Requires H100 or compatible setup)
 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: {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
 
@@ -161,6 +164,7 @@ 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)
     new_width = (new_width // 16) * 16
     new_height = (new_height // 16) * 16
     
@@ -168,33 +172,32 @@ def update_dimensions_on_upload(image):
 
 @spaces.GPU(duration=60)
 def infer(
-    image_input,
-    image_reference,
+    input_image,
     prompt,
-    style_choice,
+    lora_adapter,
     seed,
     randomize_seed,
     guidance_scale,
     steps,
     progress=gr.Progress(track_tqdm=True)
 ):
-    if image_input is None:
-        raise gr.Error("Please upload the Main Input image.")
-    if image_reference is None:
-        raise gr.Error("Please upload the Reference/Texture image.")
-    if not prompt:
-        raise gr.Error("Please enter a prompt.")
-
-    # Manage Adapters
-    # We combine the specific style LoRA with the Lightning LoRA for speed
-    if style_choice == "Texture Edit":
-        pipe.set_adapters(["texture", "lightning"], adapter_weights=[1.0, 1.0])
-    elif style_choice == "Fuse Objects":
-        pipe.set_adapters(["fusion", "lightning"], adapter_weights=[1.0, 1.0])
-    elif style_choice == "Face Swap":
-        pipe.set_adapters(["faceswap", "lightning"], adapter_weights=[1.0, 1.0])
+    if input_image is None:
+        raise gr.Error("Please upload an image to edit.")
+
+    # Map Dropdown choices to internal Adapter names
+    adapters_map = {
+        "Texture Edit": "texture-edit",
+        "Fuse-Objects": "fuse-objects",
+        "Face-Swap": "face-swap",
+    }
+    
+    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:
-        pipe.set_adapters(["lightning"], adapter_weights=[1.0])
+        pipe.set_adapters([], adapter_weights=[])
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -202,16 +205,11 @@ 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"
 
-    # Prepare Images
-    img1 = image_input.convert("RGB")
-    img2 = image_reference.convert("RGB")
-    
-    # Resize logic based on the main input image
-    width, height = update_dimensions_on_upload(img1)
+    original_image = input_image.convert("RGB")
+    width, height = update_dimensions_on_upload(original_image)
 
-    # Pass list of images [Content, Reference]
     result = pipe(
-        image=[img1, img2],
+        image=original_image,
         prompt=prompt,
         negative_prompt=negative_prompt,
         height=height,
@@ -224,29 +222,20 @@ def infer(
     return result, seed
 
 @spaces.GPU(duration=60)
-def infer_example(image_input, image_reference, prompt, style_choice):
-    if image_input is None or image_reference is None:
+def infer_example(input_image, prompt, lora_adapter):
+    if input_image is None:
         return None, 0
-    
-    guidance_scale = 1.0
-    steps = 4 # Default fast steps
-    
-    result, seed = infer(
-        image_input, 
-        image_reference, 
-        prompt, 
-        style_choice, 
-        0, # seed
-        True, # randomize
-        guidance_scale, 
-        steps
-    )
+    input_pil = input_image.convert("RGB")
+    guidance_scale = 4.0 # Slightly higher default for better adherence
+    steps = 30
+    result, seed = infer(input_pil, prompt, lora_adapter, 0, True, guidance_scale, steps)
     return result, seed
 
+
 css="""
 #col-container {
     margin: 0 auto;
-    max-width: 1100px;
+    max-width: 960px;
 }
 #main-title h1 {font-size: 2.1em !important;}
 """
@@ -254,55 +243,51 @@ 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 dual-image editing: Texture Transfer, Object Fusion, and Face Swapping using [Qwen-Image-Edit](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) LoRAs.")
+        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.")
 
         with gr.Row(equal_height=True):
-            with gr.Column(scale=1):
-                image_input = gr.Image(label="1. Main Content / Source", type="pil", height=250)
-                image_reference = gr.Image(label="2. Reference / Texture / Face", type="pil", height=250)
+            with gr.Column():
+                input_image = gr.Gallery(label="Input Images", show_label=False, type="pil", interactive=True)
                 
-                style_choice = gr.Dropdown(
-                    label="Choose Editing Style",
-                    choices=["Texture Edit", "Fuse Objects", "Face Swap"],
-                    value="Texture Edit",
-                    interactive=True
-                )
-
                 prompt = gr.Text(
-                    label="Prompt",
+                    label="Edit Prompt",
                     show_label=True,
-                    placeholder="e.g., Apply wood texture to the bottle...",
+                    placeholder="e.g., Change the material to wooden texture...",
                 )
 
-                run_button = gr.Button("Generate Fusion", variant="primary")
+                run_button = gr.Button("Edit Image", variant="primary")
 
-            with gr.Column(scale=1):
-                output_image = gr.Image(label="Output Image", interactive=False, format="png", height=550)
+            with gr.Column():
+                output_image = gr.Image(label="Output Image", interactive=False, format="png", height=350)
                 
-                with gr.Accordion("Advanced Settings", open=False, visible=True):
+                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):
                     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)
-                    steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=4)
+                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=4.0)
+                    steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=30)
         
         gr.Examples(
             examples=[
-                # Format: [Img1, Img2, Prompt, Style] - Assuming example files exist in 'examples/'
-                # You will need to ensure these files exist or remove this block
-                ["examples/coffee_mug.png", "examples/wood_texture.png", "Apply wood texture to the mug", "Texture Edit"],
-                ["examples/bg_room.jpg", "examples/cat.png", "A cat sitting in the living room", "Fuse Objects"],
-                ["examples/target_person.jpg", "examples/source_face.jpg", "Swap the face", "Face Swap"],
+                ["examples/texture_sample.jpg", "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"],
             ],
-            inputs=[image_input, image_reference, prompt, style_choice],
+            inputs=[input_image, prompt, lora_adapter],
             outputs=[output_image, seed],
             fn=infer_example,
             cache_examples=False,
-            label="Examples (Ensure local files exist)"
+            label="Examples (Ensure images exist in 'examples/' folder)"
         )
 
     run_button.click(
         fn=infer,
-        inputs=[image_input, image_reference, prompt, style_choice, seed, randomize_seed, guidance_scale, steps],
+        inputs=[input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
         outputs=[output_image, seed]
     )