Upload 5 files

Added "Prompt Enhancer" from [John6666/danbooru-tags-transformer-v2-with-wd-tagger](https://huggingface.co/spaces/John6666/danbooru-tags-transformer-v2-with-wd-tagger/blob/main/tagger/promptenhancer.py) (Thanks!) and cleared some code.

app.py

@@ -11,6 +11,7 @@ import json
 from modules.classifyTags import classify_tags,process_tags
 from modules.florence2 import process_image,single_task_list,update_task_dropdown
 from modules.llama_loader import llama_list,llama3reorganize
+from modules.tag_enhancer import prompt_enhancer
 os.environ['PYTORCH_ENABLE_MPS_FALLBACK']='1'
 
 TITLE = "Multi-Tagger"
@@ -249,9 +250,9 @@ class Predictor:
                         reverse=True,
                     )
                     sorted_general_list = [x[0] for x in sorted_general_list]
-                    #Remove values from character_list that already exist in sorted_general_list
+                    # Remove values from character_list that already exist in sorted_general_list
                     character_list = [item for item in character_list if item not in sorted_general_list]
-                    #Remove values from sorted_general_list that already exist in prepend_list or append_list
+                    # Remove values from sorted_general_list that already exist in prepend_list or append_list
                     if prepend_list:
                         sorted_general_list = [item for item in sorted_general_list if item not in prepend_list]
                     if append_list:
@@ -312,7 +313,8 @@ class Predictor:
                         "rating": rating, 
                         "character_res": character_res, 
                         "general_res": general_res, 
-                        "unclassified_tags": unclassified_tags 
+                        "unclassified_tags": unclassified_tags,
+                        "enhanced_tags": ""  # Initialize as empty string
                     }
 
                     timer.report()
@@ -348,11 +350,12 @@ def get_selection_from_gallery(gallery: list, tag_results: dict, selected_state:
         "rating": "", 
         "character_res": "", 
         "general_res": "", 
-        "unclassified_tags": "{}" 
+        "unclassified_tags": "{}", 
+        "enhanced_tags": "" 
     }
     if selected_state.value["image"]["path"] in tag_results:
         tag_result = tag_results[selected_state.value["image"]["path"]]
-    return (selected_state.value["image"]["path"], selected_state.value["caption"]), tag_result["strings"], tag_result["strings2"], tag_result["classified_tags"], tag_result["rating"], tag_result["character_res"], tag_result["general_res"], tag_result["unclassified_tags"]
+    return (selected_state.value["image"]["path"], selected_state.value["caption"]), tag_result["strings"], tag_result["strings2"], tag_result["classified_tags"], tag_result["rating"], tag_result["character_res"], tag_result["general_res"], tag_result["unclassified_tags"], tag_result["enhanced_tags"]
 def append_gallery(gallery:list,image:str):
 	if gallery is None:gallery=[]
 	if not image:return gallery,None
@@ -417,7 +420,6 @@ with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True
                         upload_button = gr.UploadButton("Upload multiple images", file_types=["image"], file_count="multiple", size="sm")
                         remove_button = gr.Button("Remove Selected Image", size="sm")
                     gallery = gr.Gallery(columns=5, rows=5, show_share_button=False, interactive=True, height="500px", label="Grid of images")
-
                 model_repo = gr.Dropdown(
                     dropdown_list,
                     value=EVA02_LARGE_MODEL_DSV3_REPO,
@@ -485,14 +487,17 @@ with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True
                         size="lg",
                     )
             with gr.Column(variant="panel"):
