app.py

import json
import os
import time
import uuid

import cv2
import gradio as gr
import numpy as np
from PIL import Image
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

from src.generation import call_generation
from src.util import upload_np_2_oss, upload_json_string_2_oss, upload_preprocess, merge_images

universal_matting = pipeline(Tasks.universal_matting, model='damo/cv_unet_universal-matting')

img = "assets/image_gallery_en/"
files = os.listdir(img)
files = [file for file in files if file.lower().endswith(('.png', '.jpg', '.jpeg'))]
files = sorted(files)
basic_usage = []
showcases = []
for idx, name in enumerate(files):
    temp = os.path.join(os.path.dirname(__file__), img, name)
    if idx < 4:
        basic_usage.append(temp)
    else:
        showcases.append(temp)

# - - - - - examples  - - - - -  #
ep = "assets/examples"
# Layout, Style, Color, Subject, Prompt，Strict Layout Edge，Layout Content Scale, Automatic Image Matting
image_examples = [
    [0, f"{ep}/00_layout.png", f"{ep}/empty.png", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0., True],
    [1, f"{ep}/01_layout.png", f"{ep}/empty.png", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, True],
    [2, f"{ep}/empty.png", f"{ep}/02_style.png", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, True],
    [3, f"{ep}/empty.png", f"{ep}/03_style.png", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, True],
    [4, f"{ep}/empty.png", f"{ep}/empty.png", f"{ep}/04_color.png", f"{ep}/empty.png", "A moose, Merry Christmas", True, 0.8, True],
    [5, f"{ep}/empty.png", f"{ep}/empty.png", f"{ep}/05_color.png", f"{ep}/empty.png", "A photo about cherry blossom",
     True, 0.8, True],
    [6, f"{ep}/06_layout.png", f"{ep}/06_style.jpeg", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, True],
    [7, f"{ep}/07_layout.jpeg", f"{ep}/07_style.jpeg", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, True],
    [8, f"{ep}/empty.png", f"{ep}/08_style.png", f"{ep}/08_color.jpeg", f"{ep}/empty.png", "", True, 0.8, True],
    [9, f"{ep}/empty.png", f"{ep}/09_style.png", f"{ep}/empty.png", f"{ep}/base_image1.jpeg", "", True, 0.8, True],
    [10, f"{ep}/10_layout.png", f"{ep}/10_style.png", f"{ep}/empty.png", f"{ep}/base_image2.png", "", True, 0.8, False],
    [11, f"{ep}/11_layout.png", f"{ep}/empty.png", f"{ep}/empty.png", f"{ep}/base_image4.png", "", False, 0.8, False],
    [12, f"{ep}/12_layout.png", f"{ep}/empty.png", f"{ep}/empty.png", f"{ep}/base_image3.png", "", False, 0.8, False],
]

example_images = [
    [0, f"{ep}/layout_image1.jpeg", None, None, None],
    [1, f"{ep}/layout_image1.jpeg", None, None, None],
    [2, None, f"{ep}/style_image1.jpeg", None, None],
    [3, None, f"{ep}/style_image1.jpeg", None, None],
    [4, None, None, f"{ep}/color_image1.jpeg", None],
    [5, None, None, f"{ep}/color_image3.jpeg", None],
    [6, f"{ep}/layout_image2.jpeg", f"{ep}/style_image2.jpeg", None, None],
    [7, f"{ep}/layout_image3.jpeg", f"{ep}/style_image3.jpeg", None, None],
    [8, None, f"{ep}/style_image4.jpeg", f"{ep}/color_image2.jpeg", None],
    [9, None, f"{ep}/style_image6.jpeg", None, f"{ep}/09_base.png"],
    [10, f"{ep}/layout_image5.jpeg", f"{ep}/style_image5.jpeg", None, f"{ep}/10_base.png"],
    [11, f"{ep}/layout_image7.jpeg", None, None, f"{ep}/11_base.png"],
    [12, f"{ep}/layout_image6.jpeg", None, None, f"{ep}/12_base.png"],
]

example_masks = [
    [0, None, None, None],
    [1, f"{ep}/layout_image1_mask.png", None, None],
    [2, None, f"{ep}/style_image1_mask.png", None],
    [3, None, f"{ep}/style_image1_mask2.png", None],
    [4, None, None, None],
    [5, None, None, f"{ep}/color_image3_mask.png"],
    [6, None, None, None],
    [7, None, None, None],
    [8, None, None, None],
    [9, None, f"{ep}/style_image6_mask.png", None],
    [10, f"{ep}/layout_image5_mask.png", None, None],
    [11, None, None, None],
    [12, f"{ep}/layout_image6_mask.png", None, None],
]


