Support CPU

app.py

@@ -1,32 +1,55 @@
-import os
-from unittest.mock import patch
-import spaces
 import gradio as gr
 from transformers import AutoProcessor, AutoModelForCausalLM
-from transformers.dynamic_module_utils import get_imports
-import torch
+import spaces
+
 import requests
-from PIL import Image, ImageDraw
-import random
-import numpy as np
+import copy
+
+from PIL import Image, ImageDraw, ImageFont 
+import io
 import matplotlib.pyplot as plt
 import matplotlib.patches as patches
-import cv2
-import io
 
-def workaround_fixed_get_imports(filename: str | os.PathLike) -> list[str]:
+import random
+import numpy as np
+
+import os
+from unittest.mock import patch
+from transformers import AutoModelForCausalLM, AutoProcessor
+from transformers.dynamic_module_utils import get_imports
+
+def fixed_get_imports(filename: str | os.PathLike) -> list[str]:
+    """Work around for https://huggingface.co/microsoft/phi-1_5/discussions/72."""
     if not str(filename).endswith("/modeling_florence2.py"):
         return get_imports(filename)
     imports = get_imports(filename)
     imports.remove("flash_attn")
     return imports
 
-with patch("transformers.dynamic_module_utils.get_imports", workaround_fixed_get_imports):
-    model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-large-ft", trust_remote_code=True).to("cuda").eval()
+
+@spaces.GPU
+def get_device_type():
+    import torch
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+model_id = 'microsoft/Florence-2-base-ft'
+
+import subprocess
+device = get_device_type()
+if (device == "cuda"):
+    subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+    model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-large-ft", trust_remote_code=True)
     processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large-ft", trust_remote_code=True)
+else:
+    #https://huggingface.co/microsoft/Florence-2-large-ft/discussions/4
+    with patch("transformers.dynamic_module_utils.get_imports", fixed_get_imports):
+        model = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-large-ft", trust_remote_code=True)
+        processor = AutoProcessor.from_pretrained("microsoft/Florence-2-large-ft", trust_remote_code=True)
+
+DESCRIPTION = "# [Florence-2 base-ft Demo with CPU inference support](https://huggingface.co/microsoft/Florence-2-base-ft)"
 
-colormap = ['blue', 'orange', 'green', 'purple', 'brown', 'pink', 'gray', 'olive', 'cyan', 'red',
-            'lime', 'indigo', 'violet', 'aqua', 'magenta', 'coral', 'gold', 'tan', 'skyblue']
+colormap = ['blue','orange','green','purple','brown','pink','gray','olive','cyan','red',
+            'lime','indigo','violet','aqua','magenta','coral','gold','tan','skyblue']
 
 def fig_to_pil(fig):
     buf = io.BytesIO()
@@ -40,21 +63,20 @@ def run_example(task_prompt, image, text_input=None):
         prompt = task_prompt
     else:
         prompt = task_prompt + text_input
-    inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda")
-    with torch.inference_mode():
-        generated_ids = model.generate(
-            input_ids=inputs["input_ids"],
-            pixel_values=inputs["pixel_values"],
-            max_new_tokens=1024,
-            early_stopping=False,
-            do_sample=False,
-            num_beams=3,
-        )
+    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,
+        early_stopping=False,
+        do_sample=False,
+        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.size[0], image.size[1])
+        image_size=(image.width, image.height)
     )
     return parsed_answer
 
@@ -65,161 +87,164 @@ def plot_bbox(image, data):
         x1, y1, x2, y2 = bbox
         rect = patches.Rectangle((x1, y1), x2-x1, y2-y1, linewidth=1, edgecolor='r', facecolor='none')
         ax.add_patch(rect)
-        plt.text(x1, y1, label, color='white', fontsize=8, bbox=dict(facecolor='indigo', alpha=0.5))
+        plt.text(x1, y1, label, color='white', fontsize=8, bbox=dict(facecolor='red', alpha=0.5))
     ax.axis('off')
-    return fig_to_pil(fig)
+    return fig
 
 def draw_polygons(image, prediction, fill_mask=False):
-    fig, ax = plt.subplots()
-    ax.imshow(image)
+    draw = ImageDraw.Draw(image)
     scale = 1
     for polygons, label in zip(prediction['polygons'], prediction['labels']):
         color = random.choice(colormap)
         fill_color = random.choice(colormap) if fill_mask else None
         for _polygon in polygons:
             _polygon = np.array(_polygon).reshape(-1, 2)
