import gradio as gr
import torch
from PIL import Image
from diffusers import QwenImageEditPipeline
import spaces

# 1. Cargar el modelo base y el adaptador (LoRA)
# Esto ocurre cuando el Space arranca, preparándolo en memoria
pipe = QwenImageEditPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2511", 
    torch_dtype=torch.bfloat16
)
pipe.load_lora_weights("WarmBloodAban/AnyoneCosplay")
pipe.to("cuda")

# 2. La función generadora con Zero-GPU
# @spaces.GPU le dice a Hugging Face: "préstame una GPU potente solo para estos segundos"
@spaces.GPU
def generate_cosplay(human_img, anime_img):
    # El modelo espera una sola imagen unida. 
    # Pegamos la foto humana (Figura 1) y el anime (Figura 2) lado a lado.
    w1, h1 = human_img.size
    w2, h2 = anime_img.size
    
    # Ajustamos la altura de la imagen anime para que coincida con la humana
    new_w2 = int((h1 / h2) * w2)
    anime_img_resized = anime_img.resize((new_w2, h1))
    
    combined_img = Image.new('RGB', (w1 + new_w2, h1))
    combined_img.paste(human_img, (0, 0))
    combined_img.paste(anime_img_resized, (w1, 0))
    
    # Palabra clave de activación requerida por el autor del modelo
    prompt = "Let the person in Figure 1 cosplay the role in Figure 2"
    
    with torch.inference_mode():
        # Ejecutamos la inferencia con los parámetros recomendados para Qwen
        output = pipe(
            image=combined_img,
            prompt=prompt,
            num_inference_steps=50,
            true_cfg_scale=4.0
        ).images[0]
        
    return output

# 3. Interfaz gráfica simple y ligera
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎭 Virtual Cosplay (Zero-GPU)")
    gr.Markdown("Sube tu foto y la del personaje de anime. El modelo transferirá la ropa manteniendo tu rostro.")
    
    with gr.Row():
        human_in = gr.Image(type="pil", label="Tu foto (Figura 1)")
        anime_in = gr.Image(type="pil", label="Personaje Anime (Figura 2)")
        
    btn = gr.Button("✨ Generar Cosplay", variant="primary")
    output_image = gr.Image(label="Resultado final")
    
    btn.click(fn=generate_cosplay, inputs=[human_in, anime_in], outputs=output_image)

# Iniciar la aplicación
demo.launch()