-                download_file = gr.File(label="Download includes: All outputs* and image(s)") # 0
-                character_res = gr.Label(label="Output (characters)") # 1
-                sorted_general_strings = gr.Textbox(label="Output (string)*", show_label=True, show_copy_button=True) # 2
-                final_categorized_output = gr.Textbox(label="Categorized (string)* - If it's too long, select an image to display tags correctly.", show_label=True, show_copy_button=True) # 3
-                categorized = gr.JSON(label="Categorized (tags)* - JSON") # 4
-                rating = gr.Label(label="Rating") # 5
-                general_res = gr.Label(label="Output (tags)") # 6
-                unclassified = gr.JSON(label="Unclassified (tags)") # 7
+                download_file = gr.File(label="Download includes: All outputs* and image(s)")  # 0
+                character_res = gr.Label(label="Output (characters)")  # 1
+                sorted_general_strings = gr.Textbox(label="Output (string)*", show_label=True, show_copy_button=True)  # 2
+                final_categorized_output = gr.Textbox(label="Categorized (string)* - If it's too long, select an image to display tags correctly.", show_label=True, show_copy_button=True)  # 3
+                pe_generate_btn = gr.Button(value="ENHANCE TAGS", size="lg", variant="primary") # 4
+                enhanced_tags = gr.Textbox(label="Enhanced Tags", show_label=True, show_copy_button=True)  # 5
+                prompt_enhancer_model = gr.Radio(["Medium", "Long", "Flux"], label="Model Choice", value="Medium", info="Enhance your prompts with Medium or Long answers") # 6
+                categorized = gr.JSON(label="Categorized (tags)* - JSON")  # 7
+                rating = gr.Label(label="Rating")  # 8
+                general_res = gr.Label(label="Output (tags)")  # 9
+                unclassified = gr.JSON(label="Unclassified (tags)")  # 10
                 clear.add(
                     [
                         download_file,
@@ -503,8 +508,10 @@ with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True
                         character_res,
                         general_res,
                         unclassified,
+                        prompt_enhancer_model,
+                        enhanced_tags,
                     ]
-                )           
+                )      
             tag_results = gr.State({})
             # Define the event listener to add the uploaded image to the gallery
             image_input.change(append_gallery, inputs=[gallery, image_input], outputs=[gallery, image_input])
@@ -512,9 +519,11 @@ with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True
             upload_button.upload(extend_gallery, inputs=[gallery, upload_button], outputs=gallery)
             # Event to update the selected image when an image is clicked in the gallery
             selected_image = gr.Textbox(label="Selected Image", visible=False)
-            gallery.select(get_selection_from_gallery, inputs=[gallery, tag_results], outputs=[selected_image, sorted_general_strings, final_categorized_output, categorized, rating, character_res, general_res, unclassified])
+            gallery.select(get_selection_from_gallery,inputs=[gallery, tag_results],outputs=[selected_image, sorted_general_strings, final_categorized_output, categorized, rating, character_res, general_res, unclassified, enhanced_tags])
             # Event to remove a selected image from the gallery
             remove_button.click(remove_image_from_gallery, inputs=[gallery, selected_image], outputs=gallery)
+            # Event to for the Prompt Enhancer Button
+            pe_generate_btn.click(lambda tags,model:prompt_enhancer('','',tags,model)[0],inputs=[final_categorized_output,prompt_enhancer_model],outputs=[enhanced_tags])
         submit.click(
             predictor.predict,
             inputs=[
@@ -543,7 +552,7 @@ with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True
                 character_mcut_enabled,
             ],
         )
-    with gr.Tab(label="Tag Categorizer"): 
+    with gr.Tab(label="Tag Categorizer + Enhancer"): 
        with gr.Row():
             with gr.Column(variant="panel"):
                 input_tags = gr.Textbox(label="Input Tags (Danbooru comma-separated)", placeholder="1girl, cat, horns, blue hair, ...")
