Upload 46 files

.gitattributes

@@ -33,3 +33,4 @@ 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/storyboard_en.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+**/__pycache__
+workdir
+**/.vscode
+*.sh
+.idea/
+.DS_Store

README.md

@@ -11,4 +11,4 @@ license: apache-2.0
 short_description: Synthesize images from text prompts and visual references
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,387 @@
+##!/usr/bin/python3
+# -*- coding: utf-8 -*-
+# @Time    : 2025-03-12
+# @Author  : Junjie He
+import os
+import time
+import uuid
+
+import gradio as gr
+import numpy as np
+from PIL import Image
+
+from src.agent import extend_prompt, translate_prompt
+from src.anystory import call_anystory
+from src.matting import ImageUniversalMatting
+from src.util import upload_pil_2_oss
+
+if not os.path.exists("models/tf_matting.pb"):
+    os.makedirs("models", exist_ok=True)
+    os.system(f"wget -O  models/tf_matting.pb {os.getenv('MATTING_PATH')}")
+universal_matting = ImageUniversalMatting("models/tf_matting.pb")
+
+
+def image_matting(pil_image):
+    if pil_image.mode == "RGBA":
+        mask = np.array(pil_image)[..., -1] > 200
+        if np.all(mask):
+            mask = ((universal_matting(pil_image)[..., -1] > 200) * 255).astype(np.uint8)
+        else:
+            mask = ((np.array(pil_image)[..., -1] > 200) * 255).astype(np.uint8)
+    else:
+        mask = ((universal_matting(pil_image.convert("RGB"))[..., -1] > 200) * 255).astype(np.uint8)
+    pil_mask = Image.fromarray(mask)
+
+    np_image = np.array(pil_image.convert("RGB"))
+    np_mask = np.array(pil_mask)[..., None] / 255.
+    pil_masked_image = Image.fromarray((np_mask * np_image + (1 - np_mask) * 255.).astype(np.uint8))
+    return pil_masked_image, pil_mask
+
+
+def process(
+        pil_subject_A_image=None,
+        pil_subject_A_mask=None,
+        pil_subject_B_image=None,
+        pil_subject_B_mask=None,
+        prompt="",
+):
+    request_id = time.strftime('%Y%m%d-', time.localtime(time.time())) + str(uuid.uuid4())
+
+    if prompt == "":
+        raise gr.Error("Please enter your prompt")
+
+    if pil_subject_A_image is None and pil_subject_B_image is None:
+        raise gr.Error("Please upload your reference image(s)")
+
+    image_urls = []
+    if pil_subject_A_image is not None:
+        if pil_subject_A_mask is not None:
+            if pil_subject_A_mask.size != pil_subject_A_image.size:
+                raise gr.Error("Subject [A] image & mask size mismatch")
+            pil_subject_A_image = pil_subject_A_image.convert("RGB")
+            pil_subject_A_mask = pil_subject_A_mask.convert("L")
+            pil_subject_A_image = Image.merge("RGBA", (*pil_subject_A_image.split(), pil_subject_A_mask))
+        image_urls.append(upload_pil_2_oss(pil_subject_A_image, name=request_id + "_A.png"))
+    if pil_subject_B_image is not None:
+        if pil_subject_B_mask is not None:
+            if pil_subject_B_mask.size != pil_subject_B_image.size:
+                raise gr.Error("Subject [B] image & mask size mismatch")
+            pil_subject_B_image = pil_subject_B_image.convert("RGB")
+            pil_subject_B_mask = pil_subject_B_mask.convert("L")
+            pil_subject_B_image = Image.merge("RGBA", (*pil_subject_B_image.split(), pil_subject_B_mask))
+        image_urls.append(upload_pil_2_oss(pil_subject_B_image, name=request_id + "_B.png"))
+
+    res = call_anystory(image_urls, prompt)[0]
+
+    return res
+
+
+def interface():
+    with gr.Row(variant="panel"):
+        gr.HTML(description + "<br>" + tips)
+    with gr.Row(variant="panel"):
+        with gr.Column(scale=2, min_width=100):
+            with gr.Row(equal_height=False):
+                with gr.Column(scale=1, min_width=100):
+                    with gr.Tab(label="Subject [A]"):
+                        with gr.Group():
+                            with gr.Row(equal_height=True):
+                                with gr.Column(min_width=100):
+                                    pil_subject_A_image = gr.Image(type="pil", label="Subject [A] Reference Image",
+                                                                   format="png", show_label=True, image_mode="RGBA")
+                                with gr.Column(min_width=100):
+                                    with gr.Group():
+                                        pil_subject_A_mask = gr.Image(type="pil",
+                                                                      label="Subject [A] Mask (upload supported)",
+                                                                      format="png", show_label=True, image_mode="L")
+                                        seg_subject_A = gr.Button(value="Segment Subject")
+
+                with gr.Column(scale=1, min_width=100):
+                    with gr.Tab(label="Subject [B]"):
+                        with gr.Group():
+                            with gr.Row(equal_height=True):
+                                with gr.Column(min_width=100):
+                                    pil_subject_B_image = gr.Image(type="pil", label="Subject [B] Reference Image",
+                                                                   format="png", show_label=True, image_mode="RGBA")
+                                with gr.Column(min_width=100):
+                                    with gr.Group():
+                                        pil_subject_B_mask = gr.Image(type="pil",
+                                                                      label="Subject [B] Mask (upload supported)",
+                                                                      format="png", show_label=True, image_mode="L")
+                                        seg_subject_B = gr.Button(value="Segment Subject")
+
+            with gr.Group():
+                prompt = gr.Textbox(value="", label='Prompt', lines=6, show_label=True)
+                en_prompt = gr.Textbox(value="", label='prompt', lines=6, show_label=True, visible=False)
+                # prompt_extend_button = gr.Button(value="提示词扩写")
+
+        with gr.Column(scale=1, min_width=100):
+            result_gallery = gr.Image(type="pil", label="Generated Image", visible=True, height=450)
+            # result_gallery = gr.Gallery(label='Generated Image', show_label=True, elem_id="gallery", preview=True,
+            #                             format="png", height=450)
+            run_button = gr.Button(value="🧑‍🎨 RUN")
+
+            generated_information = gr.Markdown(label="Generation Details", value="", visible=False)
+
+    seg_subject_A.click(
+        fn=set_image_seg_unfinished, outputs=generated_information
+    ).then(
+        fn=image_matting, inputs=[pil_subject_A_image], outputs=[pil_subject_A_image, pil_subject_A_mask]
+    ).then(
+        fn=set_image_seg_finished, outputs=generated_information
+    )
+
+    seg_subject_B.click(
+        fn=set_image_seg_unfinished, outputs=generated_information
+    ).then(
+        fn=image_matting, inputs=[pil_subject_B_image], outputs=[pil_subject_B_image, pil_subject_B_mask]
+    ).then(
+        fn=set_image_seg_finished, outputs=generated_information
+    )
+
+    # prompt_extend_button.click(
+    #     fn=set_prompt_extend_unfinished, outputs=generated_information
+    # ).then(
+    #     fn=extend_prompt, inputs=[prompt], outputs=[prompt]
+    # ).then(
+    #     fn=set_prompt_extend_finished, outputs=generated_information
+    # )
+
+    run_button.click(
+        fn=set_prompt_translate_unfinished, outputs=generated_information
+    ).then(
+        fn=translate_prompt, inputs=[prompt], outputs=[en_prompt]
+    ).then(
+        fn=set_image_generate_unfinished, outputs=generated_information
+    ).then(
+        fn=process,
+        inputs=[pil_subject_A_image, pil_subject_A_mask, pil_subject_B_image, pil_subject_B_mask, en_prompt],
+        outputs=[result_gallery]
+    ).then(
+        fn=set_image_generate_finished, outputs=generated_information
+    )
+
+    with gr.Row():
+        examples = [
+            [
+                "assets/examples/1.webp",
+                "assets/examples/1_mask.webp",
+                None,
+                None,
+                "Cartoon style. A sheep is ridough the city, holding a wooden sign that says \"TongYi\".",
+                "assets/examples/1_output.webp",
+            ],
+            [
+                "assets/examples/2.webp",
+                "assets/examples/2_mask.webp",
+                None,
+                None,
+                "Cartoon style. Sun Wukong stands on a tank, holding up an ancient wooden sign high in the air. The sign reads 'AnyStory'. The background is a cyberpunk-style city sky filled with towering buildings.",
+                "assets/examples/2_output.webp",
+            ],
+            [
+                "assets/examples/3.webp",
+                "assets/examples/3_mask.webp",
+                None,
+                None,
+                "A modern and stylish Nezha playing an electric guitar, dynamic pose, vibrant colors, fantasy atmosphere, mythical Chinese character with a rock-and-roll twist, red scarf flowing in the wind, traditional elements mixed with contemporary design, cinematic lighting, 4k resolution",
+                "assets/examples/3_output.webp",
+            ],
+            [
+                "assets/examples/4.webp",
+                "assets/examples/4_mask.webp",
+                None,
+                None,
+                "a man riding a bike on the road",
+                "assets/examples/4_output.webp",
+            ],
+            [
+                "assets/examples/7.webp",
+                "assets/examples/7_mask.webp",
+                None,
+                None,
+                "Nezha is surrounded by a mysterious purple glow, with a pair of eyes glowing with an eerie red light. Broken talismans and debris float around him, highlighting his demonic nature and authority.",
+                "assets/examples/7_output.webp",
+            ],
+            [
+                "assets/examples/8.webp",
+                "assets/examples/8_mask.webp",
+                None,
+                None,
+                "The car is driving through a cyberpunk city at night in the middle of a heavy downpour.",
+                "assets/examples/8_output.webp",
+            ],
+            [
+                "assets/examples/9.webp",
+                "assets/examples/9_mask.webp",
+                None,
+                None,
+                "This cosmetic is placed on a table covered with roses.",
+                "assets/examples/9_output.webp",
+            ],
+            [
+                "assets/examples/10.webp",
+                "assets/examples/10_mask.webp",
+                None,
+                None,
+                "A little boy model is posing for a photo.",
+                "assets/examples/10_output.webp",
+            ],
+            # [
+            #     "assets/examples/5_1.webp",
+            #     "assets/examples/5_1_mask.webp",
+            #     "assets/examples/5_2.webp",
+            #     "assets/examples/5_2_mask.webp",
+            #     "两个小孩骑着一辆炫酷的双人电动车，在热闹的菜市场中穿梭。周围是琳琅满目的蔬菜摊、水果筐和忙碌的摊主，他们表情专注又带点嬉笑，车篮里还装着几根胡萝卜和一把青菜，传统与现代元素在烟火气息中完美交融。",
+            #     "assets/examples/5_output.webp",
+            # ],
+            [
+                "assets/examples/6_1.webp",
+                "assets/examples/6_1_mask.webp",
+                "assets/examples/6_2.webp",
+                "assets/examples/6_2_mask.webp",
+                "Two men are sitting by a wooden table, which is laden with delicious food and a pot of wine. One of the men holds a wine glass, drinking heartily with a bold expression; the other smiles as he pours wine for his companion, both of them engaged in cheerful conversation. In the background is an ancient pavilion surrounded by emerald bamboo groves, with sunlight filtering through the leaves to cast dappled shadows.",
+                "assets/examples/6_output.webp",
+            ],
+        ]
+        gr.Examples(
+            label="Examples",
+            examples=examples,
+            inputs=[pil_subject_A_image, pil_subject_A_mask, pil_subject_B_image, pil_subject_B_mask, prompt, result_gallery],
+        )
+
+
+def set_image_seg_unfinished():
+    return gr.update(
+        visible=True,
+        value="<h3>(Unfinished) Extracting Subject Mask...</h3>",
+    )
+
+
+def set_image_seg_finished():
+    return gr.update(visible=True, value="<h3>Subject mask ready!</h3>")
+
+
+def set_prompt_extend_unfinished():
+    return gr.update(
+        visible=True,
+        value="<h3>(Unfinished) Rewriting your prompt... ✍️</h3>",
+    )
+
+
+def set_prompt_extend_finished():
+    return gr.update(visible=True, value="<h3>Prompt expanded successfully!</h3>")
+
+
+def set_prompt_translate_unfinished():
+    return gr.update(
+        visible=True,
+        value="<h3>(Unfinished) Preprocessing...</h3>",
+    )
+
+
+def set_image_generate_unfinished():
+    return gr.update(
+        visible=True,
+        value="<h3>(Unfinished) Generating images...</h3>",
+    )
+
+
+def set_image_generate_finished():
+    return gr.update(visible=True, value="<h3>Image generation is completed!</h3>")
+
+
+if __name__ == "__main__":
+    title = r"""
+            <div style="text-align: center;">
+                <h1> AnyStory: Towards Unified Single and Multiple Subject Personalization in Text-to-Image Generation </h1>
+                <h1> V2.0.0 </h1>
+                <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
+                    <a href="https://arxiv.org/pdf/2501.09503"><img src="https://img.shields.io/badge/arXiv-2501.09503-red"></a>
+                    &nbsp;
+                    <a href='https://aigcdesigngroup.github.io/AnyStory/'><img src='https://img.shields.io/badge/Project_Page-AnyStory-green' alt='Project Page'></a>
+                    &nbsp;
+                    <a href='https://modelscope.cn/studios/damo/studio_anytext'><img src='https://img.shields.io/badge/Demo-ModelScope-blue'></a>
+                </div>
+                </br>
+            </div>
+        """
+
+    title_description = r"""
+        Official demo of <b>AnyStory 2</b> 🤗. We will continuously update this demo.
+        For technical details, please refer to our tech report: <a href='https://arxiv.org/pdf/2501.09503' target='_blank'><b>AnyStory: Towards Unified Single and Multiple Subject Personalization in Text-to-Image Generation</b></a>. 😊
+        """
+
+    description = r"""🚀🚀🚀 Quick Start:<br>
+        1. Upload subject reference images (clean background; real human IDs unsupported for now), Add prompts (CN/EN supported), and Click "<b>RUN</b>".<br>
+        2. (Recommended) Click "<b>Segment Subject</b>" to create masks (or upload your own B&W masks) for subjects. This helps the model better reference the subject you specify (otherwise, we will perform automatic detection). 🤗<br>
+        """
+
+    tips = r"""💡💡💡 Tips:<br>
+    If the subject doesn't appear, try adding a detailed description of the subject in the prompt that matches the reference image, and avoid conflicting details (e.g., significantly altering the subject's appearance). Multi-subject referencing in AnyStory2 is still being optimized. 🤗<br>
+    """
+
+    citation = r"""
+        ---
+        📝 **Citation**
+        <br>
+        If our work is helpful for your research or applications, please cite us via:
+        ```bibtex
+        @article{he2025anystory,
+          title={AnyStory: Towards Unified Single and Multiple Subject Personalization in Text-to-Image Generation},
+          author={He, Junjie and Tuo, Yuxiang and Chen, Binghui and Zhong, Chongyang and Geng, Yifeng and Bo, Liefeng},
+          journal={arXiv preprint arXiv:2501.09503},
+          year={2025}
+        }
+        ```
+        If you have any questions, feel free to open an issue or contact us directly at <b>[email protected]</b>.
+        """
+
+    js = """
+    function createGradioAnimation() {
+        var container = document.createElement('div');
+        container.id = 'gradio-animation';
+        container.style.fontSize = '2em';
+        container.style.fontWeight = 'bold';
+        container.style.textAlign = 'center';
+        container.style.marginBottom = '20px';
+
+        var text = 'Welcome to AnyStory!';
+        for (var i = 0; i < text.length; i++) {
+            (function(i){
+                setTimeout(function(){
+                    var letter = document.createElement('span');
+                    letter.style.opacity = '0';
+                    letter.style.transition = 'opacity 0.5s';
+                    letter.innerText = text[i];
+
+                    container.appendChild(letter);
+
+                    setTimeout(function() {
+                        letter.style.opacity = '1';
+                    }, 50);
+                }, i * 250);
+            })(i);
+        }
+
+        var gradioContainer = document.querySelector('.gradio-container');
+        gradioContainer.insertBefore(container, gradioContainer.firstChild);
+
+        return 'Animation created';
+    }
+    """
+    block = gr.Blocks(title="AnyStory2", js=js, theme=gr.themes.Ocean()).queue()
+    with block:
+        gr.HTML(title)
+        gr.HTML(title_description)
+
+        interface()
+
+        gr.HTML("<br>More examples: Intelligent creation of AI story pictures integrated with Qwen Agent")
+        gr.Gallery(value=["assets/storyboard_en.png"], columns=1, object_fit="contain", show_label=False)
+
+        gr.Markdown(citation)
+
+    block.launch(share=True, max_threads=10)
+    # block.launch(server_name='0.0.0.0', share=False, server_port=9999, max_threads=3)
+    # block.launch(server_name='127.0.0.1', share=False, server_port=9999, allowed_paths=["/"])

