Upload 4 files

app.py

@@ -342,13 +342,14 @@ from stablepy import logger
 logger.setLevel(logging.CRITICAL)
 
 from v2 import (
-    V2UI,
-    parse_upsampling_output,
     V2_ALL_MODELS,
+    v2_random_prompt,
+    v2_upsampling_prompt,
 )
 from utils import (
     gradio_copy_text,
     COPY_ACTION_JS,
+    gradio_copy_prompt,
     V2_ASPECT_RATIO_OPTIONS,
     V2_RATING_OPTIONS,
     V2_LENGTH_OPTIONS,
@@ -782,6 +783,7 @@ CSS ="""
 .contain {display: flex; flex-direction: column; !important;}
 #component-0 {height: 100%; !important;}
 #gallery {flex-grow: 1; !important;}
+#model-info { text-align: center; }
 """
 ## END MOD
 
@@ -811,13 +813,13 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", elem_id="main", css=CSS) as app:
     )
     with gr.Row():
         with gr.Tab("Generation"):
-            v2b = V2UI()
             with gr.Column():
                 interface_mode_gui = gr.Radio(label="Quick settings", choices=["Simple", "Standard", "Fast", "LoRA"], value="Standard")
                 with gr.Accordion("Model and Task", open=False) as menu_model:
                     task_gui = gr.Dropdown(label="Task", choices=sdxl_task, value=task_model_list[0])
-                    model_name_gui = gr.Dropdown(label="Model", info="You can enter a huggingface model repo_id to want to use.", choices=get_tupled_model_list(model_list), value="votepurchase/animagine-xl-3.1", allow_custom_value=True)
-                    model_info_gui = gr.Markdown()
+                    with gr.Group():
+                        model_name_gui = gr.Dropdown(label="Model", info="You can enter a huggingface model repo_id to want to use.", choices=get_tupled_model_list(model_list), value="votepurchase/animagine-xl-3.1", allow_custom_value=True)
+                        model_info_gui = gr.Markdown(elem_id="model-info")
                     with gr.Row():
                         quick_model_type_gui = gr.Radio(label="Model Type", choices=["None", "Auto", "Animagine", "Pony"], value="Auto", interactive=True)
                         quick_genre_gui = gr.Radio(label="Genre", choices=["Anime", "Photo"], value="Anime", interactive=True)
@@ -828,17 +830,17 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", elem_id="main", css=CSS) as app:
                         style_selector_gui = gr.Dropdown(label="Style Preset", interactive=True, choices=list(preset_styles.keys()), value="None")
                         sampler_selector_gui = gr.Dropdown(label="Sampler Quick Settings", interactive=True, choices=list(preset_sampler_setting.keys()), value="None")
                         optimization_gui = gr.Dropdown(label="Optimization for SDXL", choices=list(optimization_list.keys()), value="None", interactive=True)
-                with gr.Accordion("Generate prompt from Image", open=False) as menu_from_image:
-                    input_image_gui = gr.Image(label="Input image", type="pil", sources=["upload", "clipboard"], height=256)
-                    with gr.Accordion(label="Advanced options", open=False):
-                        general_threshold_gui = gr.Slider(label="Threshold", minimum=0.0, maximum=1.0, value=0.3, step=0.01, interactive=True)
-                        character_threshold_gui = gr.Slider(label="Character threshold", minimum=0.0, maximum=1.0, value=0.8, step=0.01, interactive=True)
-                        tag_type_gui = gr.Radio(label="Convert tags to", info="danbooru for Animagine, e621 for Pony.", choices=["danbooru", "e621"], value="danbooru")
-                        recom_prompt_gui = gr.Radio(label="Insert reccomended prompt", choices=["None", "Animagine", "Pony"], value="None", interactive=True)
-                        keep_tags_gui = gr.Radio(label="Remove tags leaving only the following", choices=["body", "dress", "all"], value="all")
-                    image_algorithms = gr.CheckboxGroup(["Use WD Tagger", "Use Florence-2-SD3-Long-Captioner"], label="Algorithms", value=["Use WD Tagger"], visible=False)
-                    generate_from_image_btn_gui = gr.Button(value="GENERATE TAGS FROM IMAGE", size="lg", variant="primary")
                 with gr.Group():
