add video

modeling_sa2va_chat.py

@@ -485,6 +485,7 @@ class Sa2VAChatModel(PreTrainedModel):
                         objects_prompt_masks = objects_prompt_masks.reshape(n_obj, -1)
                         vp_embeds.append(tile_vit_embeds[objects_prompt_masks])
                         i_vp_img += 1
+
                 vp_embeds = torch.cat(vp_embeds, dim=0)
             else:
                 vp_embeds = None
@@ -583,6 +584,7 @@ class Sa2VAChatModel(PreTrainedModel):
     def predict_forward(
             self,
             image=None,
+            video=None,
             text=None,
             past_text='',
             mask_prompts=None,
@@ -593,29 +595,57 @@ class Sa2VAChatModel(PreTrainedModel):
             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
 
-        ori_image_size = image.size
+            input_dict['vp_overall_mask'] = None
+        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)
-        input_dict['g_pixel_values'] = g_pixel_values
+            # 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)
+            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
+            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
 
-        pixel_values = [self.transformer(image) for image in images]
-        pixel_values = torch.stack(pixel_values).to(self.torch_dtype)
+            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
-        num_image_tokens = pixel_values.shape[0] * self.patch_token
+
 
         if mask_prompts is not None:
             # reshape mask prompts to feature size
@@ -627,7 +657,7 @@ class Sa2VAChatModel(PreTrainedModel):
                 mode='nearest').squeeze(0) for item in mask_prompts]
             region_pixels = []
             for mask_prompt in mask_prompts[0]:
-                region_pixels.append(mask_prompt.to(torch.int64).sum())
+                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])):
@@ -645,6 +675,9 @@ class Sa2VAChatModel(PreTrainedModel):
         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 = []
 
@@ -695,16 +728,14 @@ class Sa2VAChatModel(PreTrainedModel):
 
         for seg_hidden_states in all_seg_hidden_states:
             seg_hidden_states = seg_hidden_states.unsqueeze(0)
-            g_pixel_values = torch.stack([
-                self.grounding_encoder.preprocess_image(pixel, dtype=self.torch_dtype)
-                for pixel in [input_dict['g_pixel_values']]])
+            g_pixel_values = input_dict['g_pixel_values']
             sam_states = self.grounding_encoder.get_sam2_embeddings(g_pixel_values)
-            pred_masks = self.grounding_encoder.inject_language_embd(sam_states, [seg_hidden_states])
+            pred_masks = self.grounding_encoder.language_embd_inference(sam_states, [seg_hidden_states] * num_frames)
             w, h = ori_image_size
             masks = F.interpolate(pred_masks, size=(h, w), mode='bilinear', align_corners=False)
             masks = masks[:, 0]
             masks = masks.sigmoid() > 0.5
-            masks = masks.int().cpu()
+            masks = masks.cpu().numpy()
             ret_masks.append(masks)
 
         return {'prediction': predict, 'prediction_masks': ret_masks,}
@@ -712,6 +743,8 @@ class Sa2VAChatModel(PreTrainedModel):
 def get_seg_hidden_states(hidden_states, output_ids, seg_id):
     seg_mask = output_ids == seg_id
     n_out = len(seg_mask)
+    if n_out == 0:
+        return hidden_states[0:0]
     return hidden_states[-n_out:][seg_mask]
 
 def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height,

sam2.py

@@ -623,8 +623,8 @@ class CXBlock(nn.Module):
         x = self.pwconv1(x)
         x = self.act(x)
         x = self.pwconv2(x)
-        if self.gamma is not None:
-            x = self.gamma * x
+        if self.g_weight is not None:
+            x = self.g_weight * x
         x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
 
         x = input + self.drop_path(x)