from typing import Dict, Optional, Tuple import os import torch import torch.distributed as dist import numpy as np from torchvision.ops.boxes import batched_nms, box_area from torch import nn, Tensor from transformers import PreTrainedModel, AutoModelForCausalLM, AutoConfig from peft import LoraConfig, get_peft_model, PeftModel from src.arguments import ModelArguments from src.model.processor import LLAVA_NEXT, QWEN2_VL, PHI3V, get_backbone_name, print_master, QWEN2_5_VL, INTERNVIDEO2, \ QWEN2_VL_TOKENSELECTION, QWEN2_5_VL_TOKENSELECTION, get_model_class, GME, VLM_IMAGE_TOKENS, LamRA, LamRA_QWEN2_5, \ COLPALI, QWEN2_5_OMNI, WAVE HF_LOCAL_FILES_ONLY = False class SelfAttentionConnector(nn.Module): def __init__(self, input_dim, output_dim, num_heads=8): super().__init__() self.norm = nn.LayerNorm(input_dim) self.attn = nn.MultiheadAttention(input_dim, num_heads, batch_first=True) self.proj = nn.Linear(input_dim, output_dim) def forward(self, x): x = self.norm(x) x, _ = self.attn(x, x, x) return self.proj(x) class SAMAudioVLM(nn.Module): """ Wrapper that holds frozen SAMAudio + frozen SAM3 and produces dual embeddings: (A) Video-segmenter embeddings from SAM3 instance_embeds [B,256,H,W] via Conv2d+MLP (B) Audio-segmenter embeddings from SAM-Audio DiT ODE [B,T,128] via Conv1d+MLP Both projected to vlm_hidden_size. """ DFLT_ODE_OPT = {"method": "midpoint", "options": {"step_size": 2 / 32}} def __init__(self, sam_audio_config): super().__init__() from src.sam_audio.model.model import SAMAudio from src.sam_audio.processor import SAMAudioProcessor from src.sam3.model_builder import build_sam3_image_model # --- Memory config --- ode_steps = sam_audio_config.get("ode_steps", 16) # default 16; try 4 or 8 to save memory self.ode_opt = {"method": "midpoint", "options": {"step_size": 2 / (2 * ode_steps)}} self.offload_to_cpu = sam_audio_config.get("offload_to_cpu", False) self.skip_sam3_video = sam_audio_config.get("skip_sam3_video", False) self.max_sam3_frames = sam_audio_config.get("max_sam3_frames", None) self.max_sam_audio_frames = sam_audio_config.get("max_sam_audio_frames", None) self.sam_audio_fps = sam_audio_config.get("fps", None) self.sam3_image_size = sam_audio_config.get("sam3_image_size", 1008) self.max_audio_seconds = sam_audio_config.get("max_audio_seconds", 30) self.audio_sample_rate = 16000 self.max_visual_tokens = sam_audio_config.get("max_visual_tokens", 256) # --- Frozen SAM-Audio model (bf16 to save memory) --- self.sam_audio_model = SAMAudio.from_pretrained(sam_audio_config["checkpoint"]) self.sam_audio_model.to(dtype=torch.bfloat16) for param in self.sam_audio_model.parameters(): param.requires_grad = False # --- Frozen SAM3 image model (kept float32) --- # Directly use the image model instead of the video predictor API to # avoid session management, disk I/O, CPU<->GPU transfers, and the # tracker pipeline overhead that are unnecessary during training. self.sam3_model = build_sam3_image_model( checkpoint_path=sam_audio_config["sam3_checkpoint"], device="cpu", eval_mode=True, enable_segmentation=True, ) for param in self.sam3_model.parameters(): param.requires_grad = False # CPU offloading: move frozen models to CPU, only bring to GPU during forward if self.offload_to_cpu: self.sam_audio_model.cpu() self.sam3_model.cpu() # --- SAM-Audio processor for creating Batch objects --- self.processor = SAMAudioProcessor.from_pretrained(sam_audio_config["checkpoint"]) vlm_hidden_size = sam_audio_config["vlm_hidden_size"] connector = sam_audio_config.get("connector", "mlp") self_attention_num_heads = sam_audio_config.get("self_attention_num_heads", 4) # --- Trainable: Visual path (SAM3 instance_embeds -> VLM) --- sam3_dim = 256 # SAM3 hidden_dim self.sam3_reducer = nn.Conv2d(sam3_dim, sam3_dim, kernel_size=4, stride=4) if connector == "self_attention": self.sam3vlm_layer = SelfAttentionConnector(sam3_dim, vlm_hidden_size, num_heads=self_attention_num_heads) elif connector == "mlp": self.sam3vlm_layer = nn.Sequential( nn.Linear(sam3_dim, sam3_dim), nn.GELU(), nn.Linear(sam3_dim, vlm_hidden_size), ) # --- Trainable: Audio path (SAM-Audio DiT output -> VLM) --- # DiT output is [B, T, 256], target half is [B, T, 128] target_channels = 128 self.audio_reducer = nn.Conv1d(target_channels, target_channels, kernel_size=4, stride=4) if connector == "self_attention": self.audio_projection = SelfAttentionConnector(target_channels, vlm_hidden_size, num_heads=self_attention_num_heads) elif connector == "mlp": self.audio_projection = nn.Sequential( nn.Linear(target_channels, target_channels), nn.GELU(), nn.Linear(target_channels, vlm_hidden_size), ) def train(self, mode: bool = True): """Keep frozen SAM3 and SAM-Audio in eval mode regardless of parent .train() calls.""" super().train(mode) self.sam3_model.eval() self.sam_audio_model.eval() return self def extract_video_frames(self, video_path, max_frames=None, fps=None, image_size=None): """ Pre-extract frames from a video file using FPS-based sampling. Returns a list of PIL images, or None if the video is corrupt/unreadable. """ from torchcodec.decoders import VideoDecoder from PIL import Image try: decoder = VideoDecoder(video_path, dimension_order="NCHW") except RuntimeError as e: print(f"\n{'='*60}\n[ERROR] Failed to decode video: {video_path}\n {e}\n{'='*60}\n", flush=True) return None video_fps = decoder.metadata.average_fps total_frames = decoder.metadata.num_frames if fps is not None: FRAME_FACTOR = 2 nframes = total_frames / video_fps * fps nframes = max(nframes, FRAME_FACTOR) if max_frames is not None: nframes = min(nframes, max_frames) nframes = min(nframes, total_frames) nframes = int(nframes // FRAME_FACTOR) * FRAME_FACTOR nframes = max(nframes, FRAME_FACTOR) elif max_frames is not None: nframes = min(max_frames, total_frames) else: nframes = total_frames idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist() decoded = decoder.get_frames_at(indices=idx) frames_tensor = decoded.data if hasattr(decoded, "data") else decoded pil_frames = [] for f in frames_tensor: arr = f.permute(1, 2, 0).numpy().astype(np.uint8) img = Image.fromarray(arr) if image_size is not None: img = img.resize((image_size, image_size)) pil_frames.append(img) return pil_frames def _preprocess_frame_for_sam3(self, frame, device): """ Convert a PIL image or tensor into the format SAM3 expects: [1, 3, image_size, image_size] float32, normalized with mean/std = 0.5. """ from PIL import Image image_size = self.sam3_image_size if isinstance(frame, Image.Image): img_np = np.array(frame.convert("RGB").resize((image_size, image_size))) img = torch.from_numpy(img_np).permute(2, 0, 1).float() / 255.0 elif isinstance(frame, torch.Tensor): img = frame.float() if img.dim() == 2: img = img.unsqueeze(0) if img.shape[0] not in (1, 3): img = img.permute(2, 0, 1) if img.max() > 1.0: img = img / 255.0 if img.shape[-2:] != (image_size, image_size): img = torch.nn.functional.interpolate( img.unsqueeze(0), size=(image_size, image_size), mode="bilinear", align_corners=False, ).squeeze(0) else: img = torch.tensor(np.array(frame), dtype=torch.float).permute(2, 0, 1) / 255.0 img = (img - 0.5) / 0.5 return img.unsqueeze(0).to(device=device, dtype=torch.float32) def generate_masks_and_embeds( self, video_path, frames, prompt_text=None, points=None, point_labels=None, bboxes=None, bbox_labels=None, ): """ Use SAM3 image model directly to generate per-frame masks and instance_embeds in a single forward_grounding call per frame. Returns: masks: list of torch.Tensor, per-frame binary masks on GPU instance_embeds: Tensor [1, 256, H, W] from the key frame's segmentation head, or None """ from src.sam3.model.data_misc import FindStage, convert_my_tensors from src.sam3.model.geometry_encoders import Prompt device = next(self.sam3_model.parameters()).device num_frames = len(frames) if frames is not None else 1 text = prompt_text if prompt_text is not None else "" all_masks = [] instance_embeds = None key_idx = num_frames // 2 for frame_index in range(num_frames): frame = frames[frame_index] if frames is not None else None if frame is None: h, w = self.sam3_image_size, self.sam3_image_size all_masks.append(torch.zeros(1, h, w, device=device)) continue img_tensor = self._preprocess_frame_for_sam3(frame, device) with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"): backbone_out = self.sam3_model.backbone(img_tensor, captions=[text]) find_input = FindStage( img_ids=[0], text_ids=[0], input_boxes=[torch.zeros(258, device=device)], input_boxes_mask=[torch.empty(0, dtype=torch.bool, device=device)], input_boxes_label=[torch.empty(0, dtype=torch.long, device=device)], input_points=[torch.empty(0, 257, device=device)], input_points_mask=[torch.empty(0, device=device)], ) convert_my_tensors(find_input) bs = 1 if bboxes is not None: # bboxes is [x1, y1, x2, y2] in original image pixel coordinates. # SAM3 geometric_prompt expects normalised [cx, cy, w, h]. from PIL import Image as _PIL if isinstance(frame, _PIL.Image): img_w, img_h = frame.size else: img_w, img_h = self.sam3_image_size, self.sam3_image_size x1, y1, x2, y2 = (float(v) for v in bboxes[:4]) cx = (x1 + x2) / 2.0 / img_w cy = (y1 + y2) / 2.0 / img_h bw = (x2 - x1) / img_w bh = (y2 - y1) / img_h box_t = torch.tensor( [[cx, cy, bw, bh]], device=device, dtype=torch.float ) # [1, 4] geometric_prompt = Prompt( box_embeddings=box_t.unsqueeze(0), # [1, 1, 4] box_mask=torch.zeros(bs, 1, device=device, dtype=torch.bool), box_labels=torch.ones(1, bs, device=device, dtype=torch.long), point_embeddings=torch.zeros(0, bs, 2, device=device), point_mask=torch.zeros(bs, 0, device=device, dtype=torch.bool), ) elif points is not None: # points is a list of [x, y] in normalized [0, 1] coordinates. # SAM3 expects point_embeddings: [N_points, bs, 2] normalized in [0,1]. pts = torch.tensor( [[float(p[0]), float(p[1])] for p in points], device=device, dtype=torch.float, ) # [N, 2] pts = pts.unsqueeze(1) # [N, 1, 2] n_pts = pts.shape[0] if point_labels is None: pt_labels = torch.ones(n_pts, bs, device=device, dtype=torch.long) else: pt_labels = torch.tensor( point_labels, device=device, dtype=torch.long ).view(n_pts, 1).expand(n_pts, bs) geometric_prompt = Prompt( box_embeddings=torch.zeros(0, bs, 4, device=device), box_mask=torch.zeros(bs, 0, device=device, dtype=torch.bool), box_labels=torch.zeros(0, bs, device=device, dtype=torch.long), point_embeddings=pts, # [N, 1, 2] point_mask=torch.zeros(bs, n_pts, device=device, dtype=torch.bool), point_labels=pt_labels, # [N, 1] ) else: geometric_prompt = Prompt( box_embeddings=torch.zeros(0, bs, 4, device=device), box_mask=torch.zeros(bs, 0, device=device, dtype=torch.bool), box_labels=torch.zeros(0, bs, device=device, dtype=torch.long), point_embeddings=torch.zeros(0, bs, 2, device=device), point_mask=torch.zeros(bs, 0, device=device, dtype=torch.bool), ) out = self.sam3_model.forward_grounding( backbone_out=backbone_out, find_input=find_input, find_target=None, geometric_prompt=geometric_prompt, ) pred_masks = out["pred_masks"] pred_logits = out["pred_logits"].squeeze(-1).sigmoid() best_idx = pred_logits.argmax(dim=1) mask = pred_masks[0, best_idx[0]] binary_mask = (mask > 0).float() all_masks.append(binary_mask.unsqueeze(0)) if frame_index == key_idx and "instance_embeds" in out: instance_embeds = out["instance_embeds"] return all_masks, instance_embeds def forward( self, audio_paths=None, video_paths=None, raw_audio=None, frames=None, descriptions=None, anchors=None, prompt_text=None, points=None, point_labels=None, bboxes=None, bbox_labels=None, expected_batch_size=None, sam_mask_bytes=None, sam_video_paths_for_mask=None, ): """ Produces dual embeddings: 1. Visual segment embeddings from SAM3 instance_embeds 2. Audio segment embeddings from SAM-Audio DiT ODE Works in two modes: (A) Conventional datasets: receives audio_paths/video_paths from the collator and provides dummy prompts (empty text, no video masking, no anchors) following the SAM-Audio paper. (B) Explicit prompts: receives raw_audio, descriptions, SAM3 prompts etc. from SAM-Audio-specific datasets. Returns: (embeds, valid_mask) tuple where: embeds: Tensor [B, N_vis + T', vlm_hidden_size] or None valid_mask: BoolTensor [B] indicating which batch items have valid audio/video data (True = valid) """ from torchdiffeq import odeint from src.sam_audio.processor import Batch # ---- Resolve audio inputs ---- # Priority: explicit raw_audio > audio_paths from collator audio_inputs = raw_audio # tensor or list of tensors valid_audio_indices = None if audio_inputs is None and audio_paths is not None: valid_audio_indices = [] filtered = [] for i, p in enumerate(audio_paths): if p and os.path.isfile(p): valid_audio_indices.append(i) filtered.append(p) audio_inputs = filtered if filtered else None if expected_batch_size is None: expected_batch_size = len(audio_paths) # ---- Resolve video path ---- # Priority: explicit video_path from SAM-Audio dataset > video_paths from collator. # Per-sample mode (SCaR): when frames is a list-of-lists (one inner list per batch # sample), each sample gets its own SAM3 run. All other callers provide a flat list # of frames for a single shared video, which is the original behaviour. _per_sample_video = ( frames is not None and isinstance(frames, list) and len(frames) > 0 and isinstance(frames[0], list) ) video_path_resolved = None if not _per_sample_video and video_paths is not None: # Take the first valid video path from the batch (original behaviour) for vp in video_paths: if vp and os.path.isfile(vp): video_path_resolved = vp break # ---- Default prompts for conventional data (SAM-Audio paper) ---- # "Not all training examples have all three prompt modalities. # In such cases, we provide dummy conditions: an empty string for text, # an all-zero vector sequence for video, and a sequence of special # null tokens () for span." if descriptions is None and audio_inputs is not None: # Empty string descriptions -- SAM-Audio's T5 encoder handles this batch_size = len(audio_inputs) if isinstance(audio_inputs, list) else 1 descriptions = [""] * batch_size if prompt_text is None: # Empty text for SAM3 -- triggers generic detection (analogous to # SAM2's uniform point grid approach) prompt_text = "" # anchors default to None → SAM-Audio uses null tokens internally # Nothing to process if audio_inputs is None and video_path_resolved is None and not _per_sample_video: return None, None # Determine the target GPU device (from trainable projection layers) gpu_device = next(self.sam3_reducer.parameters()).device visual_embeds = None audio_embeds = None target_features = None with torch.no_grad(): # ---- SAM3 path: Generate masks + instance_embeds ---- instance_embeds_raw = None masks_for_audio = None if _per_sample_video: # Per-sample SAM3: each frames[i] is a [pil_image] list for one query. # video_paths[i] is the corresponding image file path (used as # a presence check only — frames are already decoded by the collator). sam3_device = next(self.sam3_model.parameters()).device all_ie: list = [] for i, frame_list in enumerate(frames): vp = video_paths[i] if video_paths is not None and i < len(video_paths) else None bbox_i = bboxes[i] if bboxes is not None and i < len(bboxes) else None point_i = points[i] if points is not None and i < len(points) else None point_labels_i = point_labels[i] if point_labels is not None and i < len(point_labels) else None if frame_list and any(f is not None for f in frame_list): _, ie_i = self.generate_masks_and_embeds( video_path=vp or "", frames=frame_list, prompt_text=prompt_text, bboxes=bbox_i, points=point_i, point_labels=point_labels_i, ) else: ie_i = None all_ie.append(ie_i) valid_ie = [ie for ie in all_ie if ie is not None] if valid_ie: ref = valid_ie[0] instance_embeds_raw = torch.cat( [ie if ie is not None else torch.zeros_like(ref) for ie in all_ie], dim=0, ) elif video_path_resolved is not None: actual_frames = frames if actual_frames is not None and self.max_sam3_frames is not None and len(actual_frames) > self.max_sam3_frames: indices = np.linspace(0, len(actual_frames) - 1, self.max_sam3_frames, dtype=int) actual_frames = [actual_frames[i] for i in indices] if actual_frames is None: actual_frames = self.extract_video_frames( video_path_resolved, max_frames=self.max_sam3_frames, fps=self.sam_audio_fps, image_size=self.sam3_image_size, ) if actual_frames is None: video_path_resolved = None else: sam3_device = next(self.sam3_model.parameters()).device masks_list, instance_embeds_raw = self.generate_masks_and_embeds( video_path=video_path_resolved, frames=actual_frames, prompt_text=prompt_text, points=points, point_labels=point_labels, bboxes=bboxes, bbox_labels=bbox_labels, ) masks_for_audio = [m.float().to(sam3_device) for m in masks_list] # ---- SAM-Audio path: Encode and run ODE ---- if audio_inputs is not None: device = next(self.sam_audio_model.parameters()).device audio_list = audio_inputs if isinstance(audio_inputs, list) else [audio_inputs] # sam_mask_bytes is only injected by _inject_omnichoir_sam_fields # (OmniCHOIR mask_bytes interaction). For all other benchmarks this # parameter is None → masked_videos stays None, behaviour unchanged. masked_videos = None if sam_mask_bytes is not None and not all(b is not None for b in sam_mask_bytes): null_indices = [i for i, b in enumerate(sam_mask_bytes) if b is None] null_paths = [ (sam_video_paths_for_mask[i] if sam_video_paths_for_mask is not None else "?") for i in null_indices ] print_master( f"[SAMAudioVLM] mask_bytes missing for {len(null_indices)}/{len(sam_mask_bytes)} " f"samples in batch — skipping masked_videos for entire batch.\n" f" Samples without mask_bytes (index → video_path):\n" + "\n".join(f" [{i}] {p}" for i, p in zip(null_indices, null_paths)) ) if sam_mask_bytes is not None and all(b is not None for b in sam_mask_bytes): import io import numpy as np masks_list, vpaths_list = [], [] for mb, vp in zip(sam_mask_bytes, sam_video_paths_for_mask): if mb is not None and vp is not None: arr = np.load(io.BytesIO(mb))["video_masklet"] # (T, H, W) int64 t = torch.from_numpy(arr).unsqueeze(1).float() # (T, 1, H, W) masks_list.append(t) vpaths_list.append(vp) if masks_list: masked_videos = self.processor.mask_videos( videos=vpaths_list, masks=masks_list, ) desc_list = descriptions if isinstance(descriptions, list) else [descriptions] if masked_videos is not None and len(masked_videos) == 1 and len(desc_list) > 1: masked_videos = masked_videos * len(desc_list) batch = self.processor( descriptions=desc_list, audios=audio_list, anchors=anchors, masked_videos=masked_videos, ) max_samples = int(self.max_audio_seconds * self.audio_sample_rate) if batch.audios.shape[-1] > max_samples: batch.audios = batch.audios[..., :max_samples] batch = batch.to(device) forward_args = self.sam_audio_model._get_forward_args(batch) audio_features = forward_args["audio_features"] B, T, C = audio_features.shape # After audio truncation the codec output T may differ from the # processor-computed frame count used for anchor_ids etc. Align # all temporal-dimension tensors to the codec's actual T. for _k in ("anchor_ids", "anchor_alignment", "audio_pad_mask"): v = forward_args.get(_k) if v is not None and v.shape[1] != T: forward_args[_k] = v[:, :T] if forward_args.get("masked_video_features") is not None: mv = forward_args["masked_video_features"] if mv.shape[2] != T: forward_args["masked_video_features"] = mv[:, :, :T] noise = torch.randn_like(audio_features) def vector_field(t, noisy_audio): return self.sam_audio_model.forward( noisy_audio=noisy_audio, time=t.expand(noisy_audio.size(0)), **forward_args, ) states = odeint( vector_field, noise, torch.tensor([0.0, 1.0], device=device), **self.ode_opt, ) generated_features = states[-1] target_features = generated_features[:, :, :128] # [B, T, 128] # ---- Trainable projections (grad enabled) ---- # Visual path: SAM3 instance_embeds -> VLM space if instance_embeds_raw is not None: ie = instance_embeds_raw.to(gpu_device).float() if ie.dim() == 3: ie = ie.unsqueeze(0) # [1, 256, H, W] # Adaptive pooling to limit spatial tokens before projection. # SAM3 instance_embeds can be very large (e.g. 288×288 = 82944), # which would produce thousands of prefix tokens and blow up the # VLM's attention memory. max_side = int(self.max_visual_tokens ** 0.5) # e.g. 256 tokens → 16×16 if ie.shape[2] > max_side * 4 or ie.shape[3] > max_side * 4: ie = torch.nn.functional.adaptive_avg_pool2d(ie, (max_side * 4, max_side * 4)) # Conv2d spatial reduction (stride=4) ie_reduced = self.sam3_reducer(ie) # [B, 256, H', W'] # Flatten and project ie_flat = ie_reduced.flatten(2).transpose(1, 2) # [B, H'*W', 256] visual_embeds = self.sam3vlm_layer(ie_flat) # [B, N_vis, vlm_hidden_size] # Audio path: DiT target features -> VLM space if target_features is not None: tf = target_features.float() # [B, T, 128] # Conv1d temporal reduction tf_t = tf.transpose(1, 2) # [B, 128, T] # Pad short sequences so T >= kernel_size (4) k = self.audio_reducer.kernel_size[0] if tf_t.shape[2] < k: tf_t = torch.nn.functional.pad(tf_t, (0, k - tf_t.shape[2])) tf_reduced = self.audio_reducer(tf_t) # [B, 128, T'] tf_reduced = tf_reduced.transpose(1, 2) # [B, T', 128] audio_embeds = self.audio_projection(tf_reduced) # [B, T', vlm_hidden_size] # Align batch dimensions before concatenation if visual_embeds is not None and audio_embeds is not None: if visual_embeds.shape[0] == 1 and audio_embeds.shape[0] > 1: visual_embeds = visual_embeds.expand(audio_embeds.shape[0], -1, -1) elif audio_embeds.shape[0] == 1 and visual_embeds.shape[0] > 1: audio_embeds = audio_embeds.expand(visual_embeds.shape[0], -1, -1) # Concatenate both prefix embeddings parts = [p for p in [visual_embeds, audio_embeds] if p is not None] if len(parts) == 0: return None, None result = torch.cat(parts, dim=1) # [B_valid, N_vis + T', vlm_hidden_size] # When audio_paths were filtered, scatter the valid-item embeddings # back into a full-batch tensor so the batch dim matches inputs_embeds. B_out = expected_batch_size if expected_batch_size is not None else result.shape[0] valid_mask = torch.ones(B_out, dtype=torch.bool, device=result.device) if valid_audio_indices is not None and result.shape[0] < B_out: full = result.new_zeros(B_out, result.shape[1], result.shape[2]) for out_idx, orig_idx in enumerate(valid_audio_indices): full[orig_idx] = result[out_idx] valid_mask = torch.zeros(B_out, dtype=torch.bool, device=result.device) for orig_idx in valid_audio_indices: valid_mask[orig_idx] = True result = full return result, valid_mask class MMEBModel(nn.Module): TRANSFORMER_CLS = AutoModelForCausalLM def __init__(self, encoder: PreTrainedModel, segmentation_model: torch.nn.Module = None, pooling: str = 'last', normalize: bool = False, temperature: float = 0.02, layer_stride: int = 1, symmetric_loss: bool = False, num_hidden_layers: int = 0, hidden_size: int = 0, num_learnable_tokens: int = 0, token_diversity_type: str = 'none', token_diversity_weight: float = 0.0, layer_fusion: str = 'equal', ): super().__init__() self.config = encoder.config self.encoder = encoder self.segmentation_model = segmentation_model if layer_fusion not in ('equal', 'joint'): raise ValueError(f"layer_fusion must be 'equal' or 'joint', got '{layer_fusion}'.") self.pooling = pooling self.layer_fusion = layer_fusion self.normalize = normalize self.temperature = temperature self.layer_stride = layer_stride self.symmetric_loss = symmetric_loss self.num_learnable_tokens = num_learnable_tokens self.token_diversity_type = token_diversity_type self.token_diversity_weight = token_diversity_weight self._token_diversity_loss = None self.cross_entropy = nn.CrossEntropyLoss(reduction='mean') self.is_ddp = dist.is_initialized() if self.is_ddp: self.process_rank = dist.get_rank() self.world_size = dist.get_world_size() if num_learnable_tokens > 0: self.learnable_tokens = nn.Parameter(torch.randn(num_learnable_tokens, hidden_size) * 0.02) if self.pooling in ('all_layer', 'all_layer_mlp'): selected_indices = list(range(layer_stride - 1, num_hidden_layers, layer_stride)) self.num_selected_layers = len(selected_indices) if self.pooling == 'all_layer': if self.layer_fusion == 'joint': assert num_learnable_tokens > 0, "layer_fusion='joint' requires num_learnable_tokens > 0" self.layer_token_weights = nn.Parameter( torch.zeros(self.num_selected_layers, num_learnable_tokens) ) else: # 'equal' self.layer_weights = nn.Parameter(torch.zeros(self.num_selected_layers)) if self.pooling == 'all_layer_mlp': self.layer_fusion_mlp = nn.Sequential( nn.Linear(hidden_size * self.num_selected_layers, hidden_size), nn.GELU(), nn.Linear(hidden_size, hidden_size), ) @property def device(self) -> torch.device: """Return the device of the first parameter.""" return next(self.parameters()).device def print_parameter_count(self): print_master("========== Model parameter count: ========== ") """Print total and trainable parameters for the model.""" total_params = sum(p.numel() for p in self.parameters()) trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad) print_master(f"Total parameters: {total_params:,}") print_master(f"Trainable parameters: {trainable_params:,}") print_master(f"Non-trainable parameters: {total_params - trainable_params:,}") # Also print_master breakdown by components print_master("\nParameter breakdown:") # Main encoder encoder_params = sum(p.numel() for p in self.encoder.parameters()) encoder_trainable = sum(p.numel() for p in self.encoder.parameters() if p.requires_grad) print_master(f" Encoder: {encoder_params:,} total, {encoder_trainable:,} trainable") # Segmentation model (if exists) if self.segmentation_model is not None: sam_params = sum(p.numel() for p in self.segmentation_model.parameters()) sam_trainable = sum(p.numel() for p in self.segmentation_model.parameters() if p.requires_grad) print_master(f" Segmentation model: {sam_params:,} total, {sam_trainable:,} trainable") if self.num_learnable_tokens > 0: print_master(f" learnable_tokens: {self.learnable_tokens.numel():,} params (K={self.num_learnable_tokens}), trainable=True") if self.pooling == 'all_layer': if getattr(self, 'layer_fusion', 'equal') == 'joint': L, K = self.layer_token_weights.shape print_master(f" layer_token_weights: {self.layer_token_weights.numel()} scalars ({L}L x {K}K joint matrix), trainable=True") else: print_master(f" layer_weights: {self.layer_weights.numel()} scalars (one per selected layer), trainable=True") elif self.pooling == 'all_layer_mlp': mlp_params = sum(p.numel() for p in self.layer_fusion_mlp.parameters()) print_master(f" layer_fusion_mlp: {mlp_params:,} params (concat+MLP, text-only uses last-token), trainable=True") return total_params, trainable_params def encode_input(self, input): if getattr(self, "model_backbone", None) == INTERNVIDEO2: if "input_ids" in input.keys(): # text side text_output = self.encoder.get_text_encoder()( input["input_ids"], attention_mask=input["attention_mask"], return_dict=True, mode="text", ) text_embeds = text_output.last_hidden_state pooled_text_embeds = text_embeds[:, 0] pooled_output = self.encoder.text_proj(pooled_text_embeds) pooled_output /= pooled_output.norm(dim=-1, keepdim=True) return pooled_output else: _, vfeat = self.encoder.encode_vision(input["pixel_values"], test=True) vfeat = self.encoder.vision_proj(vfeat) vfeat /= vfeat.norm(dim=-1, keepdim=True) return vfeat elif getattr(self, "model_backbone", None) in [GME, LamRA, LamRA_QWEN2_5]: # pooled_output = self.encoder(**input, return_dict=True, output_hidden_states=True) texts = [text.replace(VLM_IMAGE_TOKENS[QWEN2_VL] + '\n', '') for text in input["texts"]] # we are actually passing video queries so this should not happen images = [] for imgs in input['images']: # if multi images are given, select the middle frame only if isinstance(imgs, list): imgs = imgs[len(imgs) // 2] assert not isinstance(imgs, list) # make sure we have extracted the middle frame and it is no longer a list images.append(imgs) else: images.append(imgs) pooled_output = self.encoder.get_fused_embeddings(texts=texts, images=images) return pooled_output elif getattr(self, "model_backbone", None) == COLPALI: pooled_output = self.encoder(**input, return_dict=True, output_hidden_states=True) return pooled_output elif getattr(self, "model_backbone", None) == LLAVA_NEXT: input['pixel_values'] = input['pixel_values'].squeeze(dim=1) input['image_sizes'] = input['image_sizes'].squeeze(dim=1) hidden_states = self.encoder(**input, return_dict=True, output_hidden_states=True) hidden_states = hidden_states.hidden_states[-1] pooled_output = self._pooling(hidden_states, input['attention_mask']) return pooled_output elif getattr(self, "model_backbone", None) == QWEN2_5_OMNI: encoder_input = dict(input) if self.num_learnable_tokens > 0: encoder_input = self._inject_learnable_tokens(encoder_input) outputs = self.encoder( **encoder_input, segmentation_model=self.segmentation_model, return_dict=True, output_hidden_states=True, ) if self.pooling == 'all_layer': pooled_output = self._pool_all_layers(outputs.hidden_states, encoder_input['attention_mask'], encoder_input['input_ids']) elif self.pooling == 'all_layer_mlp': pooled_output = self._pool_all_layers_mlp(outputs.hidden_states, encoder_input['attention_mask'], encoder_input['input_ids']) else: pooled_output = self._pooling(outputs.hidden_states[-1], encoder_input['attention_mask']) return pooled_output elif getattr(self, "model_backbone", None) == WAVE: encoder_input = dict(input) if self.num_learnable_tokens > 0: encoder_input = self._inject_learnable_tokens(encoder_input) outputs = self.encoder( **encoder_input, return_dict=True, output_hidden_states=True, ) if self.pooling == 'all_layer': pooled_output = self._pool_all_layers(outputs.hidden_states, encoder_input['attention_mask'], encoder_input['input_ids']) elif self.pooling == 'all_layer_mlp': pooled_output = self._pool_all_layers_mlp(outputs.hidden_states, encoder_input['attention_mask'], encoder_input['input_ids']) else: pooled_output = self._pooling(outputs.hidden_states[-1], encoder_input['attention_mask']) return pooled_output else: hidden_states = self.encoder(**input, segmentation_model=self.segmentation_model, return_dict=True, output_hidden_states=True) hidden_states = hidden_states.hidden_states[-1] pooled_output = self._pooling(hidden_states, input['attention_mask']) return pooled_output def _inject_learnable_tokens(self, encoder_input: dict) -> dict: """Append K learnable tokens to the input embeddings so they can attend to the full sequence. We keep input_ids in the dict (padded with K pad tokens at the end) alongside inputs_embeds. The model skips re-embedding when inputs_embeds is provided (line: `if inputs_embeds is None`) but still uses input_ids for multimodal special-token merging, feature_attention_mask stacking, and rope index computation. The K appended pad tokens are not any special modality token id, so masked_scatter for audio/image/video won't affect those positions. """ embed_layer = self.encoder.get_input_embeddings() text_embeds = embed_layer(encoder_input['input_ids']) # [B, L, H] B = text_embeds.shape[0] K = self.num_learnable_tokens suffix = self.learnable_tokens.unsqueeze(0).expand(B, -1, -1).to(dtype=text_embeds.dtype) inputs_embeds = torch.cat([text_embeds, suffix], dim=1) # [B, L+K, H] suffix_mask = torch.ones(B, K, dtype=encoder_input['attention_mask'].dtype, device=encoder_input['attention_mask'].device) attention_mask = torch.cat([encoder_input['attention_mask'], suffix_mask], dim=1) # [B, L+K] # Pad input_ids with K harmless pad tokens so all input_ids-gated logic in the backbone # (feature_attention_mask stacking, audio/image/video merging, rope index) still runs. pad_id = getattr(self.encoder.config, 'pad_token_id', 0) or 0 pad_ids = torch.full((B, K), pad_id, dtype=encoder_input['input_ids'].dtype, device=encoder_input['input_ids'].device) input_ids_padded = torch.cat([encoder_input['input_ids'], pad_ids], dim=1) # [B, L+K] result = dict(encoder_input) result['input_ids'] = input_ids_padded result['inputs_embeds'] = inputs_embeds result['attention_mask'] = attention_mask return result def _pooling(self, last_hidden_state, attention_mask): if self.pooling in ('last', 'eos', 'all_layer', 'all_layer_mlp'): if self.num_learnable_tokens > 0: K = self.num_learnable_tokens reps = last_hidden_state[:, -K:, :].mean(dim=1) # [B, H] else: reps = self._get_last_token(last_hidden_state, attention_mask) else: raise NotImplementedError if self.normalize: reps = torch.nn.functional.normalize(reps, p=2, dim=-1) return reps def _get_last_token(self, hidden_state, attention_mask): """Extract the last non-padding token from each sequence (no normalization).""" left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0]) batch_size = hidden_state.shape[0] if left_padding: return hidden_state[torch.arange(batch_size), -1, :] eos_indices = attention_mask.sum(dim=1) - 1 return hidden_state[torch.arange(batch_size, device=hidden_state.device), eos_indices] def _detect_multimodal(self, input_ids: Tensor) -> Tensor: """Return a boolean mask [B] indicating which samples contain non-text tokens.""" mm_token_ids = set() config = self.encoder.config for attr in ('image_token_id', 'video_token_id', 'audio_token_id'): token_id = getattr(config, attr, None) if token_id is not None: mm_token_ids.add(token_id) if mm_token_ids: mm_ids_tensor = torch.tensor(list(mm_token_ids), device=input_ids.device, dtype=input_ids.dtype) return (input_ids.unsqueeze(-1) == mm_ids_tensor.unsqueeze(0).unsqueeze(0)).any(dim=-1).any(dim=-1) return torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device) def _pool_all_layers(self, all_hidden_states, attention_mask, input_ids): """WAVE-style multi-layer fusion via learnable weighted sum. layer_fusion=equal: multimodal Fused = sum_l softmax(layer_weights)[l] * h_l (same mix for all K tail tokens). layer_fusion=joint: W = softmax(layer_token_weights.flatten()).reshape(L, K), then out_b = sum_{l,k} W[l,k] * h^(l)_{b,k} (joint distribution over layer-token pairs). Text-only samples: standard last-token pooling from the final layer. """ layers = all_hidden_states[1:] # skip embedding layer selected = [layers[i] for i in range(self.layer_stride - 1, len(layers), self.layer_stride)] if self.num_learnable_tokens > 0: K = self.num_learnable_tokens k_tokens = [h[:, -K:, :] for h in selected] # list of [B, K, H] stacked = torch.stack(k_tokens, dim=0) # [L, B, K, H] if self.layer_fusion == 'joint': dtype = stacked.dtype W = torch.softmax( self.layer_token_weights.to(dtype).flatten(), dim=0 ).reshape(self.num_selected_layers, K) # [L, K] mm_output = (W[:, None, :, None] * stacked).sum(dim=(0, 2)) # [B, H] else: w = torch.softmax(self.layer_weights.to(dtype=stacked.dtype), dim=0) # [L] fused = (w[:, None, None, None] * stacked).sum(0) # [B, K, H] mm_output = fused.mean(dim=1) # [B, H] if self.token_diversity_type != 'none' and self.token_diversity_weight > 0: final_layer_tokens = selected[-1][:, -K:, :] div_loss = self._compute_token_diversity_loss(final_layer_tokens) if self._token_diversity_loss is None: self._token_diversity_loss = div_loss else: self._token_diversity_loss = self._token_diversity_loss + div_loss # Text-only path: strip appended learnable tokens to find the real EOS final_hidden = all_hidden_states[-1][:, :-K, :] text_mask = attention_mask[:, :-K] text_output = self._get_last_token(final_hidden, text_mask) else: w = torch.softmax(self.layer_weights.to(selected[0].dtype), dim=0) # [L] last_tokens = [self._get_last_token(h, attention_mask) for h in selected] stacked = torch.stack(last_tokens, dim=0) # [L, B, H] mm_output = (w[:, None, None] * stacked).sum(0) # [B, H] text_output = self._get_last_token(all_hidden_states[-1], attention_mask) is_multimodal = self._detect_multimodal(input_ids) output = torch.where(is_multimodal.unsqueeze(-1), mm_output, text_output) if self.normalize: output = torch.nn.functional.normalize(output, p=2, dim=-1) return output def _pool_all_layers_mlp(self, all_hidden_states, attention_mask, input_ids): """WAVE-style multi-layer pooling: concat last-token from each selected layer through a 2-layer MLP for multimodal samples; standard last-token pooling for text-only. """ layers = all_hidden_states[1:] selected = [layers[i] for i in range(self.layer_stride - 1, len(layers), self.layer_stride)] if self.num_learnable_tokens > 0: K = self.num_learnable_tokens repr_per_layer = [h[:, -K:, :].mean(dim=1) for h in selected] # L x [B, H] text_output = self._get_last_token( all_hidden_states[-1][:, :-K, :], attention_mask[:, :-K] ) else: repr_per_layer = [self._get_last_token(h, attention_mask) for h in selected] text_output = self._get_last_token(all_hidden_states[-1], attention_mask) concat_tokens = torch.cat(repr_per_layer, dim=-1) # [B, H*L] mm_output = self.layer_fusion_mlp(concat_tokens) # [B, H] is_multimodal = self._detect_multimodal(input_ids) output = torch.where(is_multimodal.unsqueeze(-1), mm_output, text_output) if self.normalize: output = torch.nn.functional.normalize(output, p=2, dim=-1) return output def _compute_token_diversity_loss(self, tokens: Tensor) -> Tensor: """Compute diversity regularization on learnable token representations. Args: tokens: [B, K, H] token representations from the final selected layer. Returns: Weighted scalar loss. """ K = tokens.shape[1] if self.token_diversity_type == 'orthogonal': norm_tokens = torch.nn.functional.normalize(tokens, p=2, dim=-1) gram = torch.bmm(norm_tokens, norm_tokens.transpose(1, 2)) identity = torch.eye(K, device=gram.device, dtype=gram.dtype).unsqueeze(0) loss = ((gram - identity) ** 2).mean() elif self.token_diversity_type == 'vicreg': per_token = tokens.permute(1, 0, 2) # [K, B, H] std = torch.sqrt(per_token.var(dim=1) + 1e-4) # [K, H] var_loss = torch.clamp(1.0 - std, min=0).mean() tokens_flat = per_token.reshape(K, -1) # [K, B*H] tokens_flat = tokens_flat - tokens_flat.mean(dim=1, keepdim=True) cov = (tokens_flat @ tokens_flat.T) / (tokens_flat.shape[1] - 1) # [K, K] off_diag = ~torch.eye(K, dtype=torch.bool, device=cov.device) cov_loss = cov[off_diag].pow(2).mean() loss = var_loss + cov_loss else: return torch.tensor(0.0, device=tokens.device) return loss * self.token_diversity_weight @staticmethod def _resolve_num_hidden_layers(config) -> int: """Extract num_hidden_layers from various config formats.""" if hasattr(config, 'text_config') and hasattr(config.text_config, 'num_hidden_layers'): return config.text_config.num_hidden_layers if hasattr(config, 'num_hidden_layers'): return config.num_hidden_layers return 0 @classmethod def build(cls, model_args: ModelArguments, **kwargs): try: config = AutoConfig.from_pretrained(model_args.model_name, trust_remote_code=True, local_files_only=HF_LOCAL_FILES_ONLY) except (ValueError, KeyError): from src.model.vlm_backbone.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniThinkerConfig config = Qwen2_5OmniThinkerConfig.from_pretrained(model_args.model_name, local_files_only=HF_LOCAL_FILES_ONLY) model_backbone = get_backbone_name(hf_config=config) print_master(f'Loading backbone [{model_backbone}] from {model_args.model_name}') # Loading the base model if model_backbone == PHI3V: config._attn_implementation = "eager" config.padding_side = "right" config.use_cache = False base_model = get_model_class(PHI3V).from_pretrained( model_args.model_name, config=config, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, ) elif model_backbone == LLAVA_NEXT: config.use_cache = False config.padding_side = "left" base_model = get_model_class(LLAVA_NEXT).from_pretrained( model_args.model_name, config=config, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, ) elif model_backbone in [QWEN2_VL, QWEN2_5_VL]: config._attn_implementation = "flash_attention_2" config.padding_side = "left" config.use_cache = False base_model = get_model_class(model_backbone).from_pretrained( model_args.model_name, config=config, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, ) elif model_backbone == QWEN2_5_OMNI: # Load thinker model directly instead of full Qwen2.5-Omni model config._attn_implementation = "sdpa" config.padding_side = "left" config.use_cache = False # Handle both config formats: # - qwen2_5_omni: extract thinker_config from nested config # - qwen2_5_omni_thinker: use config directly as thinker_config if config.model_type == "qwen2_5_omni": thinker_config = config.thinker_config elif config.model_type == "qwen2_5_omni_thinker": thinker_config = config else: raise ValueError(f"Unexpected config.model_type: {config.model_type}. Expected 'qwen2_5_omni' or 'qwen2_5_omni_thinker'") thinker_config._attn_implementation = "sdpa" if hasattr(thinker_config, 'vision_config'): thinker_config.vision_config._attn_implementation = "sdpa" config.hidden_size = thinker_config.text_config.hidden_size config.num_hidden_layers = thinker_config.text_config.num_hidden_layers base_model = get_model_class(model_backbone).from_pretrained( model_args.model_name, config=thinker_config, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, local_files_only=HF_LOCAL_FILES_ONLY, ) elif model_backbone in [QWEN2_VL_TOKENSELECTION, QWEN2_5_VL_TOKENSELECTION]: config._attn_implementation = "flash_attention_2" config.padding_side = "left" config.use_cache = False from .utils import parse_layer_type lm_qwen_layer = 28 vis_qwen_layer = 32 lm_skip_layer = parse_layer_type(model_args.lm_skip_layer, lm_qwen_layer) vis_skip_layer = parse_layer_type(model_args.vis_skip_layer, vis_qwen_layer) base_model = get_model_class(model_backbone).from_pretrained( model_args.model_name, config=config, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, lm_skip_layer=lm_skip_layer, vis_skip_layer=vis_skip_layer, ) else: config.use_cache = False base_model = cls.TRANSFORMER_CLS.from_pretrained( model_args.model_name, **kwargs, config=config, attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16, trust_remote_code=True) # load SAM-Audio+SAM3 segmentation_model = None if model_args.sam_audio: print_master(f"Loading SAM-Audio + SAM3 from {model_args.sam_audio_config['checkpoint']}") model_args.sam_audio_config["vlm_hidden_size"] = config.hidden_size segmentation_model = SAMAudioVLM(model_args.sam_audio_config) hidden_size = getattr(config, 'hidden_size', 0) num_hidden_layers = cls._resolve_num_hidden_layers(config) common_kwargs = dict( segmentation_model=segmentation_model, pooling=model_args.pooling, normalize=model_args.normalize, temperature=model_args.temperature, layer_stride=model_args.layer_stride, symmetric_loss=model_args.symmetric_loss, num_hidden_layers=num_hidden_layers, hidden_size=hidden_size, num_learnable_tokens=getattr(model_args, 'num_learnable_tokens', 0), token_diversity_type=getattr(model_args, 'token_diversity_type', 'none'), token_diversity_weight=getattr(model_args, 'token_diversity_weight', 0.0), layer_fusion=getattr(model_args, 'layer_fusion', 'equal'), ) if model_args.lora: print_master(f'Loading lora adapter from {base_model}') lora_config = LoraConfig( r=model_args.lora_r, lora_alpha=model_args.lora_alpha, target_modules=model_args.lora_target_modules.split(','), lora_dropout=model_args.lora_dropout, init_lora_weights="gaussian", use_dora=True, inference_mode=False ) lora_model = get_peft_model(base_model, lora_config) model = cls(encoder=lora_model, **common_kwargs) else: model = cls(encoder=base_model, **common_kwargs) model.model_backbone = model_backbone # Print parameter count model.print_parameter_count() return model @classmethod def load(cls, model_args: ModelArguments, is_trainable=True, **kwargs): model_name_or_path = model_args.checkpoint_path if model_args.checkpoint_path else model_args.model_name if getattr(model_args, "model_backbone", None) == WAVE: from src.model.vlm_backbone.wave.model.qwen2_5_omni.configuration_qwen2_5_omni import \ Qwen2_5OmniThinkerConfig as WaveThinkerConfig config = WaveThinkerConfig.from_pretrained(model_name_or_path, local_files_only=HF_LOCAL_FILES_ONLY) else: try: config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True, local_files_only=HF_LOCAL_FILES_ONLY) except (ValueError, KeyError): from src.model.vlm_backbone.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniThinkerConfig config = Qwen2_5OmniThinkerConfig.from_pretrained(model_name_or_path, local_files_only=HF_LOCAL_FILES_ONLY) if not hasattr(model_args, "model_backbone") or not model_args.model_backbone: model_backbone = get_backbone_name(hf_config=config, model_type=model_args.model_type) setattr(model_args, 'model_backbone', model_backbone) print_master(f'Loading backbone [{model_args.model_backbone}] from {model_name_or_path}') if model_args.model_backbone in {LLAVA_NEXT, QWEN2_VL, QWEN2_5_VL, QWEN2_VL_TOKENSELECTION, QWEN2_5_VL_TOKENSELECTION}: config = AutoConfig.from_pretrained(model_args.model_name, trust_remote_code=True, local_files_only=HF_LOCAL_FILES_ONLY) config._attn_implementation = "flash_attention_2" config.vision_config._attn_implementation = "flash_attention_2" base_model = get_model_class(model_args.model_backbone).from_pretrained( model_args.model_name, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, config=config, local_files_only=HF_LOCAL_FILES_ONLY, ) elif model_args.model_backbone == QWEN2_5_OMNI: # Load thinker model directly instead of full Qwen2.5-Omni model # Reuse already loaded config above (may be ThinkerConfig) # Handle both config formats: # - qwen2_5_omni: extract thinker_config from nested config # - qwen2_5_omni_thinker: use config directly as thinker_config if config.model_type == "qwen2_5_omni": thinker_config = config.thinker_config elif config.model_type == "qwen2_5_omni_thinker": thinker_config = config else: raise ValueError(f"Unexpected config.model_type: {config.model_type}. Expected 'qwen2_5_omni' or 'qwen2_5_omni_thinker'") thinker_config._attn_implementation = "sdpa" if hasattr(thinker_config, 'vision_config'): thinker_config.vision_config._attn_implementation = "sdpa" config.hidden_size = thinker_config.text_config.hidden_size config.num_hidden_layers = thinker_config.text_config.num_hidden_layers base_model = get_model_class(model_args.model_backbone).from_pretrained( model_args.model_name, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, config=thinker_config, local_files_only=HF_LOCAL_FILES_ONLY, ) elif model_args.model_backbone == WAVE: # WAVE baseline: uses WAVE's own model code (vlm_backbone/wave/) # completely independent from our custom qwen2_5_omni implementation from src.model.vlm_backbone.wave.model.qwen2_5_omni.configuration_qwen2_5_omni import \ Qwen2_5OmniThinkerConfig as WaveThinkerConfig from src.model.vlm_backbone.wave.model.qwen2_5_omni.modeling_qwen2_5_omni import \ Qwen2_5OmniThinkerForConditionalGeneration as WaveThinker wave_config = WaveThinkerConfig.from_pretrained(model_name_or_path, local_files_only=HF_LOCAL_FILES_ONLY) wave_config._attn_implementation = "flash_attention_2" if hasattr(wave_config, 'vision_config'): wave_config.vision_config._attn_implementation = "flash_attention_2" base_model = WaveThinker.from_pretrained( model_args.model_name, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, config=wave_config, local_files_only=HF_LOCAL_FILES_ONLY, ) elif model_args.model_backbone == PHI3V: config = AutoConfig.from_pretrained(model_args.model_name, trust_remote_code=True) config.use_cache = False config.padding_side = "right" base_model = get_model_class(PHI3V).from_pretrained(model_args.model_name, **kwargs, config=config, torch_dtype=torch.bfloat16, trust_remote_code=True) base_model.padding_side = "right" elif model_args.model_backbone == INTERNVIDEO2: print_master(f'Loading backbone [{model_args.model_backbone}] from {"src/model/vlm_backbone/internvideo2/"}') config = AutoConfig.from_pretrained("src/model/vlm_backbone/internvideo2/", trust_remote_code=True) base_model = get_model_class(model_args.model_backbone).from_pretrained("src/model/vlm_backbone/internvideo2/", config=config, trust_remote_code=True) elif model_args.model_backbone == GME: from src.model.vlm_backbone.gme.gme_inference import GmeQwen2VL base_model = GmeQwen2VL(model_args.model_name, processor=kwargs['processor']) setattr(base_model, 'config', config) elif model_args.model_backbone == LamRA: from src.model.vlm_backbone.lamra.lamra_inference import LamRAQwen2VL base_model = LamRAQwen2VL(model_args.model_name) setattr(base_model, 'config', config) elif model_args.model_backbone == LamRA_QWEN2_5: from src.model.vlm_backbone.lamra.lamra_qwen25_inference import LamRAQwen25VL base_model = LamRAQwen25VL(model_args.model_name) setattr(base_model, 'config', config) elif model_args.model_backbone == COLPALI: from src.model.vlm_backbone.colpali import ColPali base_model = ColPali.from_pretrained(model_args.model_name) setattr(base_model, 'config', config) else: # Loading external base model from HF config = AutoConfig.from_pretrained(model_args.model_name, trust_remote_code=True) config.use_cache = False base_model = cls.TRANSFORMER_CLS.from_pretrained( model_name_or_path, **kwargs, config=config, torch_dtype=torch.bfloat16, trust_remote_code=True) # load SAM-Audio+SAM3 segmentation_model = None if model_args.sam_audio: print_master(f"Loading SAM-Audio + SAM3 from {model_args.sam_audio_config['checkpoint']}") model_args.sam_audio_config["vlm_hidden_size"] = config.hidden_size sam_audio_model = SAMAudioVLM(model_args.sam_audio_config) orig_sam_audio_params = { name: param.clone().detach() for name, param in sam_audio_model.named_parameters() if param.requires_grad } checkpoint = torch.load(os.path.join(model_name_or_path, 'sam_audio_reducer_and_projection.pth'), map_location='cpu') sam_audio_model.sam3_reducer.load_state_dict(checkpoint['sam3_reducer']) sam_audio_model.sam3vlm_layer.load_state_dict(checkpoint['sam3vlm_layer']) sam_audio_model.audio_reducer.load_state_dict(checkpoint['audio_reducer']) sam_audio_model.audio_projection.load_state_dict(checkpoint['audio_projection']) for name, param in sam_audio_model.named_parameters(): if not param.requires_grad: continue old = orig_sam_audio_params[name] new = param.detach() delta = (new - old).abs().mean().item() print(f"{name:40s} Δ mean abs = {delta:.6f}") print() segmentation_model = sam_audio_model # Building the model on top of the base hidden_size = getattr(config, 'hidden_size', 0) num_hidden_layers = cls._resolve_num_hidden_layers(config) common_kwargs = dict( segmentation_model=segmentation_model, pooling=model_args.pooling, normalize=model_args.normalize, temperature=model_args.temperature, layer_stride=getattr(model_args, 'layer_stride', 1), symmetric_loss=getattr(model_args, 'symmetric_loss', False), num_hidden_layers=num_hidden_layers, hidden_size=hidden_size, num_learnable_tokens=getattr(model_args, 'num_learnable_tokens', 0), token_diversity_type=getattr(model_args, 'token_diversity_type', 'none'), token_diversity_weight=getattr(model_args, 'token_diversity_weight', 0.0), layer_fusion=getattr(model_args, 'layer_fusion', 'equal'), ) if model_args.lora: print_master(f'Loading LoRA from {model_name_or_path}') lora_config = LoraConfig.from_pretrained(model_name_or_path, local_files_only=HF_LOCAL_FILES_ONLY) lora_model = PeftModel.from_pretrained( base_model, model_name_or_path, config=lora_config, is_trainable=is_trainable, local_files_only=HF_LOCAL_FILES_ONLY, ) lora_model.load_adapter(model_name_or_path, lora_model.active_adapter, is_trainable=is_trainable) if not is_trainable: lora_model = lora_model.merge_and_unload() model = cls(encoder=lora_model, **common_kwargs) else: model = cls(encoder=base_model, **common_kwargs) # Load learnable_tokens (standalone, works with any pooling) learnable_tokens_path = os.path.join(model_name_or_path, 'learnable_tokens.pth') if hasattr(model, 'learnable_tokens') and os.path.isfile(learnable_tokens_path): print_master(f"Loading learnable_tokens from {learnable_tokens_path}") with torch.no_grad(): model.learnable_tokens.copy_(torch.load(learnable_tokens_path, map_location='cpu')) print_master("Loaded learnable_tokens successfully.") elif hasattr(model, 'learnable_tokens'): # Backward compat: older checkpoints stored learnable_tokens inside layer_weights.pth lw_path = os.path.join(model_name_or_path, 'layer_weights.pth') if os.path.isfile(lw_path): ckpt = torch.load(lw_path, map_location='cpu') if 'learnable_tokens' in ckpt: with torch.no_grad(): model.learnable_tokens.copy_(ckpt['learnable_tokens']) print_master("Loaded learnable_tokens from layer_weights.pth (legacy).") # Load layer_weights layer_weights_path = os.path.join(model_name_or_path, 'layer_weights.pth') if hasattr(model, 'layer_weights') and os.path.isfile(layer_weights_path): print_master(f"Loading layer_weights from {layer_weights_path}") ckpt = torch.load(layer_weights_path, map_location='cpu') with torch.no_grad(): model.layer_weights.copy_(ckpt['layer_weights']) print_master("Loaded layer_weights successfully.") layer_token_weights_path = os.path.join(model_name_or_path, 'layer_token_weights.pth') if hasattr(model, 'layer_token_weights') and os.path.isfile(layer_token_weights_path): print_master(f"Loading layer_token_weights from {layer_token_weights_path}") ckpt = torch.load(layer_token_weights_path, map_location='cpu') with torch.no_grad(): model.layer_token_weights.copy_(ckpt['layer_token_weights']) print_master("Loaded layer_token_weights successfully.") layer_fusion_mlp_path = os.path.join(model_name_or_path, 'layer_fusion_mlp.pth') if hasattr(model, 'layer_fusion_mlp') and os.path.isfile(layer_fusion_mlp_path): print_master(f"Loading layer_fusion_mlp from {layer_fusion_mlp_path}") model.layer_fusion_mlp.load_state_dict(torch.load(layer_fusion_mlp_path, map_location='cpu')) print_master("Loaded layer_fusion_mlp successfully.") if not is_trainable: for param in model.parameters(): param.requires_grad = False model.model_backbone = model_args.model_backbone # Print parameter count model.print_parameter_count() return model def save(self, output_dir: str): self.encoder.save_pretrained(output_dir) # Save SAMAudioVLM trainable weights if self.segmentation_model is not None and isinstance(self.segmentation_model, SAMAudioVLM): torch.save({ 'sam3_reducer': self.segmentation_model.sam3_reducer.state_dict(), 'sam3vlm_layer': self.segmentation_model.sam3vlm_layer.state_dict(), 'audio_reducer': self.segmentation_model.audio_reducer.state_dict(), 'audio_projection': self.segmentation_model.audio_projection.state_dict(), }, os.path.join(output_dir, 'sam_audio_reducer_and_projection.pth')) if hasattr(self, 'learnable_tokens'): torch.save(self.learnable_tokens.data, os.path.join(output_dir, 'learnable_tokens.pth')) if hasattr(self, 'layer_weights'): torch.save({'layer_weights': self.layer_weights.data}, os.path.join(output_dir, 'layer_weights.pth')) if hasattr(self, 'layer_token_weights'): torch.save({'layer_token_weights': self.layer_token_weights.data}, os.path.join(output_dir, 'layer_token_weights.pth')) if hasattr(self, 'layer_fusion_mlp'): torch.save(self.layer_fusion_mlp.state_dict(), os.path.join(output_dir, 'layer_fusion_mlp.pth')) def forward(self, qry: Dict[str, Tensor] = None, tgt: Dict[str, Tensor] = None, *args, **kwargs): self._token_diversity_loss = None qry_reps = self.encode_input(qry) if qry else None # (bsz_per_device, dim) tgt_reps = self.encode_input(tgt) if tgt else None # (bsz_per_device, dim) if qry_reps is None or tgt_reps is None: result = {"qry_reps": qry_reps, "tgt_reps": tgt_reps} if self._token_diversity_loss is not None: result["aux_loss"] = self._token_diversity_loss return result if self.is_ddp: all_qry_reps = self._dist_gather_tensor(qry_reps) all_tgt_reps = self._dist_gather_tensor(tgt_reps) else: all_qry_reps = qry_reps all_tgt_reps = tgt_reps q2t_scores = self.compute_similarity(all_qry_reps, all_tgt_reps) q2t_scores = q2t_scores.view(all_qry_reps.size(0), -1) target = torch.arange(q2t_scores.size(0), device=q2t_scores.device, dtype=torch.long) target = target * (all_qry_reps.size(0) // all_tgt_reps.size(0)) loss_q2t = self.cross_entropy(q2t_scores / self.temperature, target) if self.symmetric_loss: t2q_scores = self.compute_similarity(all_tgt_reps, all_qry_reps) t2q_scores = t2q_scores.view(all_tgt_reps.size(0), -1) target_t2q = torch.arange(t2q_scores.size(0), device=t2q_scores.device, dtype=torch.long) loss_t2q = self.cross_entropy(t2q_scores / self.temperature, target_t2q) loss = 0.5 * (loss_q2t + loss_t2q) else: loss = loss_q2t if self._token_diversity_loss is not None: loss = loss + self._token_diversity_loss if self.is_ddp: loss = loss * self.world_size return loss def _dist_gather_tensor(self, t: Tensor): t = t.contiguous() all_tensors = [torch.empty_like(t) for _ in range(self.world_size)] dist.all_gather(all_tensors, t) all_tensors[self.process_rank] = t all_tensors = torch.cat(all_tensors, dim=0) return all_tensors def compute_similarity(self, q_reps, p_reps): return torch.matmul(q_reps, p_reps.transpose(0, 1))