File size: 4,501 Bytes
6e4a2a6
 
edfb9ce
6e4a2a6
 
 
 
49749f5
edfb9ce
6e4a2a6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
edfb9ce
 
49749f5
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
import os
import cv2
import gradio as gr
import numpy as np
import random
import base64
import time
from gradio_client import Client

# -----------------------------
# Connect to Hunyuan3D Space
# -----------------------------
mesh_client = Client("tencent/Hunyuan3D-2.1")  # has real API
MESH_API_NAME = "/process"

# -----------------------------
# Kolors VTON code (internal)
# -----------------------------
MAX_SEED = 999999

def tryon(person_img, garment_img, seed, randomize_seed):
    if person_img is None or garment_img is None:
        return None, seed, "Empty image"
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # Encode images
    encoded_person_img = cv2.imencode('.jpg', cv2.cvtColor(person_img, cv2.COLOR_RGB2BGR))[1].tobytes()
    encoded_person_img = base64.b64encode(encoded_person_img).decode('utf-8')
    encoded_garment_img = cv2.imencode('.jpg', cv2.cvtColor(garment_img, cv2.COLOR_RGB2BGR))[1].tobytes()
    encoded_garment_img = base64.b64encode(encoded_garment_img).decode('utf-8')

    # Send request to internal tryon server (Kolors VTON)
    tryon_url = "http://" + os.environ['tryon_url'] + "Submit"
    token = os.environ['token']
    cookie = os.environ['Cookie']
    referer = os.environ['referer']
    headers = {'Content-Type': 'application/json', 'token': token, 'Cookie': cookie, 'referer': referer}
    data = {"clothImage": encoded_garment_img, "humanImage": encoded_person_img, "seed": seed}

    # Post request
    import requests
    try:
        response = requests.post(tryon_url, headers=headers, data=json.dumps(data), timeout=50)
        if response.status_code != 200:
            return None, seed, f"Tryon POST failed: {response.status_code}"
        result_json = response.json()['result']
        if result_json['status'] != "success":
            return None, seed, f"Tryon error: {result_json['status']}"
        task_uuid = result_json['result']
    except Exception as e:
        return None, seed, f"Tryon exception: {str(e)}"

    # Poll for result
    result_img = None
    for _ in range(12):  # retry 12 times
        try:
            query_url = "http://" + os.environ['tryon_url'] + f"Query?taskId={task_uuid}"
            response = requests.get(query_url, headers=headers, timeout=20)
            if response.status_code == 200:
                r = response.json()['result']
                if r['status'] == "success":
                    img_bytes = base64.b64decode(r['result'])
                    result_np = np.frombuffer(img_bytes, np.uint8)
                    result_img = cv2.imdecode(result_np, cv2.IMREAD_COLOR)
                    result_img = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
                    break
            time.sleep(1)
        except:
            time.sleep(1)

    if result_img is None:
        return None, seed, "Tryon failed or timed out"

    return result_img, seed, "Success"

# -----------------------------
# Merge: Tryon + Hunyuan3D
# -----------------------------
def tryon_to_3d(person_img, garment_img, seed, randomize_seed):
    # Run 2D try-on
    tryon_img, seed_used, tryon_info = tryon(person_img, garment_img, seed, randomize_seed)
    if tryon_img is None:
        return None, None, seed_used, tryon_info

    # Run 3D mesh generation (Hunyuan3D)
    try:
        mesh_file = mesh_client.predict(tryon_img, api_name=MESH_API_NAME)
    except Exception as e:
        return tryon_img, None, seed_used, f"Hunyuan3D error: {str(e)}"

    return tryon_img, mesh_file, seed_used, "Success"

# -----------------------------
# Gradio Interface
# -----------------------------
with gr.Blocks() as demo:
    gr.Markdown("# Virtual Try-On → 3D Mesh Generator")

    with gr.Row():
        person_img = gr.Image(label="Person Image", type="numpy")
        garment_img = gr.Image(label="Garment Image", type="numpy")

    with gr.Row():
        seed = gr.Slider(0, MAX_SEED, step=1, value=0, label="Seed")
        randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)

    with gr.Row():
        tryon_out = gr.Image(label="2D Try-On Result")
        mesh_out = gr.Model3D(label="3D Mesh Output")

    seed_used_out = gr.Number(label="Seed Used")
    status_out = gr.Textbox(label="Status")

    run_btn = gr.Button("Run")
    run_btn.click(fn=tryon_to_3d,
                  inputs=[person_img, garment_img, seed, randomize_seed],
                  outputs=[tryon_out, mesh_out, seed_used_out, status_out])

if __name__ == "__main__":
    demo.launch()