import gradio as gr
from gradio_client import Client, handle_file

# Connect to Spaces using gradio_client
vton_client = Client("Kwai-Kolors/Kolors-Virtual-Try-On")
mesh_client = Client("tencent/Hunyuan3D-2.1")

def vton_to_3d(person_img, cloth_img):
    # Step 1: Virtual Try-On (2D image)
    vton_output = vton_client.predict(
        person_img, cloth_img,
        api_name="/tryon"  # verify via "Use via API" on the Space UI
    )
    tryon_img = vton_output

    # Step 2: Generate 3D mesh using the 2D try-on image
    mesh_output = mesh_client.predict(
        tryon_img,
        api_name="/predict"  # verify via "Use via API" on the Space UI
    )

    return tryon_img, mesh_output

demo = gr.Interface(
    fn=vton_to_3d,
    inputs=[
        gr.Image(type="filepath", label="Person Image"),
        gr.Image(type="filepath", label="Cloth Image")
    ],
    outputs=[
        gr.Image(label="2D Try-On Result"),
        gr.File(label="3D Mesh Output (.glb/.obj/.ply)")
    ],
    title="2D Try-On → 3D Mesh Pipeline",
    description="Uploads a person image and clothing image → performs virtual try-on → generates a 3D mesh."
)

if __name__ == "__main__":
    demo.launch(share=True)  # Set share=True to create a public link