def process_example(example_idx, pil_layout_image, pil_style_image, pil_color_image, pil_base_image_rgba,
                    prompt, strict_edge, layout_scale, preprocess_base_image):
    _, layout_image, style_image, color_image, base_image_rgba = example_images[example_idx]
    _, layout_mask, style_mask, color_mask = example_masks[example_idx]

    pil_layout_image = None if layout_image is None else pil_layout_image
    pil_style_image = None if style_image is None else pil_style_image
    pil_color_image = None if color_image is None else pil_color_image
    pil_base_image_rgba = None if base_image_rgba is None else pil_base_image_rgba

    return pil_layout_image, layout_mask, pil_style_image, style_mask, pil_color_image, color_mask, \
           pil_base_image_rgba, prompt, strict_edge, layout_scale, preprocess_base_image


def process(pil_base_image_rgba=None, preprocess_base_image=False,
            pil_layout_image_dict=None, layout_scale=1.0, edge_consistency=0.5,
            strict_edge=False,
            pil_color_image_dict=None, color_scale=1.0,
            pil_style_image_dict=None, style_scale=1.0, prompt="best quality", negative_prompt="",
            pil_layout_mask=None, pil_style_mask=None, pil_color_mask=None):
    request_id = time.strftime('%Y%m%d-', time.localtime(time.time())) + str(uuid.uuid4())

    output_aspect_ratio = 1.
    matting_flag = False

    if pil_base_image_rgba is None:
        base_image_url = ""
        pil_fg_mask = None
    else:
        if preprocess_base_image:
            matting_flag = True
            orig_image = np.array(pil_base_image_rgba)
            orig_alpha = np.array(pil_base_image_rgba)[..., -1]
            matting_alpha = universal_matting(pil_base_image_rgba)[OutputKeys.OUTPUT_IMG][..., -1]
            orig_image[..., -1] = ((matting_alpha > 200) * (orig_alpha > 200) * 255.).astype(np.uint8)
            pil_base_image_rgba = Image.fromarray(orig_image)
        pil_fg_mask = pil_base_image_rgba.split()[-1]

        w, h = pil_base_image_rgba.size
        output_aspect_ratio = max(1.0 * w / h, 1.0 * h / w)
        if output_aspect_ratio > 2:
            raise gr.Error("Input of subject images with aspect ratio exceeding 2 is not supported")
        if min(w, h) > 1536:
            raise gr.Error("Input of subject images with the shorter side exceeding 1536 pixels is not supported")
        base_image_url = upload_np_2_oss(np.array(pil_base_image_rgba), request_id + "_base.png")

    if pil_layout_image_dict is None:
        layout_image_url = ""
    else:
        np_layout_image = np.array(pil_layout_image_dict["image"].convert("RGBA"))
        np_layout_image, np_layout_alpha = np_layout_image[..., :3], np_layout_image[..., 3]
        np_layout_mask = np.array(pil_layout_image_dict["mask"].convert("L"))
        if pil_layout_mask is None:
            np_layout_mask = ((np_layout_alpha > 127) * (np_layout_mask < 127) * 255.).astype(np.uint8)
        else:
            np_layout_mask = ((np_layout_alpha > 127) * (np_layout_mask < 127) *
                              (np.array(pil_layout_mask) > 127) * 255.).astype(np.uint8)
        layout_image_url = upload_np_2_oss(
            np.concatenate([np_layout_image, np_layout_mask[..., None]], axis=-1), request_id + "_layout.png"
        )
        if pil_base_image_rgba is None:
            h, w, c = np_layout_image.shape
            output_aspect_ratio = max(1.0 * w / h, 1.0 * h / w)
            if output_aspect_ratio > 2:
                raise gr.Error("Input of layout images with aspect ratio exceeding 2 is not supported")
            if min(w, h) > 1536:
                raise gr.Error("Input of layout images with the shorter side exceeding 1536 pixels is not supported")

    if pil_style_image_dict is None:
        style_image_url = ""
    else:
        np_style_image = np.array(pil_style_image_dict["image"].convert("RGBA"))
        np_style_image, np_style_alpha = np_style_image[..., :3], np_style_image[..., 3]
        np_style_mask = np.array(pil_style_image_dict["mask"].convert("L"))
        if pil_style_mask is None:
            np_style_mask = ((np_style_alpha > 127) * (np_style_mask < 127) * 255.).astype(np.uint8)
        else:
            np_style_mask = ((np_style_alpha > 127) * (np_style_mask < 127) *
                             (np.array(pil_style_mask) > 127) * 255.).astype(np.uint8)
        style_image_url = upload_np_2_oss(
            np.concatenate([np_style_image, np_style_mask[..., None]], axis=-1), request_id + "_style.png"
        )

    if pil_color_image_dict is None:
        color_image_url = ""
    else:
        np_color_image = np.array(pil_color_image_dict["image"].convert("RGBA"))
        np_color_image, np_color_alpha = np_color_image[..., :3], np_color_image[..., 3]
        np_color_mask = np.array(pil_color_image_dict["mask"].convert("L"))
        if pil_color_mask is None:
            np_color_mask = ((np_color_alpha > 127) * (np_color_mask < 127) * 255.).astype(np.uint8)
        else:
            np_color_mask = ((np_color_alpha > 127) * (np_color_mask < 127) *
                             (np.array(pil_color_mask) > 127) * 255.).astype(np.uint8)
        color_image_url = upload_np_2_oss(
            np.concatenate([np_color_image, np_color_mask[..., None]], axis=-1), request_id + "_color.png"
        )

    res = call_generation(base_image_url=base_image_url, layout_image_url=layout_image_url,
                          color_image_url=color_image_url, style_image_url=style_image_url,
                          strict_edge=int(strict_edge), layout_scale=int(layout_scale * 10),
                          edge_consistency=int(edge_consistency * 10), color_scale=int(color_scale * 10),
                          style_scale=int(style_scale * 10), prompt=prompt, negative_prompt=negative_prompt,
                          output_aspect_ratio=output_aspect_ratio)

    for idx, r in enumerate(res):
        upload_np_2_oss(np.array(r), request_id + f"_{idx}.jpg")

    if matting_flag:
        res.append(pil_base_image_rgba)
    return res, request_id, True, pil_fg_mask


