Update app.py

app.py

@@ -10,6 +10,10 @@ from typing import Iterable
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 
+# -------------------------
+# Theme
+# -------------------------
+
 colors.orange_red = colors.Color(
     name="orange_red",
     c50="#FFF0E5",
@@ -78,6 +82,10 @@ class OrangeRedTheme(Soft):
 
 orange_red_theme = OrangeRedTheme()
 
+# -------------------------
+# Device
+# -------------------------
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
@@ -88,9 +96,12 @@ print("cuda device count:", torch.cuda.device_count())
 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)
 
+# -------------------------
+# Pipeline
+# -------------------------
+
 from diffusers import FlowMatchEulerDiscreteScheduler
 from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
@@ -109,7 +120,6 @@ pipe = QwenImageEditPlusPipeline.from_pretrained(
     torch_dtype=dtype,
 ).to(device)
 
-# Apply FA3 Optimization
 try:
     pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
     print("Flash Attention 3 Processor set successfully.")
@@ -119,7 +129,7 @@ except Exception as e:
 MAX_SEED = np.iinfo(np.int32).max
 
 # -------------------------
-# LoRA adapters + presets
+# LoRAs + presets
 # -------------------------
 
 NONE_LORA = "None"
@@ -170,6 +180,13 @@ ADAPTER_SPECS = {
         "adapter_name": "HRPortrait",
         "strength": 1.0,
     },
+    "BFS-Best-FaceSwap": {
+        "type": "single",
+        "repo": "Alissonerdx/BFS-Best-Face-Swap",
+        "weights": "bfs_head_v5_2511_merged_version_rank_16_fp16.safetensors",
+        "adapter_name": "BFS-Best-Faceswap",
+        "strength": 1.0,
+    },
     "Multiple-Angles": {
         "type": "single",
         "repo": "dx8152/Qwen-Edit-2509-Multiple-angles",
@@ -226,14 +243,13 @@ ADAPTER_SPECS = {
         "adapter_name": "upscale-image",
         "strength": 1.0,
     },
-    # ✅ We'll auto-bump requested resolution when this is selected
     "Upscale2K": {
         "type": "single",
         "repo": "valiantcat/Qwen-Image-Edit-2509-Upscale2K",
         "weights": "qwen_image_edit_2509_upscale.safetensors",
         "adapter_name": "upscale-2k",
         "strength": 1.0,
-        "target_long_edge": 2048,  # <-- add this
+        "target_long_edge": 2048,
     },
 }
 
@@ -243,16 +259,36 @@ LORA_PRESET_PROMPTS = {
     "AnyPose": "Make the person in image 1 do the exact same pose of the person in image 2. Changing the style and background of the image of the person in image 1 is undesirable, so don't do it. The new pose should be pixel accurate to the pose we are trying to copy. The position of the arms and head and legs should be the same as the pose we are trying to copy. Change the field of view and angle to match exactly image 2. Head tilt and eye gaze pose should match the person in image 2.",
     "Hyperrealistic-Portrait": "Transform the image into an ultra-realistic photorealistic portrait with strict identity preservation, facing straight to the camera. Enhance pore-level skin textures, realistic moisture effects, and natural wet hair clumping against the skin. Apply cool-toned soft-box lighting with subtle highlights and shadows, maintain realistic green-hazel eye catchlights without synthetic gloss, and preserve soft natural lip texture. Use shallow depth of field with a clean bokeh background, an 85mm macro photographic look, and raw photo grading without retouching to maintain realism and original details.",
     "Upscale2K": "Upscale this picture to 4K resolution.",
+    "BFS-Best-FaceSwap": "head_swap: start with Picture 1 as the base image, keeping its lighting, environment, and background. remove the head from Picture 1 completely and replace it with the head from Picture 2, strictly preserving the hair, eye color, and nose structure of Picture 2. copy the eye direction, head rotation, and micro-expressions from Picture 1. high quality, sharp details, 4k",
+}
+
+# Which LoRAs require 2 images (and how to label them in UI)
+# AnyPose: image1=character/base, image2=pose reference  [oai_citation:2‡huggingface.co](https://huggingface.co/lilylilith/AnyPose)
+# BFS Head V5: inverted order -> image1=body/base, image2=face/head donor  [oai_citation:3‡huggingface.co](https://huggingface.co/Alissonerdx/BFS-Best-Face-Swap)
+MULTI_INPUT_SPECS = {
+    "AnyPose": {
+        "image1_label": "Image 1 (Character / base)",
+        "image2_label": "Image 2 (Pose reference)",
+        "help": "AnyPose needs 2 images: Image 1 is the character/base, Image 2 is the pose you want to copy.",
+    },
+    "BFS-Best-FaceSwap": {
+        "image1_label": "Image 1 (Body / base)",
+        "image2_label": "Image 2 (Face/Head donor)",
+        "help": "BFS Head V5 needs 2 images (inverted order): Image 1 is the body/base, Image 2 is the face/head donor.",
+    },
 }
 
