Restore original SimToken source files

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.layers import drop_path, to_2tuple, trunc_normal_
-from timm.models import register_model
+from timm.models.layers import drop_path, to_2tuple, trunc_normal_
+from timm.models.registry 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,10 +11,7 @@ import os
 
 import torch
 import torch.distributed as dist
-from timm.models._hub import (
-    download_cached_file as timm_download_cached_file,
-    get_cache_dir as timm_get_cache_dir,
-)
+import timm.models.hub as timm_hub
 
 
 def setup_for_distributed(is_master):
@@ -127,14 +124,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_get_cache_dir(), filename)
+        cached_file = os.path.join(timm_hub.get_cache_dir(), filename)
 
         return cached_file
 
     if is_main_process():
-        timm_download_cached_file(url, check_hash, progress)
+        timm_hub.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

@@ -31,7 +31,7 @@ parser = argparse.ArgumentParser(
 
 
 
-parser.add_argument("--vision_pretrained",type=str,default='/workspace/SimToken/models/segment_anything/sam_vit_h_4b8939.pth')
+parser.add_argument("--vision_pretrained",type=str,default='path/to/segment_anything/sam_vit_h_4b8939.pth')
 parser.add_argument("--vision_tower",type=str,default='openai/clip-vit-large-patch14')
 parser.add_argument("--mllm",type=str,default='Chat-UniVi/Chat-UniVi-7B-v1.5')
 
@@ -44,9 +44,9 @@ parser.add_argument("--start",type=int,default=0)
 
 parser.add_argument("--name",type=str,default='testrun')
 # path to ref-avs dataset
-parser.add_argument("--data_dir",type=str,default='/workspace/SimToken/data',help=f"The data paranet dir. File arch should be: {file_arch}")
+parser.add_argument("--data_dir",type=str,default='data',help=f"The data paranet dir. File arch should be: {file_arch}")
 # path to pretrained checkpoints
-parser.add_argument("--saved_model",type=str,default='/workspace/SimToken/checkpoints/simtoken_pretrained.pth', help="the pretrained simtoken pth")
+parser.add_argument("--saved_model",type=str,default='trained_simtoken.pth', help="the pretrained simtoken pth")
 
 
 parser.add_argument("--log_root",type=str,default='log', help="where to save log during training")
@@ -64,9 +64,6 @@ 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.")
@@ -75,36 +72,6 @@ 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.")
-parser.add_argument("--max_steps", type=int, default=-1, help="Max optimizer steps during training; -1 = full schedule.")
-parser.add_argument("--overfit_samples", type=int, default=-1, help="Train on the first N train samples for overfit probes; -1 = full train set.")
-parser.add_argument("--log_gate_stats_every", type=int, default=-1, help="Log A-min gate/proj stats every N optimizer steps; -1 = disabled.")
-parser.add_argument("--skip_eval_after_train", action="store_true", help="Save checkpoint and exit without post-train evaluation.")
-parser.add_argument("--eval_train_only", action="store_true", help="After training, evaluate only the training subset and skip test splits.")
-parser.add_argument("--cache_root", type=str, default="/workspace/SimToken/cache_q", help="Root directory for cached q features.")
-parser.add_argument("--cache_split", type=str, default="train", help="Dataset split to cache or read cached q features from.")
-parser.add_argument("--overwrite_cache", action="store_true", help="Overwrite existing cached q feature files.")
-parser.add_argument("--eval_only", action="store_true", help="Only evaluate in cached-gate scripts; do not train.")
-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.")
 
 
 

load_model.py

@@ -208,10 +208,7 @@ def collate_fn(batch, tokenizer=None):
 
 import torch.multiprocessing as mp
 if __name__ == "__main__":
-    try:
-        mp.set_start_method("spawn")
-    except RuntimeError:
-        pass
+    mp.set_start_method("spawn")
     set_seed(42)
     tokenizer = transformers.AutoTokenizer.from_pretrained(
         args.mllm,
@@ -227,17 +224,14 @@ if __name__ == "__main__":
     print("seg_token_idx: ", seg_token_idx)
 
 
-    if args.eval_split not in {"test_s", "test_u", "test_n"}:
-        raise ValueError(f"Unsupported eval_split: {args.eval_split}")
+    val_dataset_s = REFAVS('test_s', args, tokenizer, input_type='refer')
+    # val_dataset_u = REFAVS('test_u', args, tokenizer, input_type='refer')
+    # val_dataset_n = REFAVS('test_n', args, tokenizer, input_type='refer')
 
-    val_dataset = REFAVS(args.eval_split, args, tokenizer, input_type='refer')
-    val_dataloader = DataLoader(
-        val_dataset,
-        batch_size=1,
-        shuffle=False,
-        num_workers=4,
-        collate_fn=partial(collate_fn, tokenizer=tokenizer),
-    )
+
+    val_dataloader_s = DataLoader(val_dataset_s, batch_size=1, shuffle=False, num_workers=4, collate_fn=partial(collate_fn, tokenizer=tokenizer))
+    # val_dataloader_u = DataLoader(val_dataset_u, batch_size=1, shuffle=False, num_workers=4, collate_fn=partial(collate_fn, tokenizer=tokenizer))
+    # val_dataloader_n = DataLoader(val_dataset_n, batch_size=2, shuffle=False, num_workers=4, collate_fn=partial(collate_fn, tokenizer=tokenizer))
 
 
 
@@ -343,12 +337,8 @@ if __name__ == "__main__":
     model = model.to("cuda")
     model.resize_token_embeddings(len(tokenizer))
 
-    missing, unexpected = model.load_state_dict(
-        torch.load(args.saved_model, map_location="cpu"),
-        strict=False,
-    )
-    print(f"saved model loaded: {args.saved_model}")
-    print(f"missing keys: {len(missing)} | unexpected keys: {len(unexpected)}")
+    model.load_state_dict(torch.load(args.saved_model), strict=False)
+    print("saved model loaded")
 
 
     save_root = args.visualization_root
@@ -414,9 +404,7 @@ if __name__ == "__main__":
         total_fscore = 0
         count = 0
 
-        for batch_idx, batch in enumerate(tqdm(dataloader, desc=f"Evaluating on {name}")):
-            if args.max_eval_rows > 0 and batch_idx >= args.max_eval_rows:
-                break
+        for batch in tqdm(dataloader, desc=f"Evaluating on {name}"):
             input_dict = dict_to_cuda(batch)
 
             with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
@@ -450,9 +438,6 @@ if __name__ == "__main__":
                 total_fscore += fscore * num_seg * T
                 count += num_seg * T
 
-        if count == 0:
-            raise RuntimeError(f"No samples were evaluated for {name}")
-
         print(f"\n  valuate on {name}:  miou: {total_iou/count}  fscore: {total_fscore/count}")
 
 
@@ -462,9 +447,7 @@ if __name__ == "__main__":
         total_metric = 0
         count = 0
 
-        for batch_idx, batch in enumerate(tqdm(dataloader, desc=f"Evaluating on Null")):
-            if args.max_eval_rows > 0 and batch_idx >= args.max_eval_rows:
-                break
+        for batch in tqdm(dataloader, desc=f"Evaluating on Null"):
             input_dict = dict_to_cuda(batch)
             with torch.no_grad():
                 output_dict = model.forward(images=input_dict["images"],
@@ -494,13 +477,14 @@ if __name__ == "__main__":
                 total_metric += null_metric * num_seg * T
                 count += num_seg * T
 
-        if count == 0:
-            raise RuntimeError("No samples were evaluated for test_n")
-
         print(f"\n  valuate on test_n_refer, metric: {total_metric / count}")
 
 
-    if args.eval_split == "test_n":
-        valuate_Null(model, val_dataloader)
-    else:
-        valuate(model, val_dataloader, args.eval_split)
+
+
+    valuate(model, val_dataloader_s, 'test_seen')
+
+    # valuate(model, val_dataloader_u, 'test_unseen')
+    #
+    # valuate_Null(model, val_dataloader_u)
+

models/avs_model.py

@@ -100,74 +100,6 @@ 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:
@@ -183,12 +115,6 @@ 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)
@@ -217,17 +143,6 @@ 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__(
@@ -319,104 +234,6 @@ 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(),
-        }
 
 
 
@@ -453,7 +270,6 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
             epoch: int =0,
             inference: bool = False,
             num_frames: int = 10,
-            target_frame: int = None,
             contrast: float = 0.0,
 
             **kwargs,
@@ -466,12 +282,14 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
         # audio_embeddings = torch.cat(audio_features, dim=0) # [B*10, 128]
         # audio_embeddings = audio_features  # [B, 10, 128]
 
-        if target_frame is None:
-            target_frame = self.bridge_target_frame
+        # train
+        if not inference:
+            target_frame = random.randint(0, 9)
+            target_frame = 5
+
         else:
-            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}")
+            target_frame = 5
+        # print("target_frame", 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
@@ -495,62 +313,7 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
             dim=1, )  # [batch_size, seq_len]
 
 
