import gradio as gr
from PIL import Image
import torch
from torchvision import transforms

# Aquí deberías importar tu modelo real
# Este es un ejemplo genérico
class DummyFashionGAN(torch.nn.Module):
    def forward(self, person, clothes):
        # Este es solo un dummy para no lanzar errores
        return person

# Cargar modelo (usa tu modelo real aquí)
model = DummyFashionGAN()
model.eval()

# Transformación de imágenes
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor()
])

# Función principal con depuración
def tryon(person_img, clothes_img):
    try:
        print("📤 Iniciando procesamiento de imágenes...")

        person = transform(person_img).unsqueeze(0)
        clothes = transform(clothes_img).unsqueeze(0)

        print("✅ Imágenes cargadas y transformadas.")

        with torch.no_grad():
            output = model(person, clothes)

        print("🎯 Generación de imagen completada.")

        result = transforms.ToPILImage()(output.squeeze(0).clamp(0, 1))
        return result

    except Exception as e:
        print("❌ Error durante el procesamiento:", str(e))
        # Devolver una imagen roja como error
        return Image.new("RGB", (256, 256), color="red")

# Interfaz de Gradio
demo = gr.Interface(
    fn=tryon,
    inputs=[
        gr.Image(label="👤 Tu maniquí o imagen de cuerpo", type="pil"),
        gr.Image(label="👕 Imagen de la prenda", type="pil")
    ],
    outputs=gr.Image(label="🪄 Resultado: Prueba virtual"),
    title="👗 Probador Virtual AI",
    description="Sube una imagen tuya (o maniquí) y una prenda para probarla virtualmente. Esta es una demo."
)

# Lanzar app
if __name__ == "__main__":
    demo.launch()