Upload 10 files

README.md

@@ -1,10 +1,10 @@
 ---
-title: Prompt Enhancer with WD Tagger & Florence 2 SD3 Captioner
+title: Prompt Enhancer with WD Tagger & Florence 2 Flux/SD3 Captioner
 emoji: 🏃📦
 colorFrom: blue
 colorTo: yellow
 sdk: gradio
-sdk_version: 4.39.0
+sdk_version: 4.42.0
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -16,12 +16,13 @@ from tagger import (
 )
 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 & SD3 Long Captioner
+## Prompt Enhancer with WD Tagger & Flux/SD3 Captioner
 (Image =>) Prompt => Upsampled longer prompt
 """
     )
@@ -33,8 +34,11 @@ def description_ui2():
  [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-base-PromptGen](https://huggingface.co/MiaoshouAI/Florence-2-base-PromptGen).
 """
     )
@@ -51,7 +55,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"], label="Algorithms", value=["Use WD Tagger", "Use Florence-2-SD3-Long-Captioner"])
+                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-SD3-Long-Captioner"])
                 generate_from_image_btn = gr.Button(value="GENERATE TAGS FROM IMAGE", size="lg", variant="primary")
             with gr.Group():
                 with gr.Row():
@@ -98,6 +102,8 @@ 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(

fl2basepromptgen.py

@@ -7,11 +7,15 @@ import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-fl_model = AutoModelForCausalLM.from_pretrained('MiaoshouAI/Florence-2-base-PromptGen', trust_remote_code=True).to(device).eval()
-fl_processor = AutoProcessor.from_pretrained('MiaoshouAI/Florence-2-base-PromptGen', trust_remote_code=True)
 
+try:
+    fl_model = AutoModelForCausalLM.from_pretrained('MiaoshouAI/Florence-2-base-PromptGen', trust_remote_code=True).to("cpu").eval()
+    fl_processor = AutoProcessor.from_pretrained('MiaoshouAI/Florence-2-base-PromptGen', trust_remote_code=True)
+except Exception as e:
+    print(e)
+    fl_model = fl_processor = None
 
-@spaces.GPU
+@spaces.GPU(duration=30)
 def fl_run(image):
     task_prompt = "<GENERATE_PROMPT>"
     prompt = task_prompt + "Describe this image in great detail."
@@ -20,6 +24,7 @@ def fl_run(image):
     if image.mode != "RGB":
         image = image.convert("RGB")
 
+    fl_model.to(device)
     inputs = fl_processor(text=prompt, images=image, return_tensors="pt").to(device)
     generated_ids = fl_model.generate(
         input_ids=inputs["input_ids"],
@@ -28,6 +33,7 @@ def fl_run(image):
         do_sample=False,
         num_beams=3
     )
+    fl_model.to("cpu")
     generated_text = fl_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
     parsed_answer = fl_processor.post_process_generation(generated_text, task=prompt, image_size=(image.width, image.height))
     return parsed_answer["<GENERATE_PROMPT>Describe this image in great detail."]

fl2flux.py

@@ -0,0 +1,90 @@
+from transformers import AutoProcessor, AutoModelForCausalLM
+import spaces
+import re
+from PIL import Image 
+import torch
+
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+try:
+    fl_model = AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-Flux', trust_remote_code=True).to("cpu").eval()
+    fl_processor = AutoProcessor.from_pretrained('gokaygokay/Florence-2-Flux', trust_remote_code=True)
+    fl_model_large = AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-Flux-Large', trust_remote_code=True).to("cpu").eval()
+    fl_processor_large = AutoProcessor.from_pretrained('gokaygokay/Florence-2-Flux-Large', trust_remote_code=True)
+except Exception as e:
+    fl_model = fl_processor = fl_model_large = fl_processor_large = None
+    print(e)
+
+def fl_modify_caption(caption: str) -> str:
+    """
+    Removes specific prefixes from captions if present, otherwise returns the original caption.
+    Args:
+        caption (str): A string containing a caption.
+    Returns:
+        str: The caption with the prefix removed if it was present, or the original caption.
+    """
+    # Define the prefixes to remove
+    prefix_substrings = [
+        ('captured from ', ''),
+        ('captured at ', '')
+    ]
+    
+    # Create a regex pattern to match any of the prefixes
+    pattern = '|'.join([re.escape(opening) for opening, _ in prefix_substrings])
+    replacers = {opening.lower(): replacer for opening, replacer in prefix_substrings}
+    
+    # Function to replace matched prefix with its corresponding replacement
+    def replace_fn(match):
+        return replacers[match.group(0).lower()]
+    
+    # Apply the regex to the caption
+    modified_caption = re.sub(pattern, replace_fn, caption, count=1, flags=re.IGNORECASE)
+    
+    # If the caption was modified, return the modified version; otherwise, return the original
+    return modified_caption if modified_caption != caption else caption
+
+
+@spaces.GPU(duration=30)
+def fl_run_example(image, algo):
+    task_prompt = "<DESCRIPTION>"
+    prompt = task_prompt + "Describe this image in great detail."
+    #prompt = task_prompt
+
+    # Ensure the image is in RGB mode
+    if image.mode != "RGB": image = image.convert("RGB")
+
+    if algo == "Use Florence-2-Flux-Large":
+        model = fl_model_large
+        processor = fl_processor_large
+    else:
+        model = fl_model
+        processor = fl_processor
+    model.to(device)
+    inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
+    generated_ids = model.generate(
+        input_ids=inputs["input_ids"],
+        pixel_values=inputs["pixel_values"],
+        max_new_tokens=1024,
+        num_beams=3
+    )
+    model.to("cpu")
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    parsed_answer = processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
+    return fl_modify_caption(parsed_answer["<DESCRIPTION>"])
+
+
+def predict_tags_fl2_flux(image: Image.Image, input_tags: str, algo: list[str]):
+    def to_list(s):
+        return [x.strip() for x in s.split(",") if not s == ""]
+    
+    def list_uniq(l):
+        return sorted(set(l), key=l.index)
+    
+    if "Use Florence-2-Flux" not in algo and "Use Florence-2-Flux-Large" not in algo:
+        return input_tags
+    tag_list = list_uniq(to_list(input_tags) + to_list(fl_run_example(image, algo) + ", "))
+    tag_list.remove("")
+    return ", ".join(tag_list)

fl2sd3longcap.py

@@ -8,9 +8,13 @@ import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-fl_model = AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True).to(device).eval()
-fl_processor = AutoProcessor.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True)
 
