Add files using upload-large-folder tool

cache_q_smoke/test_s/000000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f85d7cf7b83caf6fedb153a2cea2b36dd144ee3c0e34039483e20d208ea92d3
+size 2327

cache_q_smoke/test_s/index.jsonl

@@ -0,0 +1 @@
+{"sample_idx": 0, "path": "000000.pt", "vid": "-3ABOVeVmpU_136000_146000", "refs": ["the object that keeps making sound at all times"], "fids": [1], "resize": [576, 1024], "orgsize": [720, 1280], "num_seg": 1}

data/image_embed.tar

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0b0f5c8ae133bbddbfa558b2052b3aeb757492ffe310650988103d07e24135bb
+oid sha256:43e0e2002e80457512c6cdb2c171d0323335ea4bbce87ed364da22d267bb931d
 size 167486740480

models/avs_model.py

@@ -270,6 +270,7 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
             epoch: int =0,
             inference: bool = False,
             num_frames: int = 10,
+            target_frame: int = None,
             contrast: float = 0.0,
 
             **kwargs,
@@ -282,14 +283,12 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
         # audio_embeddings = torch.cat(audio_features, dim=0) # [B*10, 128]
         # audio_embeddings = audio_features  # [B, 10, 128]
 
-        # train
-        if not inference:
-            target_frame = random.randint(0, 9)
+        if target_frame is None:
             target_frame = 5
-
         else:
-            target_frame = 5
-        # print("target_frame", target_frame)
+            target_frame = int(target_frame)
+            if target_frame < 0 or target_frame >= num_frames:
+                raise ValueError(f"target_frame must be in [0, {num_frames}), got {target_frame}")
 
         input_ids, attention_masks, past_key_values, inputs_embeds, labels = super().prepare_inputs_labels_for_multimodal(
             input_ids, attention_masks, past_key_values=None, labels=labels, images=images_clip, audio_features=audio_embeddings, target_frame=target_frame, ref_ids=ref_ids
@@ -313,7 +312,8 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
             dim=1, )  # [batch_size, seq_len]
 
 
-        seg_embeddings = self.model.text_hidden_fcs[0](output_hidden_states[-1][seg_token_mask])  # [seg_num,256]
+        seg_hidden_states = output_hidden_states[-1][seg_token_mask]  # [seg_num, hidden_size]
+        seg_embeddings = self.model.text_hidden_fcs[0](seg_hidden_states)  # [seg_num,256]
 
         # print("seg_embeddings in this batch:", seg_embeddings.shape)
         # print("vids:", vids)
@@ -337,10 +337,12 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
 
 
         pred_embeddings = []
+        pred_hidden_states = []
         #--------------------------------------------------------------------------------------------
         pred_idx = 0
         for ref_num in refs_num:
             pred_embeddings.append(seg_embeddings[pred_idx:pred_idx + ref_num])
+            pred_hidden_states.append(seg_hidden_states[pred_idx:pred_idx + ref_num])
             pred_idx += ref_num
         # list[B]:[num_seg, 256]
 
@@ -397,6 +399,7 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
                 "pred_masks": pred_masks,     # list[B]:[num_seg, T, H, W]
                 "gt_masks": gt_masks,         # list[B]:[num_seg, T, H, W]
                 "seg_embeddings": pred_embeddings,  # list[B]:[num_seg, 256]
+                "seg_hidden_states": pred_hidden_states,  # list[B]:[num_seg, hidden_size]
             }
 
         model_output = output
@@ -462,4 +465,3 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
 
     def evaluate(self, *args, **kwargs):
         raise NotImplementedError("This method is not implemented.")
-

models/segment_anything/modeling/mask_decoder.py

@@ -140,7 +140,17 @@ class MaskDecoder(nn.Module):
         b, c, h, w = src.shape
 
         # Run the transformer
-        hs, src = self.transformer(src, pos_src, tokens)
+        referent_token_index = (
+            1 + self.num_mask_tokens if sparse_prompt_embeddings.shape[1] > 0 else None
+        )
+        hs, src = self.transformer(
+            src,
+            pos_src,
+            tokens,
+            mask_token_start=1,
+            num_mask_tokens=self.num_mask_tokens,
+            referent_token_index=referent_token_index,
+        )
         iou_token_out = hs[:, 0, :]
         mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :]
 
@@ -188,7 +198,17 @@ class MaskDecoder(nn.Module):
         _, c, h, w = src.shape 
 
         # Run the transformer
