from contextlib import nullcontext
import gradio as gr
import torch
from torch import autocast
from diffusers import StableDiffusionPipeline
import urllib, urllib.request
import os
from xml.etree import ElementTree
import random
import re
from typing import List
pokemon_types = ["Normal",
"Water",
"Fire",
"Ice",
"Psychic",
"Rock",
"Dark",
"Electric",
"Grass",
"Fighting",
"Poison",
"Ground",
"Flying",
"Bug",
"Ghost",
"Dragon",
"Steel",
"Fairy"
]
type_choices=["None", "Random"]
type_choices.extend(pokemon_types)
paper_name = None
device = "cuda" if torch.cuda.is_available() else "cpu"
context = autocast if device == "cuda" else nullcontext
dtype = torch.float16 if device == "cuda" else torch.float32
pipe = StableDiffusionPipeline.from_pretrained("lambdalabs/sd-pokemon-diffusers", torch_dtype=dtype)
pipe = pipe.to(device)
# Sometimes the nsfw checker is confused by the Pokémon images, you can disable
# it at your own risk here
disable_safety = True
if disable_safety:
def null_safety(images, **kwargs):
return images, False
pipe.safety_checker = null_safety
def infer(prompt, n_samples, steps, scale):
with context("cuda"):
images = pipe(n_samples*[prompt], guidance_scale=scale, num_inference_steps=steps).images
return images
def get_paper_name(url: str):
paper_id = os.path.basename(url)
paper_id = paper_id.split(".pdf")[0]
query_url = f"http://export.arxiv.org/api/query?id_list={paper_id}"
hdr = { "Content-Type" : "application/atom+xml" }
req = urllib.request.Request(query_url, headers=hdr)
response = urllib.request.urlopen(req)
tree = ElementTree.fromstring(response.read().decode("utf-8"))
paper_title = tree.find("{http://www.w3.org/2005/Atom}entry").find("{http://www.w3.org/2005/Atom}title").text
paper_title = paper_title.replace("\n", "")
paper_title = re.sub(' +', ' ', paper_title)
return paper_title
block = gr.Blocks()
examples = [
[
"https://arxiv.org/abs/1706.03762",
2,
7.5,
],
[
"https://arxiv.org/abs/1404.5997v2",
2,
7.5,
],
[
"https://arxiv.org/abs/2010.11929",
2,
7.5,
],
[
"https://arxiv.org/abs/1810.04805v2",
2,
7.5,
]
]
with block:
gr.HTML(
"""
<div style="text-align: center; max-width: 650px; margin: 50px auto;">
<div>
<h1 style="font-weight: 900; font-size: 3rem;">
Paper to Pokémon
</h1>
</div>
<p style="margin-bottom: 10px; margin-top: 30px; font-size: 94%">
Generate new Pokémon from an arXiv link. Just paste the link to the overview, the pdf or just give the ID of the paper.
It will create a prompt with the paper title, which you can then modify as you like or submit as it is.
For general better quality increase the step size. (This will also increase the processing time)
</p>
</div>
"""
)
with gr.Group():
with gr.Box():
with gr.Row().style(mobile_collapse=False, equal_height=True):
text = gr.Textbox(
label="Link or ID for paper",
show_label=False,
max_lines=1,
placeholder="Give arXiv link or ID for the paper",
).style(
border=(True, False, True, True),
rounded=(True, False, False, True),
container=False,
)
btn = gr.Button("Generate image").style(
margin=False,
rounded=(False, True, True, False),
)
poke_type = gr.Radio(choices=type_choices, value="None", label="Pokemon Type")
prompt_ideas = gr.CheckboxGroup(choices=["as a bird",
"with four legs",
"with wings",
"as a koala",
"with a beak",
"looking like a llama"],
label="Additional prompt ideas")
prompt_box = gr.Textbox(placeholder="Your prompt appears here", interactive=True, label="Prompt")
gallery = gr.Gallery(
label="Generated images", show_label=False, elem_id="gallery"
).style(grid=[2], height="auto")
with gr.Row(elem_id="advanced-options"):
samples = gr.Slider(label="Images", minimum=1, maximum=4, value=2, step=1)
steps = gr.Slider(label="Steps", minimum=5, maximum=50, value=25, step=5)
scale = gr.Slider(
label="Guidance Scale", minimum=0, maximum=50, value=7.5, step=0.1
)
ex = gr.Examples(examples=examples, fn=infer, inputs=[text, samples, scale], outputs=gallery, cache_examples=False)
ex.dataset.headers = [""]
def resolve_poke_type(pok_type: str):
if pok_type == "None":
return ""
elif pok_type == "Random":
idx = random.randint(0,len(pokemon_types)-1)
return pokemon_types[idx]
else:
return pok_type
def update_prompt_link(new_link: str, pok_type: str, prompt_ideas: List[str]):
global paper_name
paper_name = get_paper_name(new_link)
pok_type = resolve_poke_type(pok_type)
prompt_text = f"{paper_name} as {pok_type} type" if pok_type != "" else f"{paper_name}"
return build_prompt_text(paper_name, pok_type, prompt_ideas)
def update_prompt_type(paper_link: str, pok_type: str, prompt_ideas: List[str]):
global paper_name
if paper_name is None:
paper_name = get_paper_name(paper_link)
pok_type = resolve_poke_type(pok_type)
return build_prompt_text(paper_name, pok_type, prompt_ideas)
def build_prompt_text(paper_name, pok_type, add_ideas):
prompt_text = f"{paper_name} as {pok_type} type" if pok_type != "" else f"{paper_name}"
prompt_text = f"""{prompt_text} {" ".join(add_ideas)}"""
return prompt_text
text.change(update_prompt_link, inputs=[text, poke_type, prompt_ideas], outputs=prompt_box)
text.submit(update_prompt_link, inputs=[text, poke_type, prompt_ideas], outputs=prompt_box)
poke_type.change(update_prompt_type, inputs=[text, poke_type, prompt_ideas], outputs=prompt_box)
prompt_ideas.change(update_prompt_type, inputs=[text, poke_type, prompt_ideas], outputs=prompt_box)
btn.click(infer, inputs=[prompt_box, samples, steps, scale], outputs=gallery)
gr.HTML(
"""
<div class="footer" style="text-align: center; max-width: 650px; margin: 50px auto;">
<p>Inspired by and cloned from the great <a href="https://huggingface.co/spaces/lambdalabs/text-to-pokemon">
Text-to-Pokémon</a> space by Lambda labs</p>
<p> Gradio Demo by johko</p>
</div>
"""
)
block.launch()