app.py

import os
import cv2
import gradio as gr
import numpy as np
import random
import base64
import requests
import json
import time


def portrait_gen(person_img, template_img, prompt, seed, randomize_seed):
    post_start_time = time.time()
    if person_img is None:
        gr.Warning("Empty Person image")
        return None, None, "Empty image"
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    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_template_img = None
    if template_img is not None:
        # style image with fixed hidden prompt
        prompt_temp = imgpath_to_prompt(template_img)
        if prompt_temp is not None:
            prompt = prompt_temp
        encoded_template_img = cv2.imencode('.jpg', cv2.imread(template_img))[1].tobytes()
        encoded_template_img = base64.b64encode(encoded_template_img).decode('utf-8')

    url = "http://" + os.environ['avatar_url'] + "Submit"
    token = os.environ['token']
    referer = os.environ['referer']
    headers = {'Content-Type': 'application/json', 'token': token, 'referer': referer}
    data = {
        "templateImage": encoded_template_img,
        "humanImage": encoded_person_img,
        "seed": seed,
        "prompt": prompt
    }
    try:
        response = requests.post(url, headers=headers, data=json.dumps(data), timeout=50)
        # print("post response code", response.status_code)
        if response.status_code == 200:
            result = response.json()['result']
            status = result['status']
            if status == "success":
                uuid = result['result']
                # print(uuid)
    except Exception as err:
        print(f"Post Exception Error: {err}")
        raise gr.Error("Too many users, please try again later")
    post_end_time = time.time()
    print(f"post time used: {post_end_time-post_start_time}")

    get_start_time =time.time()
    time.sleep(9)
    Max_Retry = 12
    result_img = None
    info = ""
    err_log = ""
    for i in range(Max_Retry):
        try:
            url = "http://" + os.environ['avatar_url'] + "Query?taskId=" + uuid
            response = requests.get(url, headers=headers, timeout=20)
            # print("get response code", response.status_code)
            if response.status_code == 200:
                result = response.json()['result']
                status = result['status']
                if status == "success":
                    result = base64.b64decode(result['result'])
                    result_np = np.frombuffer(result, np.uint8)
                    result_img = cv2.imdecode(result_np, cv2.IMREAD_UNCHANGED)
                    result_img = cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR)
                    info = "Success"
                    break
                elif status == "error":
                    err_log = f"Status is Error"
                    info = "Error"
                    break
            else:
                # print(response.text)
                err_log = "URL error, pleace contact the admin"
                info = "URL error, pleace contact the admin"
                break
        except requests.exceptions.ReadTimeout:
            err_log = "Http Timeout"
            info = "Http Timeout, please try again later"
        except Exception as err:
            err_log = f"Get Exception Error: {err}"
        time.sleep(1)
    get_end_time = time.time()
    print(f"get time used: {get_end_time-get_start_time}")
    print(f"all time used: {get_end_time-get_start_time+post_end_time-post_start_time}")
    if info == "":
        err_log = f"No image after {Max_Retry} retries"
        info = "Too many users, please try again later"
    if info != "Success":
        print(f"Error Log: {err_log}")
        gr.Warning("Too many users, please try again later")

    return result_img, seed, info

MAX_SEED = 999999

example_path = os.path.join(os.path.dirname(__file__), 'assets')

refer_list = sorted(os.listdir(os.path.join(example_path,"refer_imgs")))
refer_list_path = [os.path.join(example_path,"refer_imgs",garm) for garm in refer_list]

human_list = sorted(os.listdir(os.path.join(example_path,"id_imgs")))
human_list_path = [os.path.join(example_path,"id_imgs",human) for human in human_list]

