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

# 1. Carga del modelo optimizado en memoria
model_id = "black-forest-labs/FLUX.1-schnell"
print("Cargando pipeline en GPU...")
pipe = FluxImg2ImgPipeline.from_pretrained(
    model_id, 
    torch_dtype=torch.bfloat16
).to("cuda")

# 2. Carga del LoRA de Cosplay
print("Inyectando LoRA...")
pipe.load_lora_weights("joyfox/Kontext-Cosplay-Lora", weight_name="kontext_cosplay_lora_20000.safetensors")

# 3. Función de generación corregida
@spaces.GPU
def generate_cosplay(image, prompt, strength, steps):
    if image is None:
        return None
        
    # Redimensionamos para mantener estabilidad en FLUX
    image = image.convert("RGB").resize((1024, 1024))
    
    # Trigger indispensable del LoRA
    full_prompt = f"{prompt}, r2cos"
    
    # Eliminamos la semilla fija estricta para permitir variedad por intento,
    # pero puedes volver a usar un generador si deseas replicabilidad exacta.
    
    # Ejecución con los nuevos parámetros dinámicos
    result = pipe(
        prompt=full_prompt,
        image=image,
        strength=strength,          # Control dinámico desde la UI
        num_inference_steps=steps,  # Elevado para dar tiempo al LoRA
        guidance_scale=0.0
    ).images[0]
    
    return result

# 4. Interfaz de Gradio interactiva
demo = gr.Interface(
    fn=generate_cosplay, 
    inputs=[
        gr.Image(type="pil", label="Foto de Referencia"), 
        gr.Textbox(label="Descripción del personaje (Ej: Naruto, Cyberpunk, Samurai)"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.05, label="Fuerza de Transformación (Strength)"),
        gr.Slider(minimum=10, maximum=30, value=16, step=1, label="Pasos de Inferencia Totales (Steps)")
    ], 
    outputs=gr.Image(label="Resultado Cosplay"),
    title="Kontext-Cosplay AI Pro",
    description="Ajusta los deslizadores si la imagen cambia muy poco."
)

if __name__ == "__main__":
    demo.launch()