Update app.py

app.py

@@ -1,42 +1,121 @@
+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 Hugging Face Spaces
-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: Run Virtual Try-On (Space 1)
-    vton_output = vton_client.predict(
-        person_img,
-        cloth_img,
-        api_name="/vton"   # confirmed API name for Kolors VTON
-    )
-    tryon_img = vton_output
-
-    # Step 2: Run 3D mesh generation (Space 2)
-    mesh_output = mesh_client.predict(
-        tryon_img,
-        api_name="/process"   # confirmed API name for Hunyuan3D
-    )
-
-    # Return both outputs
-    return tryon_img, mesh_output
-
-# Gradio UI
-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.Model3D(label="3D Mesh Output")  # shows interactive 3D model
-    ],
-    title="Virtual Try-On → 3D Avatar Generator",
-    description="Upload a person image and a clothing image. The system first performs virtual try-on using Kolors VTON, then converts the result into a 3D mesh using Hunyuan3D."
-)
+# -----------------------------
+# 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()