update demo.

app.py

@@ -219,8 +219,35 @@ if __name__ == '__main__':
     conv_mode = "llama_2"
     model_path = 'DAMO-NLP-SG/VideoLLaMA2-7B'
 
+    def find_cuda():
+        # Check if CUDA_HOME or CUDA_PATH environment variables are set
+        cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
+
+        if cuda_home and os.path.exists(cuda_home):
+            return cuda_home
+
+        # Search for the nvcc executable in the system's PATH
+        nvcc_path = shutil.which('nvcc')
+
+        if nvcc_path:
+            # Remove the 'bin/nvcc' part to get the CUDA installation path
+            cuda_path = os.path.dirname(os.path.dirname(nvcc_path))
+            return cuda_path
+
+        return None
+
+    cuda_path = find_cuda()
+
+    if cuda_path:
+        print(f"CUDA installation found at: {cuda_path}")
+    else:
+        print("CUDA installation not found")
+
+    device = torch.device("cuda")
+
     handler = Chat(model_path, conv_mode=conv_mode, load_8bit=False, load_4bit=True)
     # handler.model.to(dtype=torch.float16)
+    handler = handler.model.to(device)
 
     if not os.path.exists("temp"):
         os.makedirs("temp")