prompt_list = [
    "Half-length close-up of a beautiful girl, dressed in the Hogwarts uniform from the Harry Potter series, wearing a blue cloak with a crest that appears to be the Gryffindor house emblem, along with a white shirt and a red-and-gold striped tie. She is holding a majic wand, the hair is long and wavy, and the background consists of a wooden wall with several books and a framed picture placed on a wooden desk. This setting evokes an atmosphere similar to a cozy study room or a magical classroom.",
    "WM,Anime Magic,an anime girl holding glowing blue poker cards in her hands,poker cards flying in the sky,glowing magic effects,dynamic poses,highly detailed,ultra-high resolutions,32K UHD,best quality,masterpiece..The background features a full majic moon casting a soft glow, accentuating the ethereal majic ambiance.",
    "A serene beauty stands in a studio awash with soft, golden light. Her wavy brunette hair cascades over her shoulders, while her piercing blue eyes gaze into the distance, filled with quiet contemplation. She wears a delicate, flowing white dress that seems to shimmer in the gentle glow, exuding an aura of pure elegance and grace , the classic film filter applied to the image imparts a nostalgic and gritty texture,highlighting the enduring allure of youth culture and fashion. the background provide a rich,textured setting,further emphasizing the unique style and attitude of the subject. Shot with a Fujifilm X-Pro3 camera,HDR,UHD,8K,best quality,best quality,masterpiece,realistic,cute face,smug,naughty face,light blush,giggling,backlighting,rim_light,moody lighting,",
    "A man wear an Argentina shirt with No. 7, holds up the World Cup, surrounded by golden rain. The man is laughing with a gold medal around their neck. The background showcases a large stadium filled with spectators, illuminated by bright lights and surrounded by a celebratory shower of confetti. The overall atmosphere is one of triumph and celebration",
    "anime epic cinematic shot of dynamic, The transatlantic slave trade, dramatic lighting, Technicolor, in motion, main subject of high budget action movie, raw photo, motion blur, best quality, high resolution, aidmaMJ6.1, an angel with wings and armor holding a spear in front of him flies at a demon, the demon, spreading his arms, yells at the angel,",
    "Sci-fi style digital art image showing a knight in futuristic armor riding a motorcycle on a city street. The protagonist is located in the center right position, wearing a black, metallic tight armor. The surface of the armor is covered with raindrops and water stains, indicating that it has just experienced a heavy rain. His posture is leaning forward, holding the handlebars tightly with both hands, which makes him look very focused and full of power. The design of the motorcycle is very futuristic, with smooth body lines, bright red light from the headlights, and the right side mirror also reflects the ambient light. The tires of the motorcycle are wide and the tread pattern is clearly visible, which echoes the slippery road surface. The background is a city street scene with tall buildings, and the buildings are covered with colorful neon billboards. Various words and patterns are blurred against the sky in the distance. Some larger billboards with red and blue words can be seen on the left; on the right, there are more small billboards with various colors but not very clear. There are other pedestrians on the street. Although they are far away, their figures can be distinguished. Most of them have their backs to the camera and walk towards the distant buildings. The rain-soaked ground in the entire scene reflects the surrounding lights, adding depth and three-dimensionality to the picture. The details of the buildings in the background are slightly blurred, presenting a hazy feeling, highlighting the protagonist and his motorcycle in the foreground. The overall tone of the picture is cold, mainly black and gray, supplemented by red and blue light, forming a strong contrast effect. The high color saturation, moderate brightness, and uniform exposure give the picture a strong visual impact and a sense of futuristic technology. This image successfully creates a futuristic urban scene full of dynamic and technological atmosphere through delicate detail depiction and vivid color use.",
    # "a person wearing a light brown coat, a white shirt, a dark pleated short skirt, and a striped tie, showing off the long legs. The outfit is completed with knee-high white socks and black leather shoes. The individual is standing on a pathway that is covered with yellow fallen leaves, surrounded by trees that have yellow leaves still attached. The trees are planted along both sides of the pathway, creating a picturesque autumn scene. The trunks of the trees are painted white at the base. The overall setting appears to be a serene autumnal day, with bright and warm lighting",
    "Chinese pretty girl, with long, dark, wavy hair cascading over the shoulders, a sleeveless, knitted blue top with a high neck. The background is a gradient of various shades of blue, with a subtle texture that gives it a studio portrait appearance. The lighting is soft, illuminating the individual from the front, creating a calm and serene ambiance.",
    "A person dressed in a sleek, black turtleneck sweater and a matching black suit jacket. The outfit conveys a sophisticated and formal style. The background is a solid olive green color, which contrasts subtly with the dark clothing. The hair appears to be dark and styled in a slightly tousled manner, adding a casual touch to the otherwise formal attire.",
    "The upper body of a person wearing a plain, vibrant red, long-sleeved sweater. The sweater has a round neckline with a ribbed texture. The person has black hair that appears to be slightly tousled. The background of the image is plain and neutral, likely a light gray or white color, providing a clean and simple backdrop that contrasts with the rich color of the sweater. The image conveys a sense of casual yet stylish attire.",
]
prompt_list_examples = [
    'hold a sign writing "Kolors Portrait with Flux"',
    'The millionaire flung wads of cash in the air like confetti at the crowded nightclub, making it rain on the ecstatic revellers below.',
    'Wearing Iron Man costume, in fantastic space, with stars twinkling, ultra-high resolutions,32K UHD,best quality,masterpiece',   
    'a person sitting on the ground in a forest or park during autumn. The ground is covered with fallen leaves in shades of yellow and brown. They are wearing a cozy cream-colored knit sweater and jeans. Their hair is long and brown, reaching past their shoulders. In one hand, they hold up a yellow maple leaf, while the other hand rests on their lap, also holding a leaf. The background features trees with yellow leaves, indicating the fall season, and the sunlight filters through the foliage, casting a warm, golden light on the scene. The atmosphere is serene and picturesque, capturing the essence of autumn.',
    'A masterpiece of dreamy photography. The girl holds a glowing glass jar in her hands, which emits a soft yellow light, illuminating her face and the surrounding environment. She wears a white lace dress with blue and pink decorations and a light blue bow tied around her waist. The girl wears a white headdress dotted with blue and pink flowers and carries twinkling blue wings on her back. Behind the girl, you can see a vast sea of ​​flowers, which sparkle with pink and gold light and are scattered on the green grass. The sky is full of sparkling stars, from the upper left corner to the lower right corner, as if the Milky Way is scattered in the night sky. The overall light is soft and full of magical colors, making the whole picture look like a fairy tale world.',
]
prompt_list_single = [
    'hold a sign writing "Kolors Portrait with Flux"',
    "A beautiful girl reading book, high quality.",
    "Full body shot of young Asian face woman in sun hat and white dress standing on sunny beach with sea and mountains in background, high quality, sharp focus,",
]

