Update app.py

app.py

@@ -2,43 +2,29 @@ import spaces
 import gradio as gr
 import torch
 torch.jit.script = lambda f: f  # Avoid script error in lambda
-from t2v_metrics import VQAScore, list_all_vqascore_models
+from t2v_metrics import VQAScore
+from functools import lru_cache
 
+# Remove any global model loading or CUDA initialization
+# Do not call torch.cuda.is_available() at the global scope
 
-def update_model(model_name):
+@lru_cache()
+def get_model(model_name):
+    # This function will cache the model per process
     return VQAScore(model=model_name, device="cuda")
 
-# Use global variables for model pipe and current model name
-global model_pipe, cur_model_name
-cur_model_name = "clip-flant5-xl"
-model_pipe = update_model(cur_model_name)
-
-
-# Ensure GPU context manager is imported correctly (assuming spaces is a module you have)
-#try:
-    #from spaces import GPU # i believe this is wrong, spaces package does not have "GPU"
-#except ImportError:
-    # GPU = lambda duration: (lambda f: f)  # Dummy decorator if spaces.GPU is not available
-
-if torch.cuda.is_available():
-    model_pipe.device = "cuda"
-else:
-    print("CUDA is not available")
-    
-@spaces.GPU # a duration lower than 60 does not work, leave as is.
+@spaces.GPU  # Decorate the function to use GPU
 def generate(model_name, image, text):
-    global model_pipe, cur_model_name
-    
-    if model_name != cur_model_name:
-        cur_model_name = model_name  # Update the current model name
-        model_pipe = update_model(model_name)
+    # Load the model inside the GPU context
+    model_pipe = get_model(model_name)
     
-    print("Image:", image)  # Debug: Print image path
-    print("Text:", text)  # Debug: Print text input
+    print("Image:", image)
+    print("Text:", text)
     print("Using model:", model_name)
     
     try:
-        result = model_pipe(images=[image], texts=[text]).cpu()[0][0].item()  # Perform the model inference
+        # Perform the model inference
+        result = model_pipe(images=[image], texts=[text]).cpu()[0][0].item()
         print("Result:", result)
     except RuntimeError as e:
         print(f"RuntimeError during model inference: {e}")
@@ -46,30 +32,106 @@ def generate(model_name, image, text):
     
     return result
 
-
+@spaces.GPU  # Decorate the function to use GPU
 def rank_images(model_name, images, text):
-    global model_pipe, cur_model_name
-    
-    if model_name != cur_model_name:
-        cur_model_name = model_name  # Update the current model name
-        model_pipe = update_model(model_name)
+    # Load the model inside the GPU context
+    model_pipe = get_model(model_name)
 
     images = [image_tuple[0] for image_tuple in images]
-    print("Images:", images)  # Debug: Print image paths
-    print("Text:", text)  # Debug: Print text input
+    print("Images:", images)
+    print("Text:", text)
     print("Using model:", model_name)
     
     try:
-        results = model_pipe(images=images, texts=[text]).cpu()[:, 0].tolist()  # Perform the model inference on all images
-        print("Initial results: should be imgs x texts", results)
-        ranked_results = sorted(zip(images, results), key=lambda x: x[1], reverse=True)  # Rank results
-        ranked_images = [(img, f"Rank: {rank + 1} - Score: {score:.2f}") for rank, (img, score) in enumerate(ranked_results)]  # Pair images with their scores and rank
+        # Perform the model inference on all images
+        results = model_pipe(images=images, texts=[text]).cpu()[:, 0].tolist()
+        print("Initial results:", results)
+        # Rank results
+        ranked_results = sorted(zip(images, results), key=lambda x: x[1], reverse=True)
+        # Pair images with their scores and rank
+        ranked_images = [
+            (img, f"Rank: {rank + 1} - Score: {score:.2f}")
+            for rank, (img, score) in enumerate(ranked_results)
+        ]
         print("Ranked Results:", ranked_results)
     except RuntimeError as e:
         print(f"RuntimeError during model inference: {e}")
         raise e
     
     return ranked_images
+    
+# import spaces
+# import gradio as gr
+# import torch
+# torch.jit.script = lambda f: f  # Avoid script error in lambda
+# from t2v_metrics import VQAScore, list_all_vqascore_models
+
+
+# def update_model(model_name):
+#     return VQAScore(model=model_name, device="cuda")
+
+# # Use global variables for model pipe and current model name
+# global model_pipe, cur_model_name
+# cur_model_name = "clip-flant5-xl"
+# model_pipe = update_model(cur_model_name)
+
+
+# # Ensure GPU context manager is imported correctly (assuming spaces is a module you have)
+# #try:
+#     #from spaces import GPU # i believe this is wrong, spaces package does not have "GPU"
+# #except ImportError:
+#     # GPU = lambda duration: (lambda f: f)  # Dummy decorator if spaces.GPU is not available
+
+# if torch.cuda.is_available():
+#     model_pipe.device = "cuda"
+# else:
+#     print("CUDA is not available")
+    
+# @spaces.GPU # a duration lower than 60 does not work, leave as is.
+# def generate(model_name, image, text):
+#     global model_pipe, cur_model_name
+    
+#     if model_name != cur_model_name:
+#         cur_model_name = model_name  # Update the current model name
+#         model_pipe = update_model(model_name)
+    
+#     print("Image:", image)  # Debug: Print image path
+#     print("Text:", text)  # Debug: Print text input
+#     print("Using model:", model_name)
+    
+#     try:
+#         result = model_pipe(images=[image], texts=[text]).cpu()[0][0].item()  # Perform the model inference
+#         print("Result:", result)
+#     except RuntimeError as e:
+#         print(f"RuntimeError during model inference: {e}")
+#         raise e
+    
+#     return result
+
+
+# def rank_images(model_name, images, text):
+#     global model_pipe, cur_model_name
+    
+#     if model_name != cur_model_name:
+#         cur_model_name = model_name  # Update the current model name
+#         model_pipe = update_model(model_name)
+
+#     images = [image_tuple[0] for image_tuple in images]
+#     print("Images:", images)  # Debug: Print image paths
+#     print("Text:", text)  # Debug: Print text input
+#     print("Using model:", model_name)
+    
+#     try:
+#         results = model_pipe(images=images, texts=[text]).cpu()[:, 0].tolist()  # Perform the model inference on all images
+#         print("Initial results: should be imgs x texts", results)
+#         ranked_results = sorted(zip(images, results), key=lambda x: x[1], reverse=True)  # Rank results
+#         ranked_images = [(img, f"Rank: {rank + 1} - Score: {score:.2f}") for rank, (img, score) in enumerate(ranked_results)]  # Pair images with their scores and rank
+#         print("Ranked Results:", ranked_results)
+#     except RuntimeError as e:
+#         print(f"RuntimeError during model inference: {e}")
+#         raise e
+    
+#     return ranked_images
 
 
 ### EXAMPLES ###
@@ -190,4 +252,4 @@ with gr.Blocks() as demo_vqascore_ranking:
 
 # Launch the interface
 demo_vqascore_ranking.queue()
-demo_vqascore_ranking.launch(share=False)
+demo_vqascore_ranking.launch(share=True)