Support cuda and cpu

modeling_GOT.py

@@ -19,6 +19,8 @@ DEFAULT_IMAGE_PATCH_TOKEN = '<imgpad>'
 DEFAULT_IM_START_TOKEN = '<img>'
 DEFAULT_IM_END_TOKEN = '</img>'
 
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
 from enum import auto, Enum
 class SeparatorStyle(Enum):
     """Different separator style."""
@@ -164,7 +166,7 @@ class GOTQwenModel(Qwen2Model):
         use_im_start_end=False,
         vision_select_layer=-1,
         dtype=torch.float16,
-        device="cuda"
+        device=DEVICE
     ):
 
 
@@ -453,7 +455,7 @@ class GOTQwenForCausalLM(Qwen2ForCausalLM):
         tokenizer, 
         freeze_lm_model=False, 
         pretrained_stage1_model=None,
-        device="cuda"
+        device=DEVICE
     ):
         config = self.get_model().config
 
@@ -558,7 +560,7 @@ class GOTQwenForCausalLM(Qwen2ForCausalLM):
 
         image_tensor_1 = image_processor_high(image)
 
-        input_ids = torch.as_tensor(inputs.input_ids).cuda()
+        input_ids = torch.as_tensor(inputs.input_ids).to(DEVICE)
 
         stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
         keywords = [stop_str]
@@ -566,10 +568,10 @@ class GOTQwenForCausalLM(Qwen2ForCausalLM):
         streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
         if stream_flag:
-            with torch.autocast("cuda", dtype=torch.bfloat16):
+            with torch.autocast(DEVICE, dtype=torch.bfloat16):
                 output_ids = self.generate(
                     input_ids,
-                    images=[image_tensor_1.unsqueeze(0).half().cuda()],
+                    images=[image_tensor_1.unsqueeze(0).half().to(DEVICE)],
                     do_sample=False,
                     num_beams = 1,
                     no_repeat_ngram_size = 20,
@@ -578,10 +580,10 @@ class GOTQwenForCausalLM(Qwen2ForCausalLM):
                     stopping_criteria=[stopping_criteria]
                     )
         else:
-            with torch.autocast("cuda", dtype=torch.bfloat16):
+            with torch.autocast(DEVICE, dtype=torch.bfloat16):
                 output_ids = self.generate(
                     input_ids,
-                    images=[image_tensor_1.unsqueeze(0).half().cuda()],
+                    images=[image_tensor_1.unsqueeze(0).half().to(DEVICE)],
                     do_sample=False,
                     num_beams = 1,
                     no_repeat_ngram_size = 20,
@@ -812,7 +814,7 @@ class GOTQwenForCausalLM(Qwen2ForCausalLM):
 
         inputs = tokenizer([prompt])
 
-        input_ids = torch.as_tensor(inputs.input_ids).cuda()
+        input_ids = torch.as_tensor(inputs.input_ids).to(DEVICE)
 
         stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
         keywords = [stop_str]
@@ -820,10 +822,10 @@ class GOTQwenForCausalLM(Qwen2ForCausalLM):
         streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
         if stream_flag:
-            with torch.autocast("cuda", dtype=torch.bfloat16):
+            with torch.autocast(DEVICE, dtype=torch.bfloat16):
                 output_ids = self.generate(
                     input_ids,
-                    images=[image_list.half().cuda()],
+                    images=[image_list.half().to(DEVICE)],
                     do_sample=False,
                     num_beams = 1,
                     # no_repeat_ngram_size = 20,
@@ -832,10 +834,10 @@ class GOTQwenForCausalLM(Qwen2ForCausalLM):
                     stopping_criteria=[stopping_criteria]
                     )
         else:
-            with torch.autocast("cuda", dtype=torch.bfloat16):
+            with torch.autocast(DEVICE, dtype=torch.bfloat16):
                 output_ids = self.generate(
                     input_ids,
-                    images=[image_list.half().cuda()],
+                    images=[image_list.half().to(DEVICE)],
                     do_sample=False,
                     num_beams = 1,
                     # no_repeat_ngram_size = 20,