import gradio as gr
import torch
from PIL import Image
import spaces  # DESCOMENTADO para ZeroGPU
from diffusers import QwenImageEditPipeline

device = "cuda" if torch.cuda.is_available() else "cpu"

# Modelo Qwen-Image-Edit (no SDXL)
model_repo_id = "Qwen/Qwen-Image-Edit"

if torch.cuda.is_available():
    torch_dtype = torch.bfloat16
else:
    torch_dtype = torch.float32

print("🔄 Cargando modelo Qwen-Image-Edit...")

# Cargar pipeline de EDICIÓN (no generación)
pipe = QwenImageEditPipeline.from_pretrained(
    model_repo_id, 
    torch_dtype=torch_dtype
)

print("✅ Modelo cargado correctamente")

# DESCOMENTADO para usar ZeroGPU
@spaces.GPU(duration=90)
def editar(
    imagen_entrada,
    prompt,
    seed,
    randomize_seed,
    guidance_scale,
    num_inference_steps,
    progress=gr.Progress(track_tqdm=True),
):
    """Función de edición de imagen con Qwen-Image-Edit"""
    
    if imagen_entrada is None:
        return None, seed
    
    # Mover modelo a GPU
    pipe.to(device)
    
    # Configurar semilla
    if randomize_seed:
        import random
        seed = random.randint(0, 2147483647)
    
    generator = torch.Generator().manual_seed(seed)
    
    # Editar imagen (no generar desde cero)
    resultado = pipe(
        image=imagen_entrada,
        prompt=prompt,
        generator=generator,
        true_cfg_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        negative_prompt=""
    )
    
    return resultado.images[0], seed


# Ejemplos para edición de imágenes
examples = [
    "Change the car color to red",
    "Add text 'SALE' in large red letters",
    "Change the background to a beach sunset",
    "Rotate the object 90 degrees",
    "Remove the small object on the left",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 960px;
}
"""

with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🎨 Editor de Imágenes con Qwen-Image-Edit
            ### Edición profesional de imágenes usando ZeroGPU H200
            
            **Sube una imagen y describe qué quieres cambiar**
            """
        )

        with gr.Row():
            # Columna izquierda - Entrada
            with gr.Column(scale=1):
                # INPUT DE IMAGEN (nuevo, no estaba antes)
                imagen_input = gr.Image(
                    label="📸 Imagen Original",
                    type="pil",
                    height=400
                )
                
                prompt = gr.Textbox(
                    label="✍️ Describe la edición",
                    show_label=True,
                    max_lines=3,
                    placeholder="Ejemplo: Cambia el color del auto a rojo, agrega texto 'OFERTA'...",
                    lines=3
                )

                run_button = gr.Button("🚀 Editar Imagen", variant="primary", size="lg")

            # Columna derecha - Salida
            with gr.Column(scale=1):
                result = gr.Image(label="✨ Resultado", show_label=True, height=400)

        # Configuración avanzada
        with gr.Accordion("⚙️ Configuración Avanzada", open=False):
            
            seed = gr.Slider(
                label="Semilla (para reproducibilidad)",
                minimum=0,
                maximum=2147483647,
                step=1,
                value=42,
            )

            randomize_seed = gr.Checkbox(label="Semilla aleatoria", value=True)

            guidance_scale = gr.Slider(
                label="CFG Scale (qué tan fiel al prompt)",
                minimum=1.0,
                maximum=10.0,
                step=0.5,
                value=4.0,
            )

            num_inference_steps = gr.Slider(
                label="Pasos de inferencia (más = mejor calidad)",
                minimum=20,
                maximum=100,
                step=5,
                value=50,
            )

        # Ejemplos (ahora sin imagen, se deben subir manualmente)
        gr.Markdown("### 💡 Ejemplos de prompts para edición:")
        gr.Examples(
            examples=examples, 
            inputs=[prompt],
            label="Haz clic en un ejemplo para probarlo"
        )
        
        gr.Markdown(
            """
            ---
            ### ℹ️ Información
            - **Hardware:** GPU H200 (ZeroGPU)
            - **Plan Pro:** 25 minutos diarios incluidos
            - **Modelo:** Qwen-Image-Edit (20B parámetros)
            - **Licencia:** Apache 2.0 - Uso comercial permitido
            
            ### 🎯 Capacidades:
            - ✨ Edición semántica y de apariencia
            - 📝 Edición precisa de texto (inglés/chino)
            - 🎨 Transferencia de estilo
            - 🔄 Rotación de objetos
            - ➕➖ Agregar/eliminar elementos
            """
        )
    
    # Conectar botón y submit
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=editar,
        inputs=[
            imagen_input,  # NUEVO: entrada de imagen
            prompt,
            seed,
            randomize_seed,
            guidance_scale,
            num_inference_steps,
        ],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    demo.queue(max_size=20)
    demo.launch()