app.py

import gradio as gr
import spaces
import torch
from diffusers import QwenImageEditPlusPipeline
from diffusers.utils import load_image
from PIL import Image
import numpy as np
from typing import Optional, Tuple, List

# Model configurations
BASE_MODEL = "Qwen/Qwen-Image-Edit-2511"  # Qwen Image Edit model
BFS_LORA = "Alissonerdx/BFS-Best-Face-Swap"
BFS_LORA_FILENAME = "bfs_head_v5_2511_merged_version_rank_32_fp32.safetensors"  # Qwen-specific version
ANGLES_LORA = "dx8152/Qwen-Edit-2509-Multiple-angles"
SKIN_LORA = "tlennon-ie/qwen-edit-skin"

# Fixed prompt for head swap
FIXED_PROMPT = """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, nose structure of Picture 2. copy the direction of the eye, head rotation, micro expressions from Picture 1, high quality, sharp details, 1080p"""

DEFAULT_NEGATIVE_PROMPT = "bad quality, noise, blurry, worst quality, low resolution, blur, distortion, unnatural blending, cartoon, illustration, painting"

# Cache for loaded pipe
pipe_cache = None

# Função auxiliar para redimensionar mantendo aspect ratio
def smart_resize(image, target_long_edge=1024):
    width, height = image.size
    
    # Calcular nova proporção mantendo o aspect ratio
    if width > height:
        new_width = target_long_edge
        new_height = int(height * (target_long_edge / width))
    else:
        new_height = target_long_edge
        new_width = int(width * (target_long_edge / height))
    
    # Arredondar para múltiplos de 32 (necessário para o modelo)
    new_width = (new_width // 32) * 32
    new_height = (new_height // 32) * 32
    
    # Redimensionar usando LANCZOS para alta qualidade
    resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    return resized_image, new_width, new_height

@spaces.GPU(duration=300)
def face_swap(
    body_image,
    face_image,
    custom_prompt_addon,
    bfs_lora_scale,
    angles_lora_scale,
    skin_lora_scale,
    enable_angles_lora,
    enable_skin_lora,
    num_inference_steps,
    cfg,
    seed
):
    """
    Perform head swap using Qwen-Image-Edit with multiple LoRAs
    """
    
    # Validate inputs
    if body_image is None or face_image is None:
        raise gr.Error("Please provide both body (Picture 1) and face (Picture 2) images")
    
    # Set seed for reproducibility
    if seed != -1:
        torch.manual_seed(seed)
        generator = torch.Generator(device="cuda").manual_seed(seed)
    else:
        generator = None
    
    try:
        global pipe_cache
        
        # Load the pipeline (only once)
        if pipe_cache is None:
            print(f"Loading pipeline: {BASE_MODEL}")
            pipe_cache = QwenImageEditPlusPipeline.from_pretrained(
                BASE_MODEL,
                torch_dtype=torch.bfloat16,  # Qwen uses bfloat16
                device_map="cuda"
            )
        
        pipe = pipe_cache
        
        # Prepare the LoRA adapters list
        adapters = []
        adapter_weights = []
        
        # Unload existing adapters to start fresh for this request
        try:
            pipe.unload_lora_weights()
        except:
            pass
        
        if bfs_lora_scale > 0:
            print(f"Loading BFS Face Swap LoRA (Qwen version) with scale {bfs_lora_scale}")
            try:
                from huggingface_hub import hf_hub_download
                import safetensors.torch
                import os
                
                # Verifica se já temos uma versão convertida em cache
                cache_path = "/tmp/bfs_lora_qwen_diffusers.safetensors"
                
                if not os.path.exists(cache_path):
                    print("Converting LoRA from ComfyUI to diffusers format...")
                    
                    # Baixa o original
                    original_path = hf_hub_download(
                        repo_id=BFS_LORA,
                        filename=BFS_LORA_FILENAME
                    )
                    
                    # Carrega e converte
                    original_dict = safetensors.torch.load_file(original_path)
                    converted_dict = {}
                    
                    for key, value in original_dict.items():
                        # Pula chaves incompatíveis
                        if "img_in.alpha" in key:
                            continue
                        
                        # Mantém apenas pesos relevantes para o transformer
                        if any(x in key for x in ["transformer", "diffusion_model", "unet", "lora"]):
                            # Ajusta a nomenclatura
                            new_key = key
                            if "diffusion_model" in key:
                                new_key = key.replace("diffusion_model", "transformer")
                            converted_dict[new_key] = value
                    
                    # Salva versão convertida
                    safetensors.torch.save_file(converted_dict, cache_path)
                    print(f"✅ LoRA converted and cached")
                
                # Carrega a versão convertida
                pipe.load_lora_weights(
                    cache_path,
                    adapter_name="bfs_face_swap"
                )
                
                adapters.append("bfs_face_swap")
                adapter_weights.append(bfs_lora_scale)
                print("✅ BFS LoRA loaded successfully (converted format)")
                
            except Exception as e:
                print(f"⚠️ Failed to convert/load LoRA: {e}")
                print("Continuing without BFS LoRA...")
        
        # Load Multiple Angles LoRA if enabled
        if enable_angles_lora and angles_lora_scale > 0:
            print(f"Loading Multiple Angles LoRA with scale {angles_lora_scale}")
            try:
                pipe.load_lora_weights(
                    ANGLES_LORA,
                    adapter_name="angles"
                )
                adapters.append("angles")
                adapter_weights.append(angles_lora_scale)
            except Exception as e:
                print(f"Warning: Could not load Angles LoRA: {e}")
                gr.Warning(f"Multiple Angles LoRA could not be loaded: {e}")
        
        # Load Skin LoRA if enabled
        if enable_skin_lora and skin_lora_scale > 0:
            print(f"Loading Skin LoRA with scale {skin_lora_scale}")
            try:
                pipe.load_lora_weights(
                    SKIN_LORA,
                    adapter_name="skin"
                )
                adapters.append("skin")
                adapter_weights.append(skin_lora_scale)
            except Exception as e:
                print(f"Warning: Could not load Skin LoRA: {e}")
                gr.Warning(f"Skin LoRA could not be loaded: {e}")
        
        # Set the active adapters
        if len(adapters) > 0:
            if len(adapters) == 1:
                pipe.set_adapters(adapters[0], adapter_weights=adapter_weights[0])
            else:
                pipe.set_adapters(adapters, adapter_weights=adapter_weights)
            print(f"Active LoRAs: {adapters} with weights {adapter_weights}")
        
        # Prepare images
        body_img_pil = Image.fromarray(body_image).convert("RGB")
        face_img_pil = Image.fromarray(face_image).convert("RGB")
        
        # --- LÓGICA DE REDIMENSIONAMENTO INTELIGENTE ---
        # Define o tamanho alvo baseado no maior lado (1024 é um bom equilíbrio, pode subir para 1280)
        # Isso corrige a distorção mantendo o aspect ratio correto
        TARGET_RESOLUTION = 1024 
        
        body_resized, target_w, target_h = smart_resize(body_img_pil, target_long_edge=TARGET_RESOLUTION)
        # Opcional: redimensionar a face para não ficar gigante ou minúscula comparada ao corpo
        face_resized, _, _ = smart_resize(face_img_pil, target_long_edge=TARGET_RESOLUTION)
        
        print(f"Original size: {body_img_pil.size} | Generation Target: {target_w}x{target_h}")
        
        # Combine fixed prompt with any additional instructions
        final_prompt = FIXED_PROMPT
        if custom_prompt_addon and custom_prompt_addon.strip():
            final_prompt = f"{FIXED_PROMPT} {custom_prompt_addon}"
        
        print(f"Using prompt: {final_prompt[:100]}...")
        
        # Qwen Image Edit uses a list for inputs: [body, face]
        input_images_list = [body_resized, face_resized]
        
        # Generate the head swap
        result = pipe(
            image=input_images_list,
            prompt=final_prompt,
            negative_prompt=DEFAULT_NEGATIVE_PROMPT, 
            true_cfg_scale=cfg,
            guidance_scale=1.0,
            height=target_h,  # FORÇA A ALTURA CORRETA
            width=target_w,   # FORÇA A LARGURA CORRETA
            num_inference_steps=num_inference_steps,
            generator=generator
        ).images[0]
        
        # Create status message
        active_loras = []
        if bfs_lora_scale > 0:
            active_loras.append(f"BFS-Qwen-v5({bfs_lora_scale:.2f})")
        if enable_angles_lora and angles_lora_scale > 0:
            active_loras.append(f"Angles({angles_lora_scale:.2f})")
        if enable_skin_lora and skin_lora_scale > 0:
            active_loras.append(f"Skin({skin_lora_scale:.2f})")
        
        status = f"✅ Head swap completed ({target_w}x{target_h}) | Active LoRAs: {', '.join(active_loras) if active_loras else 'None'}"
        
        return result, status
        
    except Exception as e:
        print(f"Error: {str(e)}")
        error_img = Image.new('RGB', (512, 512), color=(200, 50, 50))
        return error_img, f"❌ Error: {str(e)}"

# Create the Gradio interface
with gr.Blocks(title="BFS-Best Face Swap with Qwen", theme=gr.themes.Soft(), css="""
    .container {max-width: 1200px; margin: auto;}
    .image-container {border-radius: 10px; border: 2px dashed #ccc;}
    .fixed-prompt {background-color: #000000; padding: 10px; border-radius: 5px; font-family: monospace; color: #00ff00;}
    .lora-info {background-color: #000000; padding: 8px; border-radius: 5px; margin: 5px 0; font-size: 0.9em; color: white;}
    .footer-link {text-decoration: none !important; color: #5865F2 !important; font-weight: bold;}
    .footer-link:hover {text-decoration: underline !important;}
""") as demo:
    gr.Markdown(
        """
        # 🎭 BFS - Best Face Swap with Qwen-Image-Edit-2511
        
        This interface uses:
        - **Base Model**: Qwen-Image-Edit-2511
        - **Primary LoRA**: BFS-Best Face Swap v5 Merged R32 FP32 (Qwen-optimized: `bfs_head_v5_2511_merged_version_rank_32_fp32.safetensors`)
        - **Enhancement LoRAs**: Multiple Angles & Skin Blending
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📥 Input Images")
            
            with gr.Row():
                body_image = gr.Image(
                    label="👤 Picture 1: Body/Base Image",
                    type="numpy",
                    height=300,
                    elem_classes="image-container"
                )
                face_image = gr.Image(
                    label="😊 Picture 2: Head/Face to Swap",
                    type="numpy",
                    height=300,
                    elem_classes="image-container"
                )
            
            gr.Markdown("### 🎯 Fixed Head Swap Prompt")
            gr.Markdown(
                f'<div class="fixed-prompt">{FIXED_PROMPT}</div>',
                elem_classes="fixed-prompt"
            )
            
            custom_prompt_addon = gr.Textbox(
                label="Additional Instructions (Optional)",
                placeholder="Add any extra details or style instructions...",
                value="",
                lines=2
            )
            
            with gr.Accordion("🎛️ LoRA Controls", open=True):
                gr.Markdown("#### BFS Face Swap LoRA (Primary)")
                gr.Markdown(
                    '<div class="lora-info">📌 Using: bfs_head_v5_2511_merged_version_rank_32_fp32.safetensors</div>',
                    elem_classes="lora-info"
                )
                bfs_lora_scale = gr.Slider(
                    minimum=0.0,
                    maximum=1.5,
                    step=0.05,
                    value=1.0,
                    label="BFS Face Swap Strength (Qwen v5)",
                    info="Main face swapping LoRA optimized for Qwen - set to 0 to disable"
                )
                
                gr.Markdown("#### Enhancement LoRAs")
                
                with gr.Row():
                    enable_angles_lora = gr.Checkbox(
                        label="Enable Multiple Angles LoRA",
                        value=False,
                        info="Improves head angle matching"
                    )
                    angles_lora_scale = gr.Slider(
                        minimum=0.0,
                        maximum=1.5,
                        step=0.05,
                        value=0.7,
                        label="Multiple Angles Strength",
                        interactive=True
                    )
                
                with gr.Row():
                    enable_skin_lora = gr.Checkbox(
                        label="Enable Skin Blending LoRA",
                        value=False,
                        info="Improves skin tone matching"
                    )
                    skin_lora_scale = gr.Slider(
                        minimum=0.0,
                        maximum=1.5,
                        step=0.05,
                        value=0.6,
                        label="Skin Blending Strength",
                        interactive=True
                    )
            
            with gr.Accordion("⚙️ Generation Settings", open=False):
                num_inference_steps = gr.Slider(
                    minimum=10,
                    maximum=100,
                    step=5,
                    value=40,
                    label="Inference Steps",
                    info="Higher = better quality but slower"
                )
                
                cfg = gr.Slider(
                    minimum=1.0,
                    maximum=20.0,
                    step=0.5,
                    value=4.0,
                    label="Guidance Scale (CFG)",
                    info="How closely to follow the prompt"
                )
                
                seed = gr.Number(
                    value=-1,
                    label="Seed",
                    info="Use -1 for random, or specific number for reproducible results",
                    precision=0
                )
            
            generate_btn = gr.Button("🎨 Generate Head Swap", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Output")
            output_image = gr.Image(
                label="Result",
                type="pil",
                interactive=False,
                height=500
            )
            status_text = gr.Textbox(
                label="Status",
                interactive=False,
                max_lines=2,
                value="Ready to process..."
            )
            
            gr.Markdown(
                """
                ### 💡 Quick Tips:
                - **Picture 1**: Body/environment to keep
                - **Picture 2**: Face/head to transplant
                - **BFS Strength**: 0.8-1.2 for best results
                - **Angles LoRA**: Helps with different head angles
                - **Skin LoRA**: Smooths skin tone transitions
                """
            )
    
    # Interaction logic for enabling/disabling LoRA controls
    def toggle_angles(enabled):
        return gr.update(interactive=enabled)
    
    def toggle_skin(enabled):
        return gr.update(interactive=enabled)
    
    enable_angles_lora.change(
        fn=toggle_angles,
        inputs=enable_angles_lora,
        outputs=angles_lora_scale
    )
    
    enable_skin_lora.change(
        fn=toggle_skin,
        inputs=enable_skin_lora,
        outputs=skin_lora_scale
    )
    
    # Event handlers
    generate_btn.click(
        fn=face_swap,
        inputs=[
            body_image,
            face_image,
            custom_prompt_addon,
            bfs_lora_scale,
            angles_lora_scale,
            skin_lora_scale,
            enable_angles_lora,
            enable_skin_lora,
            num_inference_steps,
            cfg,
            seed
        ],
        outputs=[output_image, status_text]
    )
    
    gr.Markdown(
        """
        ---
        ### 📚 Documentation
        
        **Model Chain:**
        1. **Qwen-Image-Edit-2511**: Advanced image editing base model
        2. **BFS-Best Face Swap v5**: Primary face swapping LoRA
        3. **Multiple Angles**: Improves head angle matching
        4. **Skin Blending**: Natural skin tone transitions
        
        **LoRA Settings Guide:**
        - **All at 0**: Uses only base Qwen model
        - **BFS only (1.0)**: Basic face swap
        - **BFS + Angles**: Better angle matching
        - **BFS + Skin**: Better skin blending
        - **All enabled**: Maximum quality (slower)
        
        ### 🔗 Resources:
        - [Qwen-Image-Edit-2511](https://huggingface.co/Qwen/Qwen-Image-Edit-2511)
        - [BFS-Best Face Swap](https://huggingface.co/Alissonerdx/BFS-Best-Face-Swap)
        - [Multiple Angles LoRA](https://huggingface.co/dx8152/Qwen-Edit-2509-Multiple-angles)
        - [Skin Blending LoRA](https://huggingface.co/tlennon-ie/qwen-edit-skin)
        """
    )

    gr.HTML(
        """
        <div style="text-align: center; margin-top: 40px; padding: 20px; border-top: 1px solid #ccc;">
            <a href="https://buymeacoffee.com/nrdx" target="_blank" style="display: inline-block; margin-bottom: 10px;">
                <img src="https://img.buymeacoffee.com/button-api/?text=Buy me a coffee&emoji=&slug=nrdx&button_colour=FFDD00&font_colour=000000&font_family=Cookie&outline_colour=000000&coffee_colour=ffffff" alt="Buy Me A Coffee" height="40">
            </a>
            <div style="display: flex; justify-content: center; gap: 20px; margin-top: 10px; font-weight: bold;">
                <a href="https://discord.gg/uYu3KzJcKB" target="_blank" class="footer-link">
                    🇧🇷 Discord Toca da IA
                </a>
                <span>|</span>
                <a href="https://discord.gg/ThrfwKcr3F" target="_blank" class="footer-link">
                    🇧🇷 Discord Hoje na IA
                </a>
            </div>
        </div>
        """
    )

# Launch the app
if __name__ == "__main__":
    demo.queue(max_size=10)
    demo.launch()