Kwai-Kolors/Kolors-Virtual-Try-On

1import os2import cv23import gradio as gr4import numpy as np5import random6import base647import requests8import json9import time10 11 12def tryon(person_img, garment_img, seed, randomize_seed):13    post_start_time = time.time()14    if person_img is None or garment_img is None:15        gr.Warning("Empty image")16        return None, None, "Empty image"17    if randomize_seed:18        seed = random.randint(0, MAX_SEED)19    encoded_person_img = cv2.imencode('.jpg', cv2.cvtColor(person_img, cv2.COLOR_RGB2BGR))[1].tobytes()20    encoded_person_img = base64.b64encode(encoded_person_img).decode('utf-8')21    encoded_garment_img = cv2.imencode('.jpg', cv2.cvtColor(garment_img, cv2.COLOR_RGB2BGR))[1].tobytes()22    encoded_garment_img = base64.b64encode(encoded_garment_img).decode('utf-8')23 24    url = "http://" + os.environ['tryon_url'] + "Submit"25    token = os.environ['token']26    cookie = os.environ['Cookie']27    referer = os.environ['referer']28    headers = {'Content-Type': 'application/json', 'token': token, 'Cookie': cookie, 'referer': referer}29    data = {30        "clothImage": encoded_garment_img,31        "humanImage": encoded_person_img,32        "seed": seed33    }34    try:35        response = requests.post(url, headers=headers, data=json.dumps(data), timeout=50)36        # print("post response code", response.status_code)37        if response.status_code == 200:38            result = response.json()['result']39            status = result['status']40            if status == "success":41                uuid = result['result']42                # print(uuid)43    except Exception as err:44        print(f"Post Exception Error: {err}")45        raise gr.Error("Too many users, please try again later")46    post_end_time = time.time()47    print(f"post time used: {post_end_time-post_start_time}")48 49    get_start_time =time.time()50    time.sleep(9)51    Max_Retry = 1252    result_img = None53    info = ""54    err_log = ""55    for i in range(Max_Retry):56        try:57            url = "http://" + os.environ['tryon_url'] + "Query?taskId=" + uuid58            response = requests.get(url, headers=headers, timeout=20)59            # print("get response code", response.status_code)60            if response.status_code == 200:61                result = response.json()['result']62                status = result['status']63                if status == "success":64                    result = base64.b64decode(result['result'])65                    result_np = np.frombuffer(result, np.uint8)66                    result_img = cv2.imdecode(result_np, cv2.IMREAD_UNCHANGED)67                    result_img = cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR)68                    info = "Success"69                    break70                elif status == "error":71                    err_log = f"Status is Error"72                    info = "Error"73                    break74            else:75                # print(response.text)76                err_log = "URL error, pleace contact the admin"77                info = "URL error, pleace contact the admin"78                break79        except requests.exceptions.ReadTimeout:80            err_log = "Http Timeout"81            info = "Http Timeout, please try again later"82        except Exception as err:83            err_log = f"Get Exception Error: {err}"84        time.sleep(1)85    get_end_time = time.time()86    print(f"get time used: {get_end_time-get_start_time}")87    print(f"all time used: {get_end_time-get_start_time+post_end_time-post_start_time}")88    if info == "":89        err_log = f"No image after {Max_Retry} retries"90        info = "Too many users, please try again later"91    if info != "Success":92        print(f"Error Log: {err_log}")93        gr.Warning("Too many users, please try again later")94 95    return result_img, seed, info96 97def start_tryon(person_img, garment_img, seed, randomize_seed):98    start_time = time.time()99    if person_img is None or garment_img is None:100        return None, None, "Empty image"101    if randomize_seed:102        seed = random.randint(0, MAX_SEED)103    encoded_person_img = cv2.imencode('.jpg', cv2.cvtColor(person_img, cv2.COLOR_RGB2BGR))[1].tobytes()104    encoded_person_img = base64.b64encode(encoded_person_img).decode('utf-8')105    encoded_garment_img = cv2.imencode('.jpg', cv2.cvtColor(garment_img, cv2.COLOR_RGB2BGR))[1].tobytes()106    encoded_garment_img = base64.b64encode(encoded_garment_img).decode('utf-8')107 108    url = "http://" + os.environ['tryon_url']109    token = os.environ['token']110    cookie = os.environ['Cookie']111    referer = os.environ['referer']112 113    headers = {'Content-Type': 'application/json', 'token': token, 'Cookie': cookie, 'referer': referer}114    data = {115        "clothImage": encoded_garment_img,116        "humanImage": encoded_person_img,117        "seed": seed118    }119 120    result_img = None121    try:122        session = requests.Session()123        response = session.post(url, headers=headers, data=json.dumps(data), timeout=60)124        print("response code", response.status_code)125        if response.status_code == 200:126            result = response.json()['result']127            status = result['status']128            if status == "success":129                result = base64.b64decode(result['result'])130                result_np = np.frombuffer(result, np.uint8)131                result_img = cv2.imdecode(result_np, cv2.IMREAD_UNCHANGED)132                result_img = cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR)133                info = "Success"134            else:135                info = "Try again latter"136        else:137            print(response.text)138            info = "URL error, pleace contact the admin"139    except requests.exceptions.ReadTimeout:140        print("timeout")141        info = "Too many users, please try again later"142        raise gr.Error("Too many users, please try again later")143    except Exception as err:144        print(f"其他错误: {err}")145        info = "Error, pleace contact the admin"146    end_time = time.time()147    print(f"time used: {end_time-start_time}")148 149    return result_img, seed, info150 151MAX_SEED = 999999152 153example_path = os.path.join(os.path.dirname(__file__), 'assets')154 155garm_list = os.listdir(os.path.join(example_path,"cloth"))156garm_list_path = [os.path.join(example_path,"cloth",garm) for garm in garm_list]157 158human_list = os.listdir(os.path.join(example_path,"human"))159human_list_path = [os.path.join(example_path,"human",human) for human in human_list]160 161css="""162#col-left {163    margin: 0 auto;164    max-width: 430px;165}166#col-mid {167    margin: 0 auto;168    max-width: 430px;169}170#col-right {171    margin: 0 auto;172    max-width: 430px;173}174#col-showcase {175    margin: 0 auto;176    max-width: 1100px;177}178#button {179    color: blue;180}181"""182 183def load_description(fp):184    with open(fp, 'r', encoding='utf-8') as f:185        content = f.read()186    return content187 188def change_imgs(image1, image2):189    return image1, image2190 191with gr.Blocks(css=css) as Tryon:192    gr.HTML(load_description("assets/title.md"))193    with gr.Row():194        with gr.Column(elem_id = "col-left"):195            gr.HTML("""196            <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">197                <div>198                Step 1.  Upload a person image ⬇️199                </div>200            </div>201            """)202        with gr.Column(elem_id = "col-mid"):203            gr.HTML("""204            <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">205                <div>206                Step 2. Upload a garment image ⬇️207                </div>208            </div>209            """)210        with gr.Column(elem_id = "col-right"):211            gr.HTML("""212            <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">213                <div>214                Step 3. Press “Run” to get try-on results215                </div>216            </div>217            """)218    with gr.Row():219        with gr.Column(elem_id = "col-left"):220            imgs = gr.Image(label="Person image", sources='upload', type="numpy")221            # category = gr.Dropdown(label="Garment category", choices=['upper_body', 'lower_body', 'dresses'],  value="upper_body")222            example = gr.Examples(223                inputs=imgs,224                examples_per_page=12,225                examples=human_list_path226            )227        with gr.Column(elem_id = "col-mid"):228            garm_img = gr.Image(label="Garment image", sources='upload', type="numpy")229            example = gr.Examples(230                inputs=garm_img,231                examples_per_page=12,232                examples=garm_list_path233            )234        with gr.Column(elem_id = "col-right"):235            image_out = gr.Image(label="Result", show_share_button=False)236            with gr.Row():237                seed = gr.Slider(238                    label="Seed",239                    minimum=0,240                    maximum=MAX_SEED,241                    step=1,242                    value=0,243                )244                randomize_seed = gr.Checkbox(label="Random seed", value=True)245            with gr.Row():246                seed_used = gr.Number(label="Seed used")247                result_info = gr.Text(label="Response")248            # try_button = gr.Button(value="Run", elem_id="button")249            test_button = gr.Button(value="Run", elem_id="button")250 251 252    # try_button.click(fn=start_tryon, inputs=[imgs, garm_img, seed, randomize_seed], outputs=[image_out, seed_used, result_info], api_name='tryon',concurrency_limit=10)253    test_button.click(fn=tryon, inputs=[imgs, garm_img, seed, randomize_seed], outputs=[image_out, seed_used, result_info], api_name=False, concurrency_limit=45)254 255    with gr.Column(elem_id = "col-showcase"):256        gr.HTML("""257        <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">258            <div> </div>259            <br>260            <div>261            Virtual try-on examples in pairs of person and garment images262            </div>263        </div>264        """)265        show_case = gr.Examples(266            examples=[267                ["assets/examples/model2.png", "assets/examples/garment2.png", "assets/examples/result2.png"],268                ["assets/examples/model3.png", "assets/examples/garment3.png", "assets/examples/result3.png"],269                ["assets/examples/model1.png", "assets/examples/garment1.png", "assets/examples/result1.png"],270            ],271            inputs=[imgs, garm_img, image_out],272            label=None273        )274 275Tryon.queue(api_open=False).launch(show_api=False)276