Update app.py

app.py

@@ -1,121 +1,97 @@
-import os
-import cv2
 import gradio as gr
+import cv2
 import numpy as np
-import random
-import base64
-import time
-from gradio_client import Client
 
 # -----------------------------
-# Connect to Hunyuan3D Space
+# Local "try-on" function
 # -----------------------------
-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):
+def tryon_local(person_img, garment_img, seed, randomize_seed):
+    """
+    Simulates try-on by resizing the garment to fit the person and overlaying.
+    """
     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"
+        return None, None, "Empty image"
+
+    # Resize garment to width of person image
+    h_person, w_person = person_img.shape[:2]
+    h_garment, w_garment = garment_img.shape[:2]
+
+    scale = w_person / w_garment
+    new_w = int(w_garment * scale)
+    new_h = int(h_garment * scale)
+    garment_resized = cv2.resize(garment_img, (new_w, new_h))
+
+    # Simple overlay: top of garment onto top of person
+    overlay = person_img.copy()
+    y_offset = 0
+    x_offset = max(0, (w_person - new_w) // 2)
+
+    # Handle transparency if 4-channel
+    if garment_resized.shape[2] == 4:
+        alpha = garment_resized[:, :, 3] / 255.0
+        for c in range(3):
+            overlay[y_offset:y_offset+new_h, x_offset:x_offset+new_w, c] = \
+                alpha * garment_resized[:, :, c] + (1 - alpha) * overlay[y_offset:y_offset+new_h, x_offset:x_offset+new_w, c]
+    else:
+        overlay[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = garment_resized
+
+    info = "Success (local simulation)"
+    return overlay, seed, info
 
 # -----------------------------
-# Merge: Tryon + Hunyuan3D
+# Hunyuan3D API call
 # -----------------------------
+from gradio_client import Client
+hunyuan_client = Client("tencent/Hunyuan3D-2.1")
+
 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)
+    # Step 1: Run local try-on
+    tryon_img, seed_used, tryon_info = tryon_local(person_img, garment_img, seed, randomize_seed)
     if tryon_img is None:
-        return None, None, seed_used, tryon_info
+        return None, "Try-on failed"
 
-    # Run 3D mesh generation (Hunyuan3D)
+    # Step 2: Send 2D try-on result to Hunyuan3D
     try:
-        mesh_file = mesh_client.predict(tryon_img, api_name=MESH_API_NAME)
+        result_3d = hunyuan_client.predict(
+            tryon_img,  # input image
+            api_name="/process"
+        )
+        return result_3d, tryon_info
     except Exception as e:
-        return tryon_img, None, seed_used, f"Hunyuan3D error: {str(e)}"
-
-    return tryon_img, mesh_file, seed_used, "Success"
+        return None, f"Hunyuan3D API error: {e}"
 
 # -----------------------------
-# Gradio Interface
+# Gradio UI
 # -----------------------------
-with gr.Blocks() as demo:
-    gr.Markdown("# Virtual Try-On → 3D Mesh Generator")
-
+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():
-        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()
+        with gr.Column(elem_id="col-left"):
+            person_input = gr.Image(label="Person Image", source="upload", type="numpy")
+        with gr.Column(elem_id="col-mid"):
+            garment_input = gr.Image(label="Garment Image", source="upload", 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()