Show inference time for both models

app.py

@@ -2,6 +2,7 @@ import random
 import requests
 import json
 import ast
+import time
 
 import matplotlib.pyplot as plt
 from PIL import Image, ImageDraw, ImageFont
@@ -156,6 +157,7 @@ def detect_qwen(image, prompt):
         }
     ]
 
+    t0 = time.perf_counter()
     text = processor_qwen.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     image_inputs, video_inputs = process_vision_info(messages)
     inputs = processor_qwen(
@@ -173,37 +175,41 @@ def detect_qwen(image, prompt):
     output_text = processor_qwen.batch_decode(
         generated_ids_trimmed, do_sample=True, skip_special_tokens=True, clean_up_tokenization_spaces=False
     )[0]
+    elapsed_ms = (time.perf_counter() - t0) * 1_000
 
     input_height = inputs['image_grid_thw'][0][1] * 14
     input_width = inputs['image_grid_thw'][0][2] * 14
 
     annotated_image = create_annotated_image(image, output_text, input_height, input_width)
 
-    return annotated_image, output_text
+    time_taken = f"**Inference time ({model_qwen_name}):** {elapsed_ms:.0f} ms"
+    return annotated_image, output_text, time_taken
 
 
 @GPU
 def detect_moondream(image, prompt, category_input):
+    t0 = time.perf_counter()
     if category_input in ["Object Detection", "Visual Grounding + Object Detection"]:
         output_text = model_moondream.detect(image=image, object=prompt)
     elif category_input == "Visual Grounding + Keypoint Detection":
         output_text = model_moondream.point(image=image, object=prompt)
     else:
         output_text = model_moondream.query(image=image, question=prompt, reasoning=True)
+    elapsed_ms = (time.perf_counter() - t0) * 1_000
 
     annotated_image = create_annotated_image_normalized(image=image, json_data=output_text, label="object", explicit_color=None)
-    
-    return annotated_image, output_text
 
-@GPU
+    time_taken = f"**Inference time ({model_moondream_name}):** {elapsed_ms:.0f} ms"
+    return annotated_image, output_text, time_taken
+
 def detect(image, prompt_model_1, prompt_model_2, category_input):
     STANDARD_SIZE = (1024, 1024)
     image.thumbnail(STANDARD_SIZE)
     
-    annotated_image_model_1, output_text_model_1 = detect_qwen(image, prompt_model_1)
-    annotated_image_model_2, output_text_model_2 = detect_moondream(image, prompt_model_2, category_input)
+    annotated_image_model_1, output_text_model_1, timing_1 = detect_qwen(image, prompt_model_1)
+    annotated_image_model_2, output_text_model_2, timing_2 = detect_moondream(image, prompt_model_2, category_input)
 
-    return annotated_image_model_1, output_text_model_1, annotated_image_model_2, output_text_model_2
+    return annotated_image_model_1, output_text_model_1, timing_1, annotated_image_model_2, output_text_model_2, timing_2
 
 css_hide_share = """
 button#gradio-share-link-button-0 {
@@ -253,10 +259,12 @@ with gr.Blocks(theme=Ocean(), css=css_hide_share) as demo:
         with gr.Column(scale=1):
             output_image_model_1 = gr.Image(type="pil", label=f"Annotated image for {model_qwen_name}", height=400)
             output_textbox_model_1 = gr.Textbox(label=f"Model response for {model_qwen_name}", lines=10)
+            output_time_model_1 = gr.Markdown()
         
         with gr.Column(scale=1):
             output_image_model_2 = gr.Image(type="pil", label=f"Annotated image for {model_moondream_name}", height=400)
             output_textbox_model_2 = gr.Textbox(label=f"Model response for {model_moondream_name}", lines=10)
+            output_time_model_2 = gr.Markdown()
 
     gr.Markdown("### Examples")
     example_prompts = [
@@ -276,8 +284,15 @@ with gr.Blocks(theme=Ocean(), css=css_hide_share) as demo:
         label="Click an example to populate the input"
     )
 
-    generate_btn.click(fn=detect, inputs=[image_input, prompt_input_model_1, prompt_input_model_2, category_input], outputs=[output_image_model_1, output_textbox_model_1, output_image_model_2, output_textbox_model_2])
-
+    generate_btn.click(
+        fn=detect,
+        inputs=[image_input, prompt_input_model_1, prompt_input_model_2, category_input],
+        outputs=[
+            output_image_model_1, output_textbox_model_1, output_time_model_1,
+            output_image_model_2, output_textbox_model_2, output_time_model_2
+        ]
+    )
+    
 if __name__ == "__main__":
     demo.launch()