-        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"],
-            )
+        seg_embeddings = self.model.text_hidden_fcs[0](output_hidden_states[-1][seg_token_mask])  # [seg_num,256]
 
         # print("seg_embeddings in this batch:", seg_embeddings.shape)
         # print("vids:", vids)
@@ -574,14 +337,10 @@ 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]
 
@@ -598,7 +357,7 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
                 points=None,
                 boxes=None,
                 masks=None,
-                text_embeds=prompt_embeddings[i].unsqueeze(1),  # [1, 1 ,256]
+                text_embeds=pred_embeddings[i].unsqueeze(1),  # [1, 1 ,256]
             )
             # 确保数据类型一致
             sparse_embeddings = sparse_embeddings.to(pred_embeddings[i].dtype)
@@ -634,23 +393,10 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
         gt_masks = masks_list # list[B]:[num_seg, T, H, W]
 
         if inference:
-            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]
+            return {
+                "pred_masks": pred_masks,  # list[B]:[num_seg, T, H, W]
+                "gt_masks": gt_masks,  # list[B]:[num_seg, T, H, W]
             }
-            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
@@ -701,8 +447,6 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
         else:
             loss = ce_loss + mask_loss
 
-        loss = loss + bridge_teacher_loss
-
         return {
             "loss": loss,
             "ce_loss": ce_loss,
@@ -710,12 +454,6 @@ 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,
         }
