Spaces:
Running
Running
| ##!/usr/bin/python3 | |
| # -*- coding: utf-8 -*- | |
| # @Time : 2024-01-29 | |
| # @Author : Junjie He | |
| 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., False], | |
| [1, f"{ep}/01_layout.png", f"{ep}/empty.png", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, False], | |
| [2, f"{ep}/empty.png", f"{ep}/02_style.png", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, False], | |
| [3, f"{ep}/empty.png", f"{ep}/03_style.png", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, False], | |
| [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, False], | |
| [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, False], | |
| [6, f"{ep}/06_layout.png", f"{ep}/06_style.jpeg", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, False], | |
| [7, f"{ep}/07_layout.jpeg", f"{ep}/07_style.jpeg", f"{ep}/empty.png", f"{ep}/empty.png", "", True, 0.8, False], | |
| [8, f"{ep}/empty.png", f"{ep}/08_style.png", f"{ep}/08_color.jpeg", f"{ep}/empty.png", "", True, 0.8, False], | |
| [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="TransferAnything", | |
| 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> TransferAnything: Enabling Versatile Visual Information Transfer for Creative Image Synthesis </h1> | |
| <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> | |
| <h3> TransferAnything 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 under continuous development. </h3> | |
| </br> | |
| </div> | |
| """) | |
| # <h5 style="margin: 0; color: red">如果你认为该项目有所帮助的话，不妨给我们Github点个Star以便获取最新的项目进展.</h5> | |
| 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 image with the subject you wish to preserve at the pixel level to the "Subject" component, with a transparent channel required (or upload a regular RGB image and select automatic image matting, which will automatically segment the subject area for you). | |
| - ⭐️ <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): | |
| # gr.Markdown("""❤️ 温馨提示：上传图片最大边不超过1024，宽高比为1：1，将会生成更快哟""") | |
| 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") | |
| 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 | |
| ) | |
| 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) | |