Update main.py

main.py

@@ -7,7 +7,7 @@ from PIL import Image, ImageDraw
 import base64
 import io
 import asyncio
-from utils import initialize_model, sample_frame, device
+from utils import initialize_model, sample_frame
 import torch
 
 app = FastAPI()
@@ -40,7 +40,8 @@ def draw_trace(image: np.ndarray, previous_actions: List[Tuple[str, List[int]]])
 
 # Initialize the model at the start of your application
 model = initialize_model("config_csllm.yaml", "yuntian-deng/computer-model")
-
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = model.to(device)
 def predict_next_frame(previous_frames: List[np.ndarray], previous_actions: List[Tuple[str, List[int]]]) -> np.ndarray:
     width, height = 256, 256