css="""
#col-left {
    margin: 0 auto;
    max-width: 270px;
}
#col-mid {
    margin: 0 auto;
    max-width: 270px;
}
#col-right {
    margin: 0 auto;
    max-width: 550px;
}
#col-showcase {
    margin: 0 auto;
    max-width: 1100px;
}
#col-left-left {
    margin: 0 auto;
    max-width: 400px;
}
#col-left-right {
    margin: 0 auto;
    max-width: 120px;
}
#button {
    color: blue;
}
"""

def img_to_prompt(img_path,prompt):
    if img_path is None:
        # return prompt,prompt_final
        return prompt
    else:
        return None
    
def imgpath_to_prompt(img_path):
    img_name = img_path.split("/")[-1]
    if img_name in refer_list:
        prompt_id = refer_list.index(img_name)
        prompt_final = prompt_list[prompt_id]
    else:
        prompt_final = None
    return prompt_final

def text_to_text(text1, text2):
    if text1 is not None:
        return text1
    else:
        return text2
    
def load_description(fp):
    with open(fp, 'r', encoding='utf-8') as f:
        content = f.read()
    return content

def img_close():
    return None

def change_mode(mode: str):
    if mode == "Input prompt":
        return gr.update(visible=True), gr.update(visible=False, value=None)
    elif mode == "Ref. image":
        return gr.update(visible=False, value=""), gr.update(visible=True)
    else:
        raise gr.Error("no such mode!")

change_mode_js = """
    function change_mode_js(...args){
    mode = args[0];
    console.log("change_mode");
    text_examples = document.querySelector("#text_examples");
    image_examples = document.querySelector("#image_examples");
    ref_examples = document.querySelector("#ref_examples");
    prompt_examples = document.querySelector("#prompt_examples");
    if (mode === "Input prompt"){
        text_examples.hidden = false;
        prompt_examples.hidden = false;
        image_examples.hidden = true;
        ref_examples.hidden = true;
    }else {
        text_examples.hidden = true;
        prompt_examples.hidden = true;
        image_examples.hidden = false;
        ref_examples.hidden = false;
    }
    return args
}
"""


