log time for inference and GPU used

app.py

@@ -1,4 +1,4 @@
-import os
+import time
 from PIL import Image
 import gradio as gr
 from glob import glob
@@ -56,8 +56,11 @@ def calc_probs(prompt, images):
 
 def predict(prompt, image_1, image_2):
     print(f"Starting prediction for prompt: {prompt}")
+    start_time = time.time()
     probs = calc_probs(prompt, [image_1, image_2])
-    print(f"Prediction: {probs}")
+    print(f"Prediction: {probs:.3f} ({time.time() - start_time:.2f} seconds, ) ")
+    if device == "cuda":
+        print(f"GPU mem used: {round(torch.cuda.max_memory_allocated(device) / 1024 / 1024 / 1024, 2)} GB")
     return str(round(probs[0], 3)), str(round(probs[1], 3))