@@ -551,7 +560,12 @@ with gr.Blocks(title=TITLE, css=css, theme="Werli/Multi-Tagger", fill_width=True
             with gr.Column(variant="panel"):
                 categorized_string = gr.Textbox(label="Categorized (string)", show_label=True, show_copy_button=True, lines=8)
                 categorized_json = gr.JSON(label="Categorized (tags) - JSON")
-            submit_button.click(process_tags, inputs=[input_tags], outputs=[categorized_string, categorized_json])
+                submit_button.click(process_tags, inputs=[input_tags], outputs=[categorized_string, categorized_json])
+            with gr.Column(variant="panel"):
+                pe_generate_btn = gr.Button(value="ENHANCE TAGS", size="lg", variant="primary")
+                enhanced_tags = gr.Textbox(label="Enhanced Tags", show_label=True, show_copy_button=True)
+                prompt_enhancer_model = gr.Radio(["Medium", "Long", "Flux"], label="Model Choice", value="Medium", info="Enhance your prompts with Medium or Long answers")
+                pe_generate_btn.click(lambda tags,model:prompt_enhancer('','',tags,model)[0],inputs=[categorized_string,prompt_enhancer_model],outputs=[enhanced_tags])
     with gr.Tab(label="Florence 2 Image Captioning"):
         with gr.Row():
             with gr.Column(variant="panel"):

modules/classifyTags.py

@@ -171,9 +171,6 @@ def process_tags(input_tags: str):
     categorized_string = ', '.join([tag for category in classified_tags.values() for tag in category])
     categorized_json = {category: tags for category, tags in classified_tags.items()}
 
-    return categorized_string, categorized_json
+    return categorized_string, categorized_json, ""  # Initialize enhanced_prompt as empty
 
-tags = []
-if __name__ == "__main__":
-    classify_tags (tags, True)
-    process_tags(input_tags)
+tags = []

modules/florence2.py

@@ -94,9 +94,4 @@ def update_task_dropdown(choice):
     if choice == 'Cascaded task':
         return gr.Dropdown(choices=cascaded_task_list, value='Caption + Grounding')
     else:
-        return gr.Dropdown(choices=single_task_list, value='Caption')
-
-if __name__ == "__main__":
-    process_image()
-    single_task_list
-    update_task_dropdown()
+        return gr.Dropdown(choices=single_task_list, value='Caption')

modules/llama_loader.py

@@ -182,8 +182,4 @@ class llama3reorganize:
             except Exception as e:print(traceback.format_exc());print('Error reorganize text: '+str(e))
             return result
 
-llama_list=[META_LLAMA_3_3B_REPO,META_LLAMA_3_8B_REPO]
-
-if __name__ == "__main__":
-    llama3reorganize()
-    llama_list
+llama_list=[META_LLAMA_3_3B_REPO,META_LLAMA_3_8B_REPO]

modules/tag_enhancer.py

@@ -0,0 +1,53 @@
+import gradio as gr
+from transformers import pipeline,AutoTokenizer,AutoModelForSeq2SeqLM
+import re,torch
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def load_models():
+    try:
+        enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device=device)
+        enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device=device)
+        model_checkpoint = "gokaygokay/Flux-Prompt-Enhance"
+        tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint).eval().to(device=device)
+        enhancer_flux = pipeline('text2text-generation', model=model, tokenizer=tokenizer, repetition_penalty=1.5, device=device)
+    except Exception as e:
+        print(e)
+        enhancer_medium = enhancer_long = enhancer_flux = None
+    return enhancer_medium, enhancer_long, enhancer_flux
+
+enhancer_medium, enhancer_long, enhancer_flux = load_models()
+
+def enhance_prompt(input_prompt, model_choice):
+    if model_choice == "Medium":
+        result = enhancer_medium("Enhance the description: " + input_prompt)
+        enhanced_text = result[0]['summary_text']
+        
+        pattern = r'^.*?of\s+(.*?(?:\.|$))'
+        match = re.match(pattern, enhanced_text, re.IGNORECASE | re.DOTALL)
+        
+        if match:
+            remaining_text = enhanced_text[match.end():].strip()
+            modified_sentence = match.group(1).capitalize()
+            enhanced_text = modified_sentence + ' ' + remaining_text
+    elif model_choice == "Flux":
+        result = enhancer_flux("enhance prompt: " + input_prompt, max_length=256)
+        enhanced_text = result[0]['generated_text']
+    else:  # Long
+        result = enhancer_long("Enhance the description: " + input_prompt)
+        enhanced_text = result[0]['summary_text']
+    
+    return enhanced_text
+
+def prompt_enhancer(character: str, series: str, general: str, model_choice: str):
+    characters = character.split(",") if character else []
+    serieses = series.split(",") if series else []
+    generals = general.split(",") if general else []
+    tags = characters + serieses + generals
+    cprompt = ",".join(tags) if tags else ""
+
+    output = enhance_prompt(cprompt, model_choice)
+    prompt = cprompt + ", " + output
+
+    return prompt, gr.update(interactive=True), gr.update(interactive=True)