Spaces:
Running
Running
| 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() | |