Modify introductions

app.py

@@ -30,7 +30,7 @@ from kgen.formatter import seperate_tags, apply_format
 from kgen.generate import generate
 
 from diff import load_model, encode_prompts
-from meta import DEFAULT_NEGATIVE_PROMPT
+from meta import DEFAULT_NEGATIVE_PROMPT, DEFAULT_FORMAT
 
 
 sdxl_pipe = load_model()
@@ -41,17 +41,6 @@ models.load_model(
     subfolder="dan-cc-coyo_epoch2",
 )
 generate(max_new_tokens=4)
-
-
-DEFAULT_FORMAT = """<|special|>, <|characters|>, <|copyrights|>, 
-<|artist|>, 
-
-<|general|>,
-
-<|extended|>.
-
-<|quality|>, <|meta|>, <|rating|>
-""".strip()
 DEFAULT_TAGS = """
 1girl, king halo (umamusume), umamusume,
 ningen mame, ciloranko, ogipote, misu kasumi,
@@ -103,6 +92,7 @@ def generate(
     nl_prompt,
     black_list,
     temp,
+    output_format,
     target_length,
     top_p,
     min_p,
@@ -110,6 +100,7 @@ def generate(
     seed,
     escape_brackets,
 ):
+    default_format = DEFAULT_FORMAT[output_format]
     titpop.BAN_TAGS = [t.strip() for t in black_list.split(",") if t.strip()]
     generation_setting = {
         "seed": seed,
@@ -120,11 +111,11 @@ def generate(
     }
     inputs = seperate_tags(tags.split(","))
     if nl_prompt:
-        if "<|extended|>" in DEFAULT_FORMAT:
+        if "<|extended|>" in default_format:
             inputs["extended"] = nl_prompt
-        elif "<|generated|>" in DEFAULT_FORMAT:
+        elif "<|generated|>" in default_format:
             inputs["generated"] = nl_prompt
-    input_prompt = apply_format(inputs, DEFAULT_FORMAT)
+    input_prompt = apply_format(inputs, default_format)
     if escape_brackets:
         input_prompt = re.sub(r"([()\[\]])", r"\\\1", input_prompt)
 
@@ -132,13 +123,13 @@ def generate(
         seperate_tags(tags.split(",")),
         nl_prompt,
         tag_length_target=target_length,
-        generate_extra_nl_prompt="<|generated|>" in DEFAULT_FORMAT or not nl_prompt,
+        generate_extra_nl_prompt="<|generated|>" in default_format or not nl_prompt,
     )
     t0 = time()
     for result, timing in titpop.titpop_runner_generator(
         meta, operations, general, nl_prompt, **generation_setting
     ):
-        result = apply_format(result, DEFAULT_FORMAT)
+        result = apply_format(result, default_format)
         if escape_brackets:
             result = re.sub(r"([()\[\]])", r"\\\1", result)
         timing["total"] = time() - t0
@@ -153,11 +144,11 @@ def generate_image(
     prompt2,
 ):
     torch.cuda.empty_cache()
+    set_seed(seed)
     prompt_embeds, negative_prompt_embeds, pooled_embeds2, neg_pooled_embeds2 = (
-        encode_prompts(sdxl_pipe, prompt, DEFAULT_NEGATIVE_PROMPT)
+        encode_prompts(sdxl_pipe, prompt2, DEFAULT_NEGATIVE_PROMPT)
     )
-    set_seed(seed)
-    result = sdxl_pipe(
+    result2 = sdxl_pipe(
         prompt_embeds=prompt_embeds,
         negative_prompt_embeds=negative_prompt_embeds,
         pooled_prompt_embeds=pooled_embeds2,
@@ -167,11 +158,13 @@ def generate_image(
         height=1024,
         guidance_scale=6.0,
     ).images[0]
+    yield result2, None
+
     prompt_embeds, negative_prompt_embeds, pooled_embeds2, neg_pooled_embeds2 = (
-        encode_prompts(sdxl_pipe, prompt2, DEFAULT_NEGATIVE_PROMPT)
+        encode_prompts(sdxl_pipe, prompt, DEFAULT_NEGATIVE_PROMPT)
     )
     set_seed(seed)
-    result2 = sdxl_pipe(
+    result = sdxl_pipe(
         prompt_embeds=prompt_embeds,
         negative_prompt_embeds=negative_prompt_embeds,
         pooled_prompt_embeds=pooled_embeds2,
@@ -182,19 +175,39 @@ def generate_image(
         guidance_scale=6.0,
     ).images[0]
     torch.cuda.empty_cache()
-    return result2, result
+    yield result2, result
 
 
 if __name__ == "__main__":
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
-        gr.Markdown("""# TITPOP DEMO""")
         with gr.Accordion("Introduction and Instructions", open=False):
             gr.Markdown(
                 """
-### What is this:
-TITPOP
+## TITPOP Demo
+### What is this
+TITPOP is a tool to extend, generate, refine the input prompt for T2I models.
+<br>It can work on both Danbooru tags and Natural Language. Which means you can use it on almost all the existed T2I models.
+<br>You can take it as "pro max" version of [DTG](https://huggingface.co/KBlueLeaf/DanTagGen-delta-rev2)
+
+### How to use this demo
+1. Enter your tags(optional): put the desired tags into "danboru tags" box
+2. Enter your NL Prompt(optional): put the desired natural language prompt into "Natural Language Prompt" box
+3. Enter your black list(optional): put the desired black list into "black list" box
+4. Adjust the settings: length, temp, top_p, min_p, top_k, seed ...
+4. Click "TITPOP" button: you will see refined prompt on "result" box
+5. If you like the result, click "Generate Image From Result" button
+    * You will see 2 generated images, left one is based on your prompt, right one is based on refined prompt
+    * The backend is diffusers, there are no weighting mechanism, so Escape Brackets is default to False
 
-**The implementation is a little bit inefficient, image gen may be a little bit slower than expected.**
+### Why inference code is private? When will it be open sourced?
+1. This model/tool is still under development, currently is early Alpha version.
+2. I'm doing some research and projects based on this.
+3. The model is released under CC-BY-NC-ND License currently. If you have interest, you can implement inference by yourself.
+4. Once the project/research are done, I will open source all these models/codes with Apache2 license.
+
+### Notification
+**ITPOP is NOT a T2I model. It is Prompt Gen, or, Text-to-Text model. 
+<br>The generated image is come from [Kohaku-XL-Zeta](https://huggingface.co/KBlueLeaf/Kohaku-XL-Zeta) model**
 """
             )
         with gr.Row():
