Update app.py

app.py

@@ -1,16 +1,15 @@
 import os
-import math  # Added to fix NameError
-import gradio as gr
-import numpy as np
 import random
+import sys
+from typing import Sequence, Mapping, Any, Union
 import torch
-import spaces
 from PIL import Image
-from diffusers import FlowMatchEulerDiscreteScheduler
-from optimization import optimize_pipeline_
-from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
-from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
-from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
+from huggingface_hub import hf_hub_download
+import spaces
+import subprocess
+import gradio
+import gradio_client
+import gradio as gr
 import requests  # For translation API
 
 # --- Translation Function ---
@@ -37,114 +36,133 @@ def translate_albanian_to_english(text):
                 raise gr.Error("Translation failed. Please try again.")
     raise gr.Error("Translation failed. Please try again.")
 
-# --- Model Loading ---
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# --- Monkey-patch for gradio_client ---
+import gradio_client.utils as _gc_utils
 
-# Scheduler configuration for Lightning
-scheduler_config = {
-    "base_image_seq_len": 256,
-    "base_shift": math.log(3),
-    "invert_sigmas": False,
-    "max_image_seq_len": 8192,
-    "max_shift": math.log(3),
-    "num_train_timesteps": 1000,
-    "shift": 1.0,
-    "shift_terminal": None,
-    "stochastic_sampling": False,
-    "time_shift_type": "exponential",
-    "use_beta_sigmas": False,
-    "use_dynamic_shifting": True,
-    "use_exponential_sigmas": False,
-    "use_karras_sigmas": False,
-}
+_orig_js2pt = _gc_utils._json_schema_to_python_type
+_orig_get_type = _gc_utils.get_type
 
-# Initialize scheduler with Lightning config
-scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
+def _safe_json_schema_to_python_type(schema, defs=None):
+    if isinstance(schema, bool):
+        return "Any"
+    return _orig_js2pt(schema, defs)
 
-# Load the model pipeline
-pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2509", 
-                                                 scheduler=scheduler,
-                                                 torch_dtype=dtype).to(device)
-pipe.load_lora_weights(
-    "lightx2v/Qwen-Image-Lightning", 
-    weight_name="Qwen-Image-Lightning-4steps-V2.0.safetensors"
-)
-pipe.fuse_lora()
+def _safe_get_type(schema):
+    if isinstance(schema, bool):
+        return "Any"
+    return _orig_get_type(schema)
 
-# Apply the same optimizations from the first version
-pipe.transformer.__class__ = QwenImageTransformer2DModel
-pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+_gc_utils._json_schema_to_python_type = _safe_json_schema_to_python_type
+_gc_utils.get_type = _safe_get_type
 
-# --- Ahead-of-time compilation ---
-optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")
+print("gradio version:", gradio.__version__)
+print("gradio_client version:", gradio_client.__version__)
 
-# --- UI Constants and Helpers ---
-MAX_SEED = np.iinfo(np.int32).max
-QUALITY_PROMPT = ", high quality, detailed, vibrant, professional lighting"
+# --- Model Downloads ---
+hf_hub_download(repo_id="ezioruan/inswapper_128.onnx", filename="inswapper_128.onnx", local_dir="models/insightface")
+hf_hub_download(repo_id="martintomov/comfy", filename="facerestore_models/GPEN-BFR-512.onnx", local_dir="models/facerestore_models")
+hf_hub_download(repo_id="darkeril/collection", filename="detection_Resnet50_Final.pth", local_dir="models/facedetection")
+hf_hub_download(repo_id="gmk123/GFPGAN", filename="parsing_parsenet.pth", local_dir="models/facedetection")
+hf_hub_download(repo_id="MonsterMMORPG/tools", filename="1k3d68.onnx", local_dir="models/insightface/models/buffalo_l")
+hf_hub_download(repo_id="MonsterMMORPG/tools", filename="2d106det.onnx", local_dir="models/insightface/models/buffalo_l")
+hf_hub_download(repo_id="maze/faceX", filename="det_10g.onnx", local_dir="models/insightface/models/buffalo_l")
+hf_hub_download(repo_id="typhoon01/aux_models", filename="genderage.onnx", local_dir="models/insightface/models/buffalo_l")
+hf_hub_download(repo_id="maze/faceX", filename="w600k_r50.onnx", local_dir="models/insightface/models/buffalo_l")
 
