Change the call interface and adjust the program execution logic

README.md

@@ -30,6 +30,7 @@ base_model:
 
 ## Evaluation
 
+If you just want to use this code, please refer to this sample below
 ```python
 model_path = "DeepGlint-AI/MLCD-Seg" # or use your local path
 mlcd_seg = AutoModel.from_pretrained(
@@ -40,19 +41,23 @@ mlcd_seg = AutoModel.from_pretrained(
 tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
 # Assuming you have an image named test.jpg
 seg_img = Image.open("test.jpg").convert('RGB')
-seg_prompt = "The <image> provides an overview of the picture.\nCould you provide a segmentation mask for the right giraffe in this image?"
-pred_mask = model.predict_forward(seg_img, seg_prompt, tokenizer, force_seg=False)
+seg_prompt = "Could you provide a segmentation mask for the right giraffe in this image?"
+pred_mask = model.seg(seg_img, seg_prompt, tokenizer, force_seg=False)
 ```
 
-## Tips for updating this repo in the future
-
-
-Huggingface uses cache management module code, so manual clearing of cache is required after repo update
-
-
-```bash
-cd ~/.cache/huggingface/modules/transformers_modules
-rm mlcd_seg.py vision_projector.py vision_resampler.py vision_tower.py sam.py conversation_mlcd_seg.py
+If you want to use this code measurement dataset (e.g. refcoco), then you need to use the following method
+```python
+model_path = "DeepGlint-AI/MLCD-Seg" # or use your local path
+mlcd_seg = AutoModel.from_pretrained(
+    model_path,
+    torch_dtype=torch.float16,
+    trust_remote_code=True
+).cuda()
+tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+# Assuming you have an image named test.jpg
+seg_img = Image.open("test.jpg").convert('RGB')
+seg_prompt = "Could you provide a segmentation mask for the right giraffe in this image?"
+pred_mask = model.seg(seg_img, seg_prompt, tokenizer, force_seg=True)
 ```
 
 

config.json

@@ -1,5 +1,5 @@
 {
-    "_name_or_path": "DeepGlint-AI/MLCD-Embodied-7B",
+    "_name_or_path": "DeepGlint-AI/MLCD-Seg",
     "add_faster_video": false,
     "add_time_instruction": false,
     "architectures": [

mlcd_seg.py

@@ -27,6 +27,7 @@ import random
 import ast
 import re
 import json
+import os
 import numpy as np
 import torch
 import torch.nn as nn
@@ -42,7 +43,7 @@ from .vision_tower import build_vision_tower
 from .vision_resampler import build_vision_resampler
 from .vision_projector import build_vision_projector
 from .sam import build_sam_vit_h, text2sam_projection_layer
-from .conversation_mlcd_seg import default_conversation
+from .conversation_mlcd_seg import conv_templates, default_conversation
 from .transform import ResizeLongestSide
 from typing import Optional, Any, List, Tuple, Union, Dict
 
@@ -140,7 +141,10 @@ class MLCDSegMetaModel:
 
     def dispatch_weight(self, config):
         safetensors_set = set()
-        index_file = Path(getattr(config, "name_or_path", "./")) / "model.safetensors.index.json"
+        index_folder = Path(getattr(config, "_name_or_path", "./"))
+        index_file = index_folder / "model.safetensors.index.json"
+        if not index_file.exists():
+            os.getenv("")
         with open(index_file, "r") as safetensors_index:
             safetensors_map = json.loads(safetensors_index.read())
             for key, value in safetensors_map["weight_map"].items():
@@ -152,7 +156,7 @@ class MLCDSegMetaModel:
         projector_weight = {}
         text2sam_projection_weight = {}
         for safetensors_file in safetensors_set:
-            temp_load = safetensors_load(safetensors_file)
+            temp_load = safetensors_load(index_folder / safetensors_file)
             for key, value in temp_load.items():
                 if key.startswith("model.sam."):
                     sam_weight[key.replace("model.sam.", "")] = value
@@ -170,85 +174,85 @@ class MLCDSegMetaModel:
             vision_tower = vision_tower[0]
         return vision_tower
 
-    def initialize_vision_modules(self, model_args, fsdp=None):
-        vision_tower = model_args.vision_tower
-        mm_vision_select_layer = model_args.mm_vision_select_layer
-        mm_vision_select_feature = model_args.mm_vision_select_feature
-        pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter
-        mm_patch_merge_type = model_args.mm_patch_merge_type
-
-        self.config.mm_vision_tower = vision_tower
-        self.config.vision_tower_pretrained = getattr(model_args, "vision_tower_pretrained", "")
-
-        if self.get_vision_tower() is None:
-            vision_tower = build_vision_tower(model_args)
-            vision_resampler = build_vision_resampler(model_args, vision_tower=vision_tower)
-            for k, v in vision_resampler.config.items():
-                setattr(self.config, k, v)
-
-            if fsdp is not None and len(fsdp) > 0:
-                self.vision_tower = [vision_tower]
-                self.vision_resampler = [vision_resampler]
-            else:
-                self.vision_tower = vision_tower
-                self.vision_resampler = vision_resampler
-        else:
-            if fsdp is not None and len(fsdp) > 0:
-                vision_resampler = self.vision_resampler[0]
-                vision_tower = self.vision_tower[0]
-            else:
-                vision_resampler = self.vision_resampler
-                vision_tower = self.vision_tower
-            vision_tower.load_model()
-
-            # In case it is frozen by LoRA
-            for p in self.vision_resampler.parameters():
-                p.requires_grad = True
-
-        self.config.use_mm_proj = True
-        self.config.mm_projector_type = getattr(model_args, "mm_projector_type", "linear")
-        self.config.mm_hidden_size = getattr(vision_resampler, "hidden_size", vision_tower.hidden_size)
-        self.config.mm_vision_select_layer = mm_vision_select_layer
-        self.config.mm_vision_select_feature = mm_vision_select_feature
-        self.config.mm_patch_merge_type = mm_patch_merge_type
+    # def initialize_vision_modules(self, model_args, fsdp=None):
+    #     vision_tower = model_args.vision_tower
+    #     mm_vision_select_layer = model_args.mm_vision_select_layer
+    #     mm_vision_select_feature = model_args.mm_vision_select_feature
+    #     pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter
+    #     mm_patch_merge_type = model_args.mm_patch_merge_type
+
+    #     self.config.mm_vision_tower = vision_tower
+    #     self.config.vision_tower_pretrained = getattr(model_args, "vision_tower_pretrained", "")
+
+    #     if self.get_vision_tower() is None:
+    #         vision_tower = build_vision_tower(model_args)
+    #         vision_resampler = build_vision_resampler(model_args, vision_tower=vision_tower)
+    #         for k, v in vision_resampler.config.items():
+    #             setattr(self.config, k, v)
+
+    #         if fsdp is not None and len(fsdp) > 0:
+    #             self.vision_tower = [vision_tower]
+    #             self.vision_resampler = [vision_resampler]
+    #         else:
+    #             self.vision_tower = vision_tower
+    #             self.vision_resampler = vision_resampler
+    #     else:
+    #         if fsdp is not None and len(fsdp) > 0:
+    #             vision_resampler = self.vision_resampler[0]
+    #             vision_tower = self.vision_tower[0]
+    #         else:
+    #             vision_resampler = self.vision_resampler
+    #             vision_tower = self.vision_tower
+    #         vision_tower.load_model()
+
+    #         # In case it is frozen by LoRA
+    #         for p in self.vision_resampler.parameters():
+    #             p.requires_grad = True
+
+    #     self.config.use_mm_proj = True
+    #     self.config.mm_projector_type = getattr(model_args, "mm_projector_type", "linear")
+    #     self.config.mm_hidden_size = getattr(vision_resampler, "hidden_size", vision_tower.hidden_size)
+    #     self.config.mm_vision_select_layer = mm_vision_select_layer
+    #     self.config.mm_vision_select_feature = mm_vision_select_feature
+    #     self.config.mm_patch_merge_type = mm_patch_merge_type
         
-        for key in vars(model_args):  
-            if key.startswith('sam_'):
-                setattr(self.config, key, getattr(model_args, key))
+    #     for key in vars(model_args):  
+    #         if key.startswith('sam_'):
+    #             setattr(self.config, key, getattr(model_args, key))
         
-        if not hasattr(self.config, 'add_faster_video'):
-            if model_args.add_faster_video:
-                embed_std = 1 / torch.sqrt(torch.tensor(self.config.hidden_size, dtype=self.dtype))
-                self.faster_token = nn.Parameter(
-                    torch.randn(self.config.hidden_size, dtype=self.dtype) * embed_std
-                )
-
-        if getattr(self, "mm_projector", None) is None:
-            self.mm_projector = build_vision_projector(self.config, vision_cfg=vision_tower.config)
-
-            if "unpad" in mm_patch_merge_type:
-                embed_std = 1 / torch.sqrt(torch.tensor(self.config.hidden_size, dtype=self.dtype))
-                self.image_newline = nn.Parameter(torch.randn(self.config.hidden_size, dtype=self.dtype) * embed_std)
+    #     if not hasattr(self.config, 'add_faster_video'):
+    #         if model_args.add_faster_video:
+    #             embed_std = 1 / torch.sqrt(torch.tensor(self.config.hidden_size, dtype=self.dtype))
+    #             self.faster_token = nn.Parameter(
+    #                 torch.randn(self.config.hidden_size, dtype=self.dtype) * embed_std
+    #             )
+
+    #     if getattr(self, "mm_projector", None) is None:
+    #         self.mm_projector = build_vision_projector(self.config, vision_cfg=vision_tower.config)
+
+    #         if "unpad" in mm_patch_merge_type:
+    #             embed_std = 1 / torch.sqrt(torch.tensor(self.config.hidden_size, dtype=self.dtype))
+    #             self.image_newline = nn.Parameter(torch.randn(self.config.hidden_size, dtype=self.dtype) * embed_std)
                 
-            if getattr(self.config, 'sam_path', None) is not None:
-                self.sam = build_sam_vit_h(self.config.sam_path)
-                self.text2sam_projection = text2sam_projection_layer(self.config)
-        else:
-            if getattr(self.config, 'sam_path', None) is not None and self.config.sam_path !="":
-                self.sam = build_sam_vit_h(self.config.sam_path)
-                self.text2sam_projection = text2sam_projection_layer(self.config)
-            # In case it is frozen by LoRA
-            for p in self.mm_projector.parameters():
-                p.requires_grad = True
+    #         if getattr(self.config, 'sam_path', None) is not None:
+    #             self.sam = build_sam_vit_h(self.config.sam_path)
+    #             self.text2sam_projection = text2sam_projection_layer(self.config)
+    #     else:
+    #         if getattr(self.config, 'sam_path', None) is not None and self.config.sam_path !="":
+    #             self.sam = build_sam_vit_h(self.config.sam_path)
+    #             self.text2sam_projection = text2sam_projection_layer(self.config)
+    #         # In case it is frozen by LoRA
+    #         for p in self.mm_projector.parameters():
+    #             p.requires_grad = True
 
-        if pretrain_mm_mlp_adapter is not None:
-            mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location="cpu")
+    #     if pretrain_mm_mlp_adapter is not None:
+    #         mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location="cpu")
 
-            def get_w(weights, keyword):
-                return {k.split(keyword + ".")[1]: v for k, v in weights.items() if keyword in k}
+    #         def get_w(weights, keyword):
+    #             return {k.split(keyword + ".")[1]: v for k, v in weights.items() if keyword in k}
 
-            incompatible_keys = self.mm_projector.load_state_dict(get_w(mm_projector_weights, "mm_projector"))
-            incompatible_keys = self.vision_resampler.load_state_dict(get_w(mm_projector_weights, "vision_resampler"), strict=False)
+    #         incompatible_keys = self.mm_projector.load_state_dict(get_w(mm_projector_weights, "mm_projector"))
+    #         incompatible_keys = self.vision_resampler.load_state_dict(get_w(mm_projector_weights, "vision_resampler"), strict=False)
 
 
 def unpad_image(tensor, original_size):
@@ -774,8 +778,61 @@ class MLCDSegForCausalLM(Qwen2ForCausalLM, MLCDSegMetaForCausalLM):
         image_sam_resizes: Optional[List[torch.FloatTensor]] = None,
         original_sizes: Optional[List[torch.FloatTensor]] = None,
         masks_list: Optional[List[List[torch.FloatTensor]]] = None,
-        infer: bool = False,
-        force_seg: bool = True
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        if inputs_embeds is None:
+            (
+                input_ids,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                inputs_embeds,
+                labels
+            ) = self.prepare_inputs_labels_for_multimodal(
+                input_ids,
+                position_ids,
+                attention_mask,
+                past_key_values,
+                labels,
+                images,
+                modalities,
+                image_sizes
+            )
+        output = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=True,
+            return_dict=return_dict,
+            cache_position=cache_position
+        )
+        return MLCDSegOutputWithPast(**output)
+    
+    def seg_forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        images: Optional[torch.FloatTensor] = None,
+        image_sizes: Optional[List[List[int]]] = None,
+        return_dict: Optional[bool] = None,
+        modalities: Optional[List[str]] = ["image"],
+        dpo_forward: Optional[bool] = False,
+        cache_position=None,
+        grounding_enc_imgs: Optional[List[torch.FloatTensor]] = None,
+        image_sam_resizes: Optional[List[torch.FloatTensor]] = None,
+        original_sizes: Optional[List[torch.FloatTensor]] = None,
+        masks_list: Optional[List[List[torch.FloatTensor]]] = None,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         input_ids_ = input_ids
         if inputs_embeds is None:
@@ -832,16 +889,10 @@ class MLCDSegForCausalLM(Qwen2ForCausalLM, MLCDSegMetaForCausalLM):
                 cache_position=cache_position
             )
         sam_image_embeddings = self.get_grounding_encoder_embs(grounding_enc_imgs)
-        if force_seg:
-            seg_token_mask = self.create_seg_token_mask(input_ids_, old_attention_mask, img_token_num, num_images_batch)
-        else:
-            # should be raise NotImplementedError
-            seg_token_mask = self.create_seg_token_mask(input_ids_, old_attention_mask, img_token_num, num_images_batch)
+        seg_token_mask = self.create_seg_token_mask(input_ids_, old_attention_mask, img_token_num, num_images_batch)
         seg_text_embeds_batch = self.process_hidden_states(output["hidden_states"], seg_token_mask)
         pred_masks_batch = self.generate_and_postprocess_masks(seg_text_embeds_batch, sam_image_embeddings, num_images_batch, image_sam_resizes, original_sizes)
-        if infer:
-            return {"output":output, "pred_masks":pred_masks_batch}
-        return MLCDSegOutputWithPast(**output)
+        return pred_masks_batch
 
     @torch.no_grad()
     def generate(
@@ -856,13 +907,29 @@ class MLCDSegForCausalLM(Qwen2ForCausalLM, MLCDSegMetaForCausalLM):
         attention_mask = kwargs.pop("attention_mask", None)
         if "inputs_embeds" in kwargs:
             raise NotImplementedError("`inputs_embeds` is not supported")
+        (
+            inputs,
+            position_ids,
+            attention_mask,
+            _,
+            inputs_embeds,
+            _,
+            old_attention_mask,
+            img_token_num,
+            num_images_batch
+        ) = self.prepare_inputs_labels_for_multimodal(
+            inputs,
+            position_ids,
+            attention_mask,
+            None,
+            None,
+            images,
+            image_sizes=image_sizes,
+            # batch_pboxes=all_pboxes
+        )
+        llm_out = super().generate(position_ids=position_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, output_hidden_states=True, return_dict_in_generate=True, max_length=4096, **kwargs)
+        return llm_out.sequences
 
-        if images is not None:
-            (inputs, position_ids, attention_mask, _, inputs_embeds, _) = self.prepare_inputs_labels_for_multimodal(inputs, position_ids, attention_mask, None, None, images, modalities, image_sizes=image_sizes)
-        else:
-            inputs_embeds = self.get_model().embed_tokens(inputs)
-
-        return super().generate(position_ids=position_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, **kwargs)
 
     def generate_and_postprocess_masks(self, seg_text_embeds_batch, sam_image_embeddings, num_images_batch, image_sam_resizes, original_sizes):
         assert len(seg_text_embeds_batch) == len(num_images_batch)
@@ -911,20 +978,18 @@ class MLCDSegForCausalLM(Qwen2ForCausalLM, MLCDSegMetaForCausalLM):
             mask = input_ids[i][num_images_batch[i]:] == self.seg_token_idx
             seg_token_mask.append(
                 torch.cat(
-                    [torch.zeros((1, img_token_num[i])).bool().cuda(), mask.unsqueeze(0), torch.zeros((1, max_len-(len(input_ids[i]) + img_token_num[i] - num_images_batch[i]))).bool().cuda()], dim=1
+                    [torch.zeros((1, img_token_num[i])).bool().to(device=self.device), mask.unsqueeze(0), torch.zeros((1, max_len-(len(input_ids[i]) + img_token_num[i] - num_images_batch[i]))).bool().to(device=self.device)], dim=1
                 )
             )
         return torch.cat(seg_token_mask, dim=0)
             
     def get_grounding_encoder_embs(self, batch_images: torch.FloatTensor):
-        # with torch.no_grad():
         batch_feats = []
         for images in batch_images:
             batch_feats.append(torch.cat([self._encode_single_image(img) for img in images], dim=0))
         return batch_feats
 
     def _encode_single_image(self, image):
-        # torch.cuda.empty_cache()
         return self.model.sam.image_encoder(image.unsqueeze(0))
  
     def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
@@ -937,22 +1002,34 @@ class MLCDSegForCausalLM(Qwen2ForCausalLM, MLCDSegMetaForCausalLM):
             inputs["image_sizes"] = image_sizes
         return inputs
     
-    def process_prompt(self, text, tokenizer: PreTrainedTokenizer, force_seg=True) -> Dict:
-        conv = default_conversation.copy()
-        BEGIN_SIGNAL = "### "
-        END_SIGNAL = "\n"
-        roles = {"human": conv.roles[0], "gpt": conv.roles[1]}
-        # Apply prompt templates
-        sys_prompt = default_conversation.system + "\n\n"
-        full_prompt = sys_prompt + BEGIN_SIGNAL + roles["human"] + ": " + text + END_SIGNAL
-        if force_seg:
+    def process_prompt(self, text, tokenizer: PreTrainedTokenizer, stage="gen") -> Dict:
+        if stage.lower() not in ["gen", "seg"]:
+            stage = "seg"
+        if stage.lower() == "gen":
+            conv = conv_templates['qwen_2'].copy()
+            conv.append_message(conv.roles[0], text)
+            conv.append_message(conv.roles[1], None)
+            full_prompt = conv.get_prompt()
+            input_ids = torch.stack([gen_image_token(full_prompt, tokenizer, return_tensors='pt')], dim=0)
+            return dict(
+                input_ids=input_ids,
+                labels=None,
+            )
+        else:
+            conv = default_conversation.copy()
+            BEGIN_SIGNAL = "### "
+            END_SIGNAL = "\n"
+            roles = {"human": conv.roles[0], "gpt": conv.roles[1]}
+            # Apply prompt templates
+            sys_prompt = default_conversation.system + "\n\n" + "The <image> provides an overview of the picture.\n"
+            full_prompt = sys_prompt + BEGIN_SIGNAL + roles["human"] + ": " + text + END_SIGNAL
             full_prompt += BEGIN_SIGNAL + roles["gpt"] + ": It is [SEG]." + END_SIGNAL
-        full_prompt += BEGIN_SIGNAL
-        input_ids = torch.stack([gen_image_token(full_prompt, tokenizer, return_tensors='pt')], dim=0)
-        return dict(
-            input_ids=input_ids,
-            labels=None,
-        )
+            full_prompt += BEGIN_SIGNAL
+            input_ids = torch.stack([gen_image_token(full_prompt, tokenizer, return_tensors='pt')], dim=0)
+            return dict(
+                input_ids=input_ids,
+                labels=None,
+            )
     
     def process_images(self, images, image_processor, model_cfg):
         image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)
@@ -967,11 +1044,11 @@ class MLCDSegForCausalLM(Qwen2ForCausalLM, MLCDSegMetaForCausalLM):
             new_images = torch.stack(new_images, dim=0)
         return new_images
 
-    def predict_forward(self, image, prompt, tokenizer, force_seg=True):
+    def seg(self, image, prompt, tokenizer, force_seg = False):
         self.seg_token_idx = tokenizer(DEFAULT_SEG_TOKEN, add_special_tokens=False).input_ids[0]
         image_np = np.array(image)
         image_sizes = [image.size]
-        input_ids = self.process_prompt(prompt, tokenizer, force_seg)["input_ids"].to(self.device) # 这里需要设置对应的device
+        input_ids = self.process_prompt(prompt, tokenizer, "gen")["input_ids"].to(self.device)
         image_processor = self.get_vision_tower().image_processor
         image_tensors = self.process_images([image], image_processor, self.config)
         image_np_resize = self.sam_transform.apply_image(image_np)
@@ -994,21 +1071,39 @@ class MLCDSegForCausalLM(Qwen2ForCausalLM, MLCDSegMetaForCausalLM):
                 image_tensors = [[x_.unsqueeze(dim=0).to(dtype=self.dtype, device=self.device, non_blocking=True) for x_ in image_tensors]]
         else:
             image_tensors = image_tensors.to(dtype=self.dtype, device='cuda', non_blocking=True)
-        
-        with torch.inference_mode():
-            net_out = self.forward(
-                input_ids=input_ids,
-                output_hidden_states=True,
-                images=image_tensors,
-                image_sizes=image_sizes,
-                grounding_enc_imgs=[torch.stack(grounding_enc_img_list, dim=0)],
-                image_sam_resizes=[image_sam_resize_list],
-                original_sizes=[(mask_h, mask_w)],
-                infer=True,
-                force_seg=force_seg
-            )
-            pred_mask = net_out["pred_masks"][0]
+        if not force_seg:
+            attention_mask = torch.ones(input_ids.shape).bool().to(device=self.device)
+            with torch.inference_mode():
+                llm_gen = self.generate(
+                    inputs=input_ids,
+                    attention_mask=attention_mask,
+                    images=image_tensors,
+                    image_sizes=image_sizes,
+                    grounding_enc_imgs=[torch.stack(grounding_enc_img_list, dim=0)],
+                    image_sam_resizes=[image_sam_resize_list],
+                    original_sizes=[(mask_h, mask_w)],
+                    pad_token_id=tokenizer.eos_token_id
+                )
+            seg_flag = llm_gen == self.seg_token_idx
+            seg_flag = torch.sum(seg_flag.int()).item()
+            if seg_flag > 0:
+                force_seg = True
+        if force_seg:
+            input_ids = self.process_prompt(prompt, tokenizer, "seg")["input_ids"].to(self.device)
+            with torch.inference_mode():
+                net_out = self.seg_forward(
+                    input_ids=input_ids,
+                    output_hidden_states=True,
+                    images=image_tensors,
+                    image_sizes=image_sizes,
+                    grounding_enc_imgs=[torch.stack(grounding_enc_img_list, dim=0)],
+                    image_sam_resizes=[image_sam_resize_list],
+                    original_sizes=[(mask_h, mask_w)],
+                )
+            pred_mask = net_out[0]
             return pred_mask
+        else:
+            return torch.zeros([0] + list(image_np.shape[:2]), device=self.device)
 
 
 def gen_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):