requirements.txt

@@ -0,0 +1,4 @@
+oss2
+phidata
+dashscope
+tensorflow==2.15

src/agent.py

@@ -0,0 +1,111 @@
+import os
+import time
+
+import gradio as gr
+from phi.assistant import Assistant
+from phi.llm.openai import OpenAIChat
+from pydantic import BaseModel
+
+from .log import logger
+
+qwen_model = 'qwen-plus'
+qwen_api_key = os.getenv("QWEN_DS_API_KEY")
+qwen_url = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
+
+
+class TranslateFormat(BaseModel):
+    en: str
+
+
+class PormptExtendFormat(BaseModel):
+    prompt: str
+
+
+trans_assist = Assistant(
+    llm=OpenAIChat(model=qwen_model, max_tokens=3000, temperature=0.3, api_key=qwen_api_key, base_url=qwen_url),
+    description='你是专业的英语翻译，请将用户输入的一句中文翻译成英文',
+    output_model=TranslateFormat
+)
+
+requirements = [
+    "1.对于过于简短的用户输入，在不改变原意及主体外观的前提下，合理补充具象化细节，避免使用过于抽象描述。",
+    "2.如果涉及人像，完善与用户输入相符的人物外貌、表情、种族、姿态、景别、穿着、影调、质感等方面描述。",
+    "3.匹配符合用户意图且精准详细的风格描述。如果用户未指定，则根据画面选择最恰当的风格，或使用纪实摄影照片。",
+    "4.若输入为中文，则整体中文输出；若输入为全英文，则整体英文输出；保留引号中原文以及重要的输入信息，不要改写。",
+    "5.在语义完整前提下，改写后prompt字数小于200字，避免prompt冗长。",
+    """
+    6.改写后 prompt 示例：
+    (1) 日系小清新写真照片，扎着双麻花辫的波西米亚小女孩坐在船边。女孩穿着白色方领泡泡袖连衣裙，裙子上有褶皱和纽扣装饰。她皮肤白皙，五官精致，眼泪汪汪直视镜头。她双手扶船，长发刘海遮住部分额头。背景是清澈明亮的户外场景，可见蓝天、山峦和一些干枯植物。高清写实摄影，近景中心对称构图。\n
+    (2) 二次元厚涂动漫插画，一个猫耳东亚萌妹手持文件夹，怒气冲冲走向电脑。她深紫色爆炸头，红色眼睛，头顶有一个粉色光圈。少女身穿深灰色短裙和浅灰色上衣，腰间系着白色系带，胸前佩戴名牌，居中写着黑体中文"紫阳"。室内背景中摆放了很多办公桌。粗线条的日系赛璐璐风格。近景半身略仰视视角。\n
+    (3) 美剧艺术海报风格，身穿黄色防护服的老年Walter White坐在金属折叠椅上，头顶无衬线英文写着"Breaking Bad"，周围是成堆的美元和蓝色塑料储物箱。他戴着眼镜目光直视镜头，左手拿着一支雪茄，右手放在膝盖上。背景是废弃阴暗的厂房，阳光透过窗户照射进来。画面带有明显颗粒质感纹理。长焦人物平视特写。
+    (4) CG game concept digital art featuring three giant mutant crocodiles with wide-open mouths, revealing pink tongues and sharp teeth. Their rough skin resembles grayish-white stone. On the back of the crocodile to the left, lush trees, shrubs, and some thorn-like protrusions grow. The background of the scene shows a dusk sky and a shimmering pond. The overall atmosphere is dark and cold, with a tilted composition that creates a strong sense of depth and layers. \n
+    """
+]
+
+prompt_writer = Assistant(
+    llm=OpenAIChat(model=qwen_model, max_tokens=8000, temperature=0.3, api_key=qwen_api_key, base_url=qwen_url),
+    description="你是一位prompt优化师，旨在将用户输入改写为优质prompt，在不影响原意前提下更高质量生图。",
+    instructions=requirements,
+    debug_mode=False,
+    output_model=PormptExtendFormat,
+)
+
+
+def contains_chinese(text):
+    """
+    判断文本中是否包含中文字符（基于 Unicode 范围）
+    """
+    # 定义中文的 Unicode 范围（基础汉字 + 扩展区）
+    cjk_ranges = [
+        (0x4e00, 0x9fff),  # CJK 基本汉字区（含常用汉字）
+        (0x3400, 0x4dbf),  # CJK 扩展 A
+        (0x20000, 0x2a6df),  # CJK 扩展 B
+        (0x2a700, 0x2b73f),  # CJK 扩展 C
+        (0x2b740, 0x2b81f),  # CJK 扩展 D
+        (0x2b820, 0x2ceaf),  # CJK 扩展 E
+    ]
+
+    for c in text:
+        code = ord(c)
+        for start, end in cjk_ranges:
+            if start <= code <= end:
+                return True
+    return False
+
+
+def translate_prompt(prompt, max_attempts=5):
+    attempts = 0
+    while attempts < max_attempts:
+        try:
+            if contains_chinese(prompt):
+                res = trans_assist.run(prompt)
+                logger.info(f"translate Chinese prompt into English: {prompt} -> {res.en}")
+                return res.en
+            else:
+                return prompt
+        except Exception as e:
+            print(f"尝试 {attempts + 1} 失败：{e}")
+            attempts += 1
+            if attempts < 5:
+                time.sleep(1)  # 延迟重试
+    raise gr.Error(f"提示词扩写达到最大尝试次数 {max_attempts}，任务失败。")
+
+
+def extend_prompt(prompt, max_attempts=5):
+    attempts = 0
+    while attempts < max_attempts:
+        try:
+            res = prompt_writer.run(prompt, stream=False)
+            logger.info(f"extend prompt finished: {prompt} -> {res.prompt}.")
+            return res.prompt
+        except Exception as e:
+            print(f"尝试 {attempts + 1} 失败：{e}")
+            attempts += 1
+            if attempts < 5:
+                time.sleep(1)  # 延迟重试
+    raise gr.Error(f"提示词扩写达到最大尝试次数 {max_attempts}，任务失败。")
+
+
+if __name__ == '__main__':
+    # test_trans()
+    pass

