Make infer general, so it runs on non cuda devices

modeling_deepseekocr.py

@@ -1,6 +1,9 @@
 from .modeling_deepseekv2 import DeepseekV2Model, DeepseekV2ForCausalLM
 from .configuration_deepseek_v2 import DeepseekV2Config
-from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
 from typing import List, Optional, Tuple, Union
 from transformers.cache_utils import Cache
 import requests
@@ -25,14 +28,13 @@ import time
 
 
 def load_image(image_path):
-
     try:
         image = Image.open(image_path)
-        
+
         corrected_image = ImageOps.exif_transpose(image)
-        
+
         return corrected_image
-        
+
     except Exception as e:
         print(f"error: {e}")
         try:
@@ -42,7 +44,7 @@ def load_image(image_path):
 
 
 def re_match(text):
-    pattern = r'(<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>)'
+    pattern = r"(<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>)"
     matches = re.findall(pattern, text, re.DOTALL)
 
     # pattern1 = r'<\|ref\|>.*?<\|/ref\|>\n'
@@ -51,7 +53,7 @@ def re_match(text):
     mathes_image = []
     mathes_other = []
     for a_match in matches:
-        if '<|ref|>image<|/ref|>' in a_match[0]:
+        if "<|ref|>image<|/ref|>" in a_match[0]:
             mathes_image.append(a_match[0])
         else:
             mathes_other.append(a_match[0])
@@ -59,7 +61,6 @@ def re_match(text):
 
 
 def extract_coordinates_and_label(ref_text, image_width, image_height):
-
     try:
         label_type = ref_text[1]
         cor_list = eval(ref_text[2])
@@ -71,33 +72,36 @@ def extract_coordinates_and_label(ref_text, image_width, image_height):
 
 
 def draw_bounding_boxes(image, refs, ouput_path):
-
     image_width, image_height = image.size
-    
+
     img_draw = image.copy()
     draw = ImageDraw.Draw(img_draw)
 
-    overlay = Image.new('RGBA', img_draw.size, (0, 0, 0, 0))
+    overlay = Image.new("RGBA", img_draw.size, (0, 0, 0, 0))
     draw2 = ImageDraw.Draw(overlay)
-    
+
     # try:
     # except IOError:
     #     try:
-    #         font = ImageFont.truetype("DejaVuSans.ttf", 20) 
+    #         font = ImageFont.truetype("DejaVuSans.ttf", 20)
     #     except IOError:
     font = ImageFont.load_default()
 
     img_idx = 0
-    
+
     for i, ref in enumerate(refs):
         try:
             result = extract_coordinates_and_label(ref, image_width, image_height)
             if result:
                 label_type, points_list = result
-                
-                color = (np.random.randint(0, 200), np.random.randint(0, 200), np.random.randint(0, 255))
 
-                color_a = color + (20, )
+                color = (
+                    np.random.randint(0, 200),
+                    np.random.randint(0, 200),
+                    np.random.randint(0, 255),
+                )
+
+                color_a = color + (20,)
                 for points in points_list:
                     x1, y1, x2, y2 = points
 
@@ -107,7 +111,7 @@ def draw_bounding_boxes(image, refs, ouput_path):
                     x2 = int(x2 / 999 * image_width)
                     y2 = int(y2 / 999 * image_height)
 
-                    if label_type == 'image':
+                    if label_type == "image":
                         try:
                             cropped = image.crop((x1, y1, x2, y2))
                             cropped.save(f"{ouput_path}/images/{img_idx}.jpg")
@@ -115,24 +119,35 @@ def draw_bounding_boxes(image, refs, ouput_path):
                             print(e)
                             pass
                         img_idx += 1
-                        
+
                     try:
-                        if label_type == 'title':
+                        if label_type == "title":
                             draw.rectangle([x1, y1, x2, y2], outline=color, width=4)
-                            draw2.rectangle([x1, y1, x2, y2], fill=color_a, outline=(0, 0, 0, 0), width=1)
+                            draw2.rectangle(
+                                [x1, y1, x2, y2],
+                                fill=color_a,
+                                outline=(0, 0, 0, 0),
+                                width=1,
+                            )
                         else:
                             draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
-                            draw2.rectangle([x1, y1, x2, y2], fill=color_a, outline=(0, 0, 0, 0), width=1)
+                            draw2.rectangle(
+                                [x1, y1, x2, y2],
+                                fill=color_a,
+                                outline=(0, 0, 0, 0),
+                                width=1,
+                            )
                         text_x = x1
                         text_y = max(0, y1 - 15)
-                            
-                        
+
                         text_bbox = draw.textbbox((0, 0), label_type, font=font)
                         text_width = text_bbox[2] - text_bbox[0]
                         text_height = text_bbox[3] - text_bbox[1]
