Add Info toast message w/ time taken

app.py

@@ -1,4 +1,8 @@
+import contextlib
 import subprocess
+import time
+from typing import Iterator, Callable
+
 import gradio as gr
 import spaces
 from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
@@ -20,6 +24,21 @@ MODEL_NAMES = {
 }
 
 
+@contextlib.contextmanager
+def measure_time() -> Iterator[Callable[[], float]]:
+    """
+    A context manager for measuring execution time (in seconds) within its code block.
+
+    usage:
+        with code_timer() as timer:
+            # Code snippet to be timed
+        print(f"Code took: {timer()} seconds")
+    """
+    start_time = end_time = time.perf_counter()
+    yield lambda: end_time - start_time
+    end_time = time.perf_counter()
+
+
 @spaces.GPU(duration=300)
 def run_inference(model_key, input_type, text, image, video, fps, system_prompt, add_vision_id):
     """
@@ -78,10 +97,14 @@ def run_inference(model_key, input_type, text, image, video, fps, system_prompt,
     )
     inputs = inputs.to(model.device)
 
-    gen_ids = model.generate(**inputs, max_new_tokens=512)
-    # Trim the prompt tokens
-    trimmed = [out_ids[len(inp_ids):] for inp_ids, out_ids in zip(inputs.input_ids, gen_ids)]
-    return processor.batch_decode(trimmed, skip_special_tokens=True)[0]
+    with measure_time() as timer:
+        gen_ids = model.generate(**inputs, max_new_tokens=512)
+        # Trim the prompt tokens
+        trimmed = [out_ids[len(inp_ids):] for inp_ids, out_ids in zip(inputs.input_ids, gen_ids)]
+        result = processor.batch_decode(trimmed, skip_special_tokens=True)[0]
+
+    gr.Info(f"Finished in {timer():.2f}s", title="Success", duration=5)  # green-style info toast :contentReference[oaicite:0]{index=0}
+    return result
 
 
 # Build Gradio interface