-        hs, src = self.transformer(src, pos_src, tokens)
+        referent_token_index = (
+            1 + self.num_mask_tokens if sparse_prompt_embeddings.shape[1] > 0 else None
+        )
+        hs, src = self.transformer(
+            src,
+            pos_src,
+            tokens,
+            mask_token_start=1,
+            num_mask_tokens=self.num_mask_tokens,
+            referent_token_index=referent_token_index,
+        )
         mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :]
 
         # Upscale mask embeddings and predict masks using the mask tokens

models/segment_anything/modeling/transformer.py

@@ -9,6 +9,7 @@ from typing import Tuple, Type
 
 import torch
 from torch import Tensor, nn
+from torch.nn import functional as F
 
 from .common import MLPBlock
 
@@ -64,6 +65,9 @@ class TwoWayTransformer(nn.Module):
         image_embedding: Tensor,
         image_pe: Tensor,
         point_embedding: Tensor,
+        mask_token_start: int = None,
+        num_mask_tokens: int = None,
+        referent_token_index: int = None,
     ) -> Tuple[Tensor, Tensor]:
         """
         Args:
@@ -94,6 +98,9 @@ class TwoWayTransformer(nn.Module):
                 keys=keys,
                 query_pe=point_embedding,
                 key_pe=image_pe,
+                mask_token_start=mask_token_start,
+                num_mask_tokens=num_mask_tokens,
+                referent_token_index=referent_token_index,
             )
 
         # Apply the final attention layer from the points to the image
@@ -145,11 +152,19 @@ class TwoWayAttentionBlock(nn.Module):
         self.cross_attn_image_to_token = Attention(
             embedding_dim, num_heads, downsample_rate=attention_downsample_rate
         )
+        self.referent_gate = ReferentGate(embedding_dim)
 
         self.skip_first_layer_pe = skip_first_layer_pe
 
     def forward(
-        self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor
+        self,
+        queries: Tensor,
+        keys: Tensor,
+        query_pe: Tensor,
+        key_pe: Tensor,
+        mask_token_start: int = None,
+        num_mask_tokens: int = None,
+        referent_token_index: int = None,
     ) -> Tuple[Tensor, Tensor]:
         # Self attention block
         if self.skip_first_layer_pe:
@@ -160,6 +175,17 @@ class TwoWayAttentionBlock(nn.Module):
             queries = queries + attn_out
         queries = self.norm1(queries)
 
+        if (
+            mask_token_start is not None
+            and num_mask_tokens is not None
+            and referent_token_index is not None
+        ):
+            mask_slice = slice(mask_token_start, mask_token_start + num_mask_tokens)
+            mask_tokens = queries[:, mask_slice, :]
+            referent_token = queries[:, referent_token_index : referent_token_index + 1, :]
+            queries = queries.clone()
+            queries[:, mask_slice, :] = self.referent_gate(mask_tokens, referent_token)
+
         # Cross attention block, tokens attending to image embedding
         q = queries + query_pe
         k = keys + key_pe
@@ -182,6 +208,26 @@ class TwoWayAttentionBlock(nn.Module):
         return queries, keys
 
 
+class ReferentGate(nn.Module):
+    def __init__(self, embedding_dim: int) -> None:
+        super().__init__()
+        self.gate = nn.Linear(embedding_dim * 2 + 1, embedding_dim)
+        self.proj = nn.Linear(embedding_dim, embedding_dim)
+        nn.init.zeros_(self.gate.weight)
+        nn.init.zeros_(self.gate.bias)
+        nn.init.zeros_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+        self.last_alpha = None
+
+    def forward(self, mask_tokens: Tensor, referent_token: Tensor) -> Tensor:
+        referent = referent_token.expand_as(mask_tokens)
+        cosine = F.cosine_similarity(mask_tokens, referent, dim=-1).unsqueeze(-1)
+        gate_input = torch.cat([mask_tokens, referent, cosine], dim=-1)
+        alpha = torch.sigmoid(self.gate(gate_input))
+        self.last_alpha = alpha.detach()
+        return mask_tokens + alpha * self.proj(referent)
+
+
 class Attention(nn.Module):
     """
     An attention layer that allows for downscaling the size of the embedding

upload_hf.py

@@ -27,7 +27,7 @@ IGNORE_PATTERNS = [
     "upload.log",
 ]
 
-NUM_WORKERS  = 2    # conservative; increase to 8 if no rate-limit errors
+NUM_WORKERS  = 1    # conservative; increase to 8 if no rate-limit errors
 MAX_RETRIES  = 10
 # ───────────────────────────────────────────────────────────────────────────