with gr.Blocks(css=css) as Portrait:
    gr.HTML(load_description("title.md"))
    with gr.Row():
        with gr.Column(elem_id = "col-left"):
            gr.HTML("""
            <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
                <div>
                Step 1.  Upload a portrait image ⬇️
                </div>
            </div>
            """)

             
        with gr.Column(elem_id = "col-mid"):
            gr.HTML("""
            <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
                <div>
                Step 2. Set a style reference ⬇️
                </div>
            </div>
            """)
        with gr.Column(elem_id = "col-right"):
            gr.HTML("""
            <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
                <div>
                Step 3. Press “Run” to get results ⬇️
                </div>
            </div>
            """)
    with gr.Row():
        with gr.Column(elem_id = "col-left"):
            imgs = gr.Image(label="Person Images", sources='upload', type="numpy",width=400)
            example = gr.Examples(
                inputs=imgs,
                examples_per_page=10,
                examples=human_list_path,
                label="Person Images"
            )
            
        with gr.Column(elem_id = "col-mid"):
            # gr.HTML("""
            #     <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 15px;">
            #         <div>
            #         Use image <strong>OR!!!</strong> text as reference style.
            #         </div>
            #     </div>
            #     """) 
            mode = gr.Radio(label="Set a reference image or input prompt", choices=["Ref. image", "Input prompt"], value="text")
            refer_img = gr.Image(label="Reference images", sources='upload', type="filepath")
            example = gr.Examples(
                inputs=refer_img,
                examples_per_page=9,
                examples=refer_list_path,
                label = 'Reference images',
                elem_id="ref_examples"
            )
            prompt = gr.Text(
                        label="Input prompts",
                        show_label=False,
                        max_lines=1,
                        placeholder="Enter your prompt",
                        container=False,
                    )
            examples_prompt = gr.Examples(
                inputs=prompt,
                examples=prompt_list_single,
                label = 'Input prompts',
                elem_id="prompt_examples"
            )

            Portrait.load(
                fn=lambda: gr.update(value="Input prompt"),
                outputs=[mode],
                show_progress=False
            )
            
            mode.change(
                js=change_mode_js, 
                fn=change_mode,
                inputs=[mode],
                outputs=[prompt, refer_img],
                show_progress=False
            )
            
        with gr.Column(elem_id = "col-right"):
            image_out = gr.Image(label="Results", show_share_button=False)
            with gr.Row():
                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=0,
                )
                randomize_seed = gr.Checkbox(label="Random seed", value=True)
            with gr.Row():
                seed_used = gr.Number(label="Seed used")
                result_info = gr.Text(label="Response")
            test_button = gr.Button(value="Run", elem_id="button")


    test_button.click(fn=portrait_gen, inputs=[imgs, refer_img, prompt, seed, randomize_seed], outputs=[image_out, seed_used, result_info], api_name=False, concurrency_limit=10)

    with gr.Column(elem_id = "col-showcase"):
        gr.HTML("""
        <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
            <div> </div>
            <br>
            <br>
            <div>
            Portrait examples in pairs of person images and style.
            </div>
        </div>
        """)
        
        image_show_case = gr.Examples(
            examples=[
                [human_list_path[5], refer_list_path[0], os.path.join(example_path, 'outputs/190815.jpg')],
                [human_list_path[6], refer_list_path[1], os.path.join(example_path, 'outputs/190811.jpg')],
                [human_list_path[2], refer_list_path[3], os.path.join(example_path, 'outputs/190813.jpg')],
                [human_list_path[0], refer_list_path[4], os.path.join(example_path, 'outputs/190814.jpg')],
            ],
            inputs=[imgs, refer_img, image_out],
            label=None,
            elem_id="image_examples"
        )

        text_show_case = gr.Examples(
            examples=[
                [human_list_path[6], prompt_list_examples[0], os.path.join(example_path, 'outputs/190820.jpg')],
                [human_list_path[4], prompt_list_examples[1], os.path.join(example_path, 'outputs/190812.jpg')],
                [human_list_path[3], prompt_list_examples[2], os.path.join(example_path, 'outputs/190816.jpg')],
                [human_list_path[7], prompt_list_examples[3], os.path.join(example_path, 'outputs/190818.jpg')],
                [human_list_path[8], prompt_list_examples[4], os.path.join(example_path, 'outputs/190819.jpg')],
            ],
            inputs=[imgs, prompt, image_out],
            label=None,
            elem_id="text_examples"
        )
        

Portrait.queue(api_open=False).launch(show_api=False)