-            if _polygon.shape[0] < 3:
+            if len(_polygon) < 3:
                 print('Invalid polygon:', _polygon)
                 continue
             _polygon = (_polygon * scale).reshape(-1).tolist()
-            if len(_polygon) % 2 != 0:
-                print('Invalid polygon:', _polygon)
-                continue
-            polygon_points = np.array(_polygon).reshape(-1, 2)
             if fill_mask:
-                polygon = patches.Polygon(polygon_points, edgecolor=color, facecolor=fill_color, linewidth=2)
+                draw.polygon(_polygon, outline=color, fill=fill_color)
             else:
-                polygon = patches.Polygon(polygon_points, edgecolor=color, fill=False, linewidth=2)
-            ax.add_patch(polygon)
-        plt.text(polygon_points[0, 0], polygon_points[0, 1], label, color='white', fontsize=8, bbox=dict(facecolor=color, alpha=0.5))
-    ax.axis('off')
-    return fig_to_pil(fig)
+                draw.polygon(_polygon, outline=color)
+            draw.text((_polygon[0] + 8, _polygon[1] + 2), label, fill=color)
+    return image
+
+def convert_to_od_format(data):
+    bboxes = data.get('bboxes', [])
+    labels = data.get('bboxes_labels', [])
+    od_results = {
+        'bboxes': bboxes,
+        'labels': labels
+    }
+    return od_results
 
 def draw_ocr_bboxes(image, prediction):
-    fig, ax = plt.subplots()
-    ax.imshow(image)
     scale = 1
+    draw = ImageDraw.Draw(image)
     bboxes, labels = prediction['quad_boxes'], prediction['labels']
     for box, label in zip(bboxes, labels):
         color = random.choice(colormap)
         new_box = (np.array(box) * scale).tolist()
-        polygon = patches.Polygon(new_box, edgecolor=color, fill=False, linewidth=3)
-        ax.add_patch(polygon)
-        plt.text(new_box[0], new_box[1], label, color='white', fontsize=8, bbox=dict(facecolor=color, alpha=0.5))
-    ax.axis('off')
-    return fig_to_pil(fig)
-
-
-@spaces.GPU(duration=120)
-def process_video(input_video_path, task_prompt):
-    cap = cv2.VideoCapture(input_video_path)
-    if not cap.isOpened():
-        print("Error: Can't open the video file.")
-        return
-
-    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter("output_vid.mp4", fourcc, fps, (frame_width, frame_height))
-
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        pil_image = Image.fromarray(frame_rgb)
-
-        result = run_example(task_prompt, pil_image)
-
-        if task_prompt == "<OD>":
-            processed_image = plot_bbox(pil_image, result['<OD>'])
-        elif task_prompt == "<DENSE_REGION_CAPTION>":
-            processed_image = plot_bbox(pil_image, result['<DENSE_REGION_CAPTION>'])
-        else:
-            processed_image = pil_image
-
-        processed_frame = cv2.cvtColor(np.array(processed_image), cv2.COLOR_RGB2BGR)
-        out.write(processed_frame)
-
-    cap.release()
-    out.release()
-    cv2.destroyAllWindows()
-    return "output_vid.mp4"
+        draw.polygon(new_box, width=3, outline=color)
+        draw.text((new_box[0]+8, new_box[1]+2),
+                  "{}".format(label),
+                  align="right",
+                  fill=color)
+    return image
+
+def process_image(image, task_prompt, text_input=None):
+    image = Image.fromarray(image)  # Convert NumPy array to PIL Image
+    if task_prompt == 'Caption':
+        task_prompt = '<CAPTION>'
+        result = run_example(task_prompt, image)
+        return result, None
+    elif task_prompt == 'Detailed Caption':
+        task_prompt = '<DETAILED_CAPTION>'
+        result = run_example(task_prompt, image)
+        return result, None
+    elif task_prompt == 'More Detailed Caption':
+        task_prompt = '<MORE_DETAILED_CAPTION>'
+        result = run_example(task_prompt, image)
+        return result, None
+    elif task_prompt == 'Object Detection':
+        task_prompt = '<OD>'
+        results = run_example(task_prompt, image)
+        fig = plot_bbox(image, results['<OD>'])
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'Dense Region Caption':
+        task_prompt = '<DENSE_REGION_CAPTION>'
+        results = run_example(task_prompt, image)
+        fig = plot_bbox(image, results['<DENSE_REGION_CAPTION>'])
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'Region Proposal':
+        task_prompt = '<REGION_PROPOSAL>'
+        results = run_example(task_prompt, image)
+        fig = plot_bbox(image, results['<REGION_PROPOSAL>'])
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'Caption to Phrase Grounding':
+        task_prompt = '<CAPTION_TO_PHRASE_GROUNDING>'
+        results = run_example(task_prompt, image, text_input)
+        fig = plot_bbox(image, results['<CAPTION_TO_PHRASE_GROUNDING>'])
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'Referring Expression Segmentation':
+        task_prompt = '<REFERRING_EXPRESSION_SEGMENTATION>'
+        results = run_example(task_prompt, image, text_input)
+        output_image = copy.deepcopy(image)
+        output_image = draw_polygons(output_image, results['<REFERRING_EXPRESSION_SEGMENTATION>'], fill_mask=True)
+        return results, output_image
+    elif task_prompt == 'Region to Segmentation':
+        task_prompt = '<REGION_TO_SEGMENTATION>'
+        results = run_example(task_prompt, image, text_input)
+        output_image = copy.deepcopy(image)
+        output_image = draw_polygons(output_image, results['<REGION_TO_SEGMENTATION>'], fill_mask=True)
+        return results, output_image
+    elif task_prompt == 'Open Vocabulary Detection':
+        task_prompt = '<OPEN_VOCABULARY_DETECTION>'
+        results = run_example(task_prompt, image, text_input)
+        bbox_results = convert_to_od_format(results['<OPEN_VOCABULARY_DETECTION>'])
+        fig = plot_bbox(image, bbox_results)
+        return results, fig_to_pil(fig)
+    elif task_prompt == 'Region to Category':
+        task_prompt = '<REGION_TO_CATEGORY>'
+        results = run_example(task_prompt, image, text_input)
+        return results, None
+    elif task_prompt == 'Region to Description':
+        task_prompt = '<REGION_TO_DESCRIPTION>'
+        results = run_example(task_prompt, image, text_input)
+        return results, None
+    elif task_prompt == 'OCR':
+        task_prompt = '<OCR>'
+        result = run_example(task_prompt, image)
+        return result, None
+    elif task_prompt == 'OCR with Region':
+        task_prompt = '<OCR_WITH_REGION>'
+        results = run_example(task_prompt, image)
+        output_image = copy.deepcopy(image)
+        output_image = draw_ocr_bboxes(output_image, results['<OCR_WITH_REGION>'])
+        return results, output_image
+    else:
+        return "", None  # Return empty string and None for unknown task prompts
 
 css = """
-#output {
-    min-height: 100px;
-    overflow: auto;
-    border: 1px solid #ccc;
-}
+  #output {
+    height: 500px; 
+    overflow: auto; 
+    border: 1px solid #ccc; 
+  }
 """
 
 with gr.Blocks(css=css) as demo:
