Create app.py

app.py

@@ -0,0 +1,95 @@
+import gradio as gr
+from transformers import AutoProcessor, AutoModelForCausalLM
+import re
+from PIL import Image 
+
+import subprocess
+subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+model = AutoModelForCausalLM.from_pretrained('vumichien/Florence-2-FT-Caption', trust_remote_code=True).to("cpu").eval()
+
+processor = AutoProcessor.from_pretrained('vumichien/Florence-2-FT-Caption', trust_remote_code=True)
+
+
+TITLE = "# [Florence-2 Captioner](https://huggingface.co/vumichien/Florence-2-FT-Caption)"
+
+
+def 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 run_example(image):
+    image = Image.fromarray(image)
+    task_prompt = "<CAPTION>"
+    prompt = task_prompt
+
+    # Ensure the image is in RGB mode
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    inputs = processor(text=prompt, images=image, return_tensors="pt").to("cpu")
+    generated_ids = model.generate(
+        input_ids=inputs["input_ids"],
+        pixel_values=inputs["pixel_values"],
+        max_new_tokens=1024,
+        num_beams=3
+    )
+    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 modify_caption(parsed_answer["<CAPTION>"])
+
+
+css = """
+  #output {
+    height: 500px; 
+    overflow: auto; 
+    border: 1px solid #ccc; 
+  }
+"""
+examples = ["240617143250078.JPG", "240617144124216.JPG", "240617144154631.JPG", "240617160707848.JPG"]
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown(TITLE)
+    with gr.Tab(label="Florence-2 SD3 Prompts"):
+        with gr.Row():
+            with gr.Column():
+                input_img = gr.Image(label="Input Picture",height=400)
+                submit_btn = gr.Button(value="Submit")
+            with gr.Column():
+                output_text = gr.Textbox(label="Output Text")
+
+        submit_btn.click(run_example, [input_img], [output_text])
+        examples = gr.Examples(
+            examples,
+            fn=run_example,
+            inputs=[input_img],
+            outputs=output,
+            cache_examples=True,
+        )
+
+demo.launch(debug=True)