Changed model for DocVQA and added task

app.py

@@ -16,22 +16,12 @@ import numpy as np
 import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-models = {
-    'microsoft/Florence-2-large-ft': AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-large-ft', trust_remote_code=True).to("cuda").eval(),
-    'microsoft/Florence-2-large': AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-large', trust_remote_code=True).to("cuda").eval(),
-    'microsoft/Florence-2-base-ft': AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base-ft', trust_remote_code=True).to("cuda").eval(),
-    'microsoft/Florence-2-base': AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True).to("cuda").eval(),
-}
+model = AutoModelForCausalLM.from_pretrained('HuggingFaceM4/Florence-2-DocVQA', trust_remote_code=True).to("cuda").eval()
 
-processors = {
-    'microsoft/Florence-2-large-ft': AutoProcessor.from_pretrained('microsoft/Florence-2-large-ft', trust_remote_code=True),
-    'microsoft/Florence-2-large': AutoProcessor.from_pretrained('microsoft/Florence-2-large', trust_remote_code=True),
-    'microsoft/Florence-2-base-ft': AutoProcessor.from_pretrained('microsoft/Florence-2-base-ft', trust_remote_code=True),
-    'microsoft/Florence-2-base': AutoProcessor.from_pretrained('microsoft/Florence-2-base', trust_remote_code=True),
-}
+processor = AutoProcessor.from_pretrained('HuggingFaceM4/Florence-2-DocVQA', trust_remote_code=True)
 
 
-DESCRIPTION = "# [Florence-2 Demo](https://huggingface.co/microsoft/Florence-2-large)"
+DESCRIPTION = "# [Florence-2-DocVQA Demo](https://huggingface.co/HuggingFaceM4/Florence-2-DocVQA)"
 
 colormap = ['blue','orange','green','purple','brown','pink','gray','olive','cyan','red',
             'lime','indigo','violet','aqua','magenta','coral','gold','tan','skyblue']
@@ -43,9 +33,9 @@ def fig_to_pil(fig):
     return Image.open(buf)
 
 @spaces.GPU
-def run_example(task_prompt, image, text_input=None, model_id='microsoft/Florence-2-large'):
-    model = models[model_id]
-    processor = processors[model_id]
+def run_example(task_prompt, image, text_input=None):
+    model = model
+    processor = processor
     if text_input is None:
         prompt = task_prompt
     else:
@@ -123,71 +113,75 @@ def draw_ocr_bboxes(image, prediction):
 
 def process_image(image, task_prompt, text_input=None, model_id='microsoft/Florence-2-large'):
     image = Image.fromarray(image)  # Convert NumPy array to PIL Image
-    if task_prompt == 'Caption':
+    if task_prompt == 'Document Visual Question Answering':
+        task_prompt = '<DocVQA>'
+        results = run_example(task_prompt, image)
+        return results, None
+    elif task_prompt == 'Caption':
         task_prompt = '<CAPTION>'
-        results = run_example(task_prompt, image, model_id=model_id)
+        results = run_example(task_prompt, image)
         return results, None
     elif task_prompt == 'Detailed Caption':
         task_prompt = '<DETAILED_CAPTION>'
-        results = run_example(task_prompt, image, model_id=model_id)
+        results = run_example(task_prompt, image)
         return results, None
     elif task_prompt == 'More Detailed Caption':
         task_prompt = '<MORE_DETAILED_CAPTION>'
-        results = run_example(task_prompt, image, model_id=model_id)
+        results = run_example(task_prompt, image)
         return results, None
     elif task_prompt == 'Object Detection':
         task_prompt = '<OD>'
-        results = run_example(task_prompt, image, model_id=model_id)
+        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, model_id=model_id)
+        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, model_id=model_id)
+        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, model_id)
+        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, model_id)
+        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, model_id)
+        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, model_id)
+        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, model_id)
+        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, model_id)
+        results = run_example(task_prompt, image, text_input)
         return results, None
     elif task_prompt == 'OCR':
         task_prompt = '<OCR>'
-        results = run_example(task_prompt, image, model_id=model_id)
+        results = run_example(task_prompt, image)
         return results, None
     elif task_prompt == 'OCR with Region':
         task_prompt = '<OCR_WITH_REGION>'
-        results = run_example(task_prompt, image, model_id=model_id)
+        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
@@ -208,14 +202,14 @@ with gr.Blocks(css=css) as demo:
         with gr.Row():
             with gr.Column():
                 input_img = gr.Image(label="Input Picture")
-                model_selector = gr.Dropdown(choices=list(models.keys()), label="Model", value='microsoft/Florence-2-large')
                 task_prompt = gr.Dropdown(choices=[
+                    'Document Visual Question Answering',
                     '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')
+                ], label="Task Prompt", value= 'Document Visual Question Answering')
                 text_input = gr.Textbox(label="Text Input (optional)")
                 submit_btn = gr.Button(value="Submit")
             with gr.Column():
@@ -234,6 +228,6 @@ with gr.Blocks(css=css) as demo:
             label='Try examples'
         )
 
-        submit_btn.click(process_image, [input_img, task_prompt, text_input, model_selector], [output_text, output_img])
+        submit_btn.click(process_image, [input_img, task_prompt, text_input], [output_text, output_img])
 
 demo.launch(debug=True)