-# --- Main Inference Function ---
-@spaces.GPU(duration=40)
-def infer(
-    images,
-    prompt,
-    progress=gr.Progress(track_tqdm=True),
-):
-    """
-    Generates an image using the local Qwen-Image diffusers pipeline.
-    """
-    negative_prompt = ""  # Empty as in original
-    seed = random.randint(0, MAX_SEED)  # Default: randomize_seed=True
-    true_guidance_scale = 1.0  # Default
-    num_inference_steps = 4  # Default
-    height = None  # Default
-    width = None  # Default
-    num_images_per_prompt = 1  # Default
+# --- Utility Functions ---
+def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
+    try:
+        return obj[index]
+    except KeyError:
+        return obj["result"][index]
+
+def find_path(name: str, path: str = None) -> str:
+    if path is None:
+        path = os.getcwd()
+    if name in os.listdir(path):
+        path_name = os.path.join(path, name)
+        print(f"{name} found: {path_name}")
+        return path_name
+    parent_directory = os.path.dirname(path)
+    if parent_directory == path:
+        return None
+    return find_path(name, parent_directory)
+
+def add_comfyui_directory_to_sys_path() -> None:
+    comfyui_path = find_path("ComfyUI")
+    if comfyui_path is not None and os.path.isdir(comfyui_path):
+        sys.path.append(comfyui_path)
+        print(f"'{comfyui_path}' added to sys.path")
 
-    # Translate prompt from Albanian to English
-    prompt_final = translate_albanian_to_english(prompt.strip()) + QUALITY_PROMPT
+def add_extra_model_paths() -> None:
+    try:
+        from main import load_extra_path_config
+    except ImportError:
+        print("Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead.")
+        from utils.extra_config import load_extra_path_config
+    extra_model_paths = find_path("extra_model_paths.yaml")
+    if extra_model_paths is not None:
+        load_extra_path_config(extra_model_paths)
+    else:
+        print("Could not find the extra_model_paths config file.")
 
-    # Set up the generator for reproducibility
-    generator = torch.Generator(device=device).manual_seed(seed)
-    
-    # Load input images into PIL Images
-    pil_images = []
-    if images is not None:
-        for item in images:
-            try:
-                if isinstance(item[0], Image.Image):
-                    pil_images.append(item[0].convert("RGB"))
-                elif isinstance(item[0], str):
-                    pil_images.append(Image.open(item[0]).convert("RGB"))
-                elif hasattr(item, "name"):
-                    pil_images.append(Image.open(item.name).convert("RGB"))
-            except Exception:
-                continue
+add_comfyui_directory_to_sys_path()
+add_extra_model_paths()
 
-    if height == 256 and width == 256:
-        height, width = None, None
-    print(f"Calling pipeline with prompt: '{prompt_final}'")
-    print(f"Negative Prompt: '{negative_prompt}'")
-    print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}, Size: {width}x{height}")
+def import_custom_nodes() -> None:
+    import asyncio
+    import execution
+    from nodes import init_extra_nodes
+    import server
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    server_instance = server.PromptServer(loop)
+    execution.PromptQueue(server_instance)
+    init_extra_nodes()
 
-    # Generate the image
-    image = pipe(
-        image=pil_images if len(pil_images) > 0 else None,
-        prompt=prompt_final,
-        height=height,
-        width=width,
-        negative_prompt=negative_prompt,
-        num_inference_steps=num_inference_steps,
-        generator=generator,
-        true_cfg_scale=true_guidance_scale,
-        num_images_per_prompt=num_images_per_prompt,
-    ).images
+import_custom_nodes()
+from nodes import NODE_CLASS_MAPPINGS
 
