Upload modeling_sa2va_chat.py with huggingface_hub

modeling_sa2va_chat.py

@@ -594,116 +594,137 @@ class Sa2VAChatModel(PreTrainedModel):
             assert tokenizer
             self.preparing_for_generation(tokenizer=tokenizer)
 
-        input_dict = {}
-        if video is not None:
-            pixel_values = []
-            extra_pixel_values = []
-            ori_image_size = video[0].size
-            for frame_idx, frame_image in enumerate(video):
-                assert ori_image_size == frame_image.size
-                g_image = np.array(frame_image)  # for grounding
-                g_image = self.extra_image_processor.apply_image(g_image)
-                g_image = torch.from_numpy(g_image).permute(2, 0, 1).contiguous()
-                extra_pixel_values.append(g_image)
-                if frame_idx < 5:
-                    img = self.transformer(frame_image)
-                    pixel_values.append(img)
-
-            pixel_values = torch.stack(pixel_values, dim=0).to(self.torch_dtype)  # (n_f, 3, h, w)
-            g_pixel_values = torch.stack([
-                self.grounding_encoder.preprocess_image(pixel) for pixel in extra_pixel_values
-            ]).to(self.torch_dtype)
-            num_image_tokens = self.patch_token
-            num_frames = 5
-
-            input_dict['vp_overall_mask'] = None
+        if image is None and video is None and '<image>' not in past_text:
+            text = text.replace('<image>', "")
+            input_text = ''
+            input_text += self.template['INSTRUCTION'].format(
+                input=text, round=1, bot_name=self.bot_name)
+            input_text = past_text + input_text
+            ids = self.tokenizer.encode(input_text)
+            ids = torch.tensor(ids).cuda().unsqueeze(0)
+
+            attention_mask = torch.ones_like(ids, dtype=torch.bool)
+
+            mm_inputs = {
+                'pixel_values': None,
+                'input_ids': ids,
+                'attention_mask': attention_mask,
+                'position_ids': None,
+                'past_key_values': None,
+                'labels': None,
+                'prompt_masks': None,
+                'vp_overall_mask': None,
+            }
+            ret_masks = []
         else:
-            ori_image_size = image.size
-
-            # prepare grounding images
-            g_image = np.array(image)  # for grounding
-            g_image = self.extra_image_processor.apply_image(g_image)
-            g_pixel_values = torch.from_numpy(g_image).permute(2, 0, 1).contiguous().to(self.torch_dtype)
-            extra_pixel_values = [g_pixel_values]
-            g_pixel_values = torch.stack([
-                self.grounding_encoder.preprocess_image(pixel) for pixel in extra_pixel_values
-            ]).to(self.torch_dtype)
+            input_dict = {}
+            if video is not None:
+                pixel_values = []
+                extra_pixel_values = []
+                ori_image_size = video[0].size
+                for frame_idx, frame_image in enumerate(video):
+                    assert ori_image_size == frame_image.size
+                    g_image = np.array(frame_image)  # for grounding
+                    g_image = self.extra_image_processor.apply_image(g_image)
+                    g_image = torch.from_numpy(g_image).permute(2, 0, 1).contiguous()
+                    extra_pixel_values.append(g_image)
+                    if frame_idx < 5:
+                        img = self.transformer(frame_image)
+                        pixel_values.append(img)
+
+                pixel_values = torch.stack(pixel_values, dim=0).to(self.torch_dtype)  # (n_f, 3, h, w)
+                g_pixel_values = torch.stack([
+                    self.grounding_encoder.preprocess_image(pixel) for pixel in extra_pixel_values
+                ]).to(self.torch_dtype)
+                num_image_tokens = self.patch_token
+                num_frames = len(pixel_values)
 
-            images = dynamic_preprocess(image, self.min_dynamic_patch,
-                                        self.max_dynamic_patch,
-                                        self.image_size, self.use_thumbnail)
-
-            if mask_prompts is not None:
-                vp_overall_mask = torch.Tensor([False] * (len(images) - 1) + [True])
-                input_dict['vp_overall_mask'] = vp_overall_mask
-            else:
                 input_dict['vp_overall_mask'] = None
