updated version

app.py

@@ -6,7 +6,34 @@ from diffusers import StableDiffusionPipeline
 import urllib, urllib.request
 import os
 from xml.etree import ElementTree
-
+import random
+import re
+
+
+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
@@ -26,11 +53,9 @@ if disable_safety:
   pipe.safety_checker = null_safety
 
 
-def infer(prompt, n_samples, steps, scale):
-    paper_title = get_paper_name(prompt)
-    
+def infer(prompt, n_samples, steps, scale):    
     with context("cuda"):
-        images = pipe(n_samples*[paper_title], guidance_scale=scale, num_inference_steps=steps).images
+        images = pipe(n_samples*[prompt], guidance_scale=scale, num_inference_steps=steps).images
 
     return images
 
@@ -38,12 +63,13 @@ 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' }
+    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
-    
+    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
     
 
@@ -52,10 +78,25 @@ block = gr.Blocks()
 
 examples = [
     [
-        'https://arxiv.org/abs/1706.03762',
+        "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:
@@ -77,10 +118,10 @@ with block:
         with gr.Box():
             with gr.Row().style(mobile_collapse=False, equal_height=True):
                 text = gr.Textbox(
-                    label="Enter your prompt",
+                    label="Link or ID for paper",
                     show_label=False,
                     max_lines=1,
-                    placeholder="Enter your prompt",
+                    placeholder="Give arXiv the link or ID for the paper",
                 ).style(
                     border=(True, False, True, True),
                     rounded=(True, False, False, True),
@@ -90,6 +131,17 @@ with block:
                     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"
@@ -107,9 +159,46 @@ with block:
         ex = gr.Examples(examples=examples, fn=infer, inputs=[text, samples, scale], outputs=gallery, cache_examples=False)
         ex.dataset.headers = [""]
 
-
-        text.submit(infer, inputs=[text, samples, steps, scale], outputs=gallery)
-        btn.click(infer, inputs=[text, samples, steps, scale], outputs=gallery)
+        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;">