Spaces:
Runtime error
Runtime error
import gradio as gr | |
import os | |
import argparse | |
from easydict import EasyDict as edict | |
import yaml | |
import os.path as osp | |
import random | |
import numpy.random as npr | |
import sys | |
import imageio | |
import numpy as np | |
sys.path.append('./code') | |
sys.path.append('/home/user/app/code') | |
# set up diffvg | |
# os.system('git clone https://github.com/BachiLi/diffvg.git') | |
os.system('git submodule update --init') | |
os.chdir('diffvg') | |
os.system('git submodule update --init --recursive') | |
os.system('python setup.py install --user') | |
sys.path.append("/home/user/.local/lib/python3.10/site-packages/diffvg-0.0.1-py3.8-linux-x86_64.egg") | |
os.chdir('/home/user/app') | |
import torch | |
from diffusers import StableDiffusionPipeline | |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
model = None | |
model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to(device) | |
from typing import Mapping | |
from tqdm import tqdm | |
import torch | |
from torch.optim.lr_scheduler import LambdaLR | |
import pydiffvg | |
import save_svg | |
from losses import SDSLoss, ToneLoss, ConformalLoss | |
from utils import ( | |
edict_2_dict, | |
update, | |
check_and_create_dir, | |
get_data_augs, | |
save_image, | |
preprocess, | |
learning_rate_decay, | |
combine_word) | |
import warnings | |
TITLE="""<h1 style="font-size: 42px;" align="center">Font-To-Sketch: Morphing Any Font to a Visual Representation</h1>""" | |
DESCRIPTION="""This demo builds on the [Word-As-Image for Semantic Typography](https://wordasimage.github.io/Word-As-Image-Page/) work to support **any** font and morphing whole words and phrases to a visual representation of a given semantic concept. This project started as part of an ongoing effort with the [ARBML](https://arbml.github.io/website/) community to build open-source Arabic tools using machine learning.""" | |
DESCRIPTION+="""The demo currently supports the following scripts: **Arabic**, **Simplified Chinese**, **Cyrillic**, **Greek**, **Latin**, **Tamil**. Therefore you can write the text in any language using those scripts. To add support for more fonts please check the [GitHub ReadMe](https://raw.githubusercontent.com/BKHMSI/Font-To-Sketch).""" | |
# DESCRIPTION += '\n<p>This demo is licensed under a <a rel="license" href="http://creativecommons.org/licenses/by-sa/4.0/"> Creative Commons Attribution-ShareAlike 4.0 International License</a>.</p>' | |
DESCRIPTION += '\n<p>Note: it takes about 5 minutes for 250 iterations to generate the final GIF. For faster inference without waiting in queue, you can <a href="https://colab.research.google.com/drive/1wobOAsnLpkIzaRxG5yac8NcV7iCrlycP"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a></p>' | |
if (SPACE_ID := os.getenv('SPACE_ID')) is not None: | |
DESCRIPTION = DESCRIPTION.replace("</p>", " ") | |
DESCRIPTION += f'or <a href="https://huggingface.co/spaces/{SPACE_ID}?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate the Space"/></a> and upgrade to GPU in settings.</p>' | |
else: | |
DESCRIPTION = DESCRIPTION.replace("either", "") | |
DESCRIPTION += "<img src='https://raw.githubusercontent.com/BKHMSI/Font-To-Sketch/main/images/animals_7.gif' alt='Example of Outputs'/>" | |
warnings.filterwarnings("ignore") | |
pydiffvg.set_print_timing(False) | |
gamma = 1.0 | |
def set_config(semantic_concept, word, script, prompt_suffix, font_name, num_steps, seed, is_seed_rand, dist_loss_weight, pixel_dist_kernel_blur, pixel_dist_sigma, angeles_w): | |
cfg_d = edict() | |
cfg_d.config = "code/config/base.yaml" | |
cfg_d.experiment = "default" | |
with open(cfg_d.config, 'r') as f: | |
cfg_full = yaml.load(f, Loader=yaml.FullLoader) | |
cfg_key = cfg_d.experiment | |
cfgs = [cfg_d] | |
while cfg_key: | |
cfgs.append(cfg_full[cfg_key]) | |
cfg_key = cfgs[-1].get('parent_config', 'baseline') | |
cfg = edict() | |
for options in reversed(cfgs): | |
update(cfg, options) | |
del cfgs | |
cfg.semantic_concept = semantic_concept | |
cfg.prompt_suffix = prompt_suffix | |
cfg.word = word | |
cfg.optimized_letter = word | |
cfg.script = script.lower() | |
script_path = f"code/data/fonts/{cfg.script}" | |
if cfg.font == "none": | |
cfg.font = [x for x in os.listdir(script_path) if "ttf" in x][0][:-4] | |
if is_seed_rand == "Random Seed": | |
cfg.seed = np.random.randint(10000) | |
else: | |
cfg.seed = int(seed) | |
cfg.num_iter = num_steps | |
cfg.batch_size = 1 | |
cfg.loss.tone.dist_loss_weight = int(dist_loss_weight) | |
cfg.loss.tone.pixel_dist_kernel_blur = int(pixel_dist_kernel_blur) | |
cfg.loss.tone.pixel_dist_sigma = int(pixel_dist_sigma) | |
cfg.loss.conformal.angeles_w = angeles_w | |
cfg.caption = f"a {cfg.semantic_concept}. {cfg.prompt_suffix}" | |
cfg.log_dir = f"{cfg.script}" | |
if cfg.optimized_letter in cfg.word: | |
cfg.optimized_letter = cfg.optimized_letter | |
else: | |
raise gr.Error(f'letter should be in word') | |
# if ' ' in cfg.word: | |
# cfg.optimized_letter = cfg.optimized_letter.replace(' ', '_') | |
cfg.letter = f"{cfg.font}_{cfg.optimized_letter}_scaled" | |
cfg.target = f"code/data/init/{cfg.letter}" | |
if ' ' in cfg.target: | |
cfg.target = cfg.target.replace(' ', '_') | |
# set experiment dir | |
signature = f"{cfg.word}_{cfg.semantic_concept}_{cfg.seed}" | |
cfg.experiment_dir = osp.join(cfg.log_dir, cfg.font, signature) | |
configfile = osp.join(cfg.experiment_dir, 'config.yaml') | |
# create experiment dir and save config | |
check_and_create_dir(configfile) | |
with open(osp.join(configfile), 'w') as f: | |
yaml.dump(edict_2_dict(cfg), f) | |
if cfg.seed is not None: | |
random.seed(cfg.seed) | |
npr.seed(cfg.seed) | |
torch.manual_seed(cfg.seed) | |
torch.backends.cudnn.benchmark = False | |
else: | |
assert False | |
return cfg | |
def init_shapes(svg_path, trainable: Mapping[str, bool]): | |
svg = f'{svg_path}.svg' | |
canvas_width, canvas_height, shapes_init, shape_groups_init = pydiffvg.svg_to_scene(svg) | |
parameters = edict() | |
# path points | |
if trainable.point: | |
parameters.point = [] | |
for path in shapes_init: | |
path.points.requires_grad = True | |
parameters.point.append(path.points) | |
return shapes_init, shape_groups_init, parameters | |
def run_main_ex(word, semantic_concept, script, num_steps, seed): | |
prompt_suffix = "minimal flat 2d vector. lineal color. trending on artstation" | |
is_seed_rand = "Use Set Value" | |
return list(next(run_main_app(semantic_concept, word, script, prompt_suffix, font_name, num_steps, seed, is_seed_rand, 100, 201, 30, 0.5, 0))) | |
def run_main_app(semantic_concept, word, script, prompt_suffix, font_name, num_steps, seed, is_seed_rand, dist_loss_weight, pixel_dist_kernel_blur, pixel_dist_sigma, angeles_w, example=0): | |
cfg = set_config(semantic_concept, word, script, prompt_suffix, font_name, num_steps, seed, is_seed_rand, dist_loss_weight, pixel_dist_kernel_blur, pixel_dist_sigma, angeles_w) | |
pydiffvg.set_use_gpu(torch.cuda.is_available()) | |
print("preprocessing") | |
preprocess(cfg.font, cfg.word, cfg.optimized_letter, cfg.level_of_cc) | |
filename_init = os.path.join("code/data/init/", f"{cfg.font}_{cfg.word}_scaled.svg").replace(" ", "_") | |
if not example: | |
yield gr.update(value=filename_init,visible=True),gr.update(visible=True, label='Initializing'),gr.update(visible=False),gr.update(value=cfg.caption,visible=True),gr.update(value=cfg.seed,visible=True) | |
sds_loss = SDSLoss(cfg, device, model) | |
h, w = cfg.render_size, cfg.render_size | |
data_augs = get_data_augs(cfg.cut_size) | |
render = pydiffvg.RenderFunction.apply | |
# initialize shape | |
print('initializing shape') | |
shapes, shape_groups, parameters = init_shapes(svg_path=cfg.target, trainable=cfg.trainable) | |
scene_args = pydiffvg.RenderFunction.serialize_scene(w, h, shapes, shape_groups) | |
img_init = render(w, h, 2, 2, 0, None, *scene_args) | |
img_init = img_init[:, :, 3:4] * img_init[:, :, :3] + \ | |
torch.ones(img_init.shape[0], img_init.shape[1], 3, device=device) * (1 - img_init[:, :, 3:4]) | |
img_init = img_init[:, :, :3] | |
tone_loss = ToneLoss(cfg) | |
tone_loss.set_image_init(img_init) | |
num_iter = cfg.num_iter | |
pg = [{'params': parameters["point"], 'lr': cfg.lr_base["point"]}] | |
optim = torch.optim.Adam(pg, betas=(0.9, 0.9), eps=1e-6) | |
conformal_loss = ConformalLoss(parameters, device, cfg.optimized_letter, shape_groups) | |
lr_lambda = lambda step: learning_rate_decay(step, cfg.lr.lr_init, cfg.lr.lr_final, num_iter, | |
lr_delay_steps=cfg.lr.lr_delay_steps, | |
lr_delay_mult=cfg.lr.lr_delay_mult) / cfg.lr.lr_init | |
scheduler = LambdaLR(optim, lr_lambda=lr_lambda, last_epoch=-1) # lr.base * lrlambda_f | |
print("start training") | |
# training loop | |
t_range = tqdm(range(num_iter)) | |
gif_frames = [] | |
skip = 5 | |
for step in t_range: | |
optim.zero_grad() | |
# render image | |
scene_args = pydiffvg.RenderFunction.serialize_scene(w, h, shapes, shape_groups) | |
img = render(w, h, 2, 2, step, None, *scene_args) | |
# compose image with white background | |
img = img[:, :, 3:4] * img[:, :, :3] + torch.ones(img.shape[0], img.shape[1], 3, device=device) * ( | |
1 - img[:, :, 3:4]) | |
img = img[:, :, :3] | |
if step % skip == 0: | |
gif_frames += [img.detach().cpu().numpy()*255] | |
filename = os.path.join(cfg.experiment_dir, "video-svg", f"iter{step:04d}.svg") | |
check_and_create_dir(filename) | |
save_svg.save_svg(filename, w, h, shapes, shape_groups) | |
if not example: | |
yield gr.update(visible=True),gr.update(value=filename, label=f'iters: {step} / {num_iter}', visible=True),gr.update(visible=False),gr.update(value=cfg.caption,visible=True),gr.update(value=cfg.seed,visible=True) | |
x = img.unsqueeze(0).permute(0, 3, 1, 2) # HWC -> NCHW | |
x = x.repeat(cfg.batch_size, 1, 1, 1) | |
x_aug = data_augs.forward(x) | |
# compute diffusion loss per pixel | |
loss = sds_loss(x_aug) | |
tone_loss_res = tone_loss(x, step) | |
loss = loss + tone_loss_res | |
loss_angles = conformal_loss() | |
loss_angles = cfg.loss.conformal.angeles_w * loss_angles | |
loss = loss + loss_angles | |
loss.backward() | |
optim.step() | |
scheduler.step() | |
filename = os.path.join( | |
cfg.experiment_dir, "output-svg", "output.svg") | |
check_and_create_dir(filename) | |
save_svg.save_svg( | |
filename, w, h, shapes, shape_groups) | |
# combine_word(cfg.word, cfg.optimized_letter, cfg.font, cfg.experiment_dir, device) | |
filename = os.path.join(cfg.experiment_dir, "final.gif") | |
# writer = imageio.get_writer(filename, fps=20) | |
# for frame in gif_frames: writer.append_data(frame) | |
# writer.close() | |
imageio.mimsave(filename, np.array(gif_frames).astype(np.uint8)) | |
yield gr.update(value=filename_init,visible=True),gr.update(visible=False),gr.update(value=filename,visible=True),gr.update(value=cfg.caption,visible=True),gr.update(value=cfg.seed,visible=True) | |
with gr.Blocks() as demo: | |
gr.HTML(TITLE) | |
gr.Markdown(DESCRIPTION) | |
with gr.Row(): | |
with gr.Column(): | |
word = gr.Text( | |
label='Text', | |
max_lines=1, | |
placeholder= | |
'Enter text. For example: قطة|猫|γάτα|кошка|பூனை|Cat' | |
) | |
semantic_concept = gr.Text( | |
label='Concept', | |
max_lines=1, | |
placeholder= | |
'Enter a semantic concept that you want your text to morph into (in English). For example: cat' | |
) | |
script = gr.Dropdown( | |
["Arabic", "Simplified Chinese", "Cyrillic", "Greek", "Latin", "Tamil"], | |
value="Arabic", | |
label="Font Script" | |
) | |
prompt_suffix = gr.Text( | |
label='Prompt Suffix', | |
max_lines=1, | |
value="minimal flat 2d vector. lineal color. trending on artstation" | |
) | |
with gr.Row(): | |
with gr.Accordion("Advanced Parameters", open=False, visible=True): | |
with gr.Row(): | |
is_seed_rand = gr.Radio(["Random Seed", "Use Set Value"], label="Use Random Seed", value="Random Seed") | |
seed = gr.Number( | |
label='Seed (Set Value)', | |
value=42 | |
) | |
angeles_w = gr.Number( | |
label='ACAP Deformation Loss Weight', | |
value=0.5 | |
) | |
dist_loss_weight = gr.Number( | |
label='Tone Loss: dist_loss_weight', | |
value=100 | |
) | |
pixel_dist_kernel_blur = gr.Number( | |
label='Tone Loss: pixel_dist_kernel_blur', | |
value=201 | |
) | |
pixel_dist_sigma = gr.Number( | |
label='Tone Loss: pixel_dist_sigma', | |
value=30 | |
) | |
num_steps = gr.Slider(label='Optimization Iterations', | |
minimum=0, | |
maximum=500, | |
step=10, | |
value=250) | |
font_name = gr.Text(value=None,visible=False,label="Font Name") | |
def on_select(evt: gr.SelectData): | |
return evt.value | |
font_name.value = "ArefRuqaa" | |
run = gr.Button('Generate') | |
with gr.Column(): | |
with gr.Row(): | |
prompt = gr.Text( | |
label='Prompt', | |
visible=False, | |
max_lines=1, | |
interactive=False, | |
) | |
seed_value = gr.Text( | |
label='Seed Used', | |
visible=False, | |
max_lines=1, | |
interactive=False, | |
) | |
result0 = gr.Image(type="filepath", label="Initial Word").style(height=250) | |
result1 = gr.Image(type="filepath", label="Optimization Process").style(height=300) | |
result2 = gr.Image(type="filepath", label="Final GIF",visible=False).style(height=300) | |
with gr.Row(): | |
# examples | |
examples = [ | |
["قطة", "Cat", "Arabic", 250, 42], | |
["猫", "Cat", "Simplified Chinese", 250, 42], | |
["γάτα", "Cat", "Greek", 250, 42], | |
["кошка", "Cat", "Cyrillic", 250, 42], | |
["பூனை", "Cat", "Tamil", 250, 42], | |
] | |
demo.queue(max_size=10, concurrency_count=2) | |
gr.Examples(examples=examples, | |
inputs=[ | |
word, | |
semantic_concept, | |
script, | |
num_steps, | |
seed | |
], | |
outputs=[ | |
result0, | |
result1, | |
result2 | |
], | |
fn=run_main_ex, | |
cache_examples=False) | |
# inputs | |
inputs = [ | |
semantic_concept, | |
word, | |
script, | |
prompt_suffix, | |
font_name, | |
num_steps, | |
seed, | |
is_seed_rand, | |
dist_loss_weight, | |
pixel_dist_kernel_blur, | |
pixel_dist_sigma, | |
angeles_w | |
] | |
outputs = [ | |
result0, | |
result1, | |
result2, | |
prompt, | |
seed_value | |
] | |
run.click(fn=run_main_app, inputs=inputs, outputs=outputs, queue=True) | |
demo.launch(share=False) |