import gradio as gr
from warp_design_on_dress import run_design_warp_on_dress
from PIL import Image
import os

def gradio_tryon(dress_img, design_img):
    os.makedirs("input", exist_ok=True)
    os.makedirs("output", exist_ok=True)

    dress_path = "input/dress.jpg"
    design_path = "input/design.jpg"

    dress_img.save(dress_path)
    design_img.save(design_path)

    result_path = run_design_warp_on_dress(
        dress_path, design_path,
        gmm_ckpt="checkpoints/GMM/gmm_final.pth",
        tom_ckpt="checkpoints/TOM/tom_final.pth",
        output_dir="output"
    )

    return Image.open(result_path)

gr.Interface(
    fn=gradio_tryon,
    inputs=[
        gr.Image(label="Dress Image", type="pil"),
        gr.Image(label="Design to Warp", type="pil")
    ],
    outputs=gr.Image(label="Warped Design Output"),
    title="🎨 Design Warping on Dress",
    description="Upload a dress photo and a design. The model warps and blends the design onto the dress realistically."
).launch()