fix(modeling): device-agnostic MPS/CPU support — replace .cuda() with .to(device), fix masked_scatter_ broadcast, and use device-type-aware autocast

modeling_unlimitedocr.py

@@ -579,7 +579,13 @@ 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)
+                    # Expand mask to match inputs_embeds shape and move to the same device.
+                    # masked_scatter_ requires the mask to be broadcastable to the target tensor;
+                    # unsqueeze(-1) alone leaves shape [seq_len, 1] which triggers a broadcast
+                    # error on MPS. expand_as() makes it explicit, and .to(device) replaces
+                    # the hardcoded .to(device) that crashes on Apple Silicon and CPU.
+                    _mask = images_seq_mask[idx].unsqueeze(-1).expand_as(inputs_embeds[idx]).to(inputs_embeds.device)
+                    inputs_embeds[idx].masked_scatter_(_mask, images_in_this_batch)
 
                 idx += 1
             
@@ -785,6 +791,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):
+        device = next(self.parameters()).device  # MPS / CUDA / CPU
         self.disable_torch_init()
 
         os.makedirs(output_path, exist_ok=True)
@@ -1000,9 +1007,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 +1022,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 +1032,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 +1046,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 +1063,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 +1074,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)
@@ -1137,6 +1144,7 @@ class UnlimitedOCRForCausalLM(DeepseekV2ForCausalLM):
 
 
     def infer_multi(self, tokenizer, prompt='', image_files=None, output_path='', image_size=640, save_results=False, max_length=32768, tps_interval=0, no_repeat_ngram_size=0, ngram_window=0, temperature=0.0):
+        device = next(self.parameters()).device  # MPS / CUDA / CPU
         """
         Multi-image inference. Does NOT support crop mode.
         Prompt uses a single <image> token (e.g. "<image>Multi page parsing.").
@@ -1235,12 +1243,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 +1264,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)]