@@ -723,3 +461,4 @@ class Simtoken_ForCausalLM(ChatUniViLlamaForCausalLM):
 
     def evaluate(self, *args, **kwargs):
         raise NotImplementedError("This method is not implemented.")
+

models/segment_anything/modeling/mask_decoder.py

@@ -140,17 +140,7 @@ class MaskDecoder(nn.Module):
         b, c, h, w = src.shape
 
         # Run the transformer
-        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,
-        )
+        hs, src = self.transformer(src, pos_src, tokens)
         iou_token_out = hs[:, 0, :]
         mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :]
 
@@ -198,17 +188,7 @@ class MaskDecoder(nn.Module):
         _, c, h, w = src.shape 
 
         # Run the transformer
-        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,
-        )
+        hs, src = self.transformer(src, pos_src, tokens)
         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,7 +9,6 @@ from typing import Tuple, Type
 
 import torch
 from torch import Tensor, nn
-from torch.nn import functional as F
 
 from .common import MLPBlock
 
@@ -65,9 +64,6 @@ 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:
@@ -98,9 +94,6 @@ 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
@@ -152,19 +145,11 @@ 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,
-        mask_token_start: int = None,
-        num_mask_tokens: int = None,
-        referent_token_index: int = None,
+        self, queries: Tensor, keys: Tensor, query_pe: Tensor, key_pe: Tensor
     ) -> Tuple[Tensor, Tensor]:
         # Self attention block
         if self.skip_first_layer_pe:
@@ -175,17 +160,6 @@ 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
@@ -208,26 +182,6 @@ 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

save_audio_feats.py

@@ -80,3 +80,4 @@ for vid in vids:
     # print(f"{vid}: {audio_embed.shape}")
     torch.save(audio_embed, f'{save_dir}/{vid}.pt')
     print(f'{vid} embedding saved {audio_embed.shape}')
+

train.py

@@ -1,7 +1,7 @@
 import transformers
 from datasets import REFAVS
 from configs import args
-from torch.utils.data import DataLoader, Subset
+from torch.utils.data import DataLoader
 from functools import partial
 from models.llava import conversation as conversation_lib
 # from  models.avs_model import VISAForCausalLM
@@ -21,9 +21,6 @@ import numpy as np
 import re
 import time
 import os
-import sys
-import json
-from collections import defaultdict
 
 
 import warnings
@@ -216,61 +213,10 @@ def collate_fn(batch, tokenizer=None):
     }
 
 
-def maybe_limit_dataset(dataset, max_rows, name):
-    if max_rows is None or max_rows <= 0:
-        return dataset
-    limited_n = min(max_rows, len(dataset))
-    print(f"max_eval_rows enabled: using first {limited_n} samples from {name}")
-    return Subset(dataset, list(range(limited_n)))
-
-
-def load_subset_manifest(path):
-    if not path:
-        return {}
-    with open(path, "r", encoding="utf-8") as f:
-        manifest = json.load(f)
-    if not isinstance(manifest, dict):
-        raise ValueError(f"subset_manifest must be a JSON object, got {type(manifest).__name__}")
-    if "subsets" in manifest:
-        manifest = manifest["subsets"]
-    return manifest
-
-
-def maybe_apply_manifest_subset(dataset, manifest, split_name, name):
-    if split_name not in manifest:
-        return dataset
-    indices = manifest[split_name]
-    if not isinstance(indices, list) or not all(isinstance(i, int) for i in indices):
-        raise ValueError(f"subset_manifest[{split_name!r}] must be a list of integers")
-    if not indices:
-        raise ValueError(f"subset_manifest[{split_name!r}] is empty")
-    max_index = len(dataset) - 1
-    bad_indices = [i for i in indices if i < 0 or i > max_index]
-    if bad_indices:
-        raise ValueError(
-            f"subset_manifest[{split_name!r}] contains out-of-range indices; "
-            f"dataset size={len(dataset)}, examples={bad_indices[:5]}"
-        )
-    print(f"subset_manifest enabled: using {len(indices)} fixed samples from {name} ({split_name})")
-    return Subset(dataset, indices)
-
-
-def checkpoint_requires_lora(saved_model_path):
-    if not saved_model_path or not os.path.exists(saved_model_path):
-        return False
-    state = torch.load(saved_model_path, map_location="cpu")
-    return any("lora_" in key for key in state.keys())
-
-
 import torch.multiprocessing as mp
 if __name__ == "__main__":