src/anystory.py

@@ -0,0 +1,72 @@
+import json
+import os
+import time
+
+import gradio as gr
+import requests
+
+from src.log import logger
+from src.util import download_images
+
+anystory_url = os.getenv("ANYSTORY_URL")
+anystory_api_key = os.getenv("ANYSTORY_DS_API_KEY")
+anystory_model = os.getenv("ANYSTORY_MODEL")
+
+
+def call_anystory(image_urls, prompt):
+    headers = {
+        "Content-Type": "application/json",
+        "Accept": "application/json",
+        "Authorization": f"Bearer {anystory_api_key}",
+        "X-DashScope-Async": "enable",
+        "X-DashScope-DataInspection": "enable",
+    }
+    data = {
+        "model": anystory_model,
+        "input": {
+            "image_urls": image_urls,
+            "prompt": prompt
+        },
+        "parameters": {
+        },
+    }
+
+    res = requests.post(anystory_url, data=json.dumps(data), headers=headers)
+
+    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'https://poc-dashscope.aliyuncs.com/api/v1/tasks/{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.")
+
+            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)
+        logger.info(f"task_id: {task_id}: Generate done.")
+    else:
+        logger.error(f'Fail to create Generation task: {res.content}')
+        raise gr.Error("Fail to create Generation task.")
+
+    return img_data

