|
import gradio as gr
|
|
|
|
from v2 import (
|
|
V2UI,
|
|
parse_upsampling_output,
|
|
V2_ALL_MODELS,
|
|
)
|
|
from utils import (
|
|
gradio_copy_text,
|
|
COPY_ACTION_JS,
|
|
V2_ASPECT_RATIO_OPTIONS,
|
|
V2_RATING_OPTIONS,
|
|
V2_LENGTH_OPTIONS,
|
|
V2_IDENTITY_OPTIONS
|
|
)
|
|
from tagger import (
|
|
predict_tags_wd,
|
|
convert_danbooru_to_e621_prompt,
|
|
remove_specific_prompt,
|
|
insert_recom_prompt,
|
|
compose_prompt_to_copy,
|
|
translate_prompt,
|
|
sort_tags,
|
|
select_random_character,
|
|
)
|
|
from z3de621conv import (
|
|
predict_tags_e621,
|
|
)
|
|
from fl2sd3longcap import (
|
|
predict_tags_fl2_sd3,
|
|
)
|
|
from fl2basepromptgen import (
|
|
predict_tags_fl2_base_prompt_gen,
|
|
)
|
|
from wdtagger import (
|
|
predict_tags_wd_large,
|
|
get_wdtagger_models,
|
|
SWINV2_MODEL_DSV3_REPO,
|
|
CONV_MODEL_DSV3_REPO,
|
|
VIT_MODEL_DSV3_REPO,
|
|
VIT_LARGE_MODEL_DSV3_REPO,
|
|
EVA02_LARGE_MODEL_DSV3_REPO,
|
|
MOAT_MODEL_DSV2_REPO,
|
|
SWIN_MODEL_DSV2_REPO,
|
|
CONV_MODEL_DSV2_REPO,
|
|
CONV2_MODEL_DSV2_REPO,
|
|
VIT_MODEL_DSV2_REPO,
|
|
)
|
|
|
|
def description_ui():
|
|
gr.Markdown(
|
|
"""
|
|
## Danbooru Tags Transformer V2 Demo with WD Tagger
|
|
(Image =>) Prompt => Upsampled longer prompt
|
|
"""
|
|
)
|
|
def description_ui2():
|
|
gr.Markdown(
|
|
"""
|
|
- Mod of p1atdev's [Danbooru Tags Transformer V2 Demo](https://huggingface.co/spaces/p1atdev/danbooru-tags-transformer-v2) and [WD Tagger with π€ transformers](https://huggingface.co/spaces/p1atdev/wd-tagger-transformers).
|
|
- Models: p1atdev's [wd-swinv2-tagger-v3-hf](https://huggingface.co/p1atdev/wd-swinv2-tagger-v3-hf), [dart-v2-moe-sft](https://huggingface.co/p1atdev/dart-v2-moe-sft), [dart-v2-sft](https://huggingface.co/p1atdev/dart-v2-sft)\
|
|
, toynya's [Z3D-E621-Convnext](https://huggingface.co/toynya/Z3D-E621-Convnext), gokaygokay's [Florence-2-SD3-Captioner](https://huggingface.co/gokaygokay/Florence-2-SD3-Captioner),\
|
|
MiaoshouAI's [Florence-2-base-PromptGen](https://huggingface.co/MiaoshouAI/Florence-2-base-PromptGen),\
|
|
SmilingWolf's [wd-vit-large-tagger-v3](https://huggingface.co/SmilingWolf/wd-vit-large-tagger-v3).
|
|
"""
|
|
)
|
|
|
|
|
|
def main():
|
|
|
|
v2 = V2UI()
|
|
|
|
with gr.Blocks() as ui:
|
|
description_ui()
|
|
|
|
with gr.Row():
|
|
with gr.Column(scale=2):
|
|
with gr.Group():
|
|
input_image = gr.Image(label="Input image", type="pil", sources=["upload", "clipboard"], height=256)
|
|
with gr.Accordion(label="Advanced options", open=False):
|
|
wdtagger_model = gr.Dropdown(label="Original WD Tagger Model", choices=get_wdtagger_models(), value=VIT_LARGE_MODEL_DSV3_REPO)
|
|
general_threshold = gr.Slider(label="Threshold", minimum=0.0, maximum=1.0, value=0.3, step=0.01, interactive=True)
|
|
character_threshold = gr.Slider(label="Character threshold", minimum=0.0, maximum=1.0, value=0.8, step=0.01, interactive=True)
|
|
e621_threshold = gr.Slider(label="Threshold (Z3D-E621-Convnext)", minimum=0.0, maximum=1.0, value=0.5, 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 Original WD Tagger", "Use Z3D-E621-Convnext", "Use Florence-2-SD3-Long-Captioner", "Use Florence-2-base-PromptGen"], label="Algorithms", value=["Use WD Tagger"])
|
|
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")
|
|
random_character = gr.Button(value="Pick a 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)
|
|
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")
|
|
with gr.Accordion(label="Advanced options", open=False):
|
|
input_aspect_ratio = gr.Radio(label="Aspect ratio", info="The aspect ratio of the image.", choices=list(V2_ASPECT_RATIO_OPTIONS), value="square")
|
|
input_length = gr.Radio(label="Length", info="The total length of the tags.", choices=list(V2_LENGTH_OPTIONS), value="very_long")
|
|
input_identity = gr.Radio(label="Keep identity", info="How strictly to keep the identity of the character or subject. If you specify the detail of subject in the prompt, you should choose `strict`. Otherwise, choose `none` or `lax`. `none` is very creative but sometimes ignores the input prompt.", choices=list(V2_IDENTITY_OPTIONS), value="lax")
|
|
input_ban_tags = gr.Textbox(label="Ban tags", info="Tags to ban from the output.", placeholder="alternate costumen, ...", value="censored")
|
|
model_name = gr.Dropdown(label="Model", choices=list(V2_ALL_MODELS.keys()), value=list(V2_ALL_MODELS.keys())[0])
|
|
dummy_np = gr.Textbox(label="Negative prompt", value="", visible=False)
|
|
recom_animagine = gr.Textbox(label="Animagine reccomended prompt", value="Animagine", visible=False)
|
|
recom_pony = gr.Textbox(label="Pony reccomended prompt", value="Pony", visible=False)
|
|
|
|
generate_btn = gr.Button(value="GENERATE TAGS", size="lg", variant="primary")
|
|
|
|
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.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)
|
|
description_ui2()
|
|
|
|
v2.input_components = [model_name, input_copyright, input_character, input_general,
|
|
input_rating, input_aspect_ratio, input_length, input_identity, input_ban_tags]
|
|
|
|
random_character.click(select_random_character, [input_copyright, input_character], [input_copyright, input_character], queue=False)
|
|
|
|
translate_input_prompt_button.click(translate_prompt, [input_general], [input_general], queue=False)
|
|
translate_input_prompt_button.click(translate_prompt, [input_character], [input_character], queue=False)
|
|
translate_input_prompt_button.click(translate_prompt, [input_copyright], [input_copyright], queue=False)
|
|
|
|
generate_from_image_btn.click(
|
|
predict_tags_wd,
|
|
[input_image, input_general, image_algorithms, general_threshold, character_threshold, input_copyright, input_character],
|
|
[input_copyright, input_character, input_general, copy_input_btn],
|
|
).success(
|
|
predict_tags_wd_large,
|
|
[input_image, input_general, image_algorithms, general_threshold, character_threshold, input_copyright, input_character, wdtagger_model],
|
|
[input_copyright, input_character, input_general, copy_input_btn],
|
|
).success(
|
|
predict_tags_e621, [input_image, input_general, image_algorithms, e621_threshold], [input_general],
|
|
).success(
|
|
predict_tags_fl2_sd3, [input_image, input_general, image_algorithms], [input_general],
|
|
).success(
|
|
predict_tags_fl2_base_prompt_gen, [input_image, input_general, image_algorithms], [input_general],
|
|
).success(
|
|
remove_specific_prompt, [input_general, keep_tags], [input_general], queue=False,
|
|
).success(
|
|
convert_danbooru_to_e621_prompt, [input_general, input_tag_type], [input_general], queue=False,
|
|
).success(
|
|
sort_tags, [input_general], [input_general], queue=False,
|
|
).success(
|
|
insert_recom_prompt, [input_general, dummy_np, recom_prompt], [input_general, dummy_np], queue=False,
|
|
)
|
|
copy_input_btn.click(compose_prompt_to_copy, [input_character, input_copyright, input_general], [input_tags_to_copy], queue=False)\
|
|
.success(gradio_copy_text, [input_tags_to_copy], js=COPY_ACTION_JS)
|
|
|
|
generate_btn.click(
|
|
parse_upsampling_output(v2.on_generate), [*v2.input_components],
|
|
[output_text, elapsed_time_md, copy_btn, copy_btn_pony],
|
|
).success(
|
|
sort_tags, [output_text], [output_text], queue=False,
|
|
).success(
|
|
convert_danbooru_to_e621_prompt, [output_text, tag_type], [output_text_pony], queue=False,
|
|
).success(
|
|
insert_recom_prompt, [output_text, dummy_np, recom_animagine], [output_text, dummy_np], queue=False,
|
|
).success(
|
|
insert_recom_prompt, [output_text_pony, dummy_np, recom_pony], [output_text_pony, dummy_np], queue=False,
|
|
)
|
|
copy_btn.click(gradio_copy_text, [output_text], js=COPY_ACTION_JS)
|
|
copy_btn_pony.click(gradio_copy_text, [output_text_pony], js=COPY_ACTION_JS)
|
|
|
|
ui.launch()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|