Upload 38 files

.gitattributes

@@ -33,3 +33,20 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/examples/1-output-1.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/1-output-2.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/1-output-3.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/2-output-1.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/2-output-2.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/2-output-3.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/2-output-4.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/2-stylegan2-ffhq-0100.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/2-stylegan2-ffhq-0293.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/3-output-1.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/3-output-2.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/3-output-3.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/3-output-4.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/3-style-1.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/3-stylegan2-ffhq-0293.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/3-stylegan2-ffhq-0381.png filter=lfs diff=lfs merge=lfs -text
+assets/examples/Trump-4.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,12 @@
 ---
 title: UniPortrait
-emoji: 🐠
-colorFrom: yellow
+emoji: 👩‍🎨
+colorFrom: purple
 colorTo: red
 sdk: gradio
-sdk_version: 4.41.0
+sdk_version: 4.31.2
 app_file: app.py
 pinned: false
 license: apache-2.0
 ---
-
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,382 @@
+##!/usr/bin/python3
+# -*- coding: utf-8 -*-
+# @Time    : 2024-07-31
+# @Author  : Junjie He
+import gradio as gr
+
+from src.process import (
+    text_to_single_id_generation_process,
+    text_to_multi_id_generation_process,
+    image_to_single_id_generation_process,
+)
+
+
+def text_to_single_id_generation_block():
+    gr.Markdown("## Text-to-Single-ID Generation")
+    gr.HTML(text_to_single_id_description)
+    gr.HTML(text_to_single_id_tips)
+    with gr.Row():
+        with gr.Column(scale=1, min_width=100):
+            prompt = gr.Textbox(value="", label='Prompt', lines=2)
+            negative_prompt = gr.Textbox(value="nsfw", label='Negative Prompt')
+            image_resolution = gr.Dropdown(choices=["768x512", "512x512", "512x768"], value="512x512",
+                                           label="Image Resolution (HxW)")
+            run_button = gr.Button(value="Run")
+
+            with gr.Accordion("Advanced Options", open=True):
+                seed = gr.Slider(label="Seed (-1 indicates random)", minimum=-1, maximum=2147483647, step=1, value=-1)
+
+                faceid_scale = gr.Slider(label="Face ID Scale", minimum=0.0, maximum=1.0, step=0.01, value=0.7)
+                face_structure_scale = gr.Slider(label="Face Structure Scale", minimum=0.0, maximum=1.0,
+                                                 step=0.01, value=0.1)
+
+                style_scale = gr.Slider(label="style_scale", minimum=0.0, maximum=1.0, step=0.01, value=0.7)
+
+                use_sr = gr.Checkbox(label="RealESRGAN 2x", value=True)
+
+        with gr.Column(scale=3, min_width=100):
+            with gr.Row(equal_height=False):
+                pil_faceid = gr.Image(type="pil", label="ID Image")
+                with gr.Accordion("ID Supplements", open=True):
+                    with gr.Row():
+                        pil_supp_faceids = gr.File(file_count="multiple", file_types=["image"],
+                                                   type="filepath", label="Additional ID Images")
+                    with gr.Row():
+                        with gr.Column(scale=1, min_width=100):
+                            pil_mix_faceid_1 = gr.Image(type="pil", label="Mix ID 1")
+                            mix_scale_1 = gr.Slider(label="Mix Scale 1", minimum=0.0, maximum=1.0, step=0.01, value=0.0)
+                        with gr.Column(scale=1, min_width=100):
+                            pil_mix_faceid_2 = gr.Image(type="pil", label="Mix ID 2")
+                            mix_scale_2 = gr.Slider(label="Mix Scale 2", minimum=0.0, maximum=1.0, step=0.01, value=0.0)
+                pil_style = gr.Image(type="pil", label="Style")
+
+            with gr.Row():
+                example_output = gr.Image(type="pil", label="(Example Output)", visible=False)
+                result_gallery = gr.Gallery(label='Output', show_label=True, elem_id="gallery", columns=4, preview=True,
+                                            format="png")
+    with gr.Row():
+        examples = [
+            [
+                "A young man with short black hair, wearing a black hoodie with a hood, was paired with a blue denim jacket with yellow details.",
+                "assets/examples/1-newton.jpg",
+                "assets/No-Image-Placeholder.png",
+                "assets/examples/1-output-1.png",
+            ],
+            [
+                "A young man with short black hair, wearing a black hoodie with a hood, was paired with a blue denim jacket with yellow details.",
+                "assets/examples/1-newton.jpg",
+                "assets/examples/1-style-1.jpg",
+                "assets/examples/1-output-2.png",
+            ],
+        ]
+        gr.Examples(
+            label="Examples",
+            examples=examples,
+            inputs=[prompt, pil_faceid, pil_style, example_output],
+        )
+    ips = [
+        pil_faceid, pil_supp_faceids,
+        pil_mix_faceid_1, mix_scale_1,
+        pil_mix_faceid_2, mix_scale_2,
+        faceid_scale, face_structure_scale,
+        prompt, negative_prompt,
+        pil_style, style_scale,
+        seed, image_resolution, use_sr,
+    ]
+    run_button.click(fn=text_to_single_id_generation_process, inputs=ips, outputs=[result_gallery])
+
+
+def text_to_multi_id_generation_block():
+    gr.Markdown("## Text-to-Multi-ID Generation")
+    gr.HTML(text_to_multi_id_description)
+    gr.HTML(text_to_multi_id_tips)
+    with gr.Row():
+        with gr.Column(scale=1, min_width=100):
+            prompt = gr.Textbox(value="", label='Prompt', lines=2)
+            negative_prompt = gr.Textbox(value="nsfw", label='Negative Prompt')
+            image_resolution = gr.Dropdown(choices=["768x512", "512x512", "512x768"], value="512x512",
+                                           label="Image Resolution (HxW)")
+            run_button = gr.Button(value="Run")
+
+            with gr.Accordion("Advanced Options", open=True):
+                seed = gr.Slider(label="Seed (-1 indicates random)", minimum=-1, maximum=2147483647, step=1, value=-1)
+
+                faceid_scale = gr.Slider(label="Face ID Scale", minimum=0.0, maximum=1.0, step=0.01, value=0.7)
+                face_structure_scale = gr.Slider(label="Face Structure Scale", minimum=0.0, maximum=1.0,
+                                                 step=0.01, value=0.3)
+                style_scale = gr.Slider(label="style_scale", minimum=0.0, maximum=1.0, step=0.01, value=0.7)
+
+                use_sr = gr.Checkbox(label="RealESRGAN 2x", value=True)
+
+        with gr.Column(scale=3, min_width=100):
+            with gr.Row(equal_height=False):
+                with gr.Column(scale=1, min_width=100):
+                    pil_faceid_1st = gr.Image(type="pil", label="First ID")
+                    with gr.Accordion("First ID Supplements", open=False):
+                        with gr.Row():
+                            pil_supp_faceids_1st = gr.File(file_count="multiple", file_types=["image"],
+                                                           type="filepath", label="Additional ID Images")
+                        with gr.Row():
+                            with gr.Column(scale=1, min_width=100):
+                                pil_mix_faceid_1_1st = gr.Image(type="pil", label="Mix ID 1")
+                                mix_scale_1_1st = gr.Slider(label="Mix Scale 1", minimum=0.0, maximum=1.0, step=0.01,
+                                                            value=0.0)
+                            with gr.Column(scale=1, min_width=100):
+                                pil_mix_faceid_2_1st = gr.Image(type="pil", label="Mix ID 2")
+                                mix_scale_2_1st = gr.Slider(label="Mix Scale 2", minimum=0.0, maximum=1.0, step=0.01,
+                                                            value=0.0)
+                with gr.Column(scale=1, min_width=100):
+                    pil_faceid_2nd = gr.Image(type="pil", label="Second ID")
+                    with gr.Accordion("Second ID Supplements", open=False):
+                        with gr.Row():
+                            pil_supp_faceids_2nd = gr.File(file_count="multiple", file_types=["image"],
+                                                           type="filepath", label="Additional ID Images")
+                        with gr.Row():
+                            with gr.Column(scale=1, min_width=100):
+                                pil_mix_faceid_1_2nd = gr.Image(type="pil", label="Mix ID 1")
+                                mix_scale_1_2nd = gr.Slider(label="Mix Scale 1", minimum=0.0, maximum=1.0, step=0.01,
+                                                            value=0.0)
+                            with gr.Column(scale=1, min_width=100):
+                                pil_mix_faceid_2_2nd = gr.Image(type="pil", label="Mix ID 2")
+                                mix_scale_2_2nd = gr.Slider(label="Mix Scale 2", minimum=0.0, maximum=1.0, step=0.01,
+                                                            value=0.0)
+                with gr.Column(scale=1, min_width=100):
+                    pil_style = gr.Image(type="pil", label="Style")
+
+            with gr.Row():
+                example_output = gr.Image(type="pil", label="(Example Output)", visible=False)
+                result_gallery = gr.Gallery(label='Output', show_label=True, elem_id="gallery", columns=4, preview=True,
+                                            format="png")
+    with gr.Row():
+        examples = [
+            [
+                "两个女人在欢笑和快乐中被捕捉到，他们的脸上洋溢着真挚的幸福，背景是日落时分的宁静海滩。这幅画以柔和的风格描绘，捕捉了这一刻的温暖和宁静。",
+                "assets/examples/2-stylegan2-ffhq-0100.png",
+                "assets/examples/2-stylegan2-ffhq-0293.png",
+                "assets/No-Image-Placeholder.png",
+                "assets/examples/2-output-1.png",
+            ],
+            [
+                "The two female models are drinking coffee. The background was off-white.",
+                "assets/examples/2-stylegan2-ffhq-0100.png",
+                "assets/examples/2-stylegan2-ffhq-0293.png",
+                "assets/examples/2-style-1.jpg",
+                "assets/examples/2-output-2.png",
+            ],
+        ]
+        gr.Examples(
+            label="Examples",
+            examples=examples,
+            inputs=[prompt, pil_faceid_1st, pil_faceid_2nd, pil_style, example_output],
+        )
+    with gr.Row():
+        examples = [
+            [
+                "Two men in an American poster.",
+                "assets/examples/Trump-1.jpg",
+                ["assets/examples/Trump-2.jpg", "assets/examples/Trump-3.jpg", "assets/examples/Trump-4.jpg"],
+                "assets/examples/Biden-1.jpg",
+                ["assets/examples/Biden-2.jpg", "assets/examples/Biden-3.jpg", "assets/examples/Biden-4.jpg"],
+                "assets/examples/2-output-4.png",
+            ],
+            [
+                "Two men engaged in a vigorous handshake, both wearing expressions of enthusiasm and determination, set against a backdrop of a bustling business district. The image is crafted in a sleek and modern digital art style, conveying the dynamic and competitive nature of their interaction.",
+                "assets/examples/Trump-1.jpg",
+                ["assets/examples/Trump-2.jpg", "assets/examples/Trump-3.jpg", "assets/examples/Trump-4.jpg"],
+                "assets/examples/Biden-1.jpg",
+                ["assets/examples/Biden-2.jpg", "assets/examples/Biden-3.jpg", "assets/examples/Biden-4.jpg"],
+                "assets/examples/2-output-3.png",
+            ],
+        ]
+        gr.Examples(
+            label="Examples (Multiple References)",
+            examples=examples,
+            inputs=[prompt, pil_faceid_1st, pil_supp_faceids_1st, pil_faceid_2nd, pil_supp_faceids_2nd, example_output],
+        )
+    ips = [
+        pil_faceid_1st, pil_supp_faceids_1st,
+        pil_mix_faceid_1_1st, mix_scale_1_1st,
+        pil_mix_faceid_2_1st, mix_scale_2_1st,
+        pil_faceid_2nd, pil_supp_faceids_2nd,
+        pil_mix_faceid_1_2nd, mix_scale_1_2nd,
+        pil_mix_faceid_2_2nd, mix_scale_2_2nd,
+        faceid_scale, face_structure_scale,
+        prompt, negative_prompt,
+        pil_style, style_scale,
+        seed, image_resolution, use_sr,
+    ]
+    run_button.click(fn=text_to_multi_id_generation_process, inputs=ips, outputs=[result_gallery])
+
+
+def image_to_single_id_generation_block():
+    gr.Markdown("## Image-to-Single-ID Generation")
+    gr.HTML(image_to_single_id_description)
+    gr.HTML(image_to_single_id_tips)
+    with gr.Row():
+        with gr.Column(scale=1, min_width=100):
+            image_resolution = gr.Dropdown(choices=["768x512", "512x512", "512x768"], value="512x512",
+                                           label="Image Resolution (HxW)")
+            run_button = gr.Button(value="Run")
+
+            with gr.Accordion("Advanced Options", open=True):
+                seed = gr.Slider(label="Seed (-1 indicates random)", minimum=-1, maximum=2147483647, step=1, value=-1)
+
+                style_scale = gr.Slider(label="Reference Scale", minimum=0.0, maximum=1.0, step=0.01, value=0.7)
+                faceid_scale = gr.Slider(label="Face ID Scale", minimum=0.0, maximum=1.0, step=0.01, value=0.7)
+                face_structure_scale = gr.Slider(label="Face Structure Scale", minimum=0.0, maximum=1.0, step=0.01,
+                                                 value=0.3)
+
+                use_sr = gr.Checkbox(label="RealESRGAN 2x", value=True)
+
+        with gr.Column(scale=3, min_width=100):
+            with gr.Row(equal_height=False):
+                pil_style = gr.Image(type="pil", label="Portrait Reference")
+                pil_faceid = gr.Image(type="pil", label="ID Image")
+                with gr.Accordion("ID Supplements", open=True):
+                    with gr.Row():
+                        pil_supp_faceids = gr.File(file_count="multiple", file_types=["image"],
+                                                   type="filepath", label="Additional ID Images")
+                    with gr.Row():
+                        with gr.Column(scale=1, min_width=100):
+                            pil_mix_faceid_1 = gr.Image(type="pil", label="Mix ID 1")
+                            mix_scale_1 = gr.Slider(label="Mix Scale 1", minimum=0.0, maximum=1.0, step=0.01, value=0.0)
+                        with gr.Column(scale=1, min_width=100):
+                            pil_mix_faceid_2 = gr.Image(type="pil", label="Mix ID 2")
+                            mix_scale_2 = gr.Slider(label="Mix Scale 2", minimum=0.0, maximum=1.0, step=0.01, value=0.0)
+            with gr.Row():
+                with gr.Column(scale=3, min_width=100):
+                    example_output = gr.Image(type="pil", label="(Example Output)", visible=False)
+                    result_gallery = gr.Gallery(label='Output', show_label=True, elem_id="gallery", columns=4,
+                                                preview=True, format="png")
+    with gr.Row():
+        examples = [
+            [
+                "assets/examples/3-style-1.png",
+                "assets/examples/3-stylegan2-ffhq-0293.png",
+                0.7,
+                0.3,
+                "assets/examples/3-output-1.png",
+            ],
+            [
+                "assets/examples/3-style-1.png",
+                "assets/examples/3-stylegan2-ffhq-0293.png",
+                0.6,
+                0.0,
+                "assets/examples/3-output-2.png",
+            ],
+            [
+                "assets/examples/3-style-2.jpg",
+                "assets/examples/3-stylegan2-ffhq-0381.png",
+                0.7,
+                0.3,
+                "assets/examples/3-output-3.png",
+            ],
+            [
+                "assets/examples/3-style-3.jpg",
+                "assets/examples/3-stylegan2-ffhq-0381.png",
+                0.6,
+                0.0,
+                "assets/examples/3-output-4.png",
+            ],
+        ]
+        gr.Examples(
+            label="Examples",
+            examples=examples,
+            inputs=[pil_style, pil_faceid, faceid_scale, face_structure_scale, example_output],
+        )
+    ips = [
+        pil_faceid, pil_supp_faceids,
+        pil_mix_faceid_1, mix_scale_1,
+        pil_mix_faceid_2, mix_scale_2,
+        faceid_scale, face_structure_scale,
+        pil_style, style_scale,
+        seed, image_resolution, use_sr,
+    ]
+    run_button.click(fn=image_to_single_id_generation_process, inputs=ips, outputs=[result_gallery])
+
+
+if __name__ == "__main__":
+    title = r"""
+            <div style="text-align: center;">
+                <h1> UniPortrait: A Unified Framework for Identity-Preserving Single- and Multi-Human Image Personalization </h1>
+                <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
+                    <a href="https://arxiv.org/pdf/xxxx.xxxxx"><img src="https://img.shields.io/badge/arXiv-xxxx.xxxxx-red"></a>
+                    &nbsp;
+                    <a href='https://aigcdesigngroup.github.io/UniPortrait-Page/'><img src='https://img.shields.io/badge/Project_Page-UniPortrait-green' alt='Project Page'></a>
+                    &nbsp;
+                    <a href="https://github.com/junjiehe96/UniPortrait"><img src="https://img.shields.io/badge/Github-Code-blue"></a>
+                </div>
+                </br>
+            </div>
+        """
+
+    title_description = r"""
+        This is the <b>official 🤗 Gradio demo</b> for <a href='https://arxiv.org/pdf/xxxx.xxxxx' target='_blank'><b>UniPortrait: A Unified Framework for Identity-Preserving Single- and Multi-Human Image Personalization</b></a>.<br>
+        The demo provides three capabilities: text-to-single-ID personalization, text-to-multi-ID personalization, and image-to-single-ID personalization. All of these are based on the Stable Diffusion v1-5 model. Feel free to give them a try! 😊
+        """
+
+    text_to_single_id_description = r"""🚀🚀🚀Quick start:<br>
+        1. Enter a text prompt (Chinese or English), Upload an image with a face, and Click the <b>Run</b> button.<br>
+        2. (Optional) You can also upload an image as the style reference for the results. 🤗<br>
+        """
+
+    text_to_single_id_tips = r"""💡💡💡Tips:<br>
+        1. Try to avoid creating too small faces, as this may lead to some artifacts. (Currently, the short side length of the generated image is limited to 512)<br>
+        2. It's a good idea to upload multiple reference photos of your face to improve the prompt and ID consistency. Additional references can be uploaded in the "ID supplements".<br>
+        3. The appropriate values of "Face ID Scale" and "Face Structure Scale" are important for balancing the ID and text alignment. We recommend using "Face ID Scale" (0.5~0.7) and "Face Structure Scale" (0.0~0.4).<br>
+        """
+
+    text_to_multi_id_description = r"""🚀🚀🚀Quick start:<br>
+        1. Enter a text prompt (Chinese or English), Upload an image with a face in "First ID" and "Second ID" blocks respectively, and Click the <b>Run</b> button.<br>
+        2. (Optional) You can also upload an image as the style reference for the results. 🤗<br>
+        """
+
+    text_to_multi_id_tips = r"""💡💡💡Tips:<br>
+        1. Try to avoid creating too small faces, as this may lead to some artifacts. (Currently, the short side length of the generated image is limited to 512)<br>
+        2. It's a good idea to upload multiple reference photos of your face to improve the prompt and ID consistency. Additional references can be uploaded in the "ID supplements".<br>
+        3. The appropriate values of "Face ID Scale" and "Face Structure Scale" are important for balancing the ID and text alignment. We recommend using "Face ID Scale" (0.3~0.7) and "Face Structure Scale" (0.0~0.4).<br>
+        """
+
+    image_to_single_id_description = r"""🚀🚀🚀Quick start: Upload an image as the portrait reference (can be any style), Upload a face image, and Click the <b>Run</b> button. 🤗<br>"""
+
+    image_to_single_id_tips = r"""💡💡💡Tips:<br>
+        1. Try to avoid creating too small faces, as this may lead to some artifacts. (Currently, the short side length of the generated image is limited to 512)<br>
+        2. It's a good idea to upload multiple reference photos of your face to improve ID consistency. Additional references can be uploaded in the "ID supplements".<br>
+        3. The appropriate values of "Face ID Scale" and "Face Structure Scale" are important for balancing the portrait reference and ID alignment. We recommend using "Face ID Scale" (0.5~0.7) and "Face Structure Scale" (0.0~0.4).<br>
+        """
+
+    citation = r"""
+        ---
+        📝 **Citation**
+        <br>
+        If our work is helpful for your research or applications, please cite us via:
+        ```bibtex
+        @article{wang2024instantid,
+          title={InstantID: Zero-shot Identity-Preserving Generation in Seconds},
+          author={Wang, Qixun and Bai, Xu and Wang, Haofan and Qin, Zekui and Chen, Anthony},
+          journal={arXiv preprint arXiv:2401.07519},
+          year={2024}
+        }
+        ```
+        📧 **Contact**
+        <br>
+        If you have any questions, please feel free to open an issue or directly reach us out at <b>he_junjie@zju.edu.cn</b>.
+        """
+
+    block = gr.Blocks(title="UniPortrait").queue()
+    with block:
+        gr.HTML(title)
+        gr.HTML(title_description)
+
+        with gr.TabItem("Text-to-Single-ID"):
+            text_to_single_id_generation_block()
+
+        with gr.TabItem("Text-to-Multi-ID"):
+            text_to_multi_id_generation_block()
+
+        with gr.TabItem("Image-to-Single-ID (Stylization)"):
+            image_to_single_id_generation_block()
+
+    block.launch(share=True)
+    # block.launch(server_name='0.0.0.0', share=False, server_port=9999, allowed_paths=["/"])
+    # block.launch(server_name='127.0.0.1', share=False, server_port=9999, allowed_paths=["/"])