src/log.py

@@ -0,0 +1,18 @@
+import logging
+import os
+from logging.handlers import RotatingFileHandler
+
+log_file_name = "workdir/AnyStory.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="AnyStory-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/matting.py

@@ -0,0 +1,94 @@
+import cv2
+import numpy as np
+import tensorflow as tf
+from PIL import Image
+
+if tf.__version__ >= '2.0':
+    tf = tf.compat.v1
+
+
+class ImageUniversalMatting:
+
+    def __init__(self, weight_path):
+        super().__init__()
+        config = tf.ConfigProto(allow_soft_placement=True, device_count={'GPU': 0})
+        config.gpu_options.allow_growth = True
+        self._session = tf.Session(config=config)
+        with self._session.as_default():
+            print(f'loading model from {weight_path}')
+            with tf.gfile.FastGFile(weight_path, 'rb') as f:
+                graph_def = tf.GraphDef()
+                graph_def.ParseFromString(f.read())
+                tf.import_graph_def(graph_def, name='')
+                self.output = self._session.graph.get_tensor_by_name(
+                    'output_png:0')
+                self.input_name = 'input_image:0'
+            print('load model done')
+        self._session.graph.finalize()
+
+    def __call__(self, image):
+        output = self.preprocess(image)
+        output = self.forward(output)
+        output = self.postprocess(output)
+        return output
+
+    def resize_image(self, img, limit_side_len):
+        """
+        resize image to a size multiple of 32 which is required by the network
+        args:
+            img(array): array with shape [h, w, c]
+        return(tuple):
+            img, (ratio_h, ratio_w)
+        """
+        h, w, _ = img.shape
+
+        # limit the max side
+        if max(h, w) > limit_side_len:
+            if h > w:
+                ratio = float(limit_side_len) / h
+            else:
+                ratio = float(limit_side_len) / w
+        else:
+            ratio = 1.
+        resize_h = int(h * ratio)
+        resize_w = int(w * ratio)
+
+        resize_h = int(round(resize_h / 32) * 32)
+        resize_w = int(round(resize_w / 32) * 32)
+
+        img = cv2.resize(img, (int(resize_w), int(resize_h)))
+
+        return img
+
+    @staticmethod
+    def convert_to_ndarray(img):
+        if isinstance(img, Image.Image):
+            img = np.array(img.convert('RGB'))
+        elif isinstance(img, np.ndarray):
+            if len(img.shape) == 2:
+                img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+                img = img[:, :, ::-1]  # convert to rgb
+        else:
+            raise TypeError(f'input should be either PIL.Image,'
+                            f' np.array, but got {type(img)}')
+        return img
+
+    def preprocess(self, input, limit_side_len=800):
+        img = self.convert_to_ndarray(input)  # rgb input
+        img = img.astype(float)
+        orig_h, orig_w, _ = img.shape
+        img = self.resize_image(img, limit_side_len)
+        result = {'img': img, 'orig_h': orig_h, 'orig_w': orig_w}
+        return result
+
+    def forward(self, input):
+        orig_h, orig_w = input['orig_h'], input['orig_w']
+        with self._session.as_default():
+            feed_dict = {self.input_name: input['img']}
+            output_img = self._session.run(self.output, feed_dict=feed_dict)  # RGBA
+            # output_img = cv2.cvtColor(output_img, cv2.COLOR_RGBA2BGRA)
+            output_img = cv2.resize(output_img, (int(orig_w), int(orig_h)))
+        return {"output_img": output_img}
+
+    def postprocess(self, inputs):
+        return inputs["output_img"]