-    try:
-        mp.set_start_method("spawn")
-    except RuntimeError:
-        pass
+    mp.set_start_method("spawn")
     set_seed(42)
-    if args.bridge_only and not args.use_residual_prompt_bridge:
-        raise ValueError("--bridge_only requires --use_residual_prompt_bridge")
     tokenizer = transformers.AutoTokenizer.from_pretrained(
         args.mllm,
         cache_dir=None,
@@ -283,34 +229,17 @@ if __name__ == "__main__":
     num_added_tokens = tokenizer.add_tokens("[SEG]")
     seg_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids[0]  # 32000
     print("seg_token_idx: ", seg_token_idx)
-    subset_manifest = load_subset_manifest(args.subset_manifest)
 
     train_dataset = REFAVS('train', args, tokenizer, input_type='refer')
     val_dataset_s_refer = REFAVS('test_s', args, tokenizer, input_type='refer')
     val_dataset_u_refer = REFAVS('test_u', args, tokenizer, input_type='refer')
     val_dataset_n_refer = REFAVS('test_n', args, tokenizer, input_type='refer')
 
-    train_dataset = maybe_apply_manifest_subset(train_dataset, subset_manifest, "train", "train")
-    val_dataset_s_refer = maybe_apply_manifest_subset(val_dataset_s_refer, subset_manifest, "test_s", "test_s")
-    val_dataset_u_refer = maybe_apply_manifest_subset(val_dataset_u_refer, subset_manifest, "test_u", "test_u")
-    val_dataset_n_refer = maybe_apply_manifest_subset(val_dataset_n_refer, subset_manifest, "test_n", "test_n")
-
-    if args.overfit_samples > 0:
-        overfit_n = min(args.overfit_samples, len(train_dataset))
-        train_dataset = Subset(train_dataset, list(range(overfit_n)))
-        print(f"overfit_samples enabled: using first {overfit_n} train samples")
-
-    train_eval_dataset = maybe_limit_dataset(train_dataset, args.max_eval_rows, "train_eval")
-    val_dataset_s_refer = maybe_limit_dataset(val_dataset_s_refer, args.max_eval_rows, "test_s")
-    val_dataset_u_refer = maybe_limit_dataset(val_dataset_u_refer, args.max_eval_rows, "test_u")
-    val_dataset_n_refer = maybe_limit_dataset(val_dataset_n_refer, args.max_eval_rows, "test_n")
-
 
     g = torch.Generator()
     g.manual_seed(42)
 
     train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=8, worker_init_fn=seed_worker,collate_fn=partial(collate_fn, tokenizer=tokenizer), generator=g)
-    train_eval_dataloader = DataLoader(train_eval_dataset, batch_size=4, shuffle=False, num_workers=0, collate_fn=partial(collate_fn, tokenizer=tokenizer))
 
     val_dataloader_s_refer = DataLoader(val_dataset_s_refer, batch_size=4, shuffle=False, num_workers=0, collate_fn=partial(collate_fn, tokenizer=tokenizer))
     val_dataloader_u_refer = DataLoader(val_dataset_u_refer, batch_size=4, shuffle=False, num_workers=0, collate_fn=partial(collate_fn, tokenizer=tokenizer))
@@ -320,25 +249,15 @@ if __name__ == "__main__":
     model_args = {
         "train_mask_decoder": True,
         "out_dim": 256,  # 256
-        "ce_loss_weight": args.ce_loss_weight,
-        "dice_loss_weight": args.dice_loss_weight,
-        "bce_loss_weight": args.bce_loss_weight,
+        "ce_loss_weight": 1.0,
+        "dice_loss_weight": 0.5,
+        "bce_loss_weight": 2.0,
         "seg_token_idx": seg_token_idx,
         "vision_pretrained": args.vision_pretrained,  # sam_vit_h_xxx.pth
         "vision_tower": args.vision_tower,
         "use_im_start_end": False,
         "compress": args.compress,
         "start": args.start,
-        "use_residual_prompt_bridge": args.use_residual_prompt_bridge,
-        "bridge_pm_weight": args.bridge_pm_weight,
-        "bridge_rg_weight": args.bridge_rg_weight,
-        "bridge_norm_weight": args.bridge_norm_weight,
-        "bridge_mode": args.bridge_mode,
-        "bridge_condition": args.bridge_condition,
-        "bridge_directional_alpha": args.bridge_directional_alpha,
-        "bridge_gate_bias_init": args.bridge_gate_bias_init,
-        "bridge_residual_init_std": args.bridge_residual_init_std,
-        "bridge_target_frame": args.bridge_target_frame,
     }
 
     model = Simtoken_ForCausalLM.from_pretrained(args.mllm, torch_dtype=torch.float32, low_cpu_mem_usage=True, **model_args)
@@ -374,17 +293,7 @@ if __name__ == "__main__":
     for p in model.get_model().mm_projector.parameters():
         p.requires_grad = False
 
