import gradio as gr
from gradio_client import Client, handle_file

def process_try_on(user_image_path, dress_image_path):
    print("Routing images to the massive GPU cluster...")
    try:
        # Connect to the official, GPU-powered IDM-VTON space
        client = Client("yisol/IDM-VTON")
        
        # Send the images to their API
        result = client.predict(
            dict={"background": handle_file(user_image_path), "layers": [], "composite": None},
            garm_img=handle_file(dress_image_path),
            garment_des="a clothing item",
            is_checked=True,
            is_checked_crop=False,
            denoise_steps=30,
            seed=42,
            api_name="/tryon"
        )
        
        # Their API returns a tuple (list), the first item is the image file path
        return result[0]
        
    except Exception as e:
        print(f"Failed to connect to GPU: {e}")
        return None

# Notice we changed the input type to "filepath" so it works easily with the client
dvte_interface = gr.Interface(
    fn=process_try_on,
    inputs=[
        gr.Image(type="filepath", label="User Photo"), 
        gr.Image(type="filepath", label="Dress Product Image")
    ],
    outputs=gr.Image(type="filepath", label="Final Try-On Result"),
    title="DVTE - E-commerce Try-On Engine",
    description="API Endpoint for Virtual Garment Transfer"
)

dvte_interface.launch()