src/util.py

@@ -0,0 +1,71 @@
+import concurrent.futures
+import io
+import os
+import time
+
+import oss2
+import requests
+from PIL import Image
+
+from .log import logger
+
+# oss
+oss_ak = os.getenv("OSS_AK")
+oss_sk = os.getenv("OSS_SK")
+oss_bucket = os.getenv("OSS_BUCKET")
+oss_endpoint = os.getenv("OSS_ENDPOINT")
+oss_path = os.getenv("OSS_PATH")
+
+bucket = oss2.Bucket(oss2.Auth(oss_ak, oss_sk), oss_endpoint, oss_bucket)
+
+
+def download_pil_img(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, n=1):
+    imgs_pil = [None] * n
+    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
+        to_do = []
+        for i, url in enumerate(img_urls):
+            future = executor.submit(download_pil_img, 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 resize_img_with_max_size(image, max_side_length=512):
+    width, height = image.size
+    ratio = max_side_length / max(width, height)
+    new_width = int(width * ratio)
+    new_height = int(height * ratio)
+    resized_image = image.resize((new_width, new_height))
+    return resized_image
+
+
+def upload_pil_2_oss(pil_image, name="cache.jpg"):
+    image = resize_img_with_max_size(pil_image, max_side_length=512)
+
+    imgByteArr = io.BytesIO()
+    if name.lower().endswith(".png"):
+        image.save(imgByteArr, format="PNG")
+    else:
+        image.save(imgByteArr, format="JPEG", quality=95)
+    imgByteArr = imgByteArr.getvalue()
+
+    start_time = time.perf_counter()
+    bucket.put_object(oss_path + "/" + name, imgByteArr)
+    ret = bucket.sign_url('GET', oss_path + "/" + name, 60 * 60 * 24)
+    logger.info(f"upload cost: {time.perf_counter() - start_time} s.")
+    del imgByteArr
+    return ret