+                    with gr.Accordion("Prompt from Image", open=False) as menu_from_image:
+                        input_image_gui = gr.Image(label="Input image", type="pil", sources=["upload", "clipboard"], height=256)
+                        with gr.Accordion(label="Advanced options", open=False):
+                            general_threshold_gui = gr.Slider(label="Threshold", minimum=0.0, maximum=1.0, value=0.3, step=0.01, interactive=True)
+                            character_threshold_gui = gr.Slider(label="Character threshold", minimum=0.0, maximum=1.0, value=0.8, step=0.01, interactive=True)
+                            tag_type_gui = gr.Radio(label="Convert tags to", info="danbooru for Animagine, e621 for Pony.", choices=["danbooru", "e621"], value="danbooru")
+                            recom_prompt_gui = gr.Radio(label="Insert reccomended prompt", choices=["None", "Animagine", "Pony"], value="None", interactive=True)
+                            keep_tags_gui = gr.Radio(label="Remove tags leaving only the following", choices=["body", "dress", "all"], value="all")
+                        image_algorithms = gr.CheckboxGroup(["Use WD Tagger"], label="Algorithms", value=["Use WD Tagger"], visible=False)
+                        generate_from_image_btn_gui = gr.Button(value="GENERATE TAGS FROM IMAGE")
                     prompt_gui = gr.Textbox(lines=6, placeholder="1girl, solo, ...", label="Prompt", show_copy_button=True)
                     with gr.Accordion("Negative prompt, etc.", open=False) as menu_negative:
                         neg_prompt_gui = gr.Textbox(lines=3, placeholder="lowres, (bad), ...", label="Negative prompt", show_copy_button=True)
@@ -850,21 +852,20 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", elem_id="main", css=CSS) as app:
                         with gr.Row():
                             character_dbt = gr.Textbox(lines=1, placeholder="kafuu chino, ...", label="Character names")
                             series_dbt = gr.Textbox(lines=1, placeholder="Is the order a rabbit?, ...", label="Series names")
-                            random_character_gui = gr.Button(value="Pick a random character 🎲", size="sm", variant="secondary")
+                            random_character_gui = gr.Button(value="Random character 🎲", size="sm", variant="secondary")
                             model_name_dbt = gr.Dropdown(label="Model", choices=list(V2_ALL_MODELS.keys()), value=list(V2_ALL_MODELS.keys())[0], visible=False)
                             aspect_ratio_dbt = gr.Radio(label="Aspect ratio", choices=list(V2_ASPECT_RATIO_OPTIONS), value="square", visible=False)
                             length_dbt = gr.Radio(label="Length", choices=list(V2_LENGTH_OPTIONS), value="very_long", visible=False)
                             identity_dbt = gr.Radio(label="Keep identity", choices=list(V2_IDENTITY_OPTIONS), value="lax", visible=False)                    
                             ban_tags_dbt = gr.Textbox(label="Ban tags", placeholder="alternate costumen, ...", value="futanari, censored, furry, furrification", visible=False)
-                            elapsed_time_dbt = gr.Markdown(label="Elapsed time", value="", visible=False)
                             copy_button_dbt = gr.Button(value="Copy to clipboard", visible=False)
-                        generate_db_random_button = gr.Button(value="Generate random prompt 🎲", variant="secondary")
                         rating_dbt = gr.Radio(label="Rating", choices=list(V2_RATING_OPTIONS), value="sfw")
-                        with gr.Row():
-                            set_params_gui = gr.Button(value="↙️", size="sm")
-                            clear_prompt_gui = gr.Button(value="🗑️", size="sm")
-                            set_random_seed = gr.Button(value="🎲", size="sm")
-                    
+                        generate_db_random_button = gr.Button(value="EXTEND PROMPT 🎲")
+                    with gr.Row():
+                        set_random_seed = gr.Button(value="Seed 🎲", size="sm")
+                        set_params_gui = gr.Button(value="Params ↙️", size="sm")
+                        clear_prompt_gui = gr.Button(value="Clear 🗑️", size="sm")
+
                 generate_button = gr.Button(value="GENERATE IMAGE", size="lg", variant="primary")
 
                 model_name_gui.change(
@@ -1606,16 +1607,16 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", elem_id="main", css=CSS) as app:
         ).success(
             insert_recom_prompt, [prompt_gui, neg_prompt_gui, recom_prompt_gui], [prompt_gui, neg_prompt_gui], queue=False,
         )
