Upload 2 files

app.py

@@ -1,19 +1,9 @@
 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 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
@@ -39,7 +29,7 @@ def description_ui2():
  [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-base-PromptGen-v1.5](https://huggingface.co/MiaoshouAI/Florence-2-base-PromptGen-v1.5).
+ MiaoshouAI's [Florence-2-large-PromptGen-v1.5](https://huggingface.co/MiaoshouAI/Florence-2-large-PromptGen-v1.5).
 """
     )
 
@@ -55,7 +45,7 @@ def main():
                     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-base-PromptGen", "Use Florence-2-Flux", "Use Florence-2-Flux-Large"], label="Algorithms", value=["Use WD Tagger", "Use Florence-2-Flux"])
+                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():
@@ -102,14 +92,10 @@ def main():
             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,
+        ).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,
@@ -119,13 +105,9 @@ def main():
             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,
-        )
+        ).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)
 

fl2basepromptgen.py

@@ -9,8 +9,8 @@ subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENT
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 try:
-    fl_model = AutoModelForCausalLM.from_pretrained('MiaoshouAI/Florence-2-base-PromptGen-v1.5', trust_remote_code=True).to("cpu").eval()
-    fl_processor = AutoProcessor.from_pretrained('MiaoshouAI/Florence-2-base-PromptGen-v1.5', trust_remote_code=True)
+    fl_model = AutoModelForCausalLM.from_pretrained('MiaoshouAI/Florence-2-large-PromptGen-v1.5', trust_remote_code=True).to("cpu").eval()
+    fl_processor = AutoProcessor.from_pretrained('MiaoshouAI/Florence-2-large-PromptGen-v1.5', trust_remote_code=True)
 except Exception as e:
     print(e)
     fl_model = fl_processor = None
@@ -46,7 +46,7 @@ def predict_tags_fl2_base_prompt_gen(image: Image.Image, input_tags: str, algo:
     def list_uniq(l):
         return sorted(set(l), key=l.index)
     
-    if not "Use Florence-2-base-PromptGen" in algo:
+    if not "Use Florence-2-large-PromptGen" in algo:
         return input_tags
     tag_list = list_uniq(to_list(input_tags) + to_list(fl_run(image) + ", "))
     tag_list.remove("")