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

# Initialize CUDA
zero = torch.Tensor([0]).cuda()
print(f"Initial device: {zero.device}")

# Model configurations
BASE_MODEL = "Qwen/Qwen-Image-Edit-2509"  # Qwen Image Edit model
BFS_LORA = "Alissonerdx/BFS-Best-Face-Swap"
BFS_LORA_FILENAME = "bfs_head_v3_qwen_image_edit_2509.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. ensure the head and body have correct anatomical proportions, and blend the skin tones, shadows, and lighting naturally so the final result appears as one coherent, realistic person."""

# Cache for loaded pipe
pipe_cache = None

@spaces.GPU(duration=60)
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,
    guidance_scale,
    seed
):
    """
    Perform head swap using Qwen-Image-Edit with multiple LoRAs
    """
    print(f"GPU device: {zero.device}")
    
    # 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 = DiffusionPipeline.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 = []
        
        # Always load BFS Face Swap LoRA (Qwen-specific version)
        if bfs_lora_scale > 0:
            print(f"Loading BFS Face Swap LoRA (Qwen version) with scale {bfs_lora_scale}")
            try:
                pipe.load_lora_weights(
                    BFS_LORA,
                    weight_name=BFS_LORA_FILENAME,  # Using the Qwen-specific file
                    adapter_name="bfs_face_swap"
                )
                adapters.append("bfs_face_swap")
                adapter_weights.append(bfs_lora_scale)
            except Exception as e:
                print(f"Warning: Could not load BFS LoRA: {e}")
                gr.Warning(f"BFS Face Swap LoRA could not be loaded: {e}")
        
        # 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 = Image.fromarray(body_image).convert("RGB")
        face_img = Image.fromarray(face_image).convert("RGB")
        
        # 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]}...")
        print(f"Using BFS LoRA file: {BFS_LORA_FILENAME}")
        
        # Generate the head swap
        result = pipe(
            image=body_img,  # Picture 1 - Body/Base
            prompt=final_prompt,
            input_image=face_img,  # Picture 2 - Face to swap
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            generator=generator
        ).images[0]
        
        # Create status message
        active_loras = []
        if bfs_lora_scale > 0:
            active_loras.append(f"BFS-Qwen-v3({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 | 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: #f0f0f0; padding: 10px; border-radius: 5px; font-family: monospace;}
    .lora-info {background-color: #e8f4ff; padding: 8px; border-radius: 5px; margin: 5px 0; font-size: 0.9em;}
""") as demo:
    gr.Markdown(
        """
        # 🎭 BFS-Best Face Swap with Qwen-Image-Edit-2509
        
        ## Professional Head Swap using Multiple LoRAs
        
        This interface uses:
        - **Base Model**: Qwen-Image-Edit-2509
        - **Primary LoRA**: BFS-Best Face Swap v3 (Qwen-optimized: `bfs_head_v3_qwen_image_edit_2509.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_v3_qwen_image_edit_2509.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 v3)",
                    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=True,
                        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=True,
                        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=30,
                    label="Inference Steps",
                    info="Higher = better quality but slower"
                )
                
                guidance_scale = gr.Slider(
                    minimum=1.0,
                    maximum=20.0,
                    step=0.5,
                    value=7.5,
                    label="Guidance Scale",
                    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
    )
    
    # Examples
    gr.Examples(
        examples=[
            [
                None,  # body_image
                None,  # face_image
                "",  # custom_prompt_addon
                1.0,  # bfs_lora_scale
                0.7,  # angles_lora_scale
                0.6,  # skin_lora_scale
                True,  # enable_angles_lora
                True,  # enable_skin_lora
                30,  # num_inference_steps
                7.5,  # guidance_scale
                42  # seed
            ],
            [
                None,
                None,
                "professional lighting, high resolution",
                1.2,
                0.8,
                0.5,
                True,
                True,
                40,
                8.0,
                123
            ],
            [
                None,
                None,
                "",
                0.9,
                0.0,
                0.0,
                False,
                False,
                25,
                7.0,
                -1
            ]
        ],
        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,
            guidance_scale,
            seed
        ],
        outputs=[output_image, status_text],
        fn=face_swap,
        cache_examples=False
    )
    
    # 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,
            guidance_scale,
            seed
        ],
        outputs=[output_image, status_text]
    )
    
    gr.Markdown(
        """
        ---
        ### 📚 Documentation
        
        **Model Chain:**
        1. **Qwen-Image-Edit-2509**: Advanced image editing base model
        2. **BFS-Best Face Swap v3**: 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-2509](https://huggingface.co/Qwen/Qwen-Image-Edit-2509)
        - [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)
        """
    )

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