fix: Handle CUDA availability in OmniChatServer

server.py

@@ -2,17 +2,21 @@ import flask
 import base64
 import tempfile
 import traceback
+import torch
 from flask import Flask, Response, stream_with_context
 from inference import OmniInference
 
 
 class OmniChatServer(object):
     def __init__(self, ip='0.0.0.0', port=60808, run_app=True,
-                 ckpt_dir='./checkpoint', device='cuda:0') -> None:
+                 ckpt_dir='./checkpoint', device=None) -> None:
         server = Flask(__name__)
         # CORS(server, resources=r"/*")
         # server.config["JSON_AS_ASCII"] = False
 
+        if device is None:
+            device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+        
         self.client = OmniInference(ckpt_dir, device)
         self.client.warm_up()
 
@@ -46,9 +50,8 @@ def create_app():
     return server.server
 
 
-def serve(ip='0.0.0.0', port=60808, device='cuda:0'):
-
-    OmniChatServer(ip, port=port,run_app=True, device=device)
+def serve(ip='0.0.0.0', port=60808, device=None):
+    OmniChatServer(ip, port=port, run_app=True, device=device)
 
 
 if __name__ == "__main__":