-                        draw.rectangle([text_x, text_y, text_x + text_width, text_y + text_height], 
-                                    fill=(255, 255, 255, 30))
-                        
+                        draw.rectangle(
+                            [text_x, text_y, text_x + text_width, text_y + text_height],
+                            fill=(255, 255, 255, 30),
+                        )
+
                         draw.text((text_x, text_y), label_type, font=font, fill=color)
                     except:
                         pass
@@ -143,17 +158,13 @@ def draw_bounding_boxes(image, refs, ouput_path):
 
 
 def process_image_with_refs(image, ref_texts, output_path):
-
     result_image = draw_bounding_boxes(image, ref_texts, output_path)
-    
-    return result_image
-
-
 
+    return result_image
 
 
 def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
-    best_ratio_diff = float('inf')
+    best_ratio_diff = float("inf")
     best_ratio = (1, 1)
     area = width * height
     for ratio in target_ratios:
@@ -169,20 +180,27 @@ def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_
     return best_ratio
 
 
-def dynamic_preprocess(image, min_num=2, max_num=9, image_size=640, use_thumbnail=False):
+def dynamic_preprocess(
+    image, min_num=2, max_num=9, image_size=640, use_thumbnail=False
+):
     orig_width, orig_height = image.size
     aspect_ratio = orig_width / orig_height
 
     # calculate the existing image aspect ratio
     target_ratios = set(
-        (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
-        i * j <= max_num and i * j >= min_num)
+        (i, j)
+        for n in range(min_num, max_num + 1)
+        for i in range(1, n + 1)
+        for j in range(1, n + 1)
+        if i * j <= max_num and i * j >= min_num
+    )
     # print(target_ratios)
     target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
 
     # find the closest aspect ratio to the target
     target_aspect_ratio = find_closest_aspect_ratio(
-        aspect_ratio, target_ratios, orig_width, orig_height, image_size)
+        aspect_ratio, target_ratios, orig_width, orig_height, image_size
+    )
 
     # print(target_aspect_ratio)
     # calculate the target width and height