-
-        v2b.input_components = [model_name_dbt, series_dbt, character_dbt, prompt_gui,
-            rating_dbt, aspect_ratio_dbt, length_dbt, identity_dbt, ban_tags_dbt]
         
         prompt_type_button.click(convert_danbooru_to_e621_prompt, [prompt_gui, prompt_type_gui], [prompt_gui], queue=False)
         random_character_gui.click(select_random_character, [series_dbt, character_dbt], [series_dbt, character_dbt], queue=False)
         generate_db_random_button.click(
-            parse_upsampling_output(v2b.on_generate),
-            [*v2b.input_components],
-            [prompt_gui, elapsed_time_dbt, copy_button_dbt, copy_button_dbt],
+            v2_random_prompt,
+            [prompt_gui, series_dbt, character_dbt,
+            rating_dbt, aspect_ratio_dbt, length_dbt, identity_dbt, ban_tags_dbt, model_name_dbt],
+            [prompt_gui, series_dbt, character_dbt],
+        ).success(
+            convert_danbooru_to_e621_prompt, [prompt_gui, tag_type_gui], [prompt_gui], queue=False,
         )
 
         translate_prompt_button.click(translate_prompt, [prompt_gui], [prompt_gui], queue=False)
@@ -1743,7 +1744,6 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", elem_id="main", css=CSS) as app:
         ).success(save_gallery_images, [result_images], [result_images, result_images_files, result_images_files], queue=False)
 
         with gr.Tab("Danbooru Tags Transformer with WD Tagger", render=True):
-            v2 = V2UI()
             with gr.Column(scale=2):
                 with gr.Group():
                     input_image = gr.Image(label="Input image", type="pil", sources=["upload", "clipboard"], height=256)
@@ -1752,17 +1752,20 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", elem_id="main", css=CSS) as app:
                         character_threshold = gr.Slider(label="Character threshold", minimum=0.0, maximum=1.0, value=0.8, step=0.01, interactive=True)
                         input_tag_type = gr.Radio(label="Convert tags to", info="danbooru for Animagine, e621 for Pony.", choices=["danbooru", "e621"], value="danbooru")
                         recom_prompt = gr.Radio(label="Insert reccomended prompt", choices=["None", "Animagine", "Pony"], value="None", interactive=True)
-                    image_algorithms = gr.CheckboxGroup(["Use WD Tagger", "Use Florence-2-SD3-Long-Captioner"], label="Algorithms", value=["Use WD Tagger"], visible=False)
+                    image_algorithms = gr.CheckboxGroup(["Use WD Tagger"], label="Algorithms", value=["Use WD Tagger"], visible=False)
                     keep_tags = gr.Radio(label="Remove tags leaving only the following", choices=["body", "dress", "all"], value="all")
                     generate_from_image_btn = gr.Button(value="GENERATE TAGS FROM IMAGE", size="lg", variant="primary")
                 
                 with gr.Group():
-                    input_character = gr.Textbox(label="Character tags", placeholder="hatsune miku")
-                    input_copyright = gr.Textbox(label="Copyright tags", placeholder="vocaloid")
-                    pick_random_character = gr.Button(value="Pick a random character 🎲", size="sm")
+                    with gr.Row():
+                        input_character = gr.Textbox(label="Character tags", placeholder="hatsune miku")
+                        input_copyright = gr.Textbox(label="Copyright tags", placeholder="vocaloid")
+                        pick_random_character = gr.Button(value="Random character 🎲", size="sm")
                     input_general = gr.TextArea(label="General tags", lines=4, placeholder="1girl, ...", value="")
                     input_tags_to_copy = gr.Textbox(value="", visible=False)
-                    copy_input_btn = gr.Button(value="Copy to clipboard", size="sm", interactive=False)
+                    with gr.Row():
+                        copy_input_btn = gr.Button(value="Copy to clipboard", size="sm", interactive=False)
+                        copy_prompt_btn_input = gr.Button(value="Copy to primary prompt", size="sm", interactive=False)
                     translate_input_prompt_button = gr.Button(value="Translate prompt to English", size="sm", variant="secondary")
                     tag_type = gr.Radio(label="Output tag conversion", info="danbooru for Animagine, e621 for Pony.", choices=["danbooru", "e621"], value="e621", visible=False)
                     input_rating = gr.Radio(label="Rating", choices=list(V2_RATING_OPTIONS), value="explicit")
@@ -1780,18 +1783,18 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", elem_id="main", css=CSS) as app:
 
                 with gr.Group():
                     output_text = gr.TextArea(label="Output tags", interactive=False, show_copy_button=True)
