Add files using upload-large-folder tool

ChatUniVi/model/multimodal_encoder/eva_vit.py

@@ -12,8 +12,8 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.utils.checkpoint as checkpoint
-from timm.models.layers import drop_path, to_2tuple, trunc_normal_
-from timm.models.registry import register_model
+from timm.layers import drop_path, to_2tuple, trunc_normal_
+from timm.models import register_model
 
 from .utils import download_cached_file
 
@@ -445,4 +445,4 @@ def create_eva_vit_g(img_size=224, drop_path_rate=0.4, use_checkpoint=False, pre
     if precision == "fp16":
         #         model.to("cuda")
         convert_weights_to_fp16(model)
-    return model
+    return model

ChatUniVi/model/multimodal_encoder/utils.py

@@ -11,7 +11,10 @@ import os
 
 import torch
 import torch.distributed as dist
-import timm.models.hub as timm_hub
+from timm.models._hub import (
+    download_cached_file as timm_download_cached_file,
+    get_cache_dir as timm_get_cache_dir,
+)
 
 
 def setup_for_distributed(is_master):
@@ -124,14 +127,14 @@ def download_cached_file(url, check_hash=True, progress=False):
         # a hack to sync the file path across processes
         parts = torch.hub.urlparse(url)
         filename = os.path.basename(parts.path)
-        cached_file = os.path.join(timm_hub.get_cache_dir(), filename)
+        cached_file = os.path.join(timm_get_cache_dir(), filename)
 
         return cached_file
 
     if is_main_process():
-        timm_hub.download_cached_file(url, check_hash, progress)
+        timm_download_cached_file(url, check_hash, progress)
 
     if is_dist_avail_and_initialized():
         dist.barrier()
 
-    return get_cached_file_path()
+    return get_cached_file_path()

configs/config.py

@@ -64,6 +64,9 @@ parser.add_argument("--lr", type=float, default=5e-5, help='lr to fine tuning ad
 # epochs
 parser.add_argument("--epochs", type=int, default=10, help='epochs to fine tuning adapters.')
 parser.add_argument("--batch_size", type=int, default=8)
+parser.add_argument("--ce_loss_weight", type=float, default=1.0, help="Weight for language modeling loss.")
+parser.add_argument("--dice_loss_weight", type=float, default=0.5, help="Weight for dice segmentation loss.")
+parser.add_argument("--bce_loss_weight", type=float, default=2.0, help="Weight for BCE segmentation loss.")
 
 
 parser.add_argument("--gpu_id", type=str, default="0", help="The GPU device to run generation on.")
@@ -73,6 +76,7 @@ parser.add_argument("--run", type=str, default='train', help="train, test")
 parser.add_argument("--frame_n", type=int, default=10, help="Frame num of each video. Fixed to 10.")
 parser.add_argument("--text_max_len", type=int, default=25, help="Maximum textual reference length.")
 parser.add_argument("--max_eval_rows", type=int, default=-1, help="Max samples per split during eval; -1 = all.")
+parser.add_argument("--subset_manifest", type=str, default="", help="Optional JSON file that fixes train/eval subset indices per split.")
 parser.add_argument("--eval_split", type=str, default="test_u", help="Which split to evaluate: test_s, test_u, test_n.")
 parser.add_argument("--gate_only", action="store_true", help="Train only A-min referent gate parameters.")
 parser.add_argument("--init_from_saved_model", action="store_true", help="Initialize training from --saved_model before updates.")
@@ -88,6 +92,19 @@ parser.add_argument("--eval_only", action="store_true", help="Only evaluate in c
 parser.add_argument("--disable_gate", action="store_true", help="Force A-min gate to identity for cached pipeline baseline checks.")
 parser.add_argument("--gate_checkpoint", type=str, default="", help="Optional referent-gate-only checkpoint to overlay after loading --saved_model.")
 parser.add_argument("--save_gate_only", action="store_true", help="In cached-gate training, save only referent_gate parameters.")
+parser.add_argument("--use_residual_prompt_bridge", action="store_true", help="Enable the image-conditioned residual prompt bridge before SAM prompt encoding.")
+parser.add_argument("--bridge_only", action="store_true", help="Freeze all parameters except the residual prompt bridge.")
+parser.add_argument("--bridge_pm_weight", type=float, default=0.0, help="Weight for prompt-manifold teacher loss.")
+parser.add_argument("--bridge_rg_weight", type=float, default=0.0, help="Weight for region-semantic teacher loss.")
+parser.add_argument("--bridge_norm_weight", type=float, default=0.0, help="Weight for prompt-norm preservation loss.")
+parser.add_argument("--bridge_mode", type=str, default="additive", choices=["additive", "directional"], help="Prompt bridge parameterization.")
+parser.add_argument("--bridge_condition", type=str, default="image", choices=["image", "q_only"], help="Condition source for the prompt bridge.")
+parser.add_argument("--bridge_directional_alpha", type=float, default=0.1, help="Step size used by directional bridge updates after orthogonalization.")
+parser.add_argument("--bridge_gate_bias_init", type=float, default=-4.0, help="Initial bias for bridge gate sigmoid.")
+parser.add_argument("--bridge_residual_init_std", type=float, default=1e-3, help="Std used to initialize the bridge residual projection.")
+parser.add_argument("--bridge_target_frame", type=int, default=5, help="Frame index used to build bridge teachers.")
+parser.add_argument("--bridge_sanity_only", action="store_true", help="Run only bridge sanity checks (gradient, identity, teacher norms) and exit.")
+parser.add_argument("--bridge_sanity_batches", type=int, default=3, help="How many batches to scan during bridge sanity stats collection.")
 
 
 

models/avs_model.py

@@ -100,6 +100,74 @@ def compute_alignment_loss(q: torch.Tensor, pos_feats: list, neg_feats: list, te
     return total_loss / count
 
 
+class ResidualPromptBridge(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        mode: str = "additive",
+        condition: str = "image",
+        directional_alpha: float = 0.1,
+        gate_bias_init: float = -4.0,
+        residual_init_std: float = 1e-3,
+    ) -> None:
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.mode = mode
+        self.condition = condition
+        self.directional_alpha = directional_alpha
+        self.scale = math.sqrt(float(embedding_dim))
+        self.attn_proj = nn.Linear(embedding_dim, embedding_dim, bias=False)
+        self.residual_proj = nn.Linear(embedding_dim, embedding_dim, bias=False)
+        self.gate = nn.Linear(embedding_dim * 2, embedding_dim)
+        self.reset_parameters(gate_bias_init=gate_bias_init, residual_init_std=residual_init_std)
+
+    def reset_parameters(self, gate_bias_init: float, residual_init_std: float) -> None:
+        nn.init.xavier_uniform_(self.attn_proj.weight)
+        nn.init.normal_(self.residual_proj.weight, mean=0.0, std=residual_init_std)
+        nn.init.zeros_(self.gate.weight)
+        nn.init.constant_(self.gate.bias, gate_bias_init)
+
+    def forward(self, q: torch.Tensor, image_embeddings: torch.Tensor) -> dict:
+        if self.condition == "q_only":
+            attn = None
+            region = self.attn_proj(q)
+        else:
+            if image_embeddings.dim() != 4:
+                raise ValueError(
+                    f"ResidualPromptBridge expects image_embeddings [B, C, H, W], got {tuple(image_embeddings.shape)}"
+                )
+            image_tokens = image_embeddings.flatten(2).transpose(1, 2)  # [B, HW, C]
+            q_proj = self.attn_proj(q)  # [B, C]
+            attn_logits = torch.bmm(image_tokens, q_proj.unsqueeze(-1)).squeeze(-1) / self.scale
+            attn = torch.softmax(attn_logits, dim=-1)
+            region = torch.bmm(attn.unsqueeze(1), image_tokens).squeeze(1)
+
+        gate = torch.sigmoid(self.gate(torch.cat([q, region], dim=-1)))
+        region_update = self.residual_proj(region)
+
+        if self.mode == "directional":
+            q_dir = F.normalize(q, dim=-1)
+            q_parallel = (region_update * q_dir).sum(dim=-1, keepdim=True) * q_dir
+            region_orth = region_update - q_parallel
+            region_orth_norm = region_orth.norm(dim=-1, keepdim=True).clamp_min(1e-6)
+            region_dir = region_orth / region_orth_norm
+            alpha = self.directional_alpha * gate.mean(dim=-1, keepdim=True)
+            mixed_dir = F.normalize(q_dir + alpha * region_dir, dim=-1)
+            p_hat = q.norm(dim=-1, keepdim=True) * mixed_dir
+            delta = p_hat - q
+        else:
+            delta = gate * region_update
+            p_hat = q + delta
+
+        return {
+            "p_hat": p_hat,
+            "attn": attn,
+            "region": region,
+            "gate": gate,
+            "delta": delta,
+        }
+
+
 
 
 class Simtoken_MetaModel:
@@ -115,6 +183,12 @@ class Simtoken_MetaModel:
             self.config.train_mask_decoder = kwargs["train_mask_decoder"]
             self.config.out_dim = kwargs["out_dim"]
             self.vision_pretrained = kwargs.get("vision_pretrained", None)
+            self.config.use_residual_prompt_bridge = kwargs.get("use_residual_prompt_bridge", False)
+            self.config.bridge_mode = kwargs.get("bridge_mode", "additive")
+            self.config.bridge_condition = kwargs.get("bridge_condition", "image")
+            self.config.bridge_directional_alpha = kwargs.get("bridge_directional_alpha", 0.1)
+            self.config.bridge_gate_bias_init = kwargs.get("bridge_gate_bias_init", -4.0)
+            self.config.bridge_residual_init_std = kwargs.get("bridge_residual_init_std", 1e-3)
         else:
             self.vision_pretrained = kwargs.get("vision_pretrained", None)
             self.initialize_lisa_modules(self.config)
@@ -143,6 +217,17 @@ class Simtoken_MetaModel:
         for param in self.text_hidden_fcs.parameters():
             param.requires_grad = True
 
+        self.prompt_bridge = None
+        if getattr(config, "use_residual_prompt_bridge", False):
+            self.prompt_bridge = ResidualPromptBridge(
+                embedding_dim=out_dim,
+                mode=getattr(config, "bridge_mode", "additive"),
+                condition=getattr(config, "bridge_condition", "image"),
+                directional_alpha=getattr(config, "bridge_directional_alpha", 0.1),
+                gate_bias_init=getattr(config, "bridge_gate_bias_init", -4.0),
+                residual_init_std=getattr(config, "bridge_residual_init_std", 1e-3),
+            )
+
 
 class Simtoken_Model(Simtoken_MetaModel, ChatUniViLlamaModel):
     def __init__(
@@ -234,6 +319,104 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
         self.compress = kwargs.pop("compress", True)
 
         self.start = kwargs.pop("start")
+        self.use_residual_prompt_bridge = kwargs.pop("use_residual_prompt_bridge", False)
+        self.bridge_pm_weight = kwargs.pop("bridge_pm_weight", 0.0)
+        self.bridge_rg_weight = kwargs.pop("bridge_rg_weight", 0.0)
+        self.bridge_norm_weight = kwargs.pop("bridge_norm_weight", 0.0)
+        self.bridge_target_frame = kwargs.pop("bridge_target_frame", 5)
+
+    def _expand_prompt_level_inputs(
+        self,
+        image_features: List[torch.Tensor],
+        masks_list: List[torch.FloatTensor],
+        refs_num: List[int],
+        target_frame: int,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> tuple:
+        prompt_image_embeddings = []
+        prompt_masks = []
+        prompt_mask_size = self.model.visual_model.prompt_encoder.mask_input_size
+
+        for sample_idx, ref_num in enumerate(refs_num):
+            frame_feat = image_features[sample_idx][target_frame].to(device=device, dtype=dtype)
+            for prompt_idx in range(ref_num):
+                prompt_image_embeddings.append(frame_feat)
+                mask = masks_list[sample_idx][prompt_idx, target_frame].to(
+                    device=device, dtype=torch.float32
+                )
+                mask = F.interpolate(
+                    mask.unsqueeze(0).unsqueeze(0),
+                    size=prompt_mask_size,
+                    mode="nearest",
+                ).squeeze(0).squeeze(0)
+                prompt_masks.append(mask)
+
+        return torch.stack(prompt_image_embeddings, dim=0), torch.stack(prompt_masks, dim=0)
+
+    def _compute_prompt_bridge_teachers(
+        self,
+        prompt_image_embeddings: torch.Tensor,
+        prompt_masks: torch.Tensor,
+        dtype: torch.dtype,
+    ) -> tuple:
+        mask_lowres = prompt_masks.unsqueeze(1)
+        _, dense_mask_embeddings = self.model.visual_model.prompt_encoder(
+            points=None,
+            boxes=None,
+            masks=mask_lowres.to(dtype=dtype),
+            text_embeds=None,
+        )
+        prompt_manifold_teacher = dense_mask_embeddings.mean(dim=(2, 3))
+
+        mask_64 = F.interpolate(
+            prompt_masks.unsqueeze(1),
+            size=prompt_image_embeddings.shape[-2:],
+            mode="nearest",
+        )
+        flat_feats = prompt_image_embeddings.flatten(2)
+        flat_mask = mask_64.flatten(2)
+        masked_sum = (flat_feats * flat_mask).sum(dim=-1)
+        mask_area = flat_mask.sum(dim=-1).clamp_min(1.0)
+        region_teacher = masked_sum / mask_area
+
+        return prompt_manifold_teacher, region_teacher
+
+    def _summarize_prompt_bridge(
+        self,
+        q: torch.Tensor,
+        p_hat: torch.Tensor,
+        prompt_manifold_teacher: torch.Tensor,
+        region_teacher: torch.Tensor,
+        gate: torch.Tensor,
+    ) -> dict:
+        delta = p_hat - q
+        q_norm = q.norm(dim=-1)
+        p_hat_norm = p_hat.norm(dim=-1)
+        pm_cos = F.cosine_similarity(p_hat, prompt_manifold_teacher, dim=-1)
+        rg_cos = F.cosine_similarity(p_hat, region_teacher, dim=-1)
+        qq_cos = F.cosine_similarity(p_hat, q, dim=-1)
+        teacher_cos = F.cosine_similarity(prompt_manifold_teacher, region_teacher, dim=-1)
+        delta_q_cos = F.cosine_similarity(delta, q, dim=-1)
+        delta_pm_cos = F.cosine_similarity(delta, prompt_manifold_teacher, dim=-1)
+        delta_rg_cos = F.cosine_similarity(delta, region_teacher, dim=-1)
+
+        return {
+            "q_norm_mean": q_norm.mean().item(),
+            "p_hat_norm_mean": p_hat_norm.mean().item(),
+            "delta_norm_mean": delta.norm(dim=-1).mean().item(),
+            "cos_p_hat_q_mean": qq_cos.mean().item(),
+            "cos_p_hat_p_mask_mean": pm_cos.mean().item(),
+            "cos_p_hat_z_gt_mean": rg_cos.mean().item(),
+            "cos_delta_q_mean": delta_q_cos.mean().item(),
+            "cos_delta_p_mask_mean": delta_pm_cos.mean().item(),
+            "cos_delta_z_gt_mean": delta_rg_cos.mean().item(),
+            "p_mask_norm_mean": prompt_manifold_teacher.norm(dim=-1).mean().item(),
+            "z_gt_norm_mean": region_teacher.norm(dim=-1).mean().item(),
+            "cos_p_mask_z_gt_mean": teacher_cos.mean().item(),
+            "gate_mean": gate.mean().item(),
+            "gate_std": gate.std(unbiased=False).item(),
+        }
 
 
 
@@ -284,7 +467,7 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
         # audio_embeddings = audio_features  # [B, 10, 128]
 
         if target_frame is None:
-            target_frame = 5
+            target_frame = self.bridge_target_frame
         else:
             target_frame = int(target_frame)
             if target_frame < 0 or target_frame >= num_frames:
@@ -315,6 +498,60 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
         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]
 
+        prompt_embeddings_all = seg_embeddings
+        bridge_metrics = {}
+        bridge_pm_loss = seg_embeddings.new_zeros(())
+        bridge_rg_loss = seg_embeddings.new_zeros(())
+        bridge_norm_loss = seg_embeddings.new_zeros(())
+        bridge_teacher_loss = seg_embeddings.new_zeros(())
+        bridge_teacher_loss_raw = seg_embeddings.new_zeros(())
+        prompt_image_embeddings = None
+        prompt_manifold_teacher = None
+        region_teacher = None
+
+        if self.use_residual_prompt_bridge:
+            prompt_image_embeddings, prompt_masks = self._expand_prompt_level_inputs(
+                image_features=image_features,
+                masks_list=masks_list,
+                refs_num=refs_num,
+                target_frame=target_frame,
+                dtype=seg_embeddings.dtype,
+                device=seg_embeddings.device,
+            )
+            bridge_outputs = self.model.prompt_bridge(seg_embeddings, prompt_image_embeddings)
+            prompt_embeddings_all = bridge_outputs["p_hat"]
+            prompt_manifold_teacher, region_teacher = self._compute_prompt_bridge_teachers(
+                prompt_image_embeddings=prompt_image_embeddings,
+                prompt_masks=prompt_masks,
+                dtype=seg_embeddings.dtype,
+            )
+
+            pm_l1 = F.smooth_l1_loss(prompt_embeddings_all, prompt_manifold_teacher)
+            pm_cos = 1.0 - F.cosine_similarity(
+                prompt_embeddings_all, prompt_manifold_teacher, dim=-1
+            ).mean()
+            bridge_pm_loss = pm_l1 + pm_cos
+            bridge_rg_loss = 1.0 - F.cosine_similarity(
+                prompt_embeddings_all, region_teacher, dim=-1
+            ).mean()
+            bridge_norm_loss = F.mse_loss(
+                prompt_embeddings_all.norm(dim=-1),
+                seg_embeddings.norm(dim=-1),
+            )
+            bridge_teacher_loss_raw = bridge_pm_loss + bridge_rg_loss + bridge_norm_loss
+            bridge_teacher_loss = (
+                self.bridge_pm_weight * bridge_pm_loss
+                + self.bridge_rg_weight * bridge_rg_loss
+                + self.bridge_norm_weight * bridge_norm_loss
+            )
+            bridge_metrics = self._summarize_prompt_bridge(
+                q=seg_embeddings,
+                p_hat=prompt_embeddings_all,
+                prompt_manifold_teacher=prompt_manifold_teacher,
+                region_teacher=region_teacher,
+                gate=bridge_outputs["gate"],
+            )
+
         # print("seg_embeddings in this batch:", seg_embeddings.shape)
         # print("vids:", vids)
         # print("fids:", fids)
@@ -337,11 +574,13 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
 
 
         pred_embeddings = []
+        prompt_embeddings = []
         pred_hidden_states = []
         #--------------------------------------------------------------------------------------------
         pred_idx = 0
         for ref_num in refs_num:
             pred_embeddings.append(seg_embeddings[pred_idx:pred_idx + ref_num])
+            prompt_embeddings.append(prompt_embeddings_all[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]
@@ -359,7 +598,7 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
                 points=None,
                 boxes=None,
                 masks=None,
-                text_embeds=pred_embeddings[i].unsqueeze(1),  # [1, 1 ,256]
+                text_embeds=prompt_embeddings[i].unsqueeze(1),  # [1, 1 ,256]
             )
             # 确保数据类型一致
             sparse_embeddings = sparse_embeddings.to(pred_embeddings[i].dtype)
@@ -395,12 +634,23 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
         gt_masks = masks_list # list[B]:[num_seg, T, H, W]
 
         if inference:
-            return {
+            result = {
                 "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]
+                "prompt_embeddings": prompt_embeddings,  # list[B]:[num_seg, 256]
                 "seg_hidden_states": pred_hidden_states,  # list[B]:[num_seg, hidden_size]
             }
+            if self.use_residual_prompt_bridge:
+                result.update(
+                    {
+                        "bridge_metrics": bridge_metrics,
+                        "bridge_pm_loss": bridge_pm_loss.detach(),
+                        "bridge_rg_loss": bridge_rg_loss.detach(),
+                        "bridge_norm_loss": bridge_norm_loss.detach(),
+                    }
+                )
+            return result
 
         model_output = output
         output = model_output.logits
@@ -451,6 +701,8 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
         else:
             loss = ce_loss + mask_loss
 
+        loss = loss + bridge_teacher_loss
+
         return {
             "loss": loss,
             "ce_loss": ce_loss,
@@ -458,6 +710,12 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
             "mask_dice_loss": mask_dice_loss,
             "mask_loss": mask_loss,
             "ct_loss": ct_loss,
+            "bridge_pm_loss": bridge_pm_loss,
+            "bridge_rg_loss": bridge_rg_loss,
+            "bridge_norm_loss": bridge_norm_loss,
+            "bridge_teacher_loss": bridge_teacher_loss,
+            "bridge_teacher_loss_raw": bridge_teacher_loss_raw,
+            "bridge_metrics": bridge_metrics,
             "pred_masks": pred_masks,
             "gt_masks": gt_masks,
         }