+try:
+    fl_model = AutoModelForCausalLM.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True).to("cpu").eval()
+    fl_processor = AutoProcessor.from_pretrained('gokaygokay/Florence-2-SD3-Captioner', trust_remote_code=True)
+except Exception as e:
+    print(e)
+    fl_model = fl_processor = None
 
 def fl_modify_caption(caption: str) -> str:
     """
@@ -41,7 +45,7 @@ def fl_modify_caption(caption: str) -> str:
     return modified_caption if modified_caption != caption else caption
 
 
-@spaces.GPU
+@spaces.GPU(duration=30)
 def fl_run_example(image):
     task_prompt = "<DESCRIPTION>"
     prompt = task_prompt + "Describe this image in great detail."
@@ -50,6 +54,7 @@ def fl_run_example(image):
     if image.mode != "RGB":
         image = image.convert("RGB")
 
+    fl_model.to(device)
     inputs = fl_processor(text=prompt, images=image, return_tensors="pt").to(device)
     generated_ids = fl_model.generate(
         input_ids=inputs["input_ids"],
@@ -57,6 +62,7 @@ def fl_run_example(image):
         max_new_tokens=1024,
         num_beams=3
     )
+    fl_model.to("cpu")
     generated_text = fl_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
     parsed_answer = fl_processor.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
     return fl_modify_caption(parsed_answer["<DESCRIPTION>"])

promptenhancer.py

@@ -1,22 +1,32 @@
 import spaces
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import re
 import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def load_models():
-    enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device=device)
-    enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device=device)
-    return enhancer_medium, enhancer_long
+    try:
+        enhancer_medium = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device="cpu")
+        enhancer_long = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device="cpu")
+        model_checkpoint = "gokaygokay/Flux-Prompt-Enhance"
+        tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint).eval().to(device="cpu")
+        enhancer_flux = pipeline('text2text-generation', model=model, tokenizer=tokenizer, repetition_penalty=1.5, device="cpu")
+    except Exception as e:
+        print(e)
+        enhancer_medium = enhancer_long = enhancer_flux = None
+    return enhancer_medium, enhancer_long, enhancer_flux
 
-enhancer_medium, enhancer_long = load_models()
+enhancer_medium, enhancer_long, enhancer_flux = load_models()
 
 @spaces.GPU
 def enhance_prompt(input_prompt, model_choice):
     if model_choice == "Medium":
+        enhancer_medium.to(device=device)
         result = enhancer_medium("Enhance the description: " + input_prompt)
+        enhancer_medium.to(device="cpu")
         enhanced_text = result[0]['summary_text']
         
         pattern = r'^.*?of\s+(.*?(?:\.|$))'
@@ -26,8 +36,15 @@ def enhance_prompt(input_prompt, model_choice):
             remaining_text = enhanced_text[match.end():].strip()
             modified_sentence = match.group(1).capitalize()
             enhanced_text = modified_sentence + ' ' + remaining_text
+    elif model_choice == "Flux":
+        enhancer_flux.to(device=device)
+        result = enhancer_flux("enhance prompt: " + input_prompt, max_length = 256)
+        enhancer_flux.to(device="cpu")
+        enhanced_text = result[0]['generated_text']
     else:  # Long
+        enhancer_long.to(device=device)
         result = enhancer_long("Enhance the description: " + input_prompt)
+        enhancer_long.to(device="cpu")
         enhanced_text = result[0]['summary_text']
     
     return enhanced_text

tagger.py

@@ -12,10 +12,15 @@ from pathlib import Path
 WD_MODEL_NAMES = ["p1atdev/wd-swinv2-tagger-v3-hf"]
 WD_MODEL_NAME = WD_MODEL_NAMES[0]
 
-wd_model = AutoModelForImageClassification.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
-wd_model.to("cuda" if torch.cuda.is_available() else "cpu")
-wd_processor = AutoImageProcessor.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+default_device = device
 
+try:
+    wd_model = AutoModelForImageClassification.from_pretrained(WD_MODEL_NAME, trust_remote_code=True).to(default_device).eval()
+    wd_processor = AutoImageProcessor.from_pretrained(WD_MODEL_NAME, trust_remote_code=True)
+except Exception as e:
+    print(e)
+    wd_model = wd_processor = None
 
 def _people_tag(noun: str, minimum: int = 1, maximum: int = 5):
     return (
@@ -506,7 +511,7 @@ def gen_prompt(rating: list[str], character: list[str], general: list[str]):
     return ", ".join(all_tags)
 
 
-@spaces.GPU()
+@spaces.GPU(duration=30)
 def predict_tags(image: Image.Image, general_threshold: float = 0.3, character_threshold: float = 0.8):
     inputs = wd_processor.preprocess(image, return_tensors="pt")
 
@@ -514,9 +519,11 @@ def predict_tags(image: Image.Image, general_threshold: float = 0.3, character_t
     logits = torch.sigmoid(outputs.logits[0])  # take the first logits
 
     # get probabilities
+    if device != default_device: wd_model.to(device=device)
     results = {
         wd_model.config.id2label[i]: float(logit.float()) for i, logit in enumerate(logits)
     }
+    if device != default_device: wd_model.to(device=default_device)
     # rating, character, general
     rating, character, general = postprocess_results(
         results, general_threshold, character_threshold