update app

app.py

@@ -33,6 +33,17 @@ INTRO_TEXT = """SpaceLlama3.1 demo\n\n
 # Set model location as a constant outside the function
 MODEL_LOCATION = "remyxai/SpaceLlama3.1"  # Update as needed
 
+# Global model variable
+global_model = None
+
+def load_model():
+    """Loads the model globally."""
+    global global_model
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    global_model = load(MODEL_LOCATION)
+    global_model.to(device, dtype=torch.bfloat16)
+    logging.info("Model loaded successfully.")
+
 def compute(image, prompt):
     """Runs model inference."""
     if image is None:
@@ -44,10 +55,8 @@ def compute(image, prompt):
     if isinstance(image, str):
         image = PIL.Image.open(image).convert("RGB")
 
-    # Set device and load the model
-    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-    vlm = load(MODEL_LOCATION)  # Use the constant for model location
-    vlm.to(device, dtype=torch.bfloat16)
+    # Use the globally loaded model
+    vlm = global_model
 
     # Prepare prompt
     prompt_builder = vlm.get_prompt_builder()
@@ -116,5 +125,8 @@ if __name__ == "__main__":
     for k, v in os.environ.items():
         logging.info('environ["%s"] = %r', k, v)
 
+    # Load the model once globally
+    load_model()
+
     create_app().queue().launch()