cuda->cpu

third_party/VideoLLaMA2/videollama2/__init__.py

@@ -10,11 +10,12 @@ from .model import load_pretrained_model
 from .mm_utils import process_image, process_video, tokenizer_multimodal_token, get_model_name_from_path, KeywordsStoppingCriteria, process_audio_file
 from .constants import NUM_FRAMES, DEFAULT_IMAGE_TOKEN, DEFAULT_VIDEO_TOKEN, MODAL_INDEX_MAP, DEFAULT_AUDIO_TOKEN
 
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 def model_init(model_path=None, **kwargs):
     model_path = "DAMO-NLP-SG/VideoLLaMA2-7B" if model_path is None else model_path
     model_name = get_model_name_from_path(model_path)
-    tokenizer, model, processor, context_len = load_pretrained_model(model_path, None, model_name, **kwargs)
+    tokenizer, model, processor, context_len = load_pretrained_model(model_path, None, model_name, device=DEVICE, **kwargs)
 
     if tokenizer.pad_token is None and tokenizer.unk_token is not None:
         tokenizer.pad_token = tokenizer.unk_token
@@ -60,9 +61,15 @@ def mm_infer(image_or_video, instruct, model, tokenizer, modal='video', **kwargs
         tensor = None
     else:
         if isinstance(image_or_video, dict):
-            tensor = {k: v.half().cuda() for k, v in image_or_video.items()}
+            if DEVICE == "cuda":
+                tensor = {k: v.half().cuda() for k, v in image_or_video.items()}
+            else:
+                tensor = {k: v.half().cpu() for k, v in image_or_video.items()}
         else:
-            tensor = image_or_video.half().cuda() 
+            if DEVICE == "cuda":
+                tensor = image_or_video.half().cuda()
+            else:
+                tensor = image_or_video.half().cpu()
         tensor = [(tensor, modal)]
 
     # 2. text preprocess (tag process & generate prompt).
@@ -88,8 +95,12 @@ def mm_infer(image_or_video, instruct, model, tokenizer, modal='video', **kwargs
     message = system_message + message
     prompt = tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
 
-    input_ids = tokenizer_multimodal_token(prompt, tokenizer, modal_token, return_tensors='pt').unsqueeze(0).long().cuda()
-    attention_masks = input_ids.ne(tokenizer.pad_token_id).long().cuda()
+    if DEVICE == "cuda":
+        input_ids = tokenizer_multimodal_token(prompt, tokenizer, modal_token, return_tensors='pt').unsqueeze(0).long().cuda()
+        attention_masks = input_ids.ne(tokenizer.pad_token_id).long().cuda()
+    else:
+        input_ids = tokenizer_multimodal_token(prompt, tokenizer, modal_token, return_tensors='pt').unsqueeze(0).long().cpu()
+        attention_masks = input_ids.ne(tokenizer.pad_token_id).long().cpu()
 
     # 3. generate response according to visual signals and prompts. 
     keywords = [tokenizer.eos_token]

third_party/VideoLLaMA2/videollama2/model/__init__.py

@@ -51,6 +51,7 @@ VLLMConfigs = {
 }
 
 
+
 def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, load_4bit=False, device_map="auto", device="cuda", use_flash_attn=False, **kwargs):
     
     if 'token' in kwargs:
@@ -209,6 +210,6 @@ def load_pretrained_model(model_path, model_base, model_name, load_8bit=False, l
     if hasattr(model.config, "mm_audio_tower"):
         nname = model.config.mm_audio_tower.split("/")[-1]
         model.config.mm_audio_tower = os.path.join(model_path, nname)
-        
+
 
     return tokenizer, model, processor, context_len