-                    copy_btn = gr.Button(value="Copy to clipboard", size="sm", interactive=False)
-                    elapsed_time_md = gr.Markdown(label="Elapsed time", value="", visible=False)
+                    with gr.Row():
+                        copy_btn = gr.Button(value="Copy to clipboard", size="sm", interactive=False)
+                        copy_prompt_btn = gr.Button(value="Copy to primary prompt", size="sm", interactive=False)
 
                 with gr.Group():
                     output_text_pony = gr.TextArea(label="Output tags (Pony e621 style)", interactive=False, show_copy_button=True)
-                    copy_btn_pony = gr.Button(value="Copy to clipboard", size="sm", interactive=False)
+                    with gr.Row():
+                        copy_btn_pony = gr.Button(value="Copy to clipboard", size="sm", interactive=False)
+                        copy_prompt_btn_pony = gr.Button(value="Copy to primary prompt", size="sm", interactive=False)
 
                 description_ui()
 
-        v2.input_components = [model_name, input_copyright, input_character, input_general,
-            input_rating, input_aspect_ratio, input_length, input_identity, input_ban_tags]
-
         translate_input_prompt_button.click(translate_prompt, inputs=[input_general], outputs=[input_general], queue=False)
         translate_input_prompt_button.click(translate_prompt, inputs=[input_character], outputs=[input_character], queue=False)
         translate_input_prompt_button.click(translate_prompt, inputs=[input_copyright], outputs=[input_copyright], queue=False)
@@ -1806,25 +1809,31 @@ with gr.Blocks(theme="NoCrypt/miku@>=1.2.2", elem_id="main", css=CSS) as app:
             convert_danbooru_to_e621_prompt, inputs=[input_general, input_tag_type], outputs=[input_general], queue=False,
         ).success(
             insert_recom_prompt, inputs=[input_general, dummy_np, recom_prompt], outputs=[input_general, dummy_np], queue=False,
-        )
+        ).success(lambda: gr.update(interactive=True), None, [copy_prompt_btn_input], queue=False)
         copy_input_btn.click(compose_prompt_to_copy, inputs=[input_character, input_copyright, input_general], outputs=[input_tags_to_copy])\
             .success(gradio_copy_text, inputs=[input_tags_to_copy], js=COPY_ACTION_JS)
-
+        copy_prompt_btn_input.click(compose_prompt_to_copy, inputs=[input_character, input_copyright, input_general], outputs=[input_tags_to_copy])\
+            .success(gradio_copy_prompt, inputs=[input_tags_to_copy], outputs=[prompt_gui])
+        
         pick_random_character.click(select_random_character, [input_copyright, input_character], [input_copyright, input_character])
 
         generate_btn.click(
-            parse_upsampling_output(v2.on_generate),
-            [*v2.input_components],
-            [output_text, elapsed_time_md, copy_btn, copy_btn_pony],
+            v2_upsampling_prompt,
+            [model_name, input_copyright, input_character, input_general,
+            input_rating, input_aspect_ratio, input_length, input_identity, input_ban_tags],
+            [output_text],
         ).success(
             convert_danbooru_to_e621_prompt, inputs=[output_text, tag_type], outputs=[output_text_pony], queue=False,
         ).success(
             insert_recom_prompt, inputs=[output_text, dummy_np, recom_animagine], outputs=[output_text, dummy_np], queue=False,
         ).success(
             insert_recom_prompt, inputs=[output_text_pony, dummy_np, recom_pony], outputs=[output_text_pony, dummy_np], queue=False,
-        )
+        ).success(lambda: (gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True)),
+                   None, [copy_btn, copy_btn_pony, copy_prompt_btn, copy_prompt_btn_pony], queue=False)
         copy_btn.click(gradio_copy_text, inputs=[output_text], js=COPY_ACTION_JS)
         copy_btn_pony.click(gradio_copy_text, inputs=[output_text_pony], js=COPY_ACTION_JS)