-    gr.HTML("<h1><center>Microsoft Florence-2-large-ft</center></h1>")
-    with gr.Tab(label="Image"):
+    gr.Markdown(DESCRIPTION)
+    with gr.Tab(label="Florence-2 Image Captioning"):
         with gr.Row():
             with gr.Column():
-                input_img = gr.Image(label="Input Picture", type="pil")
-                task_radio = gr.Radio(
-                    ["Caption", "Detailed Caption", "More Detailed Caption", "Caption to Phrase Grounding",
-                     "Object Detection", "Dense Region Caption", "Region Proposal", "Referring Expression Segmentation",
-                     "Region to Segmentation", "Open Vocabulary Detection", "Region to Category", "Region to Description",
-                     "OCR", "OCR with Region"],
-                    label="Task", value="Caption"
-                )
-                text_input = gr.Textbox(label="Text Input (is Optional)", visible=False)
+                input_img = gr.Image(label="Input Picture")
+                task_prompt = gr.Dropdown(choices=[
+                    'Caption', 'Detailed Caption', 'More Detailed Caption', 'Object Detection',
+                    'Dense Region Caption', 'Region Proposal', 'Caption to Phrase Grounding',
+                    'Referring Expression Segmentation', 'Region to Segmentation',
+                    'Open Vocabulary Detection', 'Region to Category', 'Region to Description',
+                    'OCR', 'OCR with Region'
+                ], label="Task Prompt", value= 'Caption')
+                text_input = gr.Textbox(label="Text Input (optional)")
                 submit_btn = gr.Button(value="Submit")
             with gr.Column():