@@ -198,7 +216,7 @@ def dynamic_preprocess(image, min_num=2, max_num=9, image_size=640, use_thumbnai
             (i % (target_width // image_size)) * image_size,
             (i // (target_width // image_size)) * image_size,
             ((i % (target_width // image_size)) + 1) * image_size,
-            ((i // (target_width // image_size)) + 1) * image_size
+            ((i // (target_width // image_size)) + 1) * image_size,
         )
         # split the image
         split_img = resized_img.crop(box)
@@ -210,15 +228,14 @@ def dynamic_preprocess(image, min_num=2, max_num=9, image_size=640, use_thumbnai
     return processed_images, target_aspect_ratio
 
 
-
 def normalize_transform(mean, std):
     if mean is None and std is None:
         transform = None
     elif mean is None and std is not None:
-        mean = [0.] * len(std)
+        mean = [0.0] * len(std)
         transform = transforms.Normalize(mean=mean, std=std)
     elif mean is not None and std is None:
-        std = [1.] * len(mean)
+        std = [1.0] * len(mean)
         transform = transforms.Normalize(mean=mean, std=std)
     else:
         transform = transforms.Normalize(mean=mean, std=std)
@@ -226,11 +243,10 @@ def normalize_transform(mean, std):
     return transform
 
 
-
 def format_messages(
-        conversations: List[Dict[str, str]],
-        sft_format: str = "deepseek",
-        system_prompt: str = "",
+    conversations: List[Dict[str, str]],
+    sft_format: str = "deepseek",
+    system_prompt: str = "",
 ):
     """
     Applies the SFT template to conversation.
@@ -264,6 +280,7 @@ def text_encode(tokenizer, text: str, bos: bool = True, eos: bool = False):
 
     return t
 
+
 def load_pil_images(conversations: List[Dict[str, str]]) -> List[Image.Image]:
     """
 
@@ -294,7 +311,7 @@ def load_pil_images(conversations: List[Dict[str, str]]) -> List[Image.Image]:
             # print(image_path)
             # print('----------------')
             # exit()
-            
+
             # pil_img = Image.open(image_path)
             pil_img = load_image(image_path)
             pil_img = pil_img.convert("RGB")
@@ -304,7 +321,6 @@ def load_pil_images(conversations: List[Dict[str, str]]) -> List[Image.Image]:
 
 
 class BaseTransform(ABC):
-
     def set_rng(self, *args, **kwargs):
         pass
 
@@ -318,32 +334,32 @@ class BaseTransform(ABC):
 
 class BasicImageTransform(BaseTransform):
     def __init__(
-        self, 
+        self,
         mean: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5),
         std: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5),
-        normalize: bool = True
+        normalize: bool = True,
     ):
         self.mean = mean
         self.std = std
-    
-        transform_pipelines = [
-            transforms.ToTensor()
-        ]
+
+        transform_pipelines = [transforms.ToTensor()]
 
         normalize = normalize_transform(mean, std) if normalize else nn.Identity()
         if normalize is not None:
             transform_pipelines.append(normalize)
 
         self.transform = transforms.Compose(transform_pipelines)
-    
+
     def __call__(self, x):
         x = self.transform(x)
         return x
 
+
 class NoEOSTextStreamer(TextStreamer):
     def on_finalized_text(self, text: str, stream_end: bool = False):
-
-        eos_text = self.tokenizer.decode([self.tokenizer.eos_token_id], skip_special_tokens=False)
+        eos_text = self.tokenizer.decode(
+            [self.tokenizer.eos_token_id], skip_special_tokens=False
+        )
         text = text.replace(eos_text, "\n")
         print(text, flush=True, end="")
 
@@ -351,6 +367,7 @@ class NoEOSTextStreamer(TextStreamer):
 class DeepseekOCRConfig(DeepseekV2Config):
     model_type = "DeepseekOCR"
 
+
 class DeepseekOCRModel(DeepseekV2Model):
     config_class = DeepseekOCRConfig
 
@@ -361,14 +378,13 @@ class DeepseekOCRModel(DeepseekV2Model):
         self.vision_model = build_clip_l()
         # self.conv_2 = nn.Conv2d(in_channels=1024, out_channels=2048, kernel_size=2, stride=2)
         n_embed = 1280
-        self.projector =  MlpProjector(Dict(projector_type="linear", input_dim=2048, n_embed=n_embed))
+        self.projector = MlpProjector(
+            Dict(projector_type="linear", input_dim=2048, n_embed=n_embed)
+        )
         embed_std = 1 / torch.sqrt(torch.tensor(n_embed, dtype=torch.float32))
         self.image_newline = nn.Parameter(torch.randn(n_embed) * embed_std)
         self.view_seperator = nn.Parameter(torch.randn(n_embed) * embed_std)
 
-
-
-    
     def forward(
         self,
         input_ids: torch.LongTensor = None,
@@ -384,28 +400,23 @@ class DeepseekOCRModel(DeepseekV2Model):
         images_spatial_crop: Optional[torch.FloatTensor] = None,
         return_dict: Optional[bool] = None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
-
-
-
-
         if inputs_embeds is None:
             # inputs_embeds = self.embed_tokens(input_ids)
             inputs_embeds = self.get_input_embeddings()(input_ids)
 
-
-
-        sam_model = getattr(self, 'sam_model', None)
+        sam_model = getattr(self, "sam_model", None)
         # sam_model = self.sam_model
-        vision_model = getattr(self, 'vision_model', None)
-
-
-
-        if sam_model is not None and (input_ids.shape[1] != 1 or self.training) and torch.sum(images[0][1]).item() != 0:
+        vision_model = getattr(self, "vision_model", None)
 
+        if (
+            sam_model is not None
+            and (input_ids.shape[1] != 1 or self.training)
+            and torch.sum(images[0][1]).item() != 0
+        ):
             idx = 0
-            
+
             # sam_model = torch.jit.script(sam_model)
-            
+
             # start_time = time.time()
             for image, crop_shape in zip(images, images_spatial_crop):
                 images_in_this_batch = []
@@ -414,53 +425,86 @@ class DeepseekOCRModel(DeepseekV2Model):
                 image_ori = image[1]
 
                 with torch.no_grad():
-                # with torch.inference_mode(): 
-                    
+                    # with torch.inference_mode():
+
                     if torch.sum(patches).item() != 0:
                         # P, C, H, W = patches.shape
                         crop_flag = 1
                         local_features_1 = sam_model(patches)
 
-                        local_features_2 = vision_model(patches, local_features_1)  
+                        local_features_2 = vision_model(patches, local_features_1)
                         # vit_time = time.time()
-                        local_features = torch.cat((local_features_2[:, 1:], local_features_1.flatten(2).permute(0, 2, 1)), dim=-1) 
+                        local_features = torch.cat(
+                            (
+                                local_features_2[:, 1:],
+                                local_features_1.flatten(2).permute(0, 2, 1),
+                            ),
+                            dim=-1,
+                        )
                         local_features = self.projector(local_features)
 
-
                         global_features_1 = sam_model(image_ori)
-                        global_features_2 = vision_model(image_ori, global_features_1) 
-                        global_features = torch.cat((global_features_2[:, 1:], global_features_1.flatten(2).permute(0, 2, 1)), dim=-1) 
+                        global_features_2 = vision_model(image_ori, global_features_1)
+                        global_features = torch.cat(
+                            (
+                                global_features_2[:, 1:],
+                                global_features_1.flatten(2).permute(0, 2, 1),
+                            ),
+                            dim=-1,
+                        )
                         global_features = self.projector(global_features)
 
-                        print('=====================')
-                        print('BASE: ', global_features.shape)
-                        print('PATCHES: ', local_features.shape)
-                        print('=====================')
+                        print("=====================")
+                        print("BASE: ", global_features.shape)
+                        print("PATCHES: ", local_features.shape)
+                        print("=====================")
 
                         _, hw, n_dim = global_features.shape
-                        h = w = int(hw ** 0.5)
+                        h = w = int(hw**0.5)
 
                         _2, hw2, n_dim2 = local_features.shape
-                        h2 = w2 = int(hw2 ** 0.5)
+                        h2 = w2 = int(hw2**0.5)
 
                         width_crop_num, height_crop_num = crop_shape[0], crop_shape[1]
 
                         global_features = global_features.view(h, w, n_dim)
 
                         global_features = torch.cat(
-                            [global_features, self.image_newline[None, None, :].expand(h, 1, n_dim)], dim=1
+                            [
+                                global_features,
+                                self.image_newline[None, None, :].expand(h, 1, n_dim),
+                            ],
+                            dim=1,
                         )
 
                         global_features = global_features.view(-1, n_dim)
 
-
-                        local_features = local_features.view(height_crop_num, width_crop_num, h2, w2, n_dim2).permute(0, 2, 1, 3, 4).reshape(height_crop_num*h2, width_crop_num*w2, n_dim2)
+                        local_features = (
+                            local_features.view(
+                                height_crop_num, width_crop_num, h2, w2, n_dim2
+                            )
+                            .permute(0, 2, 1, 3, 4)
+                            .reshape(height_crop_num * h2, width_crop_num * w2, n_dim2)
+                        )
                         local_features = torch.cat(
-                            [local_features, self.image_newline[None, None, :].expand(height_crop_num * h2, 1, n_dim2)], dim=1
+                            [
+                                local_features,
+                                self.image_newline[None, None, :].expand(
+                                    height_crop_num * h2, 1, n_dim2
+                                ),
+                            ],
+                            dim=1,
                         )
                         local_features = local_features.view(-1, n_dim2)
 
-                        global_local_features = torch.cat([local_features, global_features, self.view_seperator[None, :]], dim=0)
+                        global_local_features = torch.cat(
+                            [
+                                local_features,
+                                global_features,
+                                self.view_seperator[None, :],
+                            ],
+                            dim=0,
+                        )
 
                         # end_time = time.time()
 
@@ -469,32 +513,42 @@ class DeepseekOCRModel(DeepseekV2Model):
                         # print('all: ', end_time - start_time)
 
                         # exit()
-                   
+
                     else:
                         global_features_1 = sam_model(image_ori)
-                        global_features_2 = vision_model(image_ori, global_features_1) 
-                        global_features = torch.cat((global_features_2[:, 1:], global_features_1.flatten(2).permute(0, 2, 1)), dim=-1) 
+                        global_features_2 = vision_model(image_ori, global_features_1)
+                        global_features = torch.cat(
+                            (
+                                global_features_2[:, 1:],
+                                global_features_1.flatten(2).permute(0, 2, 1),
+                            ),
+                            dim=-1,
+                        )
                         global_features = self.projector(global_features)
-                        print('=====================')
-                        print('BASE: ', global_features.shape)
-                        print('NO PATCHES')
-                        print('=====================')
+                        print("=====================")
+                        print("BASE: ", global_features.shape)
+                        print("NO PATCHES")
+                        print("=====================")
                         _, hw, n_dim = global_features.shape
-                        h = w = int(hw ** 0.5)
-
+                        h = w = int(hw**0.5)
 
                         global_features = global_features.view(h, w, n_dim)
 
                         global_features = torch.cat(
-                            [global_features, self.image_newline[None, None, :].expand(h, 1, n_dim)], dim=1
+                            [
+                                global_features,
+                                self.image_newline[None, None, :].expand(h, 1, n_dim),
+                            ],
+                            dim=1,
                         )
 
                         global_features = global_features.view(-1, n_dim)
 
-                        global_local_features = torch.cat([global_features, self.view_seperator[None, :]], dim=0)
+                        global_local_features = torch.cat(
+                            [global_features, self.view_seperator[None, :]], dim=0
+                        )
 
                     images_in_this_batch.append(global_local_features)
-                
 
                 # print(inputs_embeds.shape)
 
@@ -502,21 +556,27 @@ class DeepseekOCRModel(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)
+                    inputs_embeds[idx].masked_scatter_(
+                        images_seq_mask[idx].unsqueeze(-1).to(self.device),
+                        images_in_this_batch,
+                    )
 
                 idx += 1
-            
 
         return super(DeepseekOCRModel, self).forward(
-            input_ids=None, attention_mask=attention_mask, past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds, use_cache=use_cache, position_ids = position_ids,
-            output_attentions=output_attentions, output_hidden_states=output_hidden_states,
-            return_dict=return_dict
+            input_ids=None,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
         )
-    
 
-class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
 
+class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
     config_class = DeepseekOCRConfig
     # supports_gradient_checkpointing = True
 
@@ -536,7 +596,6 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
     def get_model(self):
         return self.model
 
-
     def forward(
         self,
         input_ids: torch.LongTensor = None,
@@ -552,17 +611,22 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
         images_seq_mask: Optional[torch.FloatTensor] = None,
         images_spatial_crop: Optional[torch.FloatTensor] = None,
         return_dict: Optional[bool] = None,
-        
     ) -> Union[Tuple, CausalLMOutputWithPast]:
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
         output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-
 
-        outputs  = self.model(
+        outputs = self.model(
             input_ids=input_ids,
             past_key_values=past_key_values,
             attention_mask=attention_mask,
@@ -572,14 +636,11 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             images=images,
-            images_seq_mask = images_seq_mask,
-            images_spatial_crop = images_spatial_crop,
-            return_dict=return_dict
-            
+            images_seq_mask=images_seq_mask,
+            images_spatial_crop=images_spatial_crop,
+            return_dict=return_dict,
         )
 
-
-        
         # print(transformer_outputs)
 
         hidden_states = outputs[0]
@@ -613,9 +674,13 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
             attentions=outputs.attentions,
         )
 
-
     def prepare_inputs_for_generation(
-        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        **kwargs,
     ):
         # Omit tokens covered by past_key_values
         past_length = 0
@@ -632,7 +697,10 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
             # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
             # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
             # input)
-            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+            if (
+                attention_mask is not None
+                and attention_mask.shape[1] > input_ids.shape[1]
+            ):
                 input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
             # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
             # input_ids based on the past_length.
@@ -668,7 +736,11 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
 
         # TODO @gante we should only keep a `cache_position` in generate, and do +=1.
         # same goes for position ids. Could also help with continued generation.
-        cache_position = torch.arange(past_length, past_length + position_ids.shape[-1], device=position_ids.device)
+        cache_position = torch.arange(
+            past_length,
+            past_length + position_ids.shape[-1],
+            device=position_ids.device,
+        )
 
         # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
         if inputs_embeds is not None and past_key_values is None:
@@ -688,45 +760,55 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
             }
         )
         return model_inputs
-    
 
     def disable_torch_init(self):
         """
         Disable the redundant torch default initialization to accelerate model creation.
         """
         import torch
+
         setattr(torch.nn.Linear, "reset_parameters", lambda self: None)
         setattr(torch.nn.LayerNorm, "reset_parameters", lambda self: None)
 
-
-
-    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):
+    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,
+    ):
         self.disable_torch_init()
 
         os.makedirs(output_path, exist_ok=True)
-        os.makedirs(f'{output_path}/images', exist_ok=True)
+        os.makedirs(f"{output_path}/images", exist_ok=True)
 
         if prompt and image_file:
             conversation = [
                 {
                     "role": "<|User|>",
                     # "content": "<image>\n<|grounding|>Given the layout of the image. ",
-                    "content": f'{prompt}',
+                    "content": f"{prompt}",
                     # "content": "君不见黄河之水天上来的下一句是什么？",
                     # "content": "<image>\nFree OCR. ",
                     # "content": "<image>\nParse the figure. ",
                     # "content": "<image>\nExtract the text in the image. ",
-                    "images": [f'{image_file}'],
+                    "images": [f"{image_file}"],
                 },
                 {"role": "<|Assistant|>", "content": ""},
             ]
-        
+
         elif prompt:
             conversation = [
                 {
                     "role": "<|User|>",
                     # "content": "<image>\n<|grounding|>Given the layout of the image. ",
-                    "content": f'{prompt}',
+                    "content": f"{prompt}",
                     # "content": "君不见黄河之水天上来的下一句是什么？",
                     # "content": "<image>\nFree OCR. ",
                     # "content": "<image>\nParse the figure. ",
@@ -736,9 +818,11 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
                 {"role": "<|Assistant|>", "content": ""},
             ]
         else:
-            assert False, f'prompt is none!'
-        
-        prompt = format_messages(conversations=conversation, sft_format='plain', system_prompt='')
+            assert False, f"prompt is none!"
+
+        prompt = format_messages(
+            conversations=conversation, sft_format="plain", system_prompt=""
+        )
 
         patch_size = 16
         downsample_ratio = 4
@@ -749,15 +833,16 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
 
         image_draw = images[0].copy()
 
-        w,h = image_draw.size
+        w, h = image_draw.size
         # print(w, h)
         ratio = 1 - ((max(w, h) - min(w, h)) / (max(w, h)))
-    
 
-        image_transform=BasicImageTransform(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), normalize=True)
+        image_transform = BasicImageTransform(
+            mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), normalize=True
+        )
         images_seq_mask = []
 
-        image_token = '<image>'
+        image_token = "<image>"
         image_token_id = 128815
         text_splits = prompt.split(image_token)
 
@@ -765,13 +850,11 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
         tokenized_str = []
         images_spatial_crop = []
         for text_sep, image in zip(text_splits, images):
-
             tokenized_sep = text_encode(tokenizer, text_sep, bos=False, eos=False)
             tokenized_str += tokenized_sep
             images_seq_mask += [False] * len(tokenized_sep)
 
             if crop_mode:
-
                 if image.size[0] <= 640 and image.size[1] <= 640:
                     crop_ratio = [1, 1]
 
@@ -782,23 +865,22 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
                     else:
                         # best_width, best_height = self.image_size, self.image_size
                         crop_ratio = [1, 1]
-                
+
                 """process the global view"""
                 # image = image.resize((base_size, base_size))
-                global_view = ImageOps.pad(image, (base_size, base_size),
-                                        color=tuple(int(x * 255) for x in image_transform.mean))
-                
+                global_view = ImageOps.pad(
+                    image,
+                    (base_size, base_size),
+                    color=tuple(int(x * 255) for x in image_transform.mean),
+                )
+
                 if base_size == 1024:
                     valid_img_tokens += int(256 * ratio)
                 elif base_size == 1280:
                     valid_img_tokens += int(400 * ratio)
                 # elif base_size == 640:
                 #     valid_img_tokens += int(100 * ratio)
-                
-
 
-
-                
                 images_list.append(image_transform(global_view).to(torch.bfloat16))
 
                 # global_view_tensor = image_transform(global_view).to(torch.bfloat16)
@@ -806,31 +888,34 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
                 width_crop_num, height_crop_num = crop_ratio
 
                 images_spatial_crop.append([width_crop_num, height_crop_num])
-                
-                
+
                 if width_crop_num > 1 or height_crop_num > 1:
                     """process the local views"""
-                    
+
                     for i in range(len(images_crop_raw)):
-                        images_crop_list.append(image_transform(images_crop_raw[i]).to(torch.bfloat16))
-                
+                        images_crop_list.append(
+                            image_transform(images_crop_raw[i]).to(torch.bfloat16)
+                        )
+
                 if image_size == 640:
                     valid_img_tokens += len(images_crop_list) * 100
 
                 num_queries = math.ceil((image_size // patch_size) / downsample_ratio)
-                num_queries_base = math.ceil((base_size // patch_size) / downsample_ratio)
-
-
+                num_queries_base = math.ceil(
+                    (base_size // patch_size) / downsample_ratio
+                )
 
                 """add image tokens"""
 
-                
-
-                tokenized_image = ([image_token_id] * num_queries_base + [image_token_id]) * num_queries_base
+                tokenized_image = (
+                    [image_token_id] * num_queries_base + [image_token_id]
+                ) * num_queries_base
                 tokenized_image += [image_token_id]
                 if width_crop_num > 1 or height_crop_num > 1:
-                    tokenized_image += ([image_token_id] * (num_queries * width_crop_num) + [image_token_id]) * (
-                                num_queries * height_crop_num)
+                    tokenized_image += (
+                        [image_token_id] * (num_queries * width_crop_num)
+                        + [image_token_id]
+                    ) * (num_queries * height_crop_num)
                 tokenized_str += tokenized_image
                 images_seq_mask += [True] * len(tokenized_image)
                 # num_image_tokens.append(len(tokenized_image))
@@ -841,11 +926,14 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
 
                 """process the global view"""
                 if image_size <= 640:
-                    print('directly resize')
+                    print("directly resize")
                     image = image.resize((image_size, image_size))
                 # else:
-                global_view = ImageOps.pad(image, (image_size, image_size),
-                                        color=tuple(int(x * 255) for x in image_transform.mean))
+                global_view = ImageOps.pad(
+                    image,
+                    (image_size, image_size),
+                    color=tuple(int(x * 255) for x in image_transform.mean),
+                )
                 images_list.append(image_transform(global_view).to(torch.bfloat16))
 
                 if base_size == 1024:
@@ -861,18 +949,18 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
 
                 images_spatial_crop.append([width_crop_num, height_crop_num])
 
-
                 """add image tokens"""
                 num_queries = math.ceil((image_size // patch_size) / downsample_ratio)
 
-                tokenized_image = ([image_token_id] * num_queries + [image_token_id]) * num_queries
+                tokenized_image = (
+                    [image_token_id] * num_queries + [image_token_id]
+                ) * num_queries
                 tokenized_image += [image_token_id]
                 # tokenized_image += ([self.image_token_id] * (num_queries * width_crop_num) + [self.image_token_id]) * (
                 #             num_queries * height_crop_num)
                 tokenized_str += tokenized_image
                 images_seq_mask += [True] * len(tokenized_image)
                 # num_image_tokens.append(len(tokenized_image))
-        
 
         """process the last text split"""
         tokenized_sep = text_encode(tokenizer, text_splits[-1], bos=False, eos=False)
@@ -881,19 +969,13 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
 
         """add the bos tokens"""
         bos_id = 0
-        tokenized_str = [bos_id] + tokenized_str 
+        tokenized_str = [bos_id] + tokenized_str
         images_seq_mask = [False] + images_seq_mask
 
-
-
         input_ids = torch.LongTensor(tokenized_str)
 
-
-        
-
         images_seq_mask = torch.tensor(images_seq_mask, dtype=torch.bool)
 
-
         if len(images_list) == 0:
             images_ori = torch.zeros((1, 3, image_size, image_size))
             images_spatial_crop = torch.zeros((1, 2), dtype=torch.long)
@@ -907,131 +989,157 @@ class DeepseekOCRForCausalLM(DeepseekV2ForCausalLM):
             else:
                 images_crop = torch.zeros((1, 3, base_size, base_size))
 
-
-
         if not eval_mode:
-            streamer = NoEOSTextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
-            with torch.autocast("cuda", dtype=torch.bfloat16):
+            streamer = NoEOSTextStreamer(
+                tokenizer, skip_prompt=True, skip_special_tokens=False
+            )
+            with torch.autocast(self.device.type, dtype=torch.bfloat16):
                 with torch.no_grad():
                     output_ids = self.generate(
-                        input_ids.unsqueeze(0).cuda(),
-                        images=[(images_crop.cuda(), images_ori.cuda())],
-                        images_seq_mask = images_seq_mask.unsqueeze(0).cuda(),
-                        images_spatial_crop = images_spatial_crop,
+                        input_ids.unsqueeze(0).to(self.device),
+                        images=[
+                            (images_crop.to(self.device), images_ori.to(self.device))
+                        ],
+                        images_seq_mask=images_seq_mask.unsqueeze(0).to(self.device),
+                        images_spatial_crop=images_spatial_crop,
                         # do_sample=False,
                         # num_beams = 1,
                         temperature=0.0,
                         eos_token_id=tokenizer.eos_token_id,
                         streamer=streamer,
                         max_new_tokens=8192,
-                        no_repeat_ngram_size = 20,
-                        use_cache = True
-                        )
+                        no_repeat_ngram_size=20,
+                        use_cache=True,
+                    )
 
         else:
-            with torch.autocast("cuda", dtype=torch.bfloat16):
+            with torch.autocast(self.device.type, dtype=torch.bfloat16):
                 with torch.no_grad():
                     output_ids = self.generate(
-                        input_ids.unsqueeze(0).cuda(),
-                        images=[(images_crop.cuda(), images_ori.cuda())],
-                        images_seq_mask = images_seq_mask.unsqueeze(0).cuda(),
-                        images_spatial_crop = images_spatial_crop,
+                        input_ids.unsqueeze(0).to(self.device),
+                        images=[
+                            (images_crop.to(self.device), images_ori.to(self.device))
+                        ],
+                        images_seq_mask=images_seq_mask.unsqueeze(0).to(self.device),
+                        images_spatial_crop=images_spatial_crop,
                         # do_sample=False,
                         # num_beams = 1,
                         temperature=0.0,
                         eos_token_id=tokenizer.eos_token_id,
                         max_new_tokens=8192,
-                        no_repeat_ngram_size = 35,
-                        use_cache = True
-                        )
-                
-
-        if '<image>' in conversation[0]['content'] and eval_mode:
-                outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
-                stop_str = '<｜end▁of▁sentence｜>'
-                if outputs.endswith(stop_str):
-                    outputs = outputs[:-len(stop_str)]
-                # re_match
-                outputs = outputs.strip()
-
-                return outputs
-        
-        if '<image>' in conversation[0]['content'] and test_compress:
-            outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
-            pure_texts_outputs_token_length = len(text_encode(tokenizer, outputs, bos=False, eos=False))
-            print('='*50)
-            print('image size: ', (w, h))
-            print('valid image tokens: ', int(valid_img_tokens))
-            print('output texts tokens (valid): ', pure_texts_outputs_token_length)
-            print('compression ratio: ', round(pure_texts_outputs_token_length/valid_img_tokens, 2))
-            print('='*50)
-
-
-        if '<image>' in conversation[0]['content'] and save_results:
-            outputs = tokenizer.decode(output_ids[0, input_ids.unsqueeze(0).cuda().shape[1]:])
-            stop_str = '<｜end▁of▁sentence｜>'
-
-            print('='*15 + 'save results:' + '='*15)
-            
+                        no_repeat_ngram_size=35,
+                        use_cache=True,
+                    )
+
+        if "<image>" in conversation[0]["content"] and eval_mode:
+            outputs = tokenizer.decode(
+                output_ids[0, input_ids.unsqueeze(0).to(self.device).shape[1] :]
+            )
+            stop_str = "<｜end▁of▁sentence｜>"
+            if outputs.endswith(stop_str):
+                outputs = outputs[: -len(stop_str)]
+            # re_match
+            outputs = outputs.strip()
+
+            return outputs
+
+        if "<image>" in conversation[0]["content"] and test_compress:
+            outputs = tokenizer.decode(
+                output_ids[0, input_ids.unsqueeze(0).to(self.device).shape[1] :]
+            )
+            pure_texts_outputs_token_length = len(
+                text_encode(tokenizer, outputs, bos=False, eos=False)
+            )
+            print("=" * 50)
+            print("image size: ", (w, h))
+            print("valid image tokens: ", int(valid_img_tokens))
+            print("output texts tokens (valid): ", pure_texts_outputs_token_length)
+            print(
+                "compression ratio: ",
+                round(pure_texts_outputs_token_length / valid_img_tokens, 2),
+            )
+            print("=" * 50)
+
+        if "<image>" in conversation[0]["content"] and save_results:
+            outputs = tokenizer.decode(
+                output_ids[0, input_ids.unsqueeze(0).to(self.device).shape[1] :]
+            )
+            stop_str = "<｜end▁of▁sentence｜>"
+
+            print("=" * 15 + "save results:" + "=" * 15)
+
             # # # # conv.messages[-1][-1] = outputs
             if outputs.endswith(stop_str):
-                outputs = outputs[:-len(stop_str)]
+                outputs = outputs[: -len(stop_str)]
             outputs = outputs.strip()
 
             matches_ref, matches_images, mathes_other = re_match(outputs)
             # print(matches_ref)
             result = process_image_with_refs(image_draw, matches_ref, output_path)
 
-
             for idx, a_match_image in enumerate(tqdm(matches_images, desc="image")):
-                outputs = outputs.replace(a_match_image, '![](images/' + str(idx) + '.jpg)\n')
-            
-            for idx, a_match_other in enumerate(tqdm(mathes_other, desc="other")):
-                outputs = outputs.replace(a_match_other, '').replace('\\coloneqq', ':=').replace('\\eqqcolon', '=:')
+                outputs = outputs.replace(
+                    a_match_image, "![](images/" + str(idx) + ".jpg)\n"
+                )
 
+            for idx, a_match_other in enumerate(tqdm(mathes_other, desc="other")):
+                outputs = (
+                    outputs.replace(a_match_other, "")
+                    .replace("\\coloneqq", ":=")
+                    .replace("\\eqqcolon", "=:")
+                )
 
             # if 'structural formula' in conversation[0]['content']:
             #     outputs = '<smiles>' + outputs + '</smiles>'
-            with open(f'{output_path}/result.mmd', 'w', encoding = 'utf-8') as afile:
+            with open(f"{output_path}/result.mmd", "w", encoding="utf-8") as afile:
                 afile.write(outputs)
 
-            if 'line_type' in outputs:
+            if "line_type" in outputs:
                 import matplotlib.pyplot as plt
-                lines = eval(outputs)['Line']['line']
 
-                line_type = eval(outputs)['Line']['line_type']
+                lines = eval(outputs)["Line"]["line"]
+
+                line_type = eval(outputs)["Line"]["line_type"]
                 # print(lines)
 
-                endpoints = eval(outputs)['Line']['line_endpoint']
+                endpoints = eval(outputs)["Line"]["line_endpoint"]
 
-                fig, ax = plt.subplots(figsize=(3,3), dpi=200)
+                fig, ax = plt.subplots(figsize=(3, 3), dpi=200)
                 ax.set_xlim(-15, 15)
                 ax.set_ylim(-15, 15)
 
                 for idx, line in enumerate(lines):
                     try:
-                        p0 = eval(line.split(' -- ')[0])
-                        p1 = eval(line.split(' -- ')[-1])
+                        p0 = eval(line.split(" -- ")[0])
+                        p1 = eval(line.split(" -- ")[-1])
 
-                        if line_type[idx] == '--':
-                            ax.plot([p0[0], p1[0]], [p0[1], p1[1]], linewidth=0.8, color='k')
+                        if line_type[idx] == "--":
+                            ax.plot(
+                                [p0[0], p1[0]], [p0[1], p1[1]], linewidth=0.8, color="k"
+                            )
                         else:
-                            ax.plot([p0[0], p1[0]], [p0[1], p1[1]], linewidth = 0.8, color = 'k')
+                            ax.plot(
+                                [p0[0], p1[0]], [p0[1], p1[1]], linewidth=0.8, color="k"
+                            )
 
-                        ax.scatter(p0[0], p0[1], s=5, color = 'k')
-                        ax.scatter(p1[0], p1[1], s=5, color = 'k')
+                        ax.scatter(p0[0], p0[1], s=5, color="k")
+                        ax.scatter(p1[0], p1[1], s=5, color="k")
                     except:
                         pass
 
                 for endpoint in endpoints:
-
-                    label = endpoint.split(': ')[0]
-                    (x, y) = eval(endpoint.split(': ')[1])
-                    ax.annotate(label, (x, y), xytext=(1, 1), textcoords='offset points', 
-                                fontsize=5, fontweight='light')
-                
-
-                plt.savefig(f'{output_path}/geo.jpg')
+                    label = endpoint.split(": ")[0]
+                    (x, y) = eval(endpoint.split(": ")[1])
+                    ax.annotate(
+                        label,
+                        (x, y),
+                        xytext=(1, 1),
+                        textcoords="offset points",
+                        fontsize=5,
+                        fontweight="light",
+                    )
+
+                plt.savefig(f"{output_path}/geo.jpg")
                 plt.close()
 
             result.save(f"{output_path}/result_with_boxes.jpg")