requirements.txt

@@ -0,0 +1 @@
+oss2

src/generation.py

@@ -0,0 +1,79 @@
+import json
+import os
+import time
+
+import gradio as gr
+import requests
+
+from src.log import logger
+from src.util import download_images
+
+
+def call_generation(data):
+    url_task = os.getenv("URL_TASK")
+    api_key = os.getenv("API_KEY_GENERATION")
+    model_id = os.getenv("MODEL_ID")
+    url_query = os.getenv("URL_QUERY")
+
+    batch_size = 2
+    repeat_times = 2
+
+    headers = {
+        "Content-Type": "application/json",
+        "Accept": "application/json",
+        "Authorization": f"Bearer {api_key}",
+        "X-DashScope-Async": "enable",
+    }
+    data["model"] = model_id
+    data["parameters"]["n"] = batch_size
+
+    all_res_ = []
+    for i in range(repeat_times):
+        if data["parameters"]["seed"] != -1:
+            data["parameters"]["seed"] = data["parameters"]["seed"] * (i+1)
+        res_ = requests.post(url_task, data=json.dumps(data), headers=headers)
+        all_res_.append(res_)
+
+    all_image_data = []
+    for res_ in all_res_:
+        respose_code = res_.status_code
+        if 200 == respose_code:
+            res = json.loads(res_.content.decode())
+            task_id = res['output']['task_id']
+            logger.info(f"task_id: {task_id}: Create request success. Params: {data}")
+
+            # Async query
+            is_running = True
+            while is_running:
+                res_ = requests.post(f'{url_query}/{task_id}', headers=headers)
+                respose_code = res_.status_code
+                if 200 == respose_code:
+                    res = json.loads(res_.content.decode())
+                    if "SUCCEEDED" == res['output']['task_status']:
+                        logger.info(f"task_id: {task_id}: Generation task query success.")
+                        results = res['output']['results']
+                        img_urls = [x['url'] for x in results]
+                        logger.info(f"task_id: {task_id}: {res}")
+                        break
+                    elif "FAILED" != res['output']['task_status']:
+                        logger.debug(f"task_id: {task_id}: query result...")
+                        time.sleep(1)
+                    else:
+                        raise gr.Error(
+                            "Fail to get results from Generation task. Make sure all the ID images have a clear face. If it still doesn't work, you can contact us or open an issue.")
+
+                else:
+                    logger.error(f'task_id: {task_id}: Fail to query task result: {res_.content}')
+                    raise gr.Error("Fail to query task result.")
+
+            logger.info(f"task_id: {task_id}: download generated images.")
+            img_data = download_images(img_urls, batch_size)
+            logger.info(f"task_id: {task_id}: Generate done.")
+            all_image_data += img_data
+        else:
+            logger.error(f'Fail to create Generation task: {res_.content}')
+            raise gr.Error("Fail to create Generation task.")
+
+    if len(all_image_data) != repeat_times * batch_size:
+        raise gr.Error("Fail to Generation.")
+    return all_image_data