-    use_lora_checkpoint = (
-        (args.init_from_saved_model or args.gate_only)
-        and checkpoint_requires_lora(args.saved_model)
-    )
-    if args.bridge_only and use_lora_checkpoint:
-        print(
-            "bridge_only notice: saved_model contains LoRA weights, "
-            "so LoRA modules will be instantiated for checkpoint compatibility and then frozen."
-        )
-
-    lora_r = 8 if (not args.bridge_only or use_lora_checkpoint) else 0
+    lora_r = 8
     target_modules = "q_proj,v_proj"
     if lora_r > 0:
 
@@ -440,11 +349,6 @@ if __name__ == "__main__":
     model = model.to("cuda")
     model.resize_token_embeddings(len(tokenizer))
 
-    if args.init_from_saved_model or args.gate_only:
-        state = torch.load(args.saved_model, map_location="cpu")
-        missing, unexpected = model.load_state_dict(state, strict=False)
-        print(f"initialized training from saved model: {args.saved_model}")
-        print(f"missing keys: {len(missing)} | unexpected keys: {len(unexpected)}")
 
     for name, param in model.audio_feature_layer.named_parameters():
         param.requires_grad = True
@@ -452,274 +356,25 @@ if __name__ == "__main__":
     # for name, param in model.token_compressor.named_parameters():
     #     param.requires_grad = True
 
+
     for n, p in model.named_parameters():
         if any(
-            [
-                x in n
-                for x in ["lm_head", "embed_tokens", "mask_decoder", "text_hidden_fcs"]
-            ]
+                [
+                    x in n
+                    for x in ["lm_head", "embed_tokens", "mask_decoder", "text_hidden_fcs"]
+                ]
         ):
             p.requires_grad = True
 
-    if args.bridge_only:
-        for p in model.parameters():
-            p.requires_grad = False
-        trainable_names = []
-        for n, p in model.named_parameters():
-            if "prompt_bridge" in n:
-                p.requires_grad = True
-                trainable_names.append(n)
-        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
-        total = sum(p.numel() for p in model.parameters())
-        print(f"bridge_only enabled: trainable params {trainable} / {total}")
-        for name in trainable_names:
-            print(f"  bridge trainable: {name}")
-    elif args.gate_only:
-        for p in model.parameters():
-            p.requires_grad = False
-        for n, p in model.named_parameters():
-            if "referent_gate" in n:
-                p.requires_grad = True
-        trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
-        total = sum(p.numel() for p in model.parameters())
-        print(f"gate_only enabled: trainable params {trainable} / {total}")
 
     print("will save train model")
 