-# Track what is currently loaded in memory (adapter_name values)
+# Track loaded adapters by adapter_name
 LOADED_ADAPTERS = set()
 
+# -------------------------
+# Utility: dimensions / resolution
+# -------------------------
+
 def _round8(x: int) -> int:
     return max(8, (int(x) // 8) * 8)
 
 def compute_dimensions(image: Image.Image, long_edge: int) -> tuple[int, int]:
-    """Preserve aspect ratio, set the longest edge to long_edge, then round to multiples of 8."""
     w, h = image.size
     if w >= h:
         new_w = long_edge
@@ -262,21 +298,46 @@ def compute_dimensions(image: Image.Image, long_edge: int) -> tuple[int, int]:
         new_w = int(round(long_edge * (w / h)))
     return _round8(new_w), _round8(new_h)
 
-def update_dimensions_on_upload(image):
-    if image is None:
-        return 1024, 1024
-    # default behavior: fit to 1024 long edge
-    return compute_dimensions(image, 1024)
+def get_target_long_edge_for_lora(lora_adapter: str) -> int:
+    spec = ADAPTER_SPECS.get(lora_adapter, {})
+    return int(spec.get("target_long_edge", 1024))
+
+# -------------------------
+# UI helpers
+# -------------------------
 
-def on_lora_change(selected_lora, current_prompt):
-    if selected_lora == NONE_LORA:
-        return gr.update(value=current_prompt)
+def ui_on_lora_change(selected_lora: str, current_prompt: str):
+    # Prompt autofill (only if empty)
+    new_prompt = current_prompt
+    if selected_lora != NONE_LORA:
+        preset = LORA_PRESET_PROMPTS.get(selected_lora, "")
+        if preset and (current_prompt is None or str(current_prompt).strip() == ""):
+            new_prompt = preset
+
+    # Multi-image UI toggles
+    if selected_lora in MULTI_INPUT_SPECS:
+        spec = (MULTI_INPUT_SPECS[selected_lora] or {})
+        img1_label = spec.get("image1_label", "Image 1")
+        img2_label = spec.get("image2_label", "Image 2")
+        help_txt = spec.get("help", "")
+        return (
+            gr.update(value=new_prompt),
+            gr.update(visible=True, label=img1_label),
+            gr.update(visible=True, label=img2_label),
+            gr.update(visible=True, value=help_txt),
+        )
 
-    preset = LORA_PRESET_PROMPTS.get(selected_lora, "")
-    if preset and (current_prompt is None or str(current_prompt).strip() == ""):
-        return gr.update(value=preset)
+    # Default: single image only
+    return (
+        gr.update(value=new_prompt),
+        gr.update(visible=True, label="Upload Image"),
+        gr.update(visible=False),
+        gr.update(visible=False, value=""),
+    )
 
-    return gr.update(value=current_prompt)
+# -------------------------
+# LoRA loading + activation
+# -------------------------
 
 def _ensure_loaded_and_get_active_adapters(selected_lora: str):
     spec = ADAPTER_SPECS.get(selected_lora)
@@ -331,15 +392,14 @@ def _ensure_loaded_and_get_active_adapters(selected_lora: str):
 
     return adapter_names, adapter_weights
 
-def get_target_long_edge_for_lora(lora_adapter: str) -> int:
-    """Return desired long-edge resolution for this selection (default 1024)."""
-    if lora_adapter == "Upscale2K":
-        return int(ADAPTER_SPECS["Upscale2K"].get("target_long_edge", 2048))
-    return 1024
+# -------------------------
+# Inference
+# -------------------------
 
 @spaces.GPU
 def infer(
-    input_image,
+    input_image_1,
+    input_image_2,
     prompt,
     lora_adapter,
     seed,
@@ -352,10 +412,15 @@ def infer(
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
 
-    if input_image is None:
-        raise gr.Error("Please upload an image to edit.")
+    if input_image_1 is None:
+        raise gr.Error("Please upload Image 1.")
+
+    # Determine whether we need a second image
+    needs_two = lora_adapter in MULTI_INPUT_SPECS
+    if needs_two and input_image_2 is None:
+        raise gr.Error("This LoRA needs Image 2. Please upload Image 2 as well.")
 
-    # Handle "None"
+    # Handle "None" LoRA
     if lora_adapter == NONE_LORA:
         try:
             pipe.set_adapters([], adapter_weights=[])
@@ -370,20 +435,30 @@ def infer(
         seed = random.randint(0, MAX_SEED)
 
     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_image.convert("RGB")
+    img1 = input_image_1.convert("RGB")
+    img2 = input_image_2.convert("RGB") if input_image_2 is not None else None
 
-    # ✅ Automatic resolution bump for Upscale2K
+    # Resolution based on Image 1 (base/body)
     target_long_edge = get_target_long_edge_for_lora(lora_adapter)
-    width, height = compute_dimensions(original_image, target_long_edge)
+    width, height = compute_dimensions(img1, target_long_edge)
+
+    # Build image input for pipeline
+    if needs_two:
+        # AnyPose expects: [image1(character/base), image2(pose)]  [oai_citation:4‡huggingface.co](https://huggingface.co/lilylilith/AnyPose)
+        # BFS Head V5 expects: [image1(body/base), image2(face)] inverted order  [oai_citation:5‡huggingface.co](https://huggingface.co/Alissonerdx/BFS-Best-Face-Swap)
+        pipe_image = [img1, img2]
+    else:
+        pipe_image = img1
 
     try:
         result = pipe(
-            image=original_image,
+            image=pipe_image,
             prompt=prompt,
             negative_prompt=negative_prompt,
             height=height,
@@ -401,16 +476,19 @@ def infer(
             torch.cuda.empty_cache()
 
 @spaces.GPU
-def infer_example(input_image, prompt, lora_adapter):
-    if input_image is None:
+def infer_example(input_image_1, input_image_2, prompt, lora_adapter):
+    if input_image_1 is None:
         return None, 0
 
-    input_pil = input_image.convert("RGB")
     guidance_scale = 1.0
     steps = 4
-    result, seed = infer(input_pil, prompt, lora_adapter, 0, True, guidance_scale, steps)
+    result, seed = infer(input_image_1, input_image_2, prompt, lora_adapter, 0, True, guidance_scale, steps)
     return result, seed
 
+# -------------------------
+# UI
+# -------------------------
+
 css = """
 #col-container {
     margin: 0 auto;
@@ -430,7 +508,10 @@ with gr.Blocks() as demo:
 
         with gr.Row(equal_height=True):
             with gr.Column():
-                input_image = gr.Image(label="Upload Image", type="pil", height=290)
+                input_image_1 = gr.Image(label="Upload Image", type="pil", height=290)
+                input_image_2 = gr.Image(label="Image 2 (only for certain LoRAs)", type="pil", height=290, visible=False)
+
+                multi_help = gr.Markdown(value="", visible=False)
 
                 prompt = gr.Text(
                     label="Edit Prompt",
@@ -457,33 +538,35 @@ with gr.Blocks() as demo:
                     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)
 
+        # LoRA change: prompt autofill + toggle 2nd image visibility/labels
         lora_adapter.change(
-            fn=on_lora_change,
+            fn=ui_on_lora_change,
             inputs=[lora_adapter, prompt],
-            outputs=[prompt],
+            outputs=[prompt, input_image_1, input_image_2, multi_help],
         )
 
+        # Examples now include Image 2 as well (kept as duplicate for single-image examples)
         gr.Examples(
             examples=[
-                ["examples/1.jpg", "Transform into anime.", "Photo-to-Anime"],
-                ["examples/5.jpg", "Remove shadows and relight the image using soft lighting.", "Light-Restoration"],
-                ["examples/4.jpg", "Use a subtle golden-hour filter with smooth light diffusion.", "Relight"],
-                ["examples/2.jpeg", "Rotate the camera 45 degrees to the left.", "Multiple-Angles"],
-                ["examples/12.jpg", "flatcolor Desaturate the image and lower the contrast to create a flat, ungraded look similar to a camera log profile. Preserve details in the highlights and shadows.", "Flat-Log"],
-                ["examples/7.jpg", "Light source from the Right Rear", "Multi-Angle-Lighting"],
-                ["examples/10.jpeg", "Upscale the image.", "Upscale-Image"],
-                ["examples/7.jpg", "Light source from the Below", "Multi-Angle-Lighting"],
-                ["examples/2.jpeg", "Switch the camera to a top-down right corner view.", "Multiple-Angles"],
-                ["examples/9.jpg", "The camera moves slightly forward as sunlight breaks through the clouds, casting a soft glow around the character's silhouette in the mist. Realistic cinematic style, atmospheric depth.", "Next-Scene"],
-                ["examples/8.jpg", "Make the subjects skin details more prominent and natural.", "Edit-Skin"],
-                ["examples/6.jpg", "Switch the camera to a bottom-up view.", "Multiple-Angles"],
-                ["examples/6.jpg", "Rotate the camera 180 degrees upside down.", "Multiple-Angles"],
-                ["examples/4.jpg", "Rotate the camera 45 degrees to the right.", "Multiple-Angles"],
-                ["examples/4.jpg", "Switch the camera to a top-down view.", "Multiple-Angles"],
-                ["examples/4.jpg", "Switch the camera to a wide-angle lens.", "Multiple-Angles"],
-                ["examples/11.jpg", "Upscale this picture to 4K resolution.", "Upscale2K"],
+                ["examples/1.jpg", "examples/1.jpg", "Transform into anime.", "Photo-to-Anime"],
+                ["examples/5.jpg", "examples/5.jpg", "Remove shadows and relight the image using soft lighting.", "Light-Restoration"],
+                ["examples/4.jpg", "examples/4.jpg", "Use a subtle golden-hour filter with smooth light diffusion.", "Relight"],
+                ["examples/2.jpeg", "examples/2.jpeg", "Rotate the camera 45 degrees to the left.", "Multiple-Angles"],
+                ["examples/12.jpg", "examples/12.jpg", "flatcolor Desaturate the image and lower the contrast to create a flat, ungraded look similar to a camera log profile. Preserve details in the highlights and shadows.", "Flat-Log"],
+                ["examples/7.jpg", "examples/7.jpg", "Light source from the Right Rear", "Multi-Angle-Lighting"],
+                ["examples/10.jpeg", "examples/10.jpeg", "Upscale the image.", "Upscale-Image"],
+                ["examples/7.jpg", "examples/7.jpg", "Light source from the Below", "Multi-Angle-Lighting"],
+                ["examples/2.jpeg", "examples/2.jpeg", "Switch the camera to a top-down right corner view.", "Multiple-Angles"],
+                ["examples/9.jpg", "examples/9.jpg", "The camera moves slightly forward as sunlight breaks through the clouds, casting a soft glow around the character's silhouette in the mist. Realistic cinematic style, atmospheric depth.", "Next-Scene"],
+                ["examples/8.jpg", "examples/8.jpg", "Make the subjects skin details more prominent and natural.", "Edit-Skin"],
+                ["examples/6.jpg", "examples/6.jpg", "Switch the camera to a bottom-up view.", "Multiple-Angles"],
+                ["examples/6.jpg", "examples/6.jpg", "Rotate the camera 180 degrees upside down.", "Multiple-Angles"],
+                ["examples/4.jpg", "examples/4.jpg", "Rotate the camera 45 degrees to the right.", "Multiple-Angles"],
+                ["examples/4.jpg", "examples/4.jpg", "Switch the camera to a top-down view.", "Multiple-Angles"],
+                ["examples/4.jpg", "examples/4.jpg", "Switch the camera to a wide-angle lens.", "Multiple-Angles"],
+                ["examples/11.jpg", "examples/11.jpg", "Upscale this picture to 4K resolution.", "Upscale2K"],
             ],
-            inputs=[input_image, prompt, lora_adapter],
+            inputs=[input_image_1, input_image_2, prompt, lora_adapter],
             outputs=[output_image, seed],
             fn=infer_example,
             cache_examples=False,
@@ -492,7 +575,7 @@ with gr.Blocks() as demo:
 
     run_button.click(
         fn=infer,
-        inputs=[input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
+        inputs=[input_image_1, input_image_2, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
         outputs=[output_image, seed],
     )