Spaces:
Paused
Paused
| import spaces | |
| import gradio as gr | |
| from utils import gradio_copy_text, COPY_ACTION_JS | |
| from tagger import (predict_tags_wd, convert_danbooru_to_e621_prompt, remove_specific_prompt, | |
| insert_recom_prompt, compose_prompt_to_copy, translate_prompt, select_random_character) | |
| from fl2sd3longcap import predict_tags_fl2_sd3 | |
| from fl2basepromptgen import predict_tags_fl2_base_prompt_gen | |
| from fl2flux import predict_tags_fl2_flux | |
| from promptenhancer import prompt_enhancer | |
| def description_ui(): | |
| gr.Markdown( | |
| """ | |
| ## Prompt Enhancer with WD Tagger & Flux/SD3 Captioner | |
| (Image =>) Prompt => Upsampled longer prompt | |
| """ | |
| ) | |
| def description_ui2(): | |
| gr.Markdown( | |
| """ | |
| - It's a mod. Original Spaces: p1atdev's [WD Tagger with 🤗 transformers](https://huggingface.co/spaces/p1atdev/wd-tagger-transformers),\ | |
| gokaygokay's [Prompt-Enhancer](https://huggingface.co/spaces/gokaygokay/Prompt-Enhancer) /\ | |
| [Florence-2-SD3-Captioner](https://huggingface.co/spaces/gokaygokay/Florence-2-SD3-Captioner). | |
| - Models: p1atdev's [wd-swinv2-tagger-v3-hf](https://huggingface.co/p1atdev/wd-swinv2-tagger-v3-hf),\ | |
| gokaygokay's [Florence-2-SD3-Captioner](https://huggingface.co/gokaygokay/Florence-2-SD3-Captioner),\ | |
| gokaygokay's [Florence-2-Flux](https://huggingface.co/gokaygokay/Florence-2-Flux),\ | |
| gokaygokay's [Florence-2-Flux-Large](https://huggingface.co/gokaygokay/Florence-2-Flux-Large),\ | |
| [Lamini-Prompt-Enchance](https://huggingface.co/gokaygokay/Lamini-Prompt-Enchance),\ | |
| [Lamini-Prompt-Enchance-Long](https://huggingface.co/gokaygokay/Lamini-Prompt-Enchance-Long),\ | |
| [Flux-Prompt-Enhance](https://huggingface.co/gokaygokay/Flux-Prompt-Enhance),\ | |
| MiaoshouAI's [Florence-2-large-PromptGen-v1.5](https://huggingface.co/MiaoshouAI/Florence-2-large-PromptGen-v1.5). | |
| """ | |
| ) | |
| def main(): | |
| with gr.Blocks() as ui: | |
| description_ui() | |
| with gr.Column(): | |
| 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): | |
| 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) | |
| 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) | |
| keep_tags = 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", "Use Florence-2-Flux", "Use Florence-2-Flux-Large", "Use Florence-2-large-PromptGen"], label="Algorithms", value=["Use WD Tagger", "Use Florence-2-Flux"]) | |
| generate_from_image_btn = gr.Button(value="GENERATE TAGS FROM IMAGE", size="lg", variant="primary") | |
| with gr.Group(): | |
| with gr.Row(): | |
| input_character = gr.Textbox(label="Character tags", placeholder="hatsune miku") | |
| input_copyright = gr.Textbox(label="Copyright tags", placeholder="vocaloid") | |
| 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) | |
| translate_input_prompt_button = gr.Button(value="Translate prompt to English", size="sm", variant="secondary") | |
| prompt_enhancer_model = gr.Radio(["Medium", "Long", "Flux"], label="Model Choice", value="Long", info="Enhance your prompts with Medium or Long answers") | |
| with gr.Accordion(label="Advanced options", open=False, visible=False): | |
| tag_type = gr.Radio(label="Output tag conversion", info="danbooru for Animagine, e621 for Pony.", choices=["danbooru", "e621"], value="e621", visible=False) | |
| 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.Row(): | |
| 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) | |
| 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() | |
| gr.LoginButton() | |
| gr.DuplicateButton(value="Duplicate Space for private use") | |
| 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( | |
| lambda: ("", "", ""), None, [input_copyright, input_character, input_general], queue=False, | |
| ).success( | |
| predict_tags_wd, | |
| [input_image, input_general, image_algorithms, general_threshold, character_threshold], | |
| [input_copyright, input_character, input_general, copy_input_btn], | |
| ).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(predict_tags_fl2_flux, [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(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( | |
| prompt_enhancer, | |
| [input_character, input_copyright, input_general, prompt_enhancer_model], | |
| [output_text, copy_btn, copy_btn_pony], | |
| ).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() | |