Update app.py

app.py

@@ -6,6 +6,36 @@ import torch
 model = PaliGemmaForConditionalGeneration.from_pretrained("gokaygokay/sd3-long-captioner").to("cuda").eval()
 processor = PaliGemmaProcessor.from_pretrained("gokaygokay/sd3-long-captioner")
 
+import re
+
+def modify_caption(caption: str) -> str:
+    """
+    Removes specific prefixes from captions.
+
+    Args:
+        caption (str): A string containing a caption.
+
+    Returns:
+        str: The caption with the prefix removed if it was present.
+    """
+    # 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: 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)]
+    
+    # Apply the regex to the caption
+    return re.sub(pattern, replace_fn, caption, count=1, flags=re.IGNORECASE).capitalize()
+
+# Example usage in your existing function
 @spaces.GPU
 def create_captions_rich(image):   
     prompt = "caption en"
@@ -16,7 +46,10 @@ def create_captions_rich(image):
         generation = model.generate(**model_inputs, max_new_tokens=256, do_sample=False)
         generation = generation[0][input_len:]
         decoded = processor.decode(generation, skip_special_tokens=True)
-    return decoded
+
+        # Modify the caption to remove specific prefixes
+        modified_caption = modify_caption(decoded)
+    return modified_caption
 
 css = """
   #mkd {