update app

app.py

@@ -30,7 +30,10 @@ INTRO_TEXT = """SpaceLlama3.1 demo\n\n
 **This is an experimental research model.** Make sure to add appropriate guardrails when using the model for applications.
 """
 
-def compute(image, prompt, model_location):
+# Set model location as a constant outside the function
+MODEL_LOCATION = "remyxai/SpaceLlama3.1"  # Update as needed
+
+def compute(image, prompt):
     """Runs model inference."""
     if image is None:
         raise gr.Error("Image required")
@@ -43,7 +46,7 @@ def compute(image, prompt, model_location):
 
     # Set device and load the model
     device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
-    vlm = load(model_location)  # No need to pass the token again
+    vlm = load(MODEL_LOCATION)  # Use the constant for model location
     vlm.to(device, dtype=torch.bfloat16)
 
     # Prepare prompt
@@ -85,13 +88,10 @@ def create_app():
                 clear = gr.Button("Clear")
                 highlighted_text = gr.HighlightedText(value="", label="Output", visible=True)
 
-        # Model location
-        model_location = "remyxai/SpaceLlama3.1"  # Update as needed
-
         # Button event handlers
         run.click(
             fn=compute,
-            inputs=[image, prompt, model_location],
+            inputs=[image, prompt],
             outputs=highlighted_text,
         )
         clear.click(fn=reset, inputs=None, outputs=[prompt, image])
@@ -100,7 +100,7 @@ def create_app():
         status = gr.Markdown(f"Startup: {datetime.datetime.now()}")
         gpu_kind = gr.Markdown(f"GPU=?")
         demo.load(
-            fn=lambda: [f"Model `{model_location}` loaded."],
+            fn=lambda: [f"Model `{MODEL_LOCATION}` loaded."],
             inputs=None,
             outputs=model_info,
         )