app.py

import math
import random
import os
import json
import time
import argparse
import torch
import numpy as np
from torchvision import transforms

from models.region_diffusion import RegionDiffusion
from utils.attention_utils import get_token_maps
from utils.richtext_utils import seed_everything, parse_json, get_region_diffusion_input,\
    get_attention_control_input, get_gradient_guidance_input


import gradio as gr
from PIL import Image, ImageOps


help_text = """
Instructions placeholder.
"""


def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = RegionDiffusion(device)

    def generate(
        text_input: str,
        negative_text: str,
        height: int,
        width: int,
        seed: int,
        steps: int,
        guidance_weight: float,
    ):
        run_dir = 'results/'
        # Load region diffusion model.
        steps = 41 if not steps else steps
        guidance_weight = 8.5 if not guidance_weight else guidance_weight

        # parse json to span attributes
        base_text_prompt, style_text_prompts, footnote_text_prompts, footnote_target_tokens,\
            color_text_prompts, color_names, color_rgbs, size_text_prompts_and_sizes, use_grad_guidance = parse_json(
                json.loads(text_input))

        # create control input for region diffusion
        region_text_prompts, region_target_token_ids, base_tokens = get_region_diffusion_input(
            model, base_text_prompt, style_text_prompts, footnote_text_prompts,
            footnote_target_tokens, color_text_prompts, color_names)

        # create control input for cross attention
        text_format_dict = get_attention_control_input(
            model, base_tokens, size_text_prompts_and_sizes)

        # create control input for region guidance
        text_format_dict, color_target_token_ids = get_gradient_guidance_input(
            model, base_tokens, color_text_prompts, color_rgbs, text_format_dict)

        seed_everything(seed)

        # get token maps from plain text to image generation.
        begin_time = time.time()
        if model.attention_maps is None:
            model.register_evaluation_hooks()
        else:
            model.reset_attention_maps()
        plain_img = model.produce_attn_maps([base_text_prompt], [negative_text],
                                            height=height, width=width, num_inference_steps=steps,
                                            guidance_scale=guidance_weight)
        print('time lapses to get attention maps: %.4f' % (time.time()-begin_time))
        color_obj_masks, _ = get_token_maps(
            model.attention_maps, run_dir, width//8, height//8, color_target_token_ids, seed)
        model.masks, token_maps = get_token_maps(
            model.attention_maps, run_dir, width//8, height//8, region_target_token_ids, seed, base_tokens)
        color_obj_masks = [transforms.functional.resize(color_obj_mask, (height, width),
                                                        interpolation=transforms.InterpolationMode.BICUBIC,
                                                        antialias=True)
                        for color_obj_mask in color_obj_masks]
        text_format_dict['color_obj_atten'] = color_obj_masks
        model.remove_evaluation_hooks()

        # generate image from rich text
        begin_time = time.time()
        seed_everything(seed)
        rich_img = model.prompt_to_img(region_text_prompts, [negative_text],
                                    height=height, width=width, num_inference_steps=steps,
                                    guidance_scale=guidance_weight, use_grad_guidance=use_grad_guidance,
                                    text_format_dict=text_format_dict)
        print('time lapses to generate image from rich text: %.4f' %
            (time.time()-begin_time))
        cat_img = np.concatenate([plain_img[0], rich_img[0]], 1)
        return [cat_img, token_maps]

    with gr.Blocks() as demo:
        gr.HTML("""<h1 style="font-weight: 900; margin-bottom: 7px;">Expressive Text-to-Image Generation with Rich Text</h1>
                   <p> Visit our <a href="https://rich-text-to-image.github.io/rich-text-to-json.html">rich-text-to-json interface</a> to generate rich-text JSON input.<p/>""")
        with gr.Row():
            with gr.Column():
                text_input = gr.Textbox(
                    label='Rich-text JSON Input',
                    max_lines=1,
                    placeholder='Example: \'{"ops":[{"insert":"a Gothic "},{"attributes":{"color":"#b26b00"},"insert":"church"},{"insert":" in a the sunset with a beautiful landscape in the background.\n"}]}\'')
                negative_prompt = gr.Textbox(
                    label='Negative Prompt',
                    max_lines=1,
                    placeholder='')
                seed = gr.Slider(label='Seed',
                                 minimum=0,
                                 maximum=100000,
                                 step=1,
                                 value=6)
                with gr.Accordion('Other Parameters', open=False):
                    steps = gr.Slider(label='Number of Steps',
                                          minimum=0,
                                          maximum=500,
                                          step=1,
                                          value=41)
                    guidance_weight = gr.Slider(label='CFG weight',
                                               minimum=0,
                                               maximum=50,
                                               step=0.1,
                                               value=8.5)
                    width = gr.Dropdown(choices=[512, 768, 896],
                                    value=512,
                                    label='Width',
                                    visible=True)
                    height = gr.Dropdown(choices=[512, 768, 896],
                                    value=512,
                                    label='height',
                                    visible=True)
                    
                with gr.Row():
                    with gr.Column(scale=1, min_width=100):
                        generate_button = gr.Button("Generate")

            with gr.Column():
                result = gr.Image(label='Result')
                token_map = gr.Image(label='TokenMap')

        with gr.Row():
            examples = [
                [
                    '{"ops":[{"insert":"a Gothic "},{"attributes":{"color":"#b26b00"},"insert":"church"},{"insert":" in a the sunset with a beautiful landscape in the background.\n"}]}',
                    '',
                    512,
                    512,
                    6,
                ],
                [
                    '{"ops": [{"insert": "A pizza with "}, {"attributes": {"size": "50px"}, "insert": "pineapples"}, {"insert": ", pepperonis, and mushrooms on the top, 4k, photorealistic\n"}]}',
                    'blurry, art, painting, rendering, drawing, sketch, ugly, duplicate, morbid, mutilated, mutated, deformed, disfigured low quality, worst quality',
                    768,
                    896,
                    6,
                ],
                [
                    '{"ops":[{"insert":"a "},{"attributes":{"font":"mirza"},"insert":"beautiful garden"},{"insert":" with a "},{"attributes":{"font":"roboto"},"insert":"snow mountain in the background"},{"insert":"\n"}]}',
                    '',
                    512,
                    512,
                    3,
                ],
                [
                    '{"ops":[{"insert":"A close-up 4k dslr photo of a "},{"attributes":{"link":"A cat wearing sunglasses and a bandana around its neck."},"insert":"cat"},{"insert":" riding a scooter. Palm trees in the background.\n"}]}',
                    '',
                    512,
                    512,
                    6,
                ],
            ]
            gr.Examples(examples=examples,
                        inputs=[
                            text_input,
                            negative_prompt,
                            height,
                            width,
                            seed,
                        ],
                        outputs=[
                            result,
                            token_map,
                        ],
                        fn=generate,
                        # cache_examples=True,
                        examples_per_page=20)

        generate_button.click(
            fn=generate,
            inputs=[
                text_input,
                negative_prompt,
                height,
                width,
                seed,
                steps,
                guidance_weight,
            ],
            outputs=[result, token_map],
        )

    demo.queue(concurrency_count=1)
    demo.launch(share=False)


if __name__ == "__main__":
    main()