+        copy_prompt_btn.click(gradio_copy_prompt, inputs=[output_text], outputs=[prompt_gui])
+        copy_prompt_btn_pony.click(gradio_copy_prompt, inputs=[output_text_pony], outputs=[prompt_gui])
 
     gr.DuplicateButton(value="Duplicate Space for private use (This demo does not work on CPU. Requires Space with GPU available.)", elem_id="duplicate-button", visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1")
 

env.py

@@ -58,6 +58,8 @@ download_vae_list = [
     'https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/resolve/main/sdxl.vae.safetensors?download=true',
     'https://huggingface.co/nubby/blessed-sdxl-vae-fp16-fix/resolve/main/sdxl_vae-fp16fix-c-1.1-b-0.5.safetensors?download=true',
     'https://huggingface.co/nubby/blessed-sdxl-vae-fp16-fix/resolve/main/sdxl_vae-fp16fix-blessed.safetensors?download=true',
+    "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt",
+    "https://huggingface.co/stabilityai/sd-vae-ft-ema-original/resolve/main/vae-ft-ema-560000-ema-pruned.ckpt",
 ]
 
 # - **Download LoRAs**

utils.py

@@ -43,3 +43,8 @@ COPY_ACTION_JS = """\
     navigator.clipboard.writeText(inputs);
   }
 }"""
+
+
+def gradio_copy_prompt(prompt: str):
+    gr.Info("Copied!")
+    return prompt

v2.py

@@ -1,7 +1,7 @@
 import time
-import os
 import torch
 from typing import Callable
+from pathlib import Path
 
 from dartrs.v2 import (
     V2Model,
@@ -20,9 +20,12 @@ from dartrs.utils import get_generation_config
 import gradio as gr
 from gradio.components import Component
 
-from output import UpsamplingOutput
 
-HF_TOKEN = os.getenv("HF_TOKEN", None)
+try:
+    from output import UpsamplingOutput
+except:
+    from .output import UpsamplingOutput
+
 
 V2_ALL_MODELS = {
     "dart-v2-moe-sft": {
@@ -40,8 +43,8 @@ V2_ALL_MODELS = {
 
 def prepare_models(model_config: dict):
     model_name = model_config["repo"]
-    tokenizer = DartTokenizer.from_pretrained(model_name, auth_token=HF_TOKEN)
-    model = model_config["class"].from_pretrained(model_name, auth_token=HF_TOKEN)
+    tokenizer = DartTokenizer.from_pretrained(model_name)
+    model = model_config["class"].from_pretrained(model_name)
 
     return {
         "tokenizer": tokenizer,
@@ -202,3 +205,56 @@ class V2UI:
             elapsed_time=elapsed_time,
         )
 
+
+def parse_upsampling_output_simple(upsampler: UpsamplingOutput):
+    return gen_prompt_text(upsampler)
+
+
+v2 = V2UI()
+
+
+def v2_upsampling_prompt(model: str = "dart-v2-moe-sft", copyright: str = "", character: str = "",
+                          general_tags: str = "", rating: str = "nsfw", aspect_ratio: str = "square",
+                            length: str = "very_long", identity: str = "lax", ban_tags: str = "censored"):
+    raw_prompt = parse_upsampling_output_simple(v2.on_generate(model, copyright, character, general_tags,
+                                                                rating, aspect_ratio, length, identity, ban_tags))
+    return raw_prompt
+
+
+def load_dict_from_csv(filename):
+    dict = {}
+    if not Path(filename).exists():
+        if Path('./tagger/', filename).exists(): filename = str(Path('./tagger/', filename))
+        else: return dict
+    try:
+        with open(filename, 'r', encoding="utf-8") as f:
+            lines = f.readlines()
+    except Exception:
+        print(f"Failed to open dictionary file: {filename}")
+        return dict
+    for line in lines:
+        parts = line.strip().split(',')
+        dict[parts[0]] = parts[1]
+    return dict
+
+
+anime_series_dict = load_dict_from_csv('character_series_dict.csv')
+
+
+def select_random_character(series: str, character: str):
+    from random import seed, randrange
+    seed()
+    character_list = list(anime_series_dict.keys())
+    character = character_list[randrange(len(character_list) - 1)]
+    series = anime_series_dict.get(character.split(",")[0].strip(), "")
+    return series, character
+
+
+def v2_random_prompt(general_tags: str = "", copyright: str = "", character: str = "", rating: str = "nsfw",
+                      aspect_ratio: str = "square", length: str = "very_long", identity: str = "lax",
+                      ban_tags: str = "censored", model: str = "dart-v2-moe-sft"):
+    if copyright == "" and character == "":
+        copyright, character = select_random_character("", "")
+    raw_prompt = v2_upsampling_prompt(model, copyright, character, general_tags, rating,
+                                       aspect_ratio, length, identity, ban_tags)
+    return raw_prompt, copyright, character