Fix MPS inference

modeling_unlimitedocr.py

@@ -579,7 +579,8 @@ class UnlimitedOCRModel(DeepseekV2Model):
                     images_in_this_batch = torch.cat(images_in_this_batch, dim=0)
                     # exit()
 
-                    inputs_embeds[idx].masked_scatter_(images_seq_mask[idx].unsqueeze(-1).cuda(), images_in_this_batch)
+                    scatter_mask = images_seq_mask[idx].unsqueeze(-1).to(inputs_embeds.device).expand_as(inputs_embeds[idx]).contiguous()
+                    inputs_embeds[idx].masked_scatter_(scatter_mask, images_in_this_batch)
 
                 idx += 1
             
@@ -786,6 +787,7 @@ class UnlimitedOCRForCausalLM(DeepseekV2ForCausalLM):
 
     def infer(self, tokenizer, prompt='', image_file='', output_path = '', base_size=1024, image_size=640, crop_mode=True, test_compress=False, save_results=False, eval_mode=False, max_length=32768, tps_interval=0, no_repeat_ngram_size=0, ngram_window=0, temperature=0.0):
         self.disable_torch_init()
+        device = self.device
 
         os.makedirs(output_path, exist_ok=True)
         os.makedirs(f'{output_path}/images', exist_ok=True)
@@ -1000,9 +1002,9 @@ class UnlimitedOCRForCausalLM(DeepseekV2ForCausalLM):
             self.config.sliding_window = None
             # Build logits processors for ngram
             gen_kwargs = dict(
-                input_ids=input_ids.unsqueeze(0).cuda(),
-                images=[(images_crop.cuda(), images_ori.cuda())],
-                images_seq_mask=images_seq_mask.unsqueeze(0).cuda(),
+                input_ids=input_ids.unsqueeze(0).to(device),
+                images=[(images_crop.to(device), images_ori.to(device))],
+                images_seq_mask=images_seq_mask.unsqueeze(0).to(device),
                 images_spatial_crop=images_spatial_crop,
                 do_sample=temperature > 0,
                 temperature=temperature if temperature > 0 else None,
@@ -1015,7 +1017,7 @@ class UnlimitedOCRForCausalLM(DeepseekV2ForCausalLM):
                 gen_kwargs['logits_processor'] = [SlidingWindowNoRepeatNgramProcessor(no_repeat_ngram_size, ngram_window)]
             elif no_repeat_ngram_size > 0:
                 gen_kwargs['no_repeat_ngram_size'] = no_repeat_ngram_size
-            with torch.autocast("cuda", dtype=torch.bfloat16):
+            with torch.autocast(device.type, dtype=torch.bfloat16):
                 with torch.no_grad():
                     output_ids = self.generate(**gen_kwargs)
             self.config.sliding_window = _orig_sw
@@ -1025,9 +1027,9 @@ class UnlimitedOCRForCausalLM(DeepseekV2ForCausalLM):
             self.config._ring_window = _orig_sw
             self.config.sliding_window = None
             gen_kwargs = dict(
-                input_ids=input_ids.unsqueeze(0).cuda(),
-                images=[(images_crop.cuda(), images_ori.cuda())],
-                images_seq_mask=images_seq_mask.unsqueeze(0).cuda(),
+                input_ids=input_ids.unsqueeze(0).to(device),
+                images=[(images_crop.to(device), images_ori.to(device))],
+                images_seq_mask=images_seq_mask.unsqueeze(0).to(device),
                 images_spatial_crop=images_spatial_crop,
                 do_sample=temperature > 0,
                 temperature=temperature if temperature > 0 else None,
@@ -1039,14 +1041,14 @@ class UnlimitedOCRForCausalLM(DeepseekV2ForCausalLM):
                 gen_kwargs['logits_processor'] = [SlidingWindowNoRepeatNgramProcessor(no_repeat_ngram_size, ngram_window)]
             elif no_repeat_ngram_size > 0:
                 gen_kwargs['no_repeat_ngram_size'] = no_repeat_ngram_size
-            with torch.autocast("cuda", dtype=torch.bfloat16):
+            with torch.autocast(device.type, dtype=torch.bfloat16):
                 with torch.no_grad():
                     output_ids = self.generate(**gen_kwargs)
             self.config.sliding_window = _orig_sw
                 
 
         if '<image>' in conversation[0]['content'] and eval_mode:
-                outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
+                outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).to(device).shape[1]:])
                 stop_str = '<｜end▁of▁sentence｜>'
                 if outputs.endswith(stop_str):
                     outputs = outputs[:-len(stop_str)]
@@ -1056,7 +1058,7 @@ class UnlimitedOCRForCausalLM(DeepseekV2ForCausalLM):
                 return outputs
         
         if '<image>' in conversation[0]['content'] and test_compress:
-            outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
+            outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).to(device).shape[1]:])
             pure_texts_outputs_token_length = len(text_encode(tokenizer, outputs, bos=False, eos=False))
             print('='*50)
             print('image size: ', (w, h))
@@ -1067,7 +1069,7 @@ class UnlimitedOCRForCausalLM(DeepseekV2ForCausalLM):
 
 
         if '<image>' in conversation[0]['content'] and save_results:
-            outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
+            outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).to(device).shape[1]:])
             stop_str = '<｜end▁of▁sentence｜>'
 
             print('='*15 + 'save results:' + '='*15)
@@ -1150,7 +1152,8 @@ class UnlimitedOCRForCausalLM(DeepseekV2ForCausalLM):
             save_results: whether to save output to file
         """
         self.disable_torch_init()
-        
+        device = self.device
+
         if image_files is None or len(image_files) == 0:
             assert False, 'image_files must be a non-empty list for multi-image inference!'
         
@@ -1235,12 +1238,12 @@ class UnlimitedOCRForCausalLM(DeepseekV2ForCausalLM):
         _orig_sw = getattr(self.config, 'sliding_window_size', None) or getattr(self.config, 'sliding_window', None)
         self.config._ring_window = _orig_sw  # Save for ring buffer to read
         self.config.sliding_window = None
-        with torch.autocast("cuda", dtype=torch.bfloat16):
+        with torch.autocast(device.type, dtype=torch.bfloat16):
             with torch.no_grad():
                 gen_kwargs = dict(
-                    input_ids=input_ids.unsqueeze(0).cuda(),
-                    images=[(dummy_crop.cuda(), images_ori.cuda())],
-                    images_seq_mask=images_seq_mask.unsqueeze(0).cuda(),
+                    input_ids=input_ids.unsqueeze(0).to(device),
+                    images=[(dummy_crop.to(device), images_ori.to(device))],
+                    images_seq_mask=images_seq_mask.unsqueeze(0).to(device),
                     images_spatial_crop=images_spatial_crop,
                     do_sample=temperature > 0,
                     temperature=temperature if temperature > 0 else None,
@@ -1256,7 +1259,7 @@ class UnlimitedOCRForCausalLM(DeepseekV2ForCausalLM):
                 output_ids = self.generate(**gen_kwargs)
         self.config.sliding_window = _orig_sw  # Restore
 
-        outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
+        outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).to(device).shape[1]:])
         stop_str = '<｜end▁of▁sentence｜>'
         if outputs.endswith(stop_str):
             outputs = outputs[:-len(stop_str)]