-                output_text = gr.Textbox(label="Results")
-                output_image = gr.Image(label="Image", type="pil")
+                output_text = gr.Textbox(label="Output Text")
+                output_img = gr.Image(label="Output Image")
+
+        gr.Examples(
+            examples=[
+                ["image1.jpg", 'Object Detection'],
+                ["image2.jpg", 'OCR with Region']
+            ],
+            inputs=[input_img, task_prompt],
+            outputs=[output_text, output_img],
+            fn=process_image,
+            cache_examples=True,
+            label='Try examples'
+        )
 
-    with gr.Tab(label="Video"):
-        with gr.Row():
-            with gr.Column():
-                input_video = gr.Video(label="Video")
-                video_task_radio = gr.Radio(
-                    ["Object Detection", "Dense Region Caption"],
-                    label="Video Task", value="Object Detection"
-                )
-                video_submit_btn = gr.Button(value="Process Video")
-            with gr.Column():
-                output_video = gr.Video(label="Video")
-
-    def update_text_input(task):
-        return gr.update(visible=task in ["Caption to Phrase Grounding", "Referring Expression Segmentation",
-                                           "Region to Segmentation", "Open Vocabulary Detection", "Region to Category",
-                                           "Region to Description"])
-
-    task_radio.change(fn=update_text_input, inputs=task_radio, outputs=text_input)
-
-    def process_image(image, task, text):
-        task_mapping = {
-            "Caption": ("<CAPTION>", lambda result: (result['<CAPTION>'], image)),
-            "Detailed Caption": ("<DETAILED_CAPTION>", lambda result: (result['<DETAILED_CAPTION>'], image)),
-            "More Detailed Caption": ("<MORE_DETAILED_CAPTION>", lambda result: (result['<MORE_DETAILED_CAPTION>'], image)),
-            "Caption to Phrase Grounding": ("<CAPTION_TO_PHRASE_GROUNDING>", lambda result: (str(result['<CAPTION_TO_PHRASE_GROUNDING>']), plot_bbox(image, result['<CAPTION_TO_PHRASE_GROUNDING>']))),
-            "Object Detection": ("<OD>", lambda result: (str(result['<OD>']), plot_bbox(image, result['<OD>']))),
-            "Dense Region Caption": ("<DENSE_REGION_CAPTION>", lambda result: (str(result['<DENSE_REGION_CAPTION>']), plot_bbox(image, result['<DENSE_REGION_CAPTION>']))),
-            "Region Proposal": ("<REGION_PROPOSAL>", lambda result: (str(result['<REGION_PROPOSAL>']), plot_bbox(image, result['<REGION_PROPOSAL>']))),
-            "Referring Expression Segmentation": ("<REFERRING_EXPRESSION_SEGMENTATION>", lambda result: (str(result['<REFERRING_EXPRESSION_SEGMENTATION>']), draw_polygons(image, result['<REFERRING_EXPRESSION_SEGMENTATION>'], fill_mask=True))),
-            "Region to Segmentation": ("<REGION_TO_SEGMENTATION>", lambda result: (str(result['<REGION_TO_SEGMENTATION>']), draw_polygons(image, result['<REGION_TO_SEGMENTATION>'], fill_mask=True))),
-            "Open Vocabulary Detection": ("<OPEN_VOCABULARY_DETECTION>", lambda result: (str(convert_to_od_format(result['<OPEN_VOCABULARY_DETECTION>'])), plot_bbox(image, convert_to_od_format(result['<OPEN_VOCABULARY_DETECTION>'])))),
-            "Region to Category": ("<REGION_TO_CATEGORY>", lambda result: (result['<REGION_TO_CATEGORY>'], image)),
-            "Region to Description": ("<REGION_TO_DESCRIPTION>", lambda result: (result['<REGION_TO_DESCRIPTION>'], image)),
-            "OCR": ("<OCR>", lambda result: (result['<OCR>'], image)),
-            "OCR with Region": ("<OCR_WITH_REGION>", lambda result: (str(result['<OCR_WITH_REGION>']), draw_ocr_bboxes(image, result['<OCR_WITH_REGION>']))),
-        }
-
-        if task in task_mapping:
-            prompt, process_func = task_mapping[task]
-            result = run_example(prompt, image, text)
-            return process_func(result)
-        else:
-            return "", image
-
-    submit_btn.click(fn=process_image, inputs=[input_img, task_radio, text_input], outputs=[output_text, output_image])
-    video_submit_btn.click(fn=process_video, inputs=[input_video, video_task_radio], outputs=output_video)
-
-demo.launch()
+        submit_btn.click(process_image, [input_img, task_prompt, text_input], [output_text, output_img])
+
+demo.launch(debug=True)