-    return image
+# --- Main Inference Function ---
+@spaces.GPU(duration=20)
+def generate_image(source_image, target_image):
+    """
+    Performs face swapping between source and target images using ReActorFaceSwap.
+    """
+    restore_strength = 1.0  # Default
+    target_index = 0  # Default
+
+    with torch.inference_mode():
+        loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
+        loadimage_1 = loadimage.load_image(image=target_image)
+        loadimage_3 = loadimage.load_image(image=source_image)
+        reactorfaceswap = NODE_CLASS_MAPPINGS["ReActorFaceSwap"]()
+        saveimage = NODE_CLASS_MAPPINGS["SaveImage"]()
+        reactorfaceswap_2 = reactorfaceswap.execute(
+            enabled=True,
+            swap_model="inswapper_128.onnx",
+            facedetection="retinaface_resnet50",
+            face_restore_model="GPEN-BFR-512.onnx",
+            face_restore_visibility=restore_strength,
+            codeformer_weight=0.5,
+            detect_gender_input="no",
+            detect_gender_source="no",
+            input_faces_index=str(target_index),
+            source_faces_index="0",
+            console_log_level=1,
+            input_image=get_value_at_index(loadimage_1, 0),
+            source_image=get_value_at_index(loadimage_3, 0),
+        )
+        saveimage_4 = saveimage.save_images(
+            filename_prefix="ComfyUI",
+            images=get_value_at_index(reactorfaceswap_2, 0),
+        )
+        saved_path = f"output/{saveimage_4['ui']['images'][0]['filename']}"
+        return Image.open(saved_path)
 
 # --- Gradio User Interface ---
 def create_demo():
-    with gr.Blocks(css="", title="Qwen Image Editor") as demo:
+    with gr.Blocks(css="", title="Face Swapper") as app:
         gr.HTML("""
         <style>
         @import url('https://fonts.googleapis.com/css2?family=Orbitron:wght@400;600;700&display=swap');
@@ -202,12 +220,12 @@ def create_demo():
             margin: 0.75rem 0;
             width: 100%;
         }
-        .gr-gallery {
+        .gr-image {
             width: 100%;
             border: 1px solid #FFFFFF;
             border-radius: 4px;
         }
-        input, textarea {
+        input {
             background: #000000;
             color: #FFFFFF;
             border: 1px solid #FFFFFF;
@@ -216,7 +234,7 @@ def create_demo():
             width: 100%;
             box-sizing: border-box;
         }
-        input:hover, textarea:hover {
+        input:hover {
             box-shadow: 0 0 8px rgba(255, 255, 255, 0.3);
             transition: box-shadow 0.3s;
         }
@@ -279,44 +297,44 @@ def create_demo():
         """)
 
         with gr.Row(elem_id="general_items"):
-            gr.Markdown("# Qwen Image Editor")
-            gr.Markdown("Edit your images with precise instructions", elem_id="subtitle")
+            gr.Markdown("# Face Swapper")
+            gr.Markdown("Swap faces in photos with advanced AI technology", elem_id="subtitle")
             with gr.Column(elem_id="input_column"):
-                input_images = gr.Gallery(
-                    label="Input Images",
-                    show_label=True,
-                    type="pil",
-                    interactive=True,
-                    elem_classes=["gradio-component", "gr-gallery"]
-                )
-                result = gr.Gallery(
-                    label="Result",
-                    show_label=True,
-                    type="pil",
-                    elem_classes=["gradio-component", "gr-gallery"]
+                source_image = gr.Image(
+                    label="Source Image",
+                    type="filepath",
+                    sources=["upload", "clipboard"],
+                    elem_classes=["gradio-component", "gr-image"]
                 )
-                prompt = gr.Textbox(
-                    label="Prompt",
-                    placeholder="Describe the edit instruction",
-                    lines=3,
-                    elem_classes="gradio-component"
+                target_image = gr.Image(
+                    label="Target Image",
+                    type="filepath",
+                    sources=["upload", "clipboard"],
+                    elem_classes=["gradio-component", "gr-image"]
                 )
-                run_button = gr.Button(
-                    "Edit!",
+                generate_btn = gr.Button(
+                    "Generate",
                     variant="primary",
                     elem_classes="gradio-component"
                 )
+                output_image = gr.Image(
+                    label="Generated Image",
+                    show_download_button=True,
+                    show_share_button=False,
+                    interactive=False,
+                    elem_classes=["gradio-component", "gr-image"]
+                )
 
         gr.on(
-            triggers=[run_button.click, prompt.submit],
-            fn=infer,
-            inputs=[input_images, prompt],
-            outputs=[result],
+            triggers=[generate_btn.click],
+            fn=generate_image,
+            inputs=[source_image, target_image],
+            outputs=[output_image]
         )
 
-    return demo
+    return app
 
 if __name__ == "__main__":
     print(f"Gradio version: {gr.__version__}")
-    demo = create_demo()
-    demo.queue().launch(share=True)
+    app = create_demo()
+    app.queue().launch(share=True)