-    def _total_norm(values):
-        if not values:
-            return 0.0
-        return float(sum(v * v for v in values) ** 0.5)
-
-    def collect_referent_gate_stats(model):
-        gate_modules = [(n, m) for n, m in model.named_modules() if n.endswith("referent_gate")]
-        proj_norms = []
-        gate_norms = []
-        proj_grad_norms = []
-        gate_grad_norms = []
-        alpha_tensors = []
-
-        for _, module in gate_modules:
-            proj_norms.append(module.proj.weight.detach().float().norm().item())
-            gate_norms.append(module.gate.weight.detach().float().norm().item())
-            if module.proj.weight.grad is not None:
-                proj_grad_norms.append(module.proj.weight.grad.detach().float().norm().item())
-            if module.gate.weight.grad is not None:
-                gate_grad_norms.append(module.gate.weight.grad.detach().float().norm().item())
-            if module.last_alpha is not None:
-                alpha_tensors.append(module.last_alpha.detach().float().reshape(-1))
-
-        stats = {
-            "modules": len(gate_modules),
-            "proj_norm": _total_norm(proj_norms),
-            "gate_norm": _total_norm(gate_norms),
-            "proj_grad_norm": _total_norm(proj_grad_norms),
-            "gate_grad_norm": _total_norm(gate_grad_norms),
-        }
-
-        if alpha_tensors:
-            alpha = torch.cat(alpha_tensors)
-            stats.update(
-                {
-                    "alpha_mean": alpha.mean().item(),
-                    "alpha_std": alpha.std(unbiased=False).item(),
-                    "alpha_min": alpha.min().item(),
-                    "alpha_max": alpha.max().item(),
-                }
-            )
-        else:
-            stats.update(
-                {
-                    "alpha_mean": float("nan"),
-                    "alpha_std": float("nan"),
-                    "alpha_min": float("nan"),
-                    "alpha_max": float("nan"),
-                }
-            )
-
-        return stats
-
-    def print_referent_gate_optimizer_sanity(model, optimizer):
-        optimizer_param_ids = {id(p) for group in optimizer.param_groups for p in group["params"]}
-        gate_params = [(n, p) for n, p in model.named_parameters() if "referent_gate" in n]
-        trainable_gate = [(n, p) for n, p in gate_params if p.requires_grad]
-        optimizer_gate = [(n, p) for n, p in gate_params if id(p) in optimizer_param_ids]
-        optimizer_trainable_gate = [
-            (n, p) for n, p in gate_params if p.requires_grad and id(p) in optimizer_param_ids
-        ]
-        print(
-            "referent_gate sanity: "
-            f"params={sum(p.numel() for _, p in gate_params)} | "
-            f"trainable={sum(p.numel() for _, p in trainable_gate)} | "
-            f"in_optimizer={sum(p.numel() for _, p in optimizer_gate)} | "
-            f"trainable_in_optimizer={sum(p.numel() for _, p in optimizer_trainable_gate)}"
-        )
-
-        stats = collect_referent_gate_stats(model)
-        print(
-            "referent_gate init stats: "
-            f"modules={stats['modules']} | "
-            f"proj_norm={stats['proj_norm']:.6f} | "
-            f"gate_norm={stats['gate_norm']:.6f}"
-        )
-
-    def log_referent_gate_stats(global_step, loss_value):
-        stats = collect_referent_gate_stats(model)
-        message = (
-            f"gate_stats step={global_step} "
-            f"loss={loss_value:.6f} "
-            f"proj_norm={stats['proj_norm']:.6f} "
-            f"gate_norm={stats['gate_norm']:.6f} "
-            f"proj_grad_norm={stats['proj_grad_norm']:.6f} "
-            f"gate_grad_norm={stats['gate_grad_norm']:.6f} "
-            f"alpha_mean={stats['alpha_mean']:.4f} "
-            f"alpha_std={stats['alpha_std']:.4f} "
-            f"alpha_min={stats['alpha_min']:.4f} "
-            f"alpha_max={stats['alpha_max']:.4f}"
-        )
-        print(message)
-        with open(os.path.join(args.log_root, f'{args.name}.txt'), "a") as f:
-            f.write(message + "\n")
-
-    def find_prompt_bridge_module(model):
-        for _, module in model.named_modules():
-            if module.__class__.__name__ == "ResidualPromptBridge":
-                return module
-        return None
-
-    def collect_prompt_bridge_grad_norms(model):
-        module = find_prompt_bridge_module(model)
-        if module is None:
-            return {}
-
-        def grad_norm(param):
-            if param.grad is None:
-                return None
-            return float(param.grad.detach().float().norm().item())
-
-        return {
-            "W_a": grad_norm(module.attn_proj.weight),
-            "W_r": grad_norm(module.residual_proj.weight),
-            "W_g": grad_norm(module.gate.weight),
-            "b_g": grad_norm(module.gate.bias),
-        }
-
-    def print_prompt_bridge_grad_norms(label, norms):
-        parts = []
-        for key in ["W_a", "W_r", "W_g", "b_g"]:
-            value = norms.get(key)
-            if value is None:
-                parts.append(f"{key}=None")
-            else:
-                parts.append(f"{key}={value:.6e}")
-        print(f"{label}: " + " | ".join(parts))
-
-    def run_bridge_sanity_checks(model, dataloader):
-        if not args.use_residual_prompt_bridge:
-            raise ValueError("--bridge_sanity_only requires --use_residual_prompt_bridge")
-
-        model.train()
-        batch = next(iter(dataloader))
-        input_dict = dict_to_cuda(batch)
-
-        output_dict = model.forward(
-            images=input_dict["images"],
-            images_clip=input_dict["images_clip"],
-            audio_features=input_dict["audio_feats"],
-            image_features=input_dict["image_feats"],
-            input_ids=input_dict["input_ids"],
-            labels=input_dict["labels"],
-            attention_masks=input_dict["attention_masks"],
-            masks_list=input_dict["masks"],
-            resize_list=input_dict["resizes"],
-            orgsize_list=input_dict["orgsizes"],
-            conversation_list=input_dict["convs"],
-            refs_num=input_dict["refs_num"],
-            fids=input_dict["fids"],
-            vids=input_dict["vids"],
-            contrast=0.0,
-            ref_ids=input_dict["ref_ids"],
-            epoch=0,
-            inference=False,
-            target_frame=args.bridge_target_frame,
-        )
-
-        model.zero_grad(set_to_none=True)
-        output_dict["mask_loss"].backward(retain_graph=True)
-        print_prompt_bridge_grad_norms(
-            "bridge grad check | L_mask only",
-            collect_prompt_bridge_grad_norms(model),
-        )
-
-        model.zero_grad(set_to_none=True)
-        output_dict["bridge_teacher_loss_raw"].backward()
-        print_prompt_bridge_grad_norms(
-            "bridge grad check | L_teach only",
-            collect_prompt_bridge_grad_norms(model),
-        )
-
-        metrics = output_dict["bridge_metrics"]
-        print(
-            "bridge identity check: "
-            f"delta_norm_mean={metrics['delta_norm_mean']:.6f} | "
-            f"cos(p_hat,q)={metrics['cos_p_hat_q_mean']:.6f} | "
-            f"q_norm_mean={metrics['q_norm_mean']:.6f} | "
-            f"p_hat_norm_mean={metrics['p_hat_norm_mean']:.6f} | "
-            f"gate_mean={metrics['gate_mean']:.6f} | "
-            f"gate_std={metrics['gate_std']:.6f}"
-        )
-
-        teacher_pm_norms = []
-        teacher_rg_norms = []
-        teacher_cosines = []
-        scanned_batches = max(1, args.bridge_sanity_batches)
-
-        model.eval()
-        with torch.no_grad():
-            for batch_idx, batch in enumerate(dataloader):
-                if batch_idx >= scanned_batches:
-                    break
-                input_dict = dict_to_cuda(batch)
-                result = model.forward(
-                    images=input_dict["images"],
-                    images_clip=input_dict["images_clip"],
-                    audio_features=input_dict["audio_feats"],
-                    image_features=input_dict["image_feats"],
-                    input_ids=input_dict["input_ids"],
-                    labels=input_dict["labels"],
-                    attention_masks=input_dict["attention_masks"],
-                    masks_list=input_dict["masks"],
-                    resize_list=input_dict["resizes"],
-                    orgsize_list=input_dict["orgsizes"],
-                    conversation_list=input_dict["convs"],
-                    refs_num=input_dict["refs_num"],
-                    fids=input_dict["fids"],
-                    vids=input_dict["vids"],
-                    contrast=0.0,
-                    ref_ids=input_dict["ref_ids"],
-                    inference=True,
-                    target_frame=args.bridge_target_frame,
-                )
-                bridge_metrics = result["bridge_metrics"]
-                teacher_pm_norms.append(bridge_metrics["p_mask_norm_mean"])
-                teacher_rg_norms.append(bridge_metrics["z_gt_norm_mean"])
-                teacher_cosines.append(bridge_metrics["cos_p_mask_z_gt_mean"])
-
-        print(
-            "bridge teacher sanity: "
-            f"mean||p_mask||={float(np.mean(teacher_pm_norms)):.6f} | "
-            f"mean||z_gt||={float(np.mean(teacher_rg_norms)):.6f} | "
-            f"mean cos(p_mask,z_gt)={float(np.mean(teacher_cosines)):.6f}"
-        )
-
     def valuate(model, dataloader, args, name):
         model.eval()
 
         total_iou = 0
         total_fscore = 0
         count = 0
