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