@@ -203,7 +216,7 @@ TITPOP
                     with gr.Column(scale=3):
                         tags_input = gr.TextArea(
                             label="Danbooru Tags",
-                            lines=6,
+                            lines=7,
                             show_copy_button=True,
                             interactive=True,
                             value=DEFAULT_TAGS,
@@ -211,7 +224,7 @@ TITPOP
                         )
                         nl_prompt_input = gr.Textbox(
                             label="Natural Language Prompt",
-                            lines=6,
+                            lines=7,
                             show_copy_button=True,
                             interactive=True,
                             value=DEFAULT_NL,
@@ -225,6 +238,11 @@ TITPOP
                             placeholder="Enter tag/nl black list here",
                         )
                     with gr.Column(scale=2):
+                        output_format = gr.Dropdown(
+                            label="Output Format",
+                            choices=list(DEFAULT_FORMAT.keys()),
+                            value="Both, tag first (recommend)"
+                        )
                         target_length = gr.Dropdown(
                             label="Target Length",
                             choices=["very_short", "short", "long", "very_long"],
@@ -293,6 +311,7 @@ TITPOP
                 nl_prompt_input,
                 black_list,
                 temp,
+                output_format,
                 target_length,
                 top_p,
                 min_p,
@@ -308,8 +327,13 @@ TITPOP
             ],
             queue=True,
         )