src/log.py

@@ -0,0 +1,18 @@
+import logging
+import os
+from logging.handlers import RotatingFileHandler
+
+log_file_name = "workdir/UniPortrait.log"
+os.makedirs(os.path.dirname(log_file_name), exist_ok=True)
+
+format = '[%(levelname)s] %(asctime)s "%(filename)s", line %(lineno)d, %(message)s'
+logging.basicConfig(
+    format=format,
+    datefmt="%Y-%m-%d %H:%M:%S",
+    level=logging.INFO)
+logger = logging.getLogger(name="UniPortrait-Studio")
+
+fh = RotatingFileHandler(log_file_name, maxBytes=20000000, backupCount=3)
+formatter = logging.Formatter(format, datefmt="%Y-%m-%d %H:%M:%S")
+fh.setFormatter(formatter)
+logger.addHandler(fh)

src/process.py

@@ -0,0 +1,277 @@
+import time
+import uuid
+
+import gradio as gr
+import numpy as np
+from PIL import Image
+
+from .generation import call_generation
+from .util import upload_np_2_oss
+
+
+def get_image_url(request_id, pil_image, suffix):
+    np_image = np.array(pil_image.convert("RGB"))
+    image_url = upload_np_2_oss(np_image, request_id + suffix)
+    return image_url
+
+
+def is_no_image_placeholder(pil_image):
+    from PIL import Image, ImageChops
+    no_image_placeholder = Image.open("assets/No-Image-Placeholder.png").convert("RGB")
+    diff = ImageChops.difference(pil_image, no_image_placeholder)
+    if diff.getbbox() is None:
+        return True
+    else:
+        return False
+
+
+def text_to_single_id_generation_process(
+        pil_faceid=None, pil_supp_faceids=None,
+        pil_mix_faceid_1=None, mix_scale_1=0.0,
+        pil_mix_faceid_2=None, mix_scale_2=0.0,
+        faceid_scale=0.0, face_structure_scale=0.0,
+        prompt="", negative_prompt="",
+        pil_style=None, style_scale=0.0,
+        seed=-1, image_resolution="512x512", use_sr=True,
+):
+    request_id = time.strftime('%Y%m%d-', time.localtime(time.time())) + str(uuid.uuid4())
+
+    if prompt == "":
+        raise gr.Error("Please enter the prompt")
+
+    if pil_style and is_no_image_placeholder(pil_style):
+        pil_style = None
+
+    faceid_url = []
+    if pil_faceid:
+        faceid_url.append(get_image_url(request_id, pil_faceid, suffix="_faceid.png"))
+    if pil_supp_faceids and len(pil_supp_faceids) > 0:
+        for idx, pil_supp_faceid in enumerate(pil_supp_faceids):
+            pil_supp_faceid = Image.open(pil_supp_faceid)
+            faceid_url.append(get_image_url(request_id, pil_supp_faceid, suffix=f"_supp_faceid_{idx}.png"))
+
+    mix_face_url = []
+    mix_scale = []
+    if pil_mix_faceid_1 is not None:
+        mix_face_url.append(
+            get_image_url(request_id, pil_mix_faceid_1, suffix=f"_mix_faceid_1_{mix_scale_1:.2f}.png")
+        )
+        mix_scale.append(mix_scale_1)
+
+    if pil_mix_faceid_2 is not None:
+        mix_face_url.append(
+            get_image_url(request_id, pil_mix_faceid_2, suffix=f"_mix_faceid_2_{mix_scale_2:.2f}.png")
+        )
+        mix_scale.append(mix_scale_2)
+
+    # The outer parenthesis represent the ID, and the inner parenthesis represent the mix face image for each ID.
+    faceid_image_url, mix_faceid_image_url, mix_faceid_scale = [], [], []
+    if len(faceid_url) > 0:
+        faceid_image_url.append(faceid_url)
+        mix_faceid_image_url.append(mix_face_url)
+        mix_faceid_scale.append(mix_scale)
+
+    if pil_style is None:
+        style_image_url = ""
+    else:
+        style_image_url = get_image_url(request_id, pil_style, suffix="_style.png")
+
+    data = dict(
+        input=dict(
+            faceid_image_url=faceid_image_url,
+            mix_faceid_image_url=mix_faceid_image_url,
+            mix_faceid_scale=mix_faceid_scale,
+            style_image_url=style_image_url,
+            prompt=prompt,
+            negative_prompt=negative_prompt
+        ),
+        parameters=dict(
+            image_resolution=image_resolution,
+            faceid_scale=faceid_scale * 10.,
+            face_structure_scale=face_structure_scale * 10.,
+            style_scale=style_scale * 10.,
+            use_sr=int(use_sr),
+            seed=seed,
+            mode="text-to-image",
+        )
+    )
+
+    res = call_generation(data)
+
+    return res
+
+
+def text_to_multi_id_generation_process(
+        pil_faceid_1st=None, pil_supp_faceids_1st=None,
+        pil_mix_faceid_1_1st=None, mix_scale_1_1st=0.0,
+        pil_mix_faceid_2_1st=None, mix_scale_2_1st=0.0,
+        pil_faceid_2nd=None, pil_supp_faceids_2nd=None,
+        pil_mix_faceid_1_2nd=None, mix_scale_1_2nd=0.0,
+        pil_mix_faceid_2_2nd=None, mix_scale_2_2nd=0.0,
+        faceid_scale=0.0, face_structure_scale=0.0,
+        prompt="", negative_prompt="",
+        pil_style=None, style_scale=0.0,
+        seed=-1, image_resolution="512x512", use_sr=True,
+):
+    request_id = time.strftime('%Y%m%d-', time.localtime(time.time())) + str(uuid.uuid4())
+
+    if prompt == "":
+        raise gr.Error("Please enter the prompt")
+
+    if pil_style and is_no_image_placeholder(pil_style):
+        pil_style = None
+
+    first_faceid_url = []
+    if pil_faceid_1st:
+        first_faceid_url.append(get_image_url(request_id, pil_faceid_1st, suffix="_faceid_1st.png"))
+    if pil_supp_faceids_1st and len(pil_supp_faceids_1st) > 0:
+        for idx, pil_supp_faceid_1st in enumerate(pil_supp_faceids_1st):
+            pil_supp_faceid_1st = Image.open(pil_supp_faceid_1st)
+            first_faceid_url.append(get_image_url(request_id, pil_supp_faceid_1st, suffix=f"_faceid_1st_{idx}.png"))
+
+    second_faceid_url = []
+    if pil_faceid_2nd:
+        second_faceid_url.append(get_image_url(request_id, pil_faceid_2nd, suffix="_faceid_2nd.png"))
+    if pil_supp_faceids_2nd and len(pil_supp_faceids_2nd) > 0:
+        for idx, pil_supp_faceid_2nd in enumerate(pil_supp_faceids_2nd):
+            pil_supp_faceid_2nd = Image.open(pil_supp_faceid_2nd)
+            second_faceid_url.append(get_image_url(request_id, pil_supp_faceid_2nd, suffix=f"_faceid_2nd_{idx}.png"))
+
+    first_mix_face_url, first_mix_scale = [], []
+    if pil_mix_faceid_1_1st:
+        first_mix_face_url.append(
+            get_image_url(request_id, pil_mix_faceid_1_1st, suffix=f"_mix_faceid_1_1st_{mix_scale_1_1st:.2f}.png")
+        )
+        first_mix_scale.append(mix_scale_1_1st)
+
+    if pil_mix_faceid_2_1st is not None:
+        first_mix_face_url.append(
+            get_image_url(request_id, pil_mix_faceid_2_1st, suffix=f"_mix_faceid_2_1st_{mix_scale_2_1st:.2f}.png")
+        )
+        first_mix_scale.append(mix_scale_2_1st)
+
+    second_mix_face_url, second_mix_scale = [], []
+    if pil_mix_faceid_1_2nd:
+        second_mix_face_url.append(
+            get_image_url(request_id, pil_mix_faceid_1_2nd, suffix=f"_mix_faceid_1_2nd_{mix_scale_1_2nd:.2f}.png")
+        )
+        second_mix_scale.append(mix_scale_1_2nd)
+
+    if pil_mix_faceid_2_2nd is not None:
+        second_mix_face_url.append(
+            get_image_url(request_id, pil_mix_faceid_2_2nd, suffix=f"_mix_faceid_2_2nd_{mix_scale_2_2nd:.2f}.png")
+        )
+        second_mix_scale.append(mix_scale_2_2nd)
+
+    # The outer parenthesis represent the ID, and the inner parenthesis represent the mix face image for each ID.
+    faceid_image_url, mix_faceid_image_url, mix_faceid_scale = [], [], []
+    if len(first_faceid_url) > 0:
+        faceid_image_url.append(first_faceid_url)
+        mix_faceid_image_url.append(first_mix_face_url)
+        mix_faceid_scale.append(first_mix_scale)
+    if len(second_faceid_url) > 0:
+        faceid_image_url.append(second_faceid_url)
+        mix_faceid_image_url.append(second_mix_face_url)
+        mix_faceid_scale.append(second_mix_scale)
+
+    if pil_style is None:
+        style_image_url = ""
+    else:
+        style_image_url = get_image_url(request_id, pil_style, suffix="_style.png")
+
+    data = dict(
+        input=dict(
+            faceid_image_url=faceid_image_url,
+            mix_faceid_image_url=mix_faceid_image_url,
+            mix_faceid_scale=mix_faceid_scale,
+            style_image_url=style_image_url,
+            prompt=prompt,
+            negative_prompt=negative_prompt
+        ),
+        parameters=dict(
+            image_resolution=image_resolution,
+            faceid_scale=faceid_scale * 10.,
+            face_structure_scale=face_structure_scale * 10.,
+            style_scale=style_scale * 10.,
+            use_sr=int(use_sr),
+            seed=seed,
+            mode="text-to-image",
+        )
+    )
+
+    res = call_generation(data)
+
+    return res
+
+
+def image_to_single_id_generation_process(
+        pil_faceid=None, pil_supp_faceids=None,
+        pil_mix_faceid_1=None, mix_scale_1=0.0,
+        pil_mix_faceid_2=None, mix_scale_2=0.0,
+        faceid_scale=0.0, face_structure_scale=0.0,
+        pil_style=None, style_scale=1.0,
+        seed=-1, image_resolution="768x512", use_sr=True,
+):
+    request_id = time.strftime('%Y%m%d-', time.localtime(time.time())) + str(uuid.uuid4())
+
+    if pil_faceid is None:
+        raise gr.Error("Please upload an ID image")
+
+    if pil_style is None:
+        raise gr.Error("Please upload a reference image")
+
+    faceid_url = []
+    if pil_faceid:
+        faceid_url.append(get_image_url(request_id, pil_faceid, suffix="_faceid.png"))
+    if pil_supp_faceids and len(pil_supp_faceids) > 0:
+        for idx, pil_supp_faceid in enumerate(pil_supp_faceids):
+            pil_supp_faceid = Image.open(pil_supp_faceid)
+            faceid_url.append(get_image_url(request_id, pil_supp_faceid, suffix=f"_supp_faceid_{idx}.png"))
+
+    mix_face_url = []
+    mix_scale = []
+    if pil_mix_faceid_1 is not None:
+        mix_face_url.append(
+            get_image_url(request_id, pil_mix_faceid_1, suffix=f"_mix_faceid_1_{mix_scale_1:.2f}.png")
+        )
+        mix_scale.append(mix_scale_1)
+
+    if pil_mix_faceid_2 is not None:
+        mix_face_url.append(
+            get_image_url(request_id, pil_mix_faceid_2, suffix=f"_mix_faceid_2_{mix_scale_2:.2f}.png")
+        )
+        mix_scale.append(mix_scale_2)
+
+    # The outer parenthesis represent the ID, and the inner parenthesis represent the mix face image for each ID.
+    faceid_image_url, mix_faceid_image_url, mix_faceid_scale = [], [], []
+    if len(faceid_url) > 0:
+        faceid_image_url.append(faceid_url)
+        mix_faceid_image_url.append(mix_face_url)
+        mix_faceid_scale.append(mix_scale)
+
+    if pil_style is None:
+        style_image_url = ""
+    else:
+        style_image_url = get_image_url(request_id, pil_style, suffix="_style.png")
+
+    data = dict(
+        input=dict(
+            faceid_image_url=faceid_image_url,
+            mix_faceid_image_url=mix_faceid_image_url,
+            mix_faceid_scale=mix_faceid_scale,
+            style_image_url=style_image_url,
+        ),
+        parameters=dict(
+            image_resolution=image_resolution,
+            faceid_scale=faceid_scale * 10.,
+            face_structure_scale=face_structure_scale * 10.,
+            style_scale=style_scale * 10.,
+            use_sr=int(use_sr),
+            seed=seed,
+            mode="image-to-image",
+        )
+    )
+
+    res = call_generation(data)
+
+    return res