if __name__ == "__main__":

    block = gr.Blocks(
        title="ImageSynthesizer",
        css="assets/css/style.css",
        theme=gr.themes.Soft(
            radius_size=gr.themes.sizes.radius_none,
            text_size=gr.themes.sizes.text_md
        )).queue(concurrency_count=3)
    with block:
        with gr.Row():
            with gr.Column():
                gr.HTML(f"""
                        <div style="text-align: center;">
                            <h1> ImageSynthesizer: Enables a Versatile Visual Information Transfer for Creative Image Synthesis </h1>
                            </br>
                            <h3> ImageSynthesizer supports transferring various visual information from any area of any image to create new compositions, offering higher freedom and flexibility in image synthesis. Currently, it supports the transfer of layout, color, style, and pixel content, with more visual information transfer capabilities. </h3>
                            </br>
                        </div>
                """)

        #with gr.Tabs(elem_classes=["Tab"]):
        #    with gr.TabItem("Image Gallery"):
        #        gr.Gallery(label="Basic Usage", value=basic_usage, height=400, columns=4, object_fit="scale-down")
        #        gr.Gallery(label="Advanced Combinations", value=showcases, height=1200, columns=4, object_fit="scale-down")
        #    with gr.TabItem("Image Creation"):
        #        with gr.Row():
        #            with gr.Column(scale=1):
        #                ...
        #            with gr.Column(scale=3):
        #                ... #gr.Image(value="assets/banner/banner.png", width=1024, show_label=False, show_download_button=False)
        #            with gr.Column(scale=1):
        #                ...
                with gr.Accordion(label="🧭 Instructions:", open=True, elem_id="accordion"):
                    with gr.Row(equal_height=True):
                        #     with gr.Row(elem_id="ShowCase"):
                        #             gr.Image(value="assets/banner/ra.gif")
                        gr.Markdown("""
                        - ⭐️ <b>step1：</b> (Optional) Upload or select a set of images from the examples for the "Layout", "Style", and "Color" reference. Mix and match freely, no need to select all.
                        - ⭐️ <b>step2：</b> (Optional) Use brush to erase areas in the layout, style, and color reference images (if present) that you do not want transferred.
                        - ⭐️ <b>step3：</b> (Optional) Upload an RGBA image to the "Subject" tab, with the alpha channel indicating the subject you wish to preserve at the pixel level (or upload a regular RGB image and check the automatic image matting option, which will automatically segment the subject for you; <b>the default is the latter</b>).
                        - ⭐️ <b>step4：</b> Click "Run" to start the generation process.
                        - ⭐️ <b>step5：</b> (Optional) Additionally, prompt input is supported, as well as control over advanced parameters such as layout edge consistency and conditional weights. Feel free to try these features.
                        """)
                    # with gr.Row(equal_height=True):
                    
                with gr.Row():
                    with gr.Column(scale=1, min_width=160):
                        with gr.Tabs(elem_classes=["feedback"]):
                            with gr.TabItem("Layout (Optional)"):
                                pil_layout_image_dict = gr.ImageMask(source='upload', type="pil", show_label=False,
                                                                     image_mode="RGBA")
                                pil_layout_image = gr.Image(interactive=False, type="pil", visible=False, label="Layout")
                                pil_layout_mask = gr.Image(interactive=False, type="pil", visible=False, label="Layout Mask",
                                                           image_mode="L")
                                with gr.Box():
                                    with gr.Accordion(label="Layout Parameters", open=False, elem_id="accordion"):
                                        with gr.TabItem("Layout Edge"):
                                            strict_edge = gr.Checkbox(label="Strict", value=True)
                                            strict_edge_mirror = gr.Checkbox(label="Strict Layout Edge", visible=False)
                                            edge_consistency = gr.Slider(label="Degree of Consistency (If Not Strict)", minimum=0.0,
                                                                         maximum=1.0,
                                                                         step=0.1, value=0.8, interactive=True)
                                        with gr.TabItem("Layout Content"):
                                            layout_scale = gr.Slider(label="Scale", minimum=0.0, maximum=1.0, step=0.1,
                                                                     value=0.8,
                                                                     interactive=True)
                                            layout_scale_mirror = gr.Slider(label="Layout Content Scale", visible=False)

                    with gr.Column(scale=1, min_width=160):
                        with gr.Tabs(elem_classes=["feedback"]):
                            with gr.TabItem("Style (Optional)"):
                                pil_style_image_dict = gr.ImageMask(source='upload', type="pil", show_label=False,
                                                                    image_mode="RGBA")
                                pil_style_image = gr.Image(interactive=False, type="pil", visible=False, label="Style")
                                pil_style_mask = gr.Image(interactive=False, type="pil", visible=False, label="Style Mask",
                                                          image_mode="L")
                                with gr.Box():
                                    with gr.Accordion(label="Style Parameters", open=False, elem_id="accordion"):
                                        style_scale = gr.Slider(label="Scale", minimum=0.0, maximum=1.0, step=0.1,
                                                                value=0.8, interactive=True)
                    with gr.Column(scale=1, min_width=160):
                        with gr.Tabs(elem_classes=["feedback"]):
                            with gr.TabItem("Color (Optional, recommended for use with prompt)"):
                                pil_color_image_dict = gr.ImageMask(source='upload', type="pil", show_label=False,
                                                                    image_mode="RGBA")
                                pil_color_image = gr.Image(interactive=False, type="pil", visible=False, label="Color")
                                pil_color_mask = gr.Image(interactive=False, type="pil", visible=False, label="Color Mask",
                                                          image_mode="L")
                                with gr.Box():
                                    with gr.Accordion(label="Color Parameters", open=False, elem_id="accordion"):
                                        color_scale = gr.Slider(label="Scale", minimum=0.0, maximum=1.0, step=0.1,
                                                                value=0.8, interactive=True)
                with gr.Row():
                    with gr.Column(scale=1, min_width=160):
                        with gr.Tabs(elem_classes=["feedback"]):
                            with gr.TabItem("Subject (Optional)"):
                                pil_base_image_rgba = gr.Image(source='upload',
                                                               interactive=True, show_label=False,
                                                               type="pil", image_mode="RGBA", tool="editor")
                                pil_base_image_rgba_mirror = gr.Image(label="Subject", image_mode="RGBA", visible=False)
                                with gr.Box():
                                    preprocess_base_image = gr.Checkbox(label="Automatic Image Matting", value=True)
                                    pil_fg_mask = gr.Image(interactive=False, type="pil", image_mode="L", visible=False)
                        run_button = gr.Button("Run", elem_id="btn")
                        with gr.Accordion("", open=True, elem_id="accordion1"):
                            prompt = gr.Textbox(value="", label='Prompt', lines=1, interactive=True)
                            prompt_mirror = gr.Textbox(label='Prompt', visible=False)
                            negative_prompt = gr.Textbox(value="", label='Negative prompt', lines=1,
                                                         interactive=True)
                    with gr.Column(scale=2, min_width=160):
                        with gr.Tabs(elem_classes=["feedback"]):
                            with gr.TabItem("Outputs"):
                                result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery",
                                                            preview=True)
                                recommend = gr.Button("Recommend results to Image Gallery", elem_id="recBut")
                                request_id = gr.State(value="")
                                gallery_flag = gr.State(value=False)
                with gr.Row():
                    with gr.Box():
                        example_idx = gr.Slider(label="Index", visible=False, value=0)
                        example = gr.Examples(
                            label="Input Examples",
                            examples=image_examples,
                            inputs=[example_idx,
                                    pil_layout_image, pil_style_image, pil_color_image, pil_base_image_rgba_mirror,
                                    prompt_mirror, strict_edge_mirror, layout_scale_mirror, preprocess_base_image],
                            outputs=[pil_layout_image_dict, pil_layout_mask, pil_style_image_dict, pil_style_mask,
                                     pil_color_image_dict, pil_color_mask, pil_base_image_rgba,
                                     prompt, strict_edge, layout_scale, preprocess_base_image],
                            fn=process_example,
                            run_on_click=True,
                            examples_per_page=20
                        )
                    gr.HTML(f"""
          		        </br>
                        <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
                                <a href='https://aigcdesigngroup.github.io/transfer-anything/'><img src='https://img.shields.io/badge/Project_Page-TransferAnything-green' alt='Project Page'></a>
                        </div>
                        </br>
                        """)

        def upload_to_img_gallery(pil_base_image_rgba, pil_layout_image_dict, pil_style_image_dict,
                                  pil_color_image_dict, pil_fg_mask, prompt, negative_prompt, res, re_id, flag,
                                  strict_edge, edge_consistency, layout_scale, style_scale, color_scale,
                                  preprocess_base_image):
            if flag:
                np_out_base_image, np_out_layout_image, np_out_style_image, np_out_color_image = upload_preprocess(
                    pil_base_image_rgba, pil_layout_image_dict, pil_style_image_dict, pil_color_image_dict, pil_fg_mask)
                np_out_images = [np_out_base_image, np_out_layout_image, np_out_style_image, np_out_color_image]
                for idx, r in enumerate(res):
                    if idx < 4:
                        r = cv2.imread(r['name'])
                        r = cv2.cvtColor(r, cv2.COLOR_BGR2RGB)
                        upload_np_2_oss(merge_images(*np_out_images, r, prompt, negative_prompt),
                                        name=re_id + f"_merge_{idx}.jpg", gallery=True)
                config = dict(
                    strict_edge=strict_edge,
                    edge_consistency=edge_consistency,
                    layout_scale=layout_scale,
                    style_scale=style_scale,
                    color_scale=color_scale,
                    preprocess_base_image=preprocess_base_image
                )
                upload_json_string_2_oss(json.dumps(config), name=re_id + f"_config.txt", gallery=True)

                flag = False
                gr.Info("Images have been uploaded and await review.")
            else:
                gr.Info("No images to recommend, or already suggested once.")
            return flag


        recommend.click(
            upload_to_img_gallery,
            [pil_base_image_rgba, pil_layout_image_dict, pil_style_image_dict, pil_color_image_dict, pil_fg_mask,
             prompt, negative_prompt, result_gallery, request_id, gallery_flag, strict_edge, edge_consistency,
             layout_scale, style_scale, color_scale, preprocess_base_image],
            [gallery_flag]
        )

        ips = [pil_base_image_rgba, preprocess_base_image,
               pil_layout_image_dict, layout_scale, edge_consistency, strict_edge,
               pil_color_image_dict, color_scale,
               pil_style_image_dict, style_scale, prompt, negative_prompt,
               pil_layout_mask, pil_style_mask, pil_color_mask]
        run_button.click(fn=process, inputs=ips, outputs=[result_gallery, request_id, gallery_flag, pil_fg_mask])

    block.launch(share=True)