+        
+        def generate_image_wrapper(seed, result, input_prompt):
+            for img1, img2 in generate_image(seed, result, input_prompt):
+                yield img1, img2, gr.update(interactive=False)
+            yield img1, img2, gr.update(interactive=True)
         gen_img.click(
-            lambda *args: (*generate_image(*args), gr.update(interactive=True)),
+            generate_image_wrapper,
             [seed, result, input_prompt],
             [img1, img2, submit],
             queue=True,

meta.py

@@ -8,47 +8,52 @@ DEFAULT_STYLE_LIST = {
     "no style": "",
 }
 
-MODEL_DEFAULT_QUALITY_LIST = {
-    "KBlueLeaf/Kohaku-XL-Zeta": "masterpiece, newest, absurdres",
-    "KBlueLeaf/Kohaku-XL-Epsilon-rev2": "masterpiece, newest, absurdres",
-    "KBlueLeaf/Kohaku-XL-Epsilon": "masterpiece, newest, absurdres, safe",
-    "cagliostrolab/animagine-xl-3.1": "masterpiece, newest, very aesthetic, absurdres, safe",
-}
+DEFAULT_NEGATIVE_PROMPT = """
+low quality, worst quality, normal quality, text, signature, jpeg artifacts, 
+bad anatomy, old, early, mini skirt, nsfw, chibi, multiple girls, multiple boys, 
+multiple tails, multiple views, copyright name, watermark, artist name, signature
+"""
 
-MODEL_FORMAT_LIST = {
-    "KBlueLeaf/Kohaku-XL-Zeta": """<|special|>, 
-<|characters|>, <|copyrights|>, 
+DEFAULT_FORMAT = {
+    "tag only (DTG mode)":"""
+<|special|>, <|characters|>, <|copyrights|>, 
 <|artist|>, 
 
-<|general|>, 
+<|general|>,
 
-<|quality|>, <|meta|>, <|rating|>""",
-    "KBlueLeaf/Kohaku-XL-Epsilon-rev2": """<|special|>, 
-<|characters|>, <|copyrights|>, 
+<|quality|>, <|meta|>, <|rating|>
+""".strip(),
+    "NL only (Tag to NL)": """<|extended|>.""".strip(),
+    "Both, tag first (recommend)": """
+<|special|>, <|characters|>, <|copyrights|>, 
 <|artist|>, 
 
-<|general|>, 
+<|general|>,
 
-<|quality|>, <|meta|>, <|rating|>""",
-    "KBlueLeaf/Kohaku-XL-Epsilon": """<|special|>, 
-<|characters|>, <|copyrights|>, 
+<|extended|>.
+
+<|quality|>, <|meta|>, <|rating|>
+""".strip(),
+    "Both, NL first (recommend)": """
+<|special|>, <|characters|>, <|copyrights|>, 
 <|artist|>, 
 
-<|general|>, 
+<|extended|>.
+
+<|general|>,
 
-<|quality|>, <|meta|>, <|rating|>""",
-    "cagliostrolab/animagine-xl-3.1": """<|special|>, 
-<|characters|>, <|copyrights|>, 
+<|quality|>, <|meta|>, <|rating|>
+""".strip(),
+    "Both + generated NL": """
+<|special|>, <|characters|>, <|copyrights|>, 
 <|artist|>, 
 
-<|general|>, 
+<|generated|>.
 
-<|quality|>, <|meta|>, <|rating|>""",
-}
+<|general|>,
 
+<|extended|>.
 
-DEFAULT_NEGATIVE_PROMPT = """
-low quality, worst quality, normal quality, text, signature, jpeg artifacts, 
-bad anatomy, old, early, mini skirt, nsfw, chibi, multiple girls, multiple boys, 
-multiple tails, multiple views, copyright name, watermark, artist name, signature
-"""
+<|quality|>, <|meta|>, <|rating|>
+""".strip()
+}