src/util.py

@@ -0,0 +1,58 @@
+import concurrent.futures
+import io
+import os
+
+import oss2
+import requests
+from PIL import Image
+
+from .log import logger
+
+# oss
+access_key_id = os.getenv("ACCESS_KEY_ID")
+access_key_secret = os.getenv("ACCESS_KEY_SECRET")
+bucket_name = os.getenv("BUCKET_NAME")
+endpoint = os.getenv("ENDPOINT")
+
+bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)
+oss_path = os.getenv("OSS_PATH")
+
+
+def download_img_pil(index, img_url):
+    r = requests.get(img_url, stream=True)
+    if r.status_code == 200:
+        img = Image.open(io.BytesIO(r.content))
+        return (index, img)
+    else:
+        logger.error(f"Fail to download: {img_url}")
+
+
+def download_images(img_urls, batch_size):
+    imgs_pil = [None] * batch_size
+    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
+        to_do = []
+        for i, url in enumerate(img_urls):
+            future = executor.submit(download_img_pil, i, url)
+            to_do.append(future)
+
+        for future in concurrent.futures.as_completed(to_do):
+            ret = future.result()
+            index, img_pil = ret
+            imgs_pil[index] = img_pil
+
+    return imgs_pil
+
+
+def upload_np_2_oss(input_image, name="cache.png"):
+    assert name.lower().endswith((".png", ".jpg")), name
+    imgByteArr = io.BytesIO()
+    if name.lower().endswith(".png"):
+        Image.fromarray(input_image).save(imgByteArr, format="PNG")
+    else:
+        Image.fromarray(input_image).save(imgByteArr, format="JPEG", quality=95)
+    imgByteArr = imgByteArr.getvalue()
+
+    bucket.put_object(oss_path + "/" + name, imgByteArr)
+    ret = bucket.sign_url('GET', oss_path + "/" + name, 60 * 60 * 24)
+    del imgByteArr
+    return ret