+            else:
+                ori_image_size = image.size
 
-            pixel_values = [self.transformer(image) for image in images]
-            pixel_values = torch.stack(pixel_values).to(self.torch_dtype)
-            num_image_tokens = pixel_values.shape[0] * self.patch_token
-            num_frames = 1
-        input_dict['g_pixel_values'] = g_pixel_values
-        input_dict['pixel_values'] = pixel_values
-
-
-        if mask_prompts is not None:
-            # reshape mask prompts to feature size
-            mask_prompts = [torch.Tensor(item).to(pixel_values.device) for item in mask_prompts]
-            mask_prompts = [F.interpolate(
-                item.unsqueeze(0),
-                size=(int(self.image_size // self.patch_size * self.downsample_ratio),
-                      int(self.image_size // self.patch_size * self.downsample_ratio)),
-                mode='nearest').squeeze(0) for item in mask_prompts]
-            region_pixels = []
-            for mask_prompt in mask_prompts[0]:
-                region_pixels.append(mask_prompt.bool().to(torch.int64).sum())
-
-            vp_token_str = '\nThere are {} part regions in the picture: '.format(len(mask_prompts[0]))
-            for i in range(len(mask_prompts[0])):
-                vp_token_str = vp_token_str + \
-                               f"region{i + 1}" + self.VP_START_TOKEN + \
-                               self.IMG_CONTEXT_TOKEN * region_pixels[i] + \
-                               self.VP_END_TOKEN
-                if i == len(mask_prompts[0]) - 1:
-                    vp_token_str = vp_token_str + '.\n'
+                # prepare grounding images
+                g_image = np.array(image)  # for grounding
+                g_image = self.extra_image_processor.apply_image(g_image)
+                g_pixel_values = torch.from_numpy(g_image).permute(2, 0, 1).contiguous().to(self.torch_dtype)
+                extra_pixel_values = [g_pixel_values]
+                g_pixel_values = torch.stack([
+                    self.grounding_encoder.preprocess_image(pixel) for pixel in extra_pixel_values
+                ]).to(self.torch_dtype)
+
+                images = dynamic_preprocess(image, self.min_dynamic_patch,
+                                            self.max_dynamic_patch,
+                                            self.image_size, self.use_thumbnail)
+
+                if mask_prompts is not None:
+                    vp_overall_mask = torch.Tensor([False] * (len(images) - 1) + [True])
+                    input_dict['vp_overall_mask'] = vp_overall_mask
                 else:
-                    vp_token_str = vp_token_str + ', '
-        else:
-            vp_token_str = ''
-
-        image_token_str = f'{self.IMG_START_TOKEN}' \
-                          f'{self.IMG_CONTEXT_TOKEN * num_image_tokens}' \
-                          f'{self.IMG_END_TOKEN}'
-        image_token_str = image_token_str + '\n'
-        image_token_str = image_token_str * num_frames
-        image_token_str = image_token_str.strip()
-
-        ret_masks = []
-
-        if '<image>' in text or mask_prompts is not None:
-            assert past_text is None or len(past_text) == 0
-        text = text.replace('<image>', image_token_str + vp_token_str)
-        input_text = ''
-        input_text += self.template['INSTRUCTION'].format(
-                input=text, round=1, bot_name=self.bot_name)
-        input_text = past_text + input_text
-        ids = self.tokenizer.encode(input_text)
-        ret_past_text = self.tokenizer.decode(ids)
-        ids = torch.tensor(ids).cuda().unsqueeze(0)
+                    input_dict['vp_overall_mask'] = None
 
-        attention_mask = torch.ones_like(ids, dtype=torch.bool)
+                pixel_values = [self.transformer(image) for image in images]
+                pixel_values = torch.stack(pixel_values).to(self.torch_dtype)
+                num_image_tokens = pixel_values.shape[0] * self.patch_token
+                num_frames = 1
+            input_dict['g_pixel_values'] = g_pixel_values
+            input_dict['pixel_values'] = pixel_values
 
-        mm_inputs = {
-            'pixel_values': input_dict['pixel_values'],
-            'input_ids': ids,
-            'attention_mask': attention_mask,
-            'position_ids': None,
-            'past_key_values': None,
-            'labels': None,
-            'prompt_masks': mask_prompts,
-            'vp_overall_mask': input_dict['vp_overall_mask'],
-        }
+            if mask_prompts is not None:
+                # reshape mask prompts to feature size
+                mask_prompts = [torch.Tensor(item).to(pixel_values.device) for item in mask_prompts]
+                mask_prompts = [F.interpolate(
+                    item.unsqueeze(0),
+                    size=(int(self.image_size // self.patch_size * self.downsample_ratio),
+                          int(self.image_size // self.patch_size * self.downsample_ratio)),
+                    mode='nearest').squeeze(0) for item in mask_prompts]
+                region_pixels = []
+                for mask_prompt in mask_prompts[0]:
+                    region_pixels.append(mask_prompt.bool().to(torch.int64).sum())
+
+                vp_token_str = '\nThere are {} part regions in the picture: '.format(len(mask_prompts[0]))
+                for i in range(len(mask_prompts[0])):
+                    vp_token_str = vp_token_str + \
+                                   f"region{i + 1}" + self.VP_START_TOKEN + \
+                                   self.IMG_CONTEXT_TOKEN * region_pixels[i] + \
+                                   self.VP_END_TOKEN
+                    if i == len(mask_prompts[0]) - 1:
+                        vp_token_str = vp_token_str + '.\n'
+                    else:
+                        vp_token_str = vp_token_str + ', '
+            else:
+                vp_token_str = ''
+
+            image_token_str = f'{self.IMG_START_TOKEN}' \
+                              f'{self.IMG_CONTEXT_TOKEN * num_image_tokens}' \
+                              f'{self.IMG_END_TOKEN}'
+            image_token_str = image_token_str + '\n'
+            image_token_str = image_token_str * num_frames
+            image_token_str = image_token_str.strip()
+
+            ret_masks = []
+
+            if '<image>' in text or mask_prompts is not None:
+                assert past_text is None or len(past_text) == 0
+            text = text.replace('<image>', image_token_str + vp_token_str)
+            input_text = ''
+            input_text += self.template['INSTRUCTION'].format(
+                input=text, round=1, bot_name=self.bot_name)
+            input_text = past_text + input_text
+            ids = self.tokenizer.encode(input_text)
+            ids = torch.tensor(ids).cuda().unsqueeze(0)
+
+            attention_mask = torch.ones_like(ids, dtype=torch.bool)
+
+            mm_inputs = {
+                'pixel_values': input_dict['pixel_values'],
+                'input_ids': ids,
+                'attention_mask': attention_mask,
+                'position_ids': None,
+                'past_key_values': None,
+                'labels': None,
+                'prompt_masks': mask_prompts,
+                'vp_overall_mask': input_dict['vp_overall_mask'],
+            }
 
         generate_output = self.generate(
             **mm_inputs,
@@ -716,8 +737,10 @@ class Sa2VAChatModel(PreTrainedModel):
         )
         predict = self.tokenizer.decode(
             generate_output.sequences[0], skip_special_tokens=False).strip()
-        ret_past_text = ret_past_text + self.tokenizer.decode(
-            generate_output.sequences[0], skip_special_tokens=False)
+
+        if image is None and video is None and '<image>' not in past_text:
+            return {'prediction': predict, 'prediction_masks': ret_masks, }
+
         # if have seg result, find the seg hidden states
         hidden_states = generate_output.hidden_states
         last_hidden_states = [item[-1][0] for item in hidden_states]
@@ -739,7 +762,8 @@ class Sa2VAChatModel(PreTrainedModel):
             masks = masks.sigmoid() > 0.5
             masks = masks.cpu().numpy()
             ret_masks.append(masks)
-        return {'prediction': predict, 'prediction_masks': ret_masks, "past_text": ret_past_text}
+
+        return {'prediction': predict, 'prediction_masks': ret_masks,}
 
 def get_seg_hidden_states(hidden_states, output_ids, seg_id):
     seg_mask = output_ids == seg_id