first upload
5fe8ec0 4 days ago
73.8 kB
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 (<null>) 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))