import gradio as gr
import cv2
import numpy as np
from gradio_client import Client
import io
from PIL import Image

# -----------------------------
# Hunyuan3D API setup
# -----------------------------
API_KEY = "HT_Token"  # Replace with your valid key
try:
    hunyuan_client = Client("tencent/Hunyuan3D-2.1", hf_token=API_KEY)
    # Grab first endpoint index for the 3D try-on function
    fn_index_default = hunyuan_client.endpoints[0].fn_index if hunyuan_client.endpoints else 0
    client_loaded = True
except Exception as e:
    print(f"Failed to load Hunyuan3D client: {e}")
    hunyuan_client = None
    fn_index_default = 0
    client_loaded = False

# -----------------------------
# Convert numpy image to bytes
# -----------------------------
def np_to_bytes(img_np):
    img_pil = Image.fromarray(cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB))
    buf = io.BytesIO()
    img_pil.save(buf, format="PNG")
    return buf.getvalue()

# -----------------------------
# Try-on function using Hunyuan3D API
# -----------------------------
def tryon_to_3d(person_img, garment_img, seed, randomize_seed):
    if not client_loaded:
        return None, "Hunyuan3D client not loaded. Please check your API key."

    if person_img is None or garment_img is None:
        return None, "Empty image provided."

    try:
        # Convert to bytes
        person_bytes = np_to_bytes(person_img)
        garment_bytes = np_to_bytes(garment_img)

        # Call the 3D try-on endpoint
        result = hunyuan_client.predict(
            person_bytes,
            garment_bytes,
            fn_index=fn_index_default
        )

        # Convert returned bytes to image
        result_img = np.array(Image.open(io.BytesIO(result)))
        return cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR), "3D try-on successful!"
    except Exception as e:
        return None, f"Hunyuan3D API error: {e}"

# -----------------------------
# Gradio UI
# -----------------------------
css = """
#col-left, #col-mid, #col-right {
    margin: 0 auto;
    max-width: 430px;
}
#col-showcase {
    margin: 0 auto;
    max-width: 1100px;
}
#button { color: blue; }
"""

with gr.Blocks(css=css) as app:
    with gr.Row():
        with gr.Column(elem_id="col-left"):
            person_input = gr.Image(label="Person Image", type="numpy")
        with gr.Column(elem_id="col-mid"):
            garment_input = gr.Image(label="Garment Image", type="numpy")
        with gr.Column(elem_id="col-right"):
            output_img = gr.Image(label="3D Result")
            result_info = gr.Text(label="Info")
            seed = gr.Slider(0, 999999, value=0, step=1, label="Seed")
            randomize_seed = gr.Checkbox(label="Random seed", value=True)
            run_btn = gr.Button("Run")

    run_btn.click(
        fn=tryon_to_3d,
        inputs=[person_input, garment_input, seed, randomize_seed],
        outputs=[output_img, result_info]
    )

app.launch()