-        bridge_accumulators = defaultdict(float)
-        bridge_count = 0
 
         for batch in tqdm(dataloader, desc=f"Evaluating on {name}"):
             input_dict = dict_to_cuda(batch)
@@ -740,8 +395,7 @@ if __name__ == "__main__":
                                             vids=input_dict["vids"],
                                             contrast=args.ct_weight,
                                             ref_ids=input_dict["ref_ids"],
-                                            inference=True,
-                                            target_frame=args.bridge_target_frame)
+                                            inference=True)
             pred_masks = output_dict["pred_masks"]  # list[B]:[num_seg, T, H, W]
             gt_masks = output_dict["gt_masks"]  # list[B]:[num_seg, T, H, W]
             for i in range(len(pred_masks)):
@@ -754,46 +408,23 @@ if __name__ == "__main__":
                 total_fscore += fscore * num_seg * T
                 count += num_seg * T
 
-            if args.use_residual_prompt_bridge and "bridge_metrics" in output_dict:
-                for key, value in output_dict["bridge_metrics"].items():
-                    bridge_accumulators[key] += float(value)
-                bridge_count += 1
-
         print(f"\n  valuate on {name}:  miou: {total_iou/count}  fscore: {total_fscore/count}")
 
         with open(os.path.join(args.log_root, f'{args.name}.txt'), "a") as f:
             f.write(f"valuate on {name}:  miou {total_iou/count}  true fscore {total_fscore/count} \n")
-            if bridge_count > 0:
-                bridge_summary = " | ".join(
-                    f"{key}={bridge_accumulators[key] / bridge_count:.6f}"
-                    for key in sorted(bridge_accumulators.keys())
-                )
-                print(f"  bridge on {name}: {bridge_summary}")
-                f.write(f"bridge on {name}: {bridge_summary}\n")
-
 
-    if args.bridge_sanity_only:
-        run_bridge_sanity_checks(model, train_eval_dataloader)
-        sys.exit(0)
 
     # ---------------train------------------------------------------
 
     model.train()
     epochs = args.epochs
     print("init lr:", args.lr)
-    trainable_params = [p for p in model.parameters() if p.requires_grad]
-    optimizer = AdamW(trainable_params, lr=args.lr, betas=(0.9, 0.95), weight_decay=0.01)
-    print_referent_gate_optimizer_sanity(model, optimizer)
-
-    gradient_accumulation_steps = max(1, int(16 // args.batch_size))
-    step_per_epoch = max(1, len(train_dataloader) // gradient_accumulation_steps)
-    full_total_steps = epochs * step_per_epoch
-    total_steps = min(args.max_steps, full_total_steps) if args.max_steps > 0 else full_total_steps
+    optimizer = AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.01)
+
+    gradient_accumulation_steps = int(16 // args.batch_size)
+    step_per_epoch = len(train_dataloader) // gradient_accumulation_steps
+    total_steps = epochs * step_per_epoch
     warmup_steps = int(total_steps * 0.1)
-    print(
-        f"training schedule: grad_accum={gradient_accumulation_steps} | "
-        f"step_per_epoch={step_per_epoch} | total_optimizer_steps={total_steps}"
-    )
 
     scheduler = get_cosine_schedule_with_warmup(
         optimizer,
@@ -802,9 +433,6 @@ if __name__ == "__main__":
     )
 
 
-    optimizer_step_count = 0
-    stop_training = False
-
     for epoch in range(epochs):
 
         model.train()
@@ -813,9 +441,6 @@ if __name__ == "__main__":
 
         loop = tqdm(train_dataloader, desc=f"Training Epoch {epoch + 1}/{epochs}")
         for step, batch in enumerate(loop):
-            if args.max_steps > 0 and optimizer_step_count >= args.max_steps:
-                stop_training = True
-                break
             input_dict = dict_to_cuda(batch)
             output_dict = model.forward(images=input_dict["images"],
                                         images_clip=input_dict["images_clip"],
@@ -834,7 +459,6 @@ if __name__ == "__main__":
                                         contrast=args.ct_weight,
                                         ref_ids=input_dict["ref_ids"],
                                         epoch=epoch,
-                                        gate_only=args.gate_only,
                                         inference=False)
 
             loss = output_dict["loss"]
@@ -844,57 +468,23 @@ if __name__ == "__main__":
 
 
             if (step + 1) % gradient_accumulation_steps == 0:
-                optimizer_step_count += 1
-                if (
-                    args.log_gate_stats_every > 0
-                    and optimizer_step_count % args.log_gate_stats_every == 0
-                ):
-                    log_referent_gate_stats(
-                        optimizer_step_count,
-                        loss.item() * gradient_accumulation_steps,
-                    )
                 optimizer.step()
                 scheduler.step()
                 optimizer.zero_grad()
 
                 current_lr = scheduler.get_lr()[0]
-                postfix = {
-                    "lr": current_lr,
-                    "loss": running_loss / ((step + 1) / gradient_accumulation_steps),
-                }
-                if args.use_residual_prompt_bridge:
-                    postfix["bridge"] = float(output_dict["bridge_teacher_loss"].item())
-                    postfix["pm"] = float(output_dict["bridge_pm_loss"].item())
-                    postfix["rg"] = float(output_dict["bridge_rg_loss"].item())
-                loop.set_postfix(**postfix)
-
-                if args.max_steps > 0 and optimizer_step_count >= args.max_steps:
-                    stop_training = True
-                    break
+                loop.set_postfix(lr=current_lr, loss=running_loss / ((step + 1) / gradient_accumulation_steps))
 
-        denom = max(1, optimizer_step_count)
-        print(f"  Epoch {epoch + 1}, Loss:{running_loss / denom :.4f}, Learning Rate:{scheduler.get_last_lr()[0]:.6f}")
+        print(f"  Epoch {epoch + 1}, Loss:{running_loss / ((step + 1) / gradient_accumulation_steps) :.4f}, Learning Rate:{scheduler.get_last_lr()[0]:.6f}")
 
 
         with open(os.path.join(args.log_root, f'{args.name}.txt'), "a") as f:
-            f.write(f"Epoch {epoch}: running_loss {running_loss / denom}  Learning Rate:{scheduler.get_last_lr()[0]:.6f}\n")
-
-        if stop_training:
-            print(f"stopped early at optimizer step {optimizer_step_count}")
-            break
+            f.write(f"Epoch {epoch}: running_loss {running_loss / len(train_dataloader) * gradient_accumulation_steps}  Learning Rate:{scheduler.get_last_lr()[0]:.6f}\n")
 
 
     torch.save(model.state_dict(), os.path.join(args.checkpoint_root, f"{args.name}.pth"))
     print(f"trained model saved as {args.name}.pth")
 
-    if args.skip_eval_after_train:
-        print("skip_eval_after_train enabled: exiting after checkpoint save")
-        sys.exit(0)
-
-    if args.eval_train_only:
-        valuate(model, train_eval_dataloader, args, 'train_overfit')
-        sys.exit(0)
-
     # ---------------test on seen & unseen ------------------------------------------
     model.eval()
 
@@ -941,4 +531,4 @@ if __name__ == "__main__":
     print(f"\n  valuate on test_n_refer, metric: {total_metric/count}")
 
     with open(os.path.join(args.log_root, f'{args.name}.txt'), "a") as f:
-        f.write(f"\n valuate on  test_n_refer:   metric {total_metric/count} \n")
+        f.write(f"\n valuate on  test_n_refer:   metric {total_metric/count} \n")