Update app.py

app.py

@@ -15,34 +15,13 @@ import uuid
 import subprocess
 subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
-def load_model():
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = AutoModelForCausalLM.from_pretrained(
-        "microsoft/Florence-2-large-ft", 
-        trust_remote_code=True
-    ).to(device).eval()
-    processor = AutoProcessor.from_pretrained(
-        "microsoft/Florence-2-large-ft", 
-        trust_remote_code=True
-    )
-    return model, processor, device
-
-model = None
-processor = None
-device = None
+model = AutoModelForCausalLM.from_pretrained('microsoft/Florence-2-large-ft', trust_remote_code=True).to("cuda").eval()
+processor = AutoProcessor.from_pretrained('microsoft/Florence-2-large-ft', trust_remote_code=True)
 
 @spaces.GPU
-def initialize_model():
-    global model, processor, device
-    model, processor, device = load_model()
-
 def run_example(task_prompt, image, text_input=None):
-    global model, processor, device
-    if model is None or processor is None:
-        initialize_model()
-    
     prompt = task_prompt if text_input is None else task_prompt + text_input
-    inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
+    inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda")
     with torch.inference_mode():
         generated_ids = model.generate(**inputs, 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]
@@ -60,20 +39,17 @@ def fig_to_pil(fig):
 def plot_bbox_img(image, data):
     fig, ax = plt.subplots(figsize=(10, 10))
     ax.imshow(image)
-    
     if 'bboxes' in data and 'labels' in data:
         bboxes, labels = data['bboxes'], data['labels']
     elif 'bboxes' in data and 'bboxes_labels' in data:
         bboxes, labels = data['bboxes'], data['bboxes_labels']
     else:
         return fig_to_pil(fig)
-
     for bbox, label in zip(bboxes, labels):
         x1, y1, x2, y2 = bbox
         rect = patches.Rectangle((x1, y1), x2-x1, y2-y1, linewidth=2, edgecolor='indigo', facecolor='none')
         ax.add_patch(rect)
         plt.text(x1, y1, label, color='white', fontsize=10, bbox=dict(facecolor='indigo', alpha=0.8))
-    
     ax.axis('off')
     return fig_to_pil(fig)
 
@@ -199,6 +175,7 @@ def draw_vid_polygons(image, prediction, get_color):
             draw.text(polygon[0], label, fill="white")
     return np.array(img_draw)
 
+@spaces.GPU
 def process_image(image, task, text):
     task_mapping = {
         "Caption": ("<CAPTION>", lambda result: (result['<CAPTION>'], image)),
@@ -242,10 +219,6 @@ def process_video_p(input_video, task, text_input):
         return None, "Error: Video processing failed. Check logs above for info.", str(frame_results)
     return result, result, str(frame_results)
 
-@spaces.GPU
-def process_image_with_gpu(image, task, text):
-    return process_image(image, task, text)
-
 with gr.Blocks() as demo:
     gr.HTML("<h1><center>Microsoft Florence-2-large-ft</center></h1>")
     
@@ -317,16 +290,7 @@ with gr.Blocks() as demo:
 
     video_task_dropdown.change(fn=update_video_text_input, inputs=video_task_dropdown, outputs=video_text_input)
 
-    submit_btn.click(
-        fn=process_image_with_gpu,
-        inputs=[input_img, task_dropdown, text_input],
-        outputs=[output_text, output_image]
-    )
-    
-    video_submit_btn.click(
-        fn=process_video_p,
-        inputs=[input_video, video_task_dropdown, video_text_input],
-        outputs=[output_video, output_video, frame_results_output]
-    )
+    submit_btn.click(fn=process_image, inputs=[input_img, task_dropdown, text_input], outputs=[output_text, output_image])
+    video_submit_btn.click(fn=process_video_p, inputs=[input_video, video_task_dropdown, video_text_input], outputs=[output_video, output_video, frame_results_output])
 
 demo.launch()