Update app.py

app.py

@@ -9,6 +9,7 @@ import cv2
 from datetime import datetime, timedelta
 from threading import Thread
 
+import torch
 import gradio as gr
 import spaces
 import numpy as np
@@ -88,9 +89,12 @@ SYSTEM_PROMPT = load_system_prompt(MODEL_ID, "SYSTEM_PROMPT.txt")
 # If you prefer a hardcoded system prompt, you can use:
 # SYSTEM_PROMPT = "You are a conversational agent that always answers straight to the point, and ends with an ASCII cat."
 
+# Set the device explicitly (vLLM requires an explicit device specification)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 # Initialize the Mistral LLM via vllm.
 # Note: Running this model on GPU may require very high VRAM.
-llm = LLM(model=MODEL_ID, tokenizer_mode="mistral")
+llm = LLM(model=MODEL_ID, tokenizer_mode="mistral", device=device)
 
 # -----------------------------------------------------------------------------
 # Main Generation Function