Upload fl2basepromptgen.py

fl2basepromptgen.py

@@ -1,6 +1,5 @@
 from transformers import AutoProcessor, AutoModelForCausalLM
 import spaces
-import re
 from PIL import Image 
 
 #import subprocess
@@ -10,37 +9,8 @@ fl_model = AutoModelForCausalLM.from_pretrained('MiaoshouAI/Florence-2-base-Prom
 fl_processor = AutoProcessor.from_pretrained('MiaoshouAI/Florence-2-base-PromptGen', trust_remote_code=True)
 
 
-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
-def fl_run_example(image):
+def fl_run(image):
     task_prompt = "<GENERATE_PROMPT>"
     prompt = task_prompt + "Describe this image in great detail."
 
@@ -58,8 +28,7 @@ def fl_run_example(image):
     )
     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))
-    print(parsed_answer)
-    return fl_modify_caption(parsed_answer["<PROMPT>"])
+    return parsed_answer["<GENERATE_PROMPT>Describe this image in great detail."]
 
 
 def predict_tags_fl2_base_prompt_gen(image: Image.Image, input_tags: str, algo: list[str]):
@@ -71,6 +40,6 @@ def predict_tags_fl2_base_prompt_gen(image: Image.Image, input_tags: str, algo:
     
     if not "Use Florence-2-base-PromptGen" in algo:
         return input_tags
-    tag_list = list_uniq(to_list(input_tags) + to_list(fl_run_example(image) + ", "))
+    tag_list = list_uniq(to_list(input_tags) + to_list(fl_run(image) + ", "))
     tag_list.remove("")
     return ", ".join(tag_list)