Spaces:
Running
on
A10G
Running
on
A10G
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_xl import RegionDiffusionXL | |
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 | |
from share_btn import community_icon_html, loading_icon_html, share_js, css | |
help_text = """ | |
If you are encountering an error or not achieving your desired outcome, here are some potential reasons and recommendations to consider: | |
1. If you format only a portion of a word rather than the complete word, an error may occur. | |
2. If you use font color and get completely corrupted results, you may consider decrease the color weight lambda. | |
3. Consider using a different seed. | |
""" | |
canvas_html = """<iframe id='rich-text-root' style='width:100%' height='360px' src='file=rich-text-to-json-iframe.html' frameborder='0' scrolling='no'></iframe>""" | |
get_js_data = """ | |
async (text_input, negative_prompt, num_segments, segment_threshold, inject_interval, inject_background, seed, color_guidance_weight, rich_text_input, height, width, steps, guidance_weights) => { | |
const richEl = document.getElementById("rich-text-root"); | |
const data = richEl? richEl.contentDocument.body._data : {}; | |
return [text_input, negative_prompt, num_segments, segment_threshold, inject_interval, inject_background, seed, color_guidance_weight, JSON.stringify(data), height, width, steps, guidance_weights]; | |
} | |
""" | |
set_js_data = """ | |
async (text_input) => { | |
const richEl = document.getElementById("rich-text-root"); | |
const data = text_input ? JSON.parse(text_input) : null; | |
if (richEl && data) richEl.contentDocument.body.setQuillContents(data); | |
} | |
""" | |
get_window_url_params = """ | |
async (url_params) => { | |
const params = new URLSearchParams(window.location.search); | |
url_params = Object.fromEntries(params); | |
return [url_params]; | |
} | |
""" | |
def load_url_params(url_params): | |
if 'prompt' in url_params: | |
return gr.update(visible=True), url_params | |
else: | |
return gr.update(visible=False), url_params | |
def main(): | |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
model = RegionDiffusionXL() | |
def generate( | |
text_input: str, | |
negative_text: str, | |
num_segments: int, | |
segment_threshold: float, | |
inject_interval: float, | |
inject_background: float, | |
seed: int, | |
color_guidance_weight: float, | |
rich_text_input: str, | |
height: int, | |
width: int, | |
steps: int, | |
guidance_weight: float, | |
): | |
run_dir = 'results/' | |
os.makedirs(run_dir, exist_ok=True) | |
# Load region diffusion model. | |
height = int(height) if height else 1024 | |
width = int(width) if width else 1024 | |
steps = 41 if not steps else steps | |
guidance_weight = 8.5 if not guidance_weight else guidance_weight | |
text_input = rich_text_input if rich_text_input != '' and rich_text_input != None else text_input | |
print('text_input', text_input, width, height, steps, guidance_weight, num_segments, segment_threshold, inject_interval, inject_background, color_guidance_weight, negative_text) | |
if (text_input == '' or rich_text_input == ''): | |
raise gr.Error("Please enter some text.") | |
# 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, color_guidance_weight=color_guidance_weight) | |
seed_everything(seed) | |
# get token maps from plain text to image generation. | |
begin_time = time.time() | |
if model.selfattn_maps is None and model.crossattn_maps is None: | |
model.remove_tokenmap_hooks() | |
model.register_tokenmap_hooks() | |
else: | |
model.reset_attention_maps() | |
model.remove_tokenmap_hooks() | |
plain_img = model.sample([base_text_prompt], negative_prompt=[negative_text], | |
height=height, width=width, num_inference_steps=steps, | |
guidance_scale=guidance_weight, run_rich_text=False) | |
print('time lapses to get attention maps: %.4f' % | |
(time.time()-begin_time)) | |
seed_everything(seed) | |
color_obj_masks, segments_vis, token_maps = get_token_maps(model.selfattn_maps, model.crossattn_maps, model.n_maps, run_dir, | |
1024//8, 1024//8, color_target_token_ids[:-1], seed, | |
base_tokens, segment_threshold=segment_threshold, num_segments=num_segments, | |
return_vis=True) | |
seed_everything(seed) | |
model.masks, segments_vis, token_maps = get_token_maps(model.selfattn_maps, model.crossattn_maps, model.n_maps, run_dir, | |
1024//8, 1024//8, region_target_token_ids[:-1], seed, | |
base_tokens, segment_threshold=segment_threshold, num_segments=num_segments, | |
return_vis=True) | |
color_obj_atten_all = torch.zeros_like(color_obj_masks[-1]) | |
for obj_mask in color_obj_masks[:-1]: | |
color_obj_atten_all += obj_mask | |
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 | |
text_format_dict['color_obj_atten_all'] = color_obj_atten_all | |
model.remove_tokenmap_hooks() | |
# generate image from rich text | |
begin_time = time.time() | |
seed_everything(seed) | |
rich_img = model.sample(region_text_prompts, negative_prompt=[negative_text], | |
height=height, width=width, num_inference_steps=steps, | |
guidance_scale=guidance_weight, use_guidance=use_grad_guidance, | |
text_format_dict=text_format_dict, inject_selfattn=inject_interval, | |
inject_background=inject_background, run_rich_text=True) | |
print('time lapses to generate image from rich text: %.4f' % | |
(time.time()-begin_time)) | |
return [plain_img.images[0], rich_img.images[0], segments_vis, token_maps] | |
with gr.Blocks(css=css) as demo: | |
url_params = gr.JSON({}, visible=False, label="URL Params") | |
gr.HTML("""<h1 style="font-weight: 900; margin-bottom: 7px;">Expressive Text-to-Image Generation with Rich Text</h1> | |
<p> <a href="https://songweige.github.io/">Songwei Ge</a>, <a href="https://taesung.me/">Taesung Park</a>, <a href="https://www.cs.cmu.edu/~junyanz/">Jun-Yan Zhu</a>, <a href="https://jbhuang0604.github.io/">Jia-Bin Huang</a> <p/> | |
<p> UMD, Adobe, CMU <p/> | |
<p> <a href="https://huggingface.co/spaces/songweig/rich-text-to-image?duplicate=true"><img src="https://bit.ly/3gLdBN6" style="display:inline;"alt="Duplicate Space"></a> | <a href="https://rich-text-to-image.github.io">[Website]</a> | <a href="https://github.com/SongweiGe/rich-text-to-image">[Code]</a> | <a href="https://arxiv.org/abs/2304.06720">[Paper]</a><p/> | |
<p> For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings.""") | |
with gr.Row(): | |
with gr.Column(): | |
rich_text_el = gr.HTML(canvas_html, elem_id="canvas_html") | |
rich_text_input = gr.Textbox(value="", visible=False) | |
text_input = gr.Textbox( | |
label='Rich-text JSON Input', | |
visible=False, | |
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"}]}\'', | |
elem_id="text_input" | |
) | |
negative_prompt = gr.Textbox( | |
label='Negative Prompt', | |
max_lines=1, | |
placeholder='Example: poor quality, blurry, dark, low resolution, low quality, worst quality', | |
elem_id="negative_prompt" | |
) | |
segment_threshold = gr.Slider(label='Token map threshold', | |
info='(See less area in token maps? Decrease this. See too much area? Increase this.)', | |
minimum=0, | |
maximum=1, | |
step=0.01, | |
value=0.25) | |
inject_interval = gr.Slider(label='Detail preservation', | |
info='(To preserve more structure from plain-text generation, increase this. To see more rich-text attributes, decrease this.)', | |
minimum=0, | |
maximum=1, | |
step=0.01, | |
value=0.) | |
inject_background = gr.Slider(label='Unformatted token preservation', | |
info='(To affect less the tokens without any rich-text attributes, increase this.)', | |
minimum=0, | |
maximum=1, | |
step=0.01, | |
value=0.3) | |
color_guidance_weight = gr.Slider(label='Color weight', | |
info='(To obtain more precise color, increase this, while too large value may cause artifacts.)', | |
minimum=0, | |
maximum=2, | |
step=0.1, | |
value=0.5) | |
num_segments = gr.Slider(label='Number of segments', | |
minimum=2, | |
maximum=20, | |
step=1, | |
value=9) | |
seed = gr.Slider(label='Seed', | |
minimum=0, | |
maximum=100000, | |
step=1, | |
value=6, | |
elem_id="seed" | |
) | |
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=[1024], | |
value=1024, | |
label='Width', | |
visible=True) | |
height = gr.Dropdown(choices=[1024], | |
value=1024, | |
label='height', | |
visible=True) | |
with gr.Row(): | |
with gr.Column(scale=1, min_width=100): | |
generate_button = gr.Button("Generate") | |
load_params_button = gr.Button( | |
"Load from URL Params", visible=True) | |
with gr.Column(): | |
richtext_result = gr.Image( | |
label='Rich-text', elem_id="rich-text-image") | |
richtext_result.style(height=784) | |
with gr.Row(): | |
plaintext_result = gr.Image( | |
label='Plain-text', elem_id="plain-text-image") | |
segments = gr.Image(label='Segmentation') | |
with gr.Row(): | |
token_map = gr.Image(label='Token Maps') | |
with gr.Row(visible=False) as share_row: | |
with gr.Group(elem_id="share-btn-container"): | |
community_icon = gr.HTML(community_icon_html) | |
loading_icon = gr.HTML(loading_icon_html) | |
share_button = gr.Button( | |
"Share to community", elem_id="share-btn") | |
share_button.click(None, [], [], _js=share_js) | |
with gr.Row(): | |
gr.Markdown(help_text) | |
# with gr.Row(): | |
# footnote_examples = [ | |
# [ | |
# '{"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."}]}', | |
# '', | |
# 9, | |
# 0.3, | |
# 0.3, | |
# 0.5, | |
# 3, | |
# 0, | |
# None, | |
# ], | |
# [ | |
# '{"ops":[{"insert":"A cozy "},{"attributes":{"link":"A charming wooden cabin with Christmas decoration, warm light coming out from the windows."},"insert":"cabin"},{"insert":" nestled in a "},{"attributes":{"link":"Towering evergreen trees covered in a thick layer of pristine snow."},"insert":"snowy forest"},{"insert":", and a "},{"attributes":{"link":"A cute snowman wearing a carrot nose, coal eyes, and a colorful scarf, welcoming visitors with a cheerful vibe."},"insert":"snowman"},{"insert":" stands in the yard."}]}', | |
# '', | |
# 12, | |
# 0.4, | |
# 0.3, | |
# 0.5, | |
# 3, | |
# 0, | |
# None, | |
# ], | |
# [ | |
# '{"ops":[{"insert":"A "},{"attributes":{"link":"Happy Kung fu panda art, elder, asian art, volumetric lighting, dramatic scene, ultra detailed, realism, chinese"},"insert":"panda"},{"insert":" standing on a cliff by a waterfall, wildlife photography, photograph, high quality, wildlife, f 1.8, soft focus, 8k, national geographic, award - winning photograph by nick nichols"}]}', | |
# '', | |
# 5, | |
# 0.3, | |
# 0, | |
# 0.1, | |
# 4, | |
# 0, | |
# None, | |
# ], | |
# ] | |
# gr.Examples(examples=footnote_examples, | |
# label='Footnote examples', | |
# inputs=[ | |
# text_input, | |
# negative_prompt, | |
# num_segments, | |
# segment_threshold, | |
# inject_interval, | |
# inject_background, | |
# seed, | |
# color_guidance_weight, | |
# rich_text_input, | |
# ], | |
# outputs=[ | |
# plaintext_result, | |
# richtext_result, | |
# segments, | |
# token_map, | |
# ], | |
# fn=generate, | |
# cache_examples=True, | |
# examples_per_page=20) | |
with gr.Row(): | |
color_examples = [ | |
# [ | |
# '{"ops":[{"insert":"a beautifule girl with big eye, skin, and long "},{"attributes":{"color":"#04a704"},"insert":"hair"},{"insert":", t-shirt, bursting with vivid color, intricate, elegant, highly detailed, photorealistic, digital painting, artstation, illustration, concept art."}]}', | |
# 'lowres, had anatomy, bad hands, cropped, worst quality', | |
# 11, | |
# 0.5, | |
# 0.3, | |
# 0.3, | |
# 6, | |
# 0.5, | |
# None, | |
# ], | |
# [ | |
# '{"ops":[{"insert":"a beautifule girl with big eye, skin, and long "},{"attributes":{"color":"#ff5df1"},"insert":"hair"},{"insert":", t-shirt, bursting with vivid color, intricate, elegant, highly detailed, photorealistic, digital painting, artstation, illustration, concept art."}]}', | |
# 'lowres, had anatomy, bad hands, cropped, worst quality', | |
# 11, | |
# 0.5, | |
# 0.3, | |
# 0.3, | |
# 6, | |
# 0.5, | |
# None, | |
# ], | |
[ | |
'{"ops":[{"insert":"a beautifule girl with big eye, skin, and long "},{"attributes":{"color":"#999999"},"insert":"hair"},{"insert":", t-shirt, bursting with vivid color, intricate, elegant, highly detailed, photorealistic, digital painting, artstation, illustration, concept art."}]}', | |
'lowres, had anatomy, bad hands, cropped, worst quality', | |
11, | |
0.5, | |
0.3, | |
0.3, | |
6, | |
0.5, | |
None, | |
], | |
[ | |
'{"ops":[{"insert":"a Gothic "},{"attributes":{"color":"#FD6C9E"},"insert":"church"},{"insert":" in a the sunset with a beautiful landscape in the background."}]}', | |
'', | |
10, | |
0.5, | |
0.5, | |
0.3, | |
7, | |
0.5, | |
None, | |
], | |
] | |
gr.Examples(examples=color_examples, | |
label='Font color examples', | |
inputs=[ | |
text_input, | |
negative_prompt, | |
num_segments, | |
segment_threshold, | |
inject_interval, | |
inject_background, | |
seed, | |
color_guidance_weight, | |
rich_text_input, | |
], | |
outputs=[ | |
plaintext_result, | |
richtext_result, | |
segments, | |
token_map, | |
], | |
fn=generate, | |
cache_examples=True, | |
examples_per_page=20) | |
# with gr.Row(): | |
# style_examples = [ | |
# [ | |
# '{"ops":[{"insert":"a beautiful"},{"attributes":{"font":"mirza"},"insert":" garden"},{"insert":" with a "},{"attributes":{"font":"roboto"},"insert":"snow mountain"},{"insert":" in the background"}]}', | |
# '', | |
# 10, | |
# 0.6, | |
# 0, | |
# 0.4, | |
# 5, | |
# 0, | |
# None, | |
# ], | |
# [ | |
# '{"ops":[{"insert":"a night"},{"attributes":{"font":"slabo"},"insert":" sky"},{"insert":" filled with stars above a turbulent"},{"attributes":{"font":"roboto"},"insert":" sea"},{"insert":" with giant waves"}]}', | |
# '', | |
# 2, | |
# 0.6, | |
# 0, | |
# 0, | |
# 6, | |
# 0.5, | |
# None, | |
# ], | |
# ] | |
# gr.Examples(examples=style_examples, | |
# label='Font style examples', | |
# inputs=[ | |
# text_input, | |
# negative_prompt, | |
# num_segments, | |
# segment_threshold, | |
# inject_interval, | |
# inject_background, | |
# seed, | |
# color_guidance_weight, | |
# rich_text_input, | |
# ], | |
# outputs=[ | |
# plaintext_result, | |
# richtext_result, | |
# segments, | |
# token_map, | |
# ], | |
# fn=generate, | |
# cache_examples=True, | |
# examples_per_page=20) | |
# with gr.Row(): | |
# size_examples = [ | |
# [ | |
# '{"ops": [{"insert": "A pizza with "}, {"attributes": {"size": "60px"}, "insert": "pineapple"}, {"insert": " pepperoni, and mushroom on the top"}]}', | |
# '', | |
# 5, | |
# 0.3, | |
# 0, | |
# 0, | |
# 3, | |
# 1, | |
# None, | |
# ], | |
# [ | |
# '{"ops": [{"insert": "A pizza with pineapple, "}, {"attributes": {"size": "60px"}, "insert": "pepperoni"}, {"insert": ", and mushroom on the top"}]}', | |
# '', | |
# 5, | |
# 0.3, | |
# 0, | |
# 0, | |
# 3, | |
# 1, | |
# None, | |
# ], | |
# [ | |
# '{"ops": [{"insert": "A pizza with pineapple, pepperoni, and "}, {"attributes": {"size": "60px"}, "insert": "mushroom"}, {"insert": " on the top"}]}', | |
# '', | |
# 5, | |
# 0.3, | |
# 0, | |
# 0, | |
# 3, | |
# 1, | |
# None, | |
# ], | |
# ] | |
# gr.Examples(examples=size_examples, | |
# label='Font size examples', | |
# inputs=[ | |
# text_input, | |
# negative_prompt, | |
# num_segments, | |
# segment_threshold, | |
# inject_interval, | |
# inject_background, | |
# seed, | |
# color_guidance_weight, | |
# rich_text_input, | |
# ], | |
# outputs=[ | |
# plaintext_result, | |
# richtext_result, | |
# segments, | |
# token_map, | |
# ], | |
# fn=generate, | |
# cache_examples=True, | |
# examples_per_page=20) | |
generate_button.click(fn=lambda: gr.update(visible=False), inputs=None, outputs=share_row, queue=False).then( | |
fn=generate, | |
inputs=[ | |
text_input, | |
negative_prompt, | |
num_segments, | |
segment_threshold, | |
inject_interval, | |
inject_background, | |
seed, | |
color_guidance_weight, | |
rich_text_input, | |
height, | |
width, | |
steps, | |
guidance_weight, | |
], | |
outputs=[plaintext_result, richtext_result, segments, token_map], | |
_js=get_js_data | |
).then( | |
fn=lambda: gr.update(visible=True), inputs=None, outputs=share_row, queue=False) | |
text_input.change( | |
fn=None, inputs=[text_input], outputs=None, _js=set_js_data, queue=False) | |
# load url param prompt to textinput | |
load_params_button.click(fn=lambda x: x['prompt'], inputs=[ | |
url_params], outputs=[text_input], queue=False) | |
demo.load( | |
fn=load_url_params, | |
inputs=[url_params], | |
outputs=[load_params_button, url_params], | |
_js=get_window_url_params | |
) | |
demo.queue(concurrency_count=1) | |
demo.launch(share=False) | |
if __name__ == "__main__": | |
main() | |