app.py

'''
AnyText: Multilingual Visual Text Generation And Editing
Paper: https://arxiv.org/abs/2311.03054
Code: https://github.com/tyxsspa/AnyText
Copyright (c) Alibaba, Inc. and its affiliates.
'''
import os
import time
import uuid
import cv2
import gradio as gr
import re
import json
import argparse
import random
from src.generation import call_generation
from src.util import upload_np_2_oss

def count_lines(prompt):
    prompt = prompt.replace('“', '"')
    prompt = prompt.replace('”', '"')
    p = '"(.*?)"'
    strs = re.findall(p, prompt)
    if len(strs) == 0:
        strs = [' ']
    return len(strs)


def process(mode, prompt, texts, texts2, layout_radio, sort_radio, revise_pos, base_model_path, lora_path_ratio, show_debug, img_count, ddim_steps, w, h, strength, cfg_scale, seed, eta, a_prompt, n_prompt):
    texts = '"' + texts + '"'
    n_lines = count_lines(texts)
    if texts2 != "":
        texts2 = '"' + texts2 + '"'
        n_lines += count_lines(texts2)
        texts += texts2

    layout_dict = {"top": "example_images/templates/top_512_512.png",
                   "down": "example_images/templates/bottom_512_512.png",
                   "left": "example_images/templates/left_512_512.png",
                   "right": "example_images/templates/right_512_512.png",
                   "top two": "example_images/templates/top2_512_512.png",
                   "down two": "example_images/templates/bottom2_512_512.png",
                   "left and right": "example_images/templates/left_right_512_512.png",
                   "top and down": "example_images/templates/top_bottom_512_512.png"}


    temp_path = layout_dict[layout_radio]

    pos_imgs = cv2.imread(temp_path)
    pos_imgs = cv2.resize(pos_imgs, (w, h), interpolation=cv2.INTER_NEAREST)

    pos_imgs_url = upload_np_2_oss(pos_imgs, 'pos_imgs.png')

    print(pos_imgs_url)
    
    
    results = call_generation(prompt + texts, pos_imgs_url, lora_path_ratio, w, h, img_count)

    return results



is_t2i = 'true'
block = gr.Blocks(css='css/style.css', theme=gr.themes.Soft()).queue()

#with open('javascript/bboxHint.js', 'r') as file:
#    value = file.read()
#escaped_value = json.dumps(value)

with block:
    block.load(fn=None)
    gr.HTML('<div style="text-align: center; margin: 20px auto;"> \
            <img id="banner" src="https://huggingface.co/spaces/martinxm/MemeMaster/resolve/main/example_images/banner2.png" alt="anytext_meme"> <br>  \
            [<a href="https://arxiv.org/abs/2311.03054" style="color:blue; font-size:18px;">arXiv</a>] \
            [<a href="https://github.com/tyxsspa/AnyText" style="color:blue; font-size:18px;">Code</a>] \
            [<a href="https://modelscope.cn/studios/iic/MemeMaster/summary" style="color:blue; font-size:18px;">ModelScope</a>]\
            [<a href="https://huggingface.co/spaces/modelscope/MemeMaster" style="color:blue; font-size:18px;">HuggingFace</a>]\
            version: 0.1.0 </div>')
    with gr.Row(variant='compact'):
        with gr.Accordion('🕹Instructions', open=True,):
                with gr.Tabs():
                    gr.Markdown('<span style="color:black;font-size:16px">Step1: Choose meme style</span>')
                    gr.Markdown('<span style="color:black;font-size:16px">Step2: Write texts (support 1-2 lines/Chinese and English)</span>')
                    gr.Markdown('<span style="color:black;font-size:16px">Step3: Choose texts layout</span>')
                    gr.Markdown('<span style="color:black;font-size:16px">Step4: Generate meme</span>')             
    
    
    with gr.Tab("🖼Meme Gallery"):
        gallery = gr.Gallery(
            label="Generated images", value = [os.path.join('example_images/gallery/',path) for path in os.listdir('example_images/gallery/')],show_label=False, elem_id="gallery"
                , columns=[4], rows=[1], object_fit="contain", height="auto")     
    
    with gr.Tab("🔧Meme Generation"):
        with gr.Row(variant='compact'):
            with gr.Column() as left_part:
                pass
            with gr.Column():
                result_gallery = gr.Gallery(label='results', show_label=True, preview=True, columns=2, allow_preview=True, height=600)
                # result_info = gr.Markdown('', visible=False)
            with left_part:
                with gr.Accordion('🛠Parameters', open=False):
                    # with gr.Row(variant='compact'):
                    prompt = gr.Textbox(label="Prompt", elem_id='t2i_prompt', value='Cute Chinese Dragon, (Best quality: 1.1), (Realistic: 1.1), (Photography: 1.1), (highly details: 1.1) with the words ', visible=False)
                    with gr.Row(variant='compact'):
                        img_count = gr.Slider(label="image number", minimum=1, maximum=4, value=4, step=1)
                        ddim_steps = gr.Slider(label="steps", minimum=1, maximum=100, value=20, step=1, visible=False)
                    with gr.Row(variant='compact'):
                        image_width = gr.Slider(label="width", minimum=256, maximum=768, value=512, step=64)
                        image_height = gr.Slider(label="height", minimum=256, maximum=768, value=512, step=64)
                        strength = gr.Slider(label="Strength", minimum=0.0, maximum=2.0, value=1.0, step=0.01, visible=False)
                        cfg_scale = gr.Slider(label="CFG scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1, visible=False)
                        seed = gr.Slider(label="seed", minimum=-1, maximum=99999999, step=1, randomize=False, value=-1, visible=False)
                        eta = gr.Number(label="eta (DDIM)", value=0.0, visible=False)
                        show_debug = gr.Checkbox(label='debug', value=False, visible=False)
                    a_prompt = gr.Textbox(label="Added Prompt", value='best quality, extremely detailed,4k, HD, supper legible text,  clear text edges,  clear strokes, neat writing, no watermarks', visible=False)
                    n_prompt = gr.Textbox(label="Negative Prompt", value='low-res, bad anatomy, extra digit, fewer digits, cropped, worst quality, low quality, watermark, unreadable text, messy words, distorted text, disorganized writing, advertising picture', visible=False)
                base_model_path = gr.Textbox(label='Base Model Path', visible=False, value='./')
                lora_path_ratio = gr.Textbox(label='LoRA Path and Ratio', visible=False, value='9 0.5')
                
                with gr.Column():
                    with gr.Row(variant='compact'):
                        style_radio = gr.Radio(["Q版3D小龙(3D dragon)", "红火财神爷(Red wealthgod)", "2D手绘小龙(2D dragon)","黄金财神爷(Gold wealthgod)","喜庆灯笼(Lantern)","可爱贴纸(Sticker)","可爱喷绘(Splash)","卡通门神(Doorgod)","可爱盲盒(Blindbox)","蜡笔女孩(Crayon)","肥龙在天(Fat dragon)"], value='Q版3D小龙(3D dragon)', label="Meme Style", visible=True)
                        example_style1 = gr.Image(value='example_images/example_style/3dcute_dragon.jpg',visible=True, label="Q版3D小龙(3D dragon)", height=320, width=320, elem_id="example_style1")
                        example_style2 = gr.Image(value='example_images/example_style/caishenye1.jpg',visible=False, label="红火财神爷(Red wealthgod)", height=320, width=320, elem_id="example_style2")
                        example_style3 = gr.Image(value='example_images/example_style/2dcartoon_dragon.jpg',visible=False, label="2D手绘小龙(2D dragon)", height=320, width=320, elem_id="example_style3")
                        example_style4 = gr.Image(value='example_images/example_style/caishenye2.jpg',visible=False, label="黄金财神爷(Gold wealthgod)", height=320, width=320, elem_id="example_style4")
                        example_style5 = gr.Image(value='example_images/example_style/cute_lantern_girl.jpg',visible=False, label="喜庆灯笼(Lantern)", height=320, width=320, elem_id="example_style5")
                        example_style6 = gr.Image(value='example_images/example_style/cute_baby_girl.jpg',visible=False, label="可爱贴纸(Sticker)", height=320, width=320, elem_id="example_style6")
                        example_style7 = gr.Image(value='example_images/example_style/cutegirl1.jpg',visible=False, label="可爱喷绘(Splash)", height=320, width=320, elem_id="example_style7")
                        example_style8 = gr.Image(value='example_images/example_style/cartoon_doorgod.jpg',visible=False, label="卡通门神(Doorgod)", height=320, width=320, elem_id="example_style8")
                        example_style9 = gr.Image(value='example_images/example_style/blindbox.jpg',visible=False, label="可爱盲盒(Blindbox)", height=320, width=320, elem_id="example_style9")
                        example_style10 = gr.Image(value='example_images/example_style/crayon.jpg',visible=False, label="蜡笔女孩(Crayon)", height=320, width=320, elem_id="example_style10")
                        example_style11 = gr.Image(value='example_images/example_style/fatdragon.jpg',visible=False, label="肥龙在天(Fat dragon)", height=320, width=320, elem_id="example_style11")
          

                def change_options2(selected_option):
                    if selected_option == 'Q版3D小龙(3D dragon)':
                        _text = 'Cute Chinese Dragon, (Best quality: 1.1), (Realistic: 1.1), (Photography: 1.1), (highly details: 1.1), with a sign written '
                        _lora = '9 0.5'
                    elif selected_option == '红火财神爷(Red wealthgod)':
                        _text = '1boy,chinese clothes,facial hair,solo,full body,red_background,chibi, with a sign written '
                        _lora = '13 1.0'
                    elif selected_option == '2D手绘小龙(2D dragon)':
                        _text = 'Cute Chinese Dragon, no humans,white background,solo,full body,looking at viewer,blush stickers,claws, (Best quality: 1.1), (Realistic: 1.1), (Photography: 1.1), (highly details: 1.1), written, with a sign written '
                        _lora = '9 0.5'
                    elif selected_option == '黄金财神爷(Gold wealthgod)':
                        _text = 'dafengcaishen,solo,blush,simple background,long sleeves,1boy,hat,white background,standing,full body,Smiling eyes,male focus,wide sleeves,chibi,facial hair,chinese clothes,crown,facing viewer,mustache, on the sign reads, with a sign written '
                        _lora = '10 0.8'
                    elif selected_option == '可爱贴纸(Sticker)':
                        _text = 'masterpiece,cartoon, best quality,high quality, sticker, 1girl, chinese new year, happy, smiling, with a sign written '
                        _lora = '1 0.8'
                    elif selected_option == '喜庆灯笼(Lantern)':
                        _text = "children's animation, masterpiece,best quality,chinese new year, cartoon drawing,  1girl, cute face, fish, red Chinese dress, half-body,  lantern, fireworks, cartoon style, 2d-drawing, with a sign written "
                        _lora = '3 0.65'
                    elif selected_option == '可爱喷绘(Splash)':
                        _text = 'masterpiece,best quality,chinese new year,1girl,close-up portrait, solo, joyful festive expression, cute face, looking at viewer, golden chinese dragon,short hair,blue hair,red background,dress,chinese clothes,red dress,china dress,simple background,cnd style,pose, cndstyle, on the wall that reads, with a sign written '
                        _lora = '15 0.8'
                    elif selected_option == '卡通门神(Doorgod)':
                        _text = 'chinese door gods, a painting of god with a sword and armor, full body, 1boy, male focus, solo, drawing, with a sign written '
                        _lora = '2 0.5'
                    elif selected_option == '可爱盲盒(Blindbox)':
                        _text = 'chinese new year, (masterpiece),(best quality),(ultra-detailed), (full body:1.2),1girl,chibi,cute, smile, open mouth,flower, outdoors, beret, jacket, blush, tree, :3, shirt, short hair, cherry blossoms, red headwear, blurry, brown hair, blush stickers, long sleeves, bangs, headphones, black hair, pink flower,(beautiful detailed face), (beautiful detailed eyes), with a sign written '
                        _lora = '6 0.8'
                    elif selected_option == '蜡笔女孩(Crayon)':
                        _text = 'masterpiece,cartoon, best quality,high quality, 1girl, chinese new year, happy, smiling, with a sign written '
                        _lora = '8 0.8'
                    elif selected_option == '肥龙在天(Fat dragon)':
                        _text = 'Cute yellow fat Dragon, flying in the cloud, simple background, 2D cartoon, 2D sticker, with a sign written '
                        _lora = '9 0.5'
                        
                    return [
                            gr.Textbox(value=_text),
                            gr.Textbox(value=_lora),
                            gr.Image(visible=selected_option == 'Q版3D小龙(3D dragon)'),
                            gr.Image(visible=selected_option == '红火财神爷(Red wealthgod)'),
                            gr.Image(visible=selected_option == '2D手绘小龙(2D dragon)'),
                            gr.Image(visible=selected_option == '黄金财神爷(Gold wealthgod)'),
                            gr.Image(visible=selected_option == '喜庆灯笼(Lantern)'),
                            gr.Image(visible=selected_option == '可爱贴纸(Sticker)'),
                            gr.Image(visible=selected_option == '可爱喷绘(Splash)'),
                            gr.Image(visible=selected_option == '卡通门神(Doorgod)'),
                            gr.Image(visible=selected_option == '可爱盲盒(Blindbox)'),
                            gr.Image(visible=selected_option == '蜡笔女孩(Crayon)'),
                            gr.Image(visible=selected_option == '肥龙在天(Fat dragon)'),
                            ]
                style_radio.change(change_options2, style_radio, [prompt, lora_path_ratio, example_style1, example_style2, example_style3,example_style4,example_style5,example_style6,example_style7,example_style8,example_style9, example_style10, example_style11], show_progress=False, queue=False)




                def change_buttons(selected_option):
                    if selected_option == 'Add a line':
                        return [gr.Button(value="Delete a line", scale=0.3, elem_classes='add_row'),
                                gr.Textbox(value='恭喜发财', visible=True),
                                gr.Radio(["top two","down two","left and right","top and down"], value='left right')]
                    elif selected_option == 'Delete a line':
                        return [gr.Button(value="Add a line", scale=0.3, elem_classes='add_row'),
                                gr.Textbox(value="", visible=False),
                                gr.Radio(["top","down","left","right"], value='down')]




                texts = gr.Textbox(label="Input Texts(first line)", value='新春快乐')
                texts2 = gr.Textbox(label="Input Texts(second line)", value='', visible=False)
                layout_radio = gr.Radio(["top","down","left","right"], value='down', label="Text layout", visible=True)
                
                def text_change(in_texts):
                    if len(in_texts)> 10:
                        gr.Warning("The text length is too long, please input a shorter text for better quality.")
                    
                    return


                texts.change(fn=text_change, inputs=texts, outputs=None)
                texts2.change(fn=text_change, inputs=texts2, outputs=None)

                with gr.Row():
                    gr.Markdown("")
                    add_row_bt = gr.Button(value="Add a line", scale=0.3, elem_classes='add_row')
                    gr.Markdown("")
                
                add_row_bt.click(change_buttons, add_row_bt, [add_row_bt,texts2,layout_radio], show_progress=False, queue=False)

                layout_radio.change(fn=None, inputs=layout_radio, outputs=layout_radio, show_progress=False, queue=False)


                def exp_gen_click():
                    return [gr.Slider(value=512), gr.Slider(value=512)]  # all examples are 512x512, refresh draw_img
                with gr.Tab("Examples"):
                            with gr.Row(variant='compact'):
                                exp_gen_en = gr.Examples(
                                    [
                                        ['新年快乐'],
                                        ['Happy New Year'],
                                        ['恭贺新年'],
                                        ['新春快乐'],
                                        ['龙年吉祥'],
                                        ['万事如意'],
                                        ['大吉大利'],
                                        ['恭贺新禧'],
                                        ['初五迎财神'],
                                        ['恭喜发财']                             
                                    ],
                                    [texts],
                                    examples_per_page=15,
                                    label=''
                                )
                                exp_gen_en.dataset.click(exp_gen_click, None, [image_width, image_height])
                

                with gr.Row():
                    sort_radio = gr.Radio(["↕", "↔"], value='↕', label="Sort Position", visible=False)
                    revise_pos = gr.Checkbox(label='Revise Position', value=False, visible=False)

                with gr.Row():
                    gr.Markdown("")
                    run_gen = gr.Button(value="Generate meme!", scale=0.3, elem_classes='run')
                    gr.Markdown("")
                            

  

    ips = [prompt, texts, texts2, layout_radio, sort_radio, revise_pos, base_model_path, lora_path_ratio, show_debug,img_count, ddim_steps, image_width, image_height, strength, cfg_scale, seed, eta, a_prompt, n_prompt]
    run_gen.click(fn=process, inputs=[gr.State('gen')] + ips, outputs=[result_gallery])

block.launch(share=True)