Update app.py

app.py

@@ -25,7 +25,8 @@ overwrite_config["mm_resampler_type"] = "dynamic_compressor"
 overwrite_config["patchify_video_feature"] = False
 overwrite_config["attn_implementation"] = "sdpa" if torch.__version__ >= "2.1.2" else "eager"
 tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name, device_map="cpu", overwrite_config=overwrite_config)
-model.to("cuda").eval()
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model.to(device).eval()
 
 def preprocess_qwen(sources, tokenizer: transformers.PreTrainedTokenizer, has_image: bool = False, max_len=2048, system_message: str = "You are a helpful assistant.") -> Dict:
     roles = {"human": "<|im_start|>user", "gpt": "<|im_start|>assistant"}
@@ -100,7 +101,7 @@ def oryx_inference(video, text):
     conv.append_message(conv.roles[1], None)
     prompt = conv.get_prompt()
 
-    input_ids = preprocess_qwen([{'from': 'human','value': question},{'from': 'gpt','value': None}], tokenizer, has_image=True).to("cuda")
+    input_ids = preprocess_qwen([{'from': 'human','value': question},{'from': 'gpt','value': None}], tokenizer, has_image=True).to(device)
 
     video_processed = []
     for idx, frame in enumerate(video):
@@ -116,7 +117,7 @@ def oryx_inference(video, text):
     if frame_idx is None:
         frame_idx = np.arange(0, len(video_processed), dtype=int).tolist()
     
-    video_processed = torch.cat(video_processed, dim=0).bfloat16().to("cuda")
+    video_processed = torch.cat(video_processed, dim=0).bfloat16().to(device)
     video_processed = (video_processed, video_processed)
 
     video_data = (video_processed, (384, 384), "video")