Think Before You Segment: How TGS-Agent Teaches AI to Reason About Sound Before Picking Up a Brush

Complete TGS-Agent Model Code (PyTorch)

The implementation below covers every component described in the TGS-Agent paper: the audio and visual Q-Former encoders with MLP projectors, the LoRA-enhanced LLaMA-2 backbone for Ref-Thinker, the think-answer output parsing, the Grounding-DINO integration for bounding box generation, the frozen SAM2 segmentation pipeline, the full TGSAgent end-to-end inference loop, loss functions for autoregressive cross-entropy training, and a smoke-test that runs on synthetic data without any pretrained weights.

# ==============================================================================
# TGS-Agent: Think-Ground-Segment for Referring Audio-Visual Segmentation
# Paper: "Think Before You Segment", AAAI 2026
# Authors: Zhou, Zhou, Han, Wang, Chang, Cholakkal, Anwer — MBZUAI / NUS / USTC
# ==============================================================================
# Sections:
#   1.  Imports & Configuration
#   2.  Q-Former: Compressing Audio/Visual Features into Tokens
#   3.  MLP Projector: Aligning Modality Tokens with LLM Space
#   4.  LoRA Adapter: Parameter-Efficient LLM Fine-Tuning
#   5.  RefThinker: Full Multimodal LLM with Think-Answer Reasoning
#   6.  Output Parser: Extracting F-Object and S-Object Descriptions
#   7.  GroundingDINOWrapper: Open-Set Bounding Box Generation
#   8.  SAM2Wrapper: Frozen Mask Generation from Bounding Boxes
#   9.  TGSAgent: Complete Think-Ground-Segment Pipeline
#  10.  Loss Functions: Autoregressive Cross-Entropy
#  11.  Instruction Tuning Dataset Builder
#  12.  Training Loop
#  13.  Inference & Evaluation Utilities
#  14.  R2-AVSBench Reference Transformer
#  15.  Smoke Test
# ==============================================================================

from __future__ import annotations

import re
import math
import warnings
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.utils.data import DataLoader, Dataset

warnings.filterwarnings("ignore")


# ─── SECTION 1: Configuration ──────────────────────────────────────────────────

@dataclass
class RefThinkerConfig:
    """
    Configuration for the Ref-Thinker audio-visual LLM.

    Attributes
    ----------
    audio_feat_dim   : output dimension of BEATs audio encoder (768)
    visual_feat_dim  : output dimension of CLIP-ViT-L/14 visual encoder (1024)
    llm_hidden_dim   : LLaMA-2-7b hidden dimension (4096)
    num_query_tokens : number of learnable Q-Former query tokens (32 per modality)
    qformer_layers   : number of cross-attention layers in Q-Former (6)
    qformer_heads    : number of attention heads in Q-Former (8)
    lora_rank        : LoRA decomposition rank (8, paper default)
    lora_alpha       : LoRA scaling factor (16, paper default)
    lora_dropout     : dropout inside LoRA adapter (0.1)
    max_output_len   : max tokens to generate during inference
    vocab_size       : LLM vocabulary size (32000 for LLaMA-2)
    """
    audio_feat_dim: int   = 768
    visual_feat_dim: int  = 1024
    llm_hidden_dim: int   = 4096
    num_query_tokens: int = 32
    qformer_layers: int   = 6
    qformer_heads: int    = 8
    lora_rank: int        = 8
    lora_alpha: int       = 16
    lora_dropout: float   = 0.1
    max_output_len: int   = 512
    vocab_size: int       = 32000


@dataclass
class GroundingConfig:
    """
    Configuration for the Grounding-DINO detection stage.

    Attributes
    ----------
    tau_bbox   : minimum confidence score to accept a bounding box (0.1)
    tau_text   : minimum text similarity score to accept a match (0.25)
    box_threshold : alias for tau_bbox (0.1)
    text_threshold: alias for tau_text (0.25)
    max_boxes  : maximum boxes to return per frame
    """
    tau_bbox: float        = 0.10
    tau_text: float        = 0.25
    box_threshold: float   = 0.10
    text_threshold: float  = 0.25
    max_boxes: int         = 5


# ─── SECTION 2: Q-Former — Compressing Modality Features into Tokens ──────────

class CrossAttentionLayer(nn.Module):
    """
    A single cross-attention layer used inside the Q-Former.

    Learnable query tokens attend to the modality feature sequence
    (audio or visual), compressing variable-length input into a
    fixed-length set of query embeddings.

    Q-Former design follows BLIP-2 (Li et al., ICML 2023).
    """

    def __init__(self, query_dim: int, feat_dim: int, num_heads: int,
                 dropout: float = 0.1):
        super().__init__()
        self.norm_q   = nn.LayerNorm(query_dim)
        self.norm_kv  = nn.LayerNorm(feat_dim)
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=query_dim, num_heads=num_heads,
            kdim=feat_dim, vdim=feat_dim,
            dropout=dropout, batch_first=True
        )
        self.norm_ff = nn.LayerNorm(query_dim)
        self.ff = nn.Sequential(
            nn.Linear(query_dim, query_dim * 4),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(query_dim * 4, query_dim),
            nn.Dropout(dropout),
        )

    def forward(self, queries: Tensor, features: Tensor) -> Tensor:
        """
        Parameters
        ----------
        queries  : (B, Nq, D_q) learnable query tokens
        features : (B, Ns, D_f) encoder output (audio or visual features)

        Returns
        -------
        queries  : (B, Nq, D_q) updated query tokens
        """
        q = self.norm_q(queries)
        kv = self.norm_kv(features)
        attn_out, _ = self.cross_attn(query=q, key=kv, value=kv)
        queries = queries + attn_out
        queries = queries + self.ff(self.norm_ff(queries))
        return queries


class QFormer(nn.Module):
    """
    Q-Former: compresses audio or visual encoder outputs into a fixed
    number of query token embeddings for consumption by the LLM.

    The Q-Former bridges the modality encoder and the language model
    by cross-attending to the full encoder output and distilling it
    into `num_query_tokens` dense embeddings.

    Paper: "Q-Formers along with their corresponding projectors are
    independently pretrained on domain-specific caption datasets."

    Architecture:
        - Learnable query tokens (initialized from N(0, 0.02))
        - `n_layers` cross-attention + FFN blocks
        - Output: (B, num_query_tokens, query_dim)
    """

    def __init__(self, feat_dim: int, query_dim: int,
                 num_query_tokens: int, n_layers: int, n_heads: int):
        super().__init__()
        self.query_tokens = nn.Parameter(
            torch.zeros(1, num_query_tokens, query_dim)
        )
        nn.init.trunc_normal_(self.query_tokens, std=0.02)
        self.layers = nn.ModuleList([
            CrossAttentionLayer(query_dim, feat_dim, n_heads)
            for _ in range(n_layers)
        ])

    def forward(self, features: Tensor) -> Tensor:
        """
        features : (B, T, feat_dim) — audio or visual encoder output
        Returns  : (B, num_query_tokens, query_dim)
        """
        B = features.shape[0]
        queries = self.query_tokens.expand(B, -1, -1)
        for layer in self.layers:
            queries = layer(queries, features)
        return queries


# ─── SECTION 3: MLP Projector — Aligning Tokens with LLM Space ────────────────

class MLPProjector(nn.Module):
    """
    Two-layer MLP projector that maps Q-Former output tokens into
    the LLM's token embedding space.

    Paper: "A projector module, implemented as two MLP layers,
    is further applied to align the audio/visual feature with
    the textual features processed by the LLM."
    """

    def __init__(self, in_dim: int, out_dim: int, dropout: float = 0.1):
        super().__init__()
        hidden = (in_dim + out_dim) // 2
        self.net = nn.Sequential(
            nn.LayerNorm(in_dim),
            nn.Linear(in_dim, hidden),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden, out_dim),
        )

    def forward(self, x: Tensor) -> Tensor:
        """x: (B, Nq, in_dim)  →  (B, Nq, out_dim)"""
        return self.net(x)


# ─── SECTION 4: LoRA Adapter — Parameter-Efficient LLM Fine-Tuning ────────────

class LoRALinear(nn.Module):
    """
    LoRA (Low-Rank Adaptation) wrapper around a frozen nn.Linear layer.

    Implements: W_adapted = W_frozen + (alpha/r) * B @ A
    where A ∈ R^{r×in}, B ∈ R^{out×r}, r = rank.

    Paper: "We apply the LoRA technique for parameter-efficient tuning
    of the LLM. LoRA is applied with a rank 8 and a scaling factor 16."

    Only A and B are trainable. The original weight W_frozen is unchanged.
    """

    def __init__(self, linear: nn.Linear, rank: int = 8,
                 alpha: int = 16, dropout: float = 0.1):
        super().__init__()
        in_features, out_features = linear.in_features, linear.out_features
        self.linear = linear
        for p in self.linear.parameters():
            p.requires_grad = False

        self.lora_A = nn.Parameter(torch.empty(rank, in_features))
        self.lora_B = nn.Parameter(torch.zeros(out_features, rank))
        self.scale   = alpha / rank
        self.dropout = nn.Dropout(dropout)

        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))

    def forward(self, x: Tensor) -> Tensor:
        base_out = self.linear(x)
        lora_out = self.dropout(x) @ self.lora_A.T @ self.lora_B.T
        return base_out + self.scale * lora_out


def apply_lora_to_attention(model: nn.Module, rank: int, alpha: int,
                             dropout: float) -> nn.Module:
    """
    Replace all q_proj and v_proj Linear layers in transformer
    attention blocks with LoRALinear wrappers.

    This is the standard practice: LoRA is applied to Q and V projections
    of each self-attention layer. K and output projections are left frozen.
    """
    for name, module in model.named_modules():
        if isinstance(module, nn.MultiheadAttention):
            if hasattr(module, 'in_proj_weight') and module.in_proj_weight is not None:
                # Wrap the combined QKV projection as a LoRA linear
                original = nn.Linear(
                    module.embed_dim, 3 * module.embed_dim, bias=False
                )
                original.weight = nn.Parameter(module.in_proj_weight)
                lora_layer = LoRALinear(original, rank, alpha, dropout)
                module.in_proj_weight = None
                module._lora_qkv = lora_layer
    return model


# ─── SECTION 5: RefThinker — Full Multimodal LLM ──────────────────────────────

class RefThinker(nn.Module):
    """
    Ref-Thinker: a reasoning-enhanced Multimodal Large Language Model.

    Given a reference expression, audio stream features, and video frame
    features, Ref-Thinker generates an explicit think-answer chain:

        
          The referential expression is "xxx".
          The video shows xxx (video analysis).
          The audio contains xxx (audio analysis).
          The reference relates to xxx (modality analysis).
        
        
           fine-grained object description 
           category name 
        

    Architecture:
        - Audio encoder  → QFormer (32 tokens) → MLP projector → LLM space
        - Visual encoder → QFormer (32 tokens) → MLP projector → LLM space
        - Text tokenizer → token embeddings → LLM space
        - Concatenate [audio_tokens | visual_tokens | text_tokens] as input
        - LLaMA-2-7b backbone with LoRA applied to attention QKV projections
        - Autoregressive generation using cross-entropy loss

    Training:
        Phase 1 (pretraining): Freeze LLM, train Q-Formers + projectors
                               using audio/visual captioning data.
        Phase 2 (instruction tuning): Apply LoRA, train on think-answer
                                       chains generated by Gemini-1.5-Pro.
    """

    def __init__(self, cfg: RefThinkerConfig):
        super().__init__()
        self.cfg = cfg
        D = cfg.llm_hidden_dim

        # ── Audio branch ──────────────────────────────────────────────────────
        self.audio_qformer = QFormer(
            feat_dim=cfg.audio_feat_dim,
            query_dim=D // 2,
            num_query_tokens=cfg.num_query_tokens,
            n_layers=cfg.qformer_layers,
            n_heads=cfg.qformer_heads,
        )
        self.audio_projector = MLPProjector(D // 2, D)

        # ── Visual branch ─────────────────────────────────────────────────────
        self.visual_qformer = QFormer(
            feat_dim=cfg.visual_feat_dim,
            query_dim=D // 2,
            num_query_tokens=cfg.num_query_tokens,
            n_layers=cfg.qformer_layers,
            n_heads=cfg.qformer_heads,
        )
        self.visual_projector = MLPProjector(D // 2, D)

        # ── LLM backbone (simplified Transformer decoder — replace with
        #    LlamaForCausalLM when using real pretrained weights) ─────────────
        self.token_embedding = nn.Embedding(cfg.vocab_size, D)
        self.pos_embedding    = nn.Embedding(2048, D)

        encoder_layer = nn.TransformerDecoderLayer(
            d_model=D, nhead=16, dim_feedforward=D * 4,
            dropout=0.1, batch_first=True, norm_first=True,
        )
        self.llm_backbone = nn.TransformerDecoder(encoder_layer, num_layers=4)
        self.lm_head = nn.Linear(D, cfg.vocab_size, bias=False)
        self.lm_head.weight = self.token_embedding.weight  # tie weights

        # ── Apply LoRA to LLM attention projections ───────────────────────────
        apply_lora_to_attention(
            self.llm_backbone, cfg.lora_rank, cfg.lora_alpha, cfg.lora_dropout
        )

        # Freeze LLM backbone weights (LoRA parameters remain trainable)
        self._freeze_llm_base()

    def _freeze_llm_base(self):
        """Freeze all LLM parameters except LoRA adapters and lm_head."""
        for name, param in self.llm_backbone.named_parameters():
            if 'lora_A' not in name and 'lora_B' not in name:
                param.requires_grad = False

    def encode_audio(self, audio_feats: Tensor) -> Tensor:
        """
        audio_feats : (B, T_a, audio_feat_dim) — BEATs segment features
        Returns     : (B, num_query_tokens, llm_hidden_dim)
        """
        qformer_out = self.audio_qformer(audio_feats)
        return self.audio_projector(qformer_out)

    def encode_visual(self, visual_feats: Tensor) -> Tensor:
        """
        visual_feats : (B, T_v, visual_feat_dim) — CLIP-ViT-L/14 frame features
        Returns      : (B, num_query_tokens, llm_hidden_dim)
        """
        qformer_out = self.visual_qformer(visual_feats)
        return self.visual_projector(qformer_out)

    def forward(
        self,
        audio_feats: Tensor,
        visual_feats: Tensor,
        input_ids: Tensor,
        attention_mask: Optional[Tensor] = None,
        labels: Optional[Tensor] = None,
    ) -> Dict:
        """
        Full RefThinker forward pass for training.

        The multimodal input sequence is constructed as:
            [audio_tokens (32)] [visual_tokens (32)] [text_tokens (L)]
        This concatenated sequence is fed to the LLM decoder.
        Cross-entropy loss is computed only on the text generation portion
        (labels for the multimodal prefix are set to -100 / ignored).

        Parameters
        ----------
        audio_feats   : (B, T_a, audio_feat_dim)
        visual_feats  : (B, T_v, visual_feat_dim)
        input_ids     : (B, L) tokenized input prompt + expected output
        attention_mask: (B, L) text attention mask
        labels        : (B, L) target token IDs; -100 where loss is masked

        Returns
        -------
        dict with 'loss', 'logits', 'audio_tokens', 'visual_tokens'
        """
        B, L = input_ids.shape
        Nq = self.cfg.num_query_tokens
        D  = self.cfg.llm_hidden_dim

        # Encode modalities → LLM-space tokens
        audio_tokens  = self.encode_audio(audio_feats)   # (B, Nq, D)
        visual_tokens = self.encode_visual(visual_feats)  # (B, Nq, D)

        # Embed text tokens
        pos_ids    = torch.arange(L, device=input_ids.device).unsqueeze(0)
        text_embs  = self.token_embedding(input_ids) + self.pos_embedding(pos_ids)

        # Concatenate: audio | visual | text
        full_embs  = torch.cat([audio_tokens, visual_tokens, text_embs], dim=1)
        total_len  = 2 * Nq + L

        # Build attention mask for full sequence
        prefix_mask = torch.ones(B, 2 * Nq, device=input_ids.device)
        if attention_mask is not None:
            full_mask = torch.cat([prefix_mask, attention_mask.float()], dim=1)
        else:
            full_mask = torch.ones(B, total_len, device=input_ids.device)

        # Build causal mask
        causal_mask = torch.triu(
            torch.ones(total_len, total_len, device=input_ids.device), diagonal=1
        ).bool()

        # LLM forward (decoder mode: memory = full sequence itself)
        decoder_out = self.llm_backbone(
            tgt=full_embs,
            memory=full_embs,
            tgt_mask=causal_mask,
            tgt_key_padding_mask=(full_mask == 0),
        )

        logits = self.lm_head(decoder_out)  # (B, total_len, vocab_size)

        loss = None
        if labels is not None:
            # Pad labels for the multimodal prefix (no loss on those positions)
            prefix_labels = torch.full(
                (B, 2 * Nq), fill_value=-100, device=labels.device
            )
            full_labels = torch.cat([prefix_labels, labels], dim=1)

            # Shift logits and labels for next-token prediction
            shift_logits = logits[:, :-1].contiguous()
            shift_labels = full_labels[:, 1:].contiguous()

            loss = F.cross_entropy(
                shift_logits.view(-1, self.cfg.vocab_size),
                shift_labels.view(-1),
                ignore_index=-100,
            )

        return {
            "loss": loss,
            "logits": logits,
            "audio_tokens": audio_tokens,
            "visual_tokens": visual_tokens,
        }

    @torch.no_grad()
    def generate(
        self,
        audio_feats: Tensor,
        visual_feats: Tensor,
        prompt_ids: Tensor,
        max_new_tokens: int = 512,
        eos_token_id: int = 2,
    ) -> Tensor:
        """
        Greedy autoregressive generation for Ref-Thinker inference.

        Produces the think-answer chain token-by-token, stopping at
        EOS or when max_new_tokens is reached.

        Parameters
        ----------
        audio_feats   : (1, T_a, audio_feat_dim) — single video
        visual_feats  : (1, T_v, visual_feat_dim)
        prompt_ids    : (1, L_p) tokenized user prompt
        max_new_tokens: maximum tokens to generate
        eos_token_id  : token ID for end-of-sequence

        Returns
        -------
        generated_ids : (1, L_p + num_new_tokens) full sequence
        """
        B = audio_feats.shape[0]
        device = audio_feats.device

        audio_tokens  = self.encode_audio(audio_feats)
        visual_tokens = self.encode_visual(visual_feats)

        generated = prompt_ids.clone()

        for _ in range(max_new_tokens):
            L = generated.shape[1]
            Nq = self.cfg.num_query_tokens

            pos_ids   = torch.arange(L, device=device).unsqueeze(0)
            text_embs = self.token_embedding(generated) + self.pos_embedding(pos_ids)
            full_embs = torch.cat([audio_tokens, visual_tokens, text_embs], dim=1)
            total_len = 2 * Nq + L

            causal_mask = torch.triu(
                torch.ones(total_len, total_len, device=device), diagonal=1
            ).bool()

            out = self.llm_backbone(
                tgt=full_embs, memory=full_embs, tgt_mask=causal_mask
            )
            next_token_logits = self.lm_head(out[:, -1, :])
            next_token = next_token_logits.argmax(dim=-1, keepdim=True)
            generated = torch.cat([generated, next_token], dim=1)

            if (next_token == eos_token_id).all():
                break

        return generated


# ─── SECTION 6: Output Parser — Extracting Object Descriptions ─────────────────

class RefThinkerOutputParser:
    """
    Parses the structured think-answer output of Ref-Thinker into
    the fine-grained (f_object) and simplified (s_object) descriptions
    used as prompts for Grounding-DINO in the Ground step.

    Expected output format:
         ... reasoning chain ... 
        
           fine-grained description 
           category name 
        
    """

    THINK_PATTERN    = re.compile(r'(.*?)', re.DOTALL)
    F_OBJECT_PATTERN = re.compile(r'(.*?)', re.DOTALL)
    S_OBJECT_PATTERN = re.compile(r'(.*?)', re.DOTALL)

    @classmethod
    def parse(cls, text: str) -> Dict[str, str]:
        """
        Parse raw model output text into structured components.

        Parameters
        ----------
        text : raw generated string from RefThinker.generate()

        Returns
        -------
        dict with keys: 'think', 'f_object', 's_object', 'raw'
        Values are empty strings if the pattern is not found.
        """
        think_match    = cls.THINK_PATTERN.search(text)
        f_object_match = cls.F_OBJECT_PATTERN.search(text)
        s_object_match = cls.S_OBJECT_PATTERN.search(text)

        return {
            "think":    think_match.group(1).strip() if think_match else "",
            "f_object": f_object_match.group(1).strip() if f_object_match else "",
            "s_object": s_object_match.group(1).strip() if s_object_match else "",
            "raw":      text,
        }

    @classmethod
    def get_detection_prompt(cls, parsed: Dict[str, str],
                               use_fine_grained: bool = False) -> str:
        """
        Return the appropriate object description string to pass to
        Grounding-DINO.

        Paper ablation finding: simplified (s_object) category description
        generally works better because Grounding-DINO prefers short,
        specific prompts. Fine-grained (f_object) is more useful when
        multiple instances of the same category appear in the scene.

        Parameters
        ----------
        parsed           : output of parse()
        use_fine_grained : if True, return f_object; else return s_object
        """
        if use_fine_grained:
            return parsed.get("f_object", "") or parsed.get("s_object", "")
        return parsed.get("s_object", "") or parsed.get("f_object", "")


# ─── SECTION 7: GroundingDINOWrapper — Open-Set Bounding Box Generation ────────

class BoundingBox:
    """Simple container for a bounding box prediction."""
    def __init__(self, x1: float, y1: float, x2: float, y2: float, score: float):
        self.x1 = x1; self.y1 = y1; self.x2 = x2; self.y2 = y2
        self.score = score

    def to_tuple(self) -> Tuple[float, float, float, float]:
        return self.x1, self.y1, self.x2, self.y2

    def __repr__(self):
        return f"BBox([{self.x1:.2f},{self.y1:.2f},{self.x2:.2f},{self.y2:.2f}] s={self.score:.3f})"


class GroundingDINOWrapper:
    """
    Wrapper for Grounding-DINO open-set object detection.

    Grounding-DINO combines the DINO detector with large-scale grounded
    pretraining, enabling detection from arbitrary text descriptions
    (category names, referring expressions).

    In TGS-Agent, this takes the T_f or T_s description from Ref-Thinker
    and produces bounding boxes in each video frame.

    Real usage:
        from groundingdino.util.inference import load_model, predict
        self.model = load_model(config_path, checkpoint_path)

    This wrapper provides a mock implementation for demonstration,
    which returns synthetic boxes to enable end-to-end smoke testing
    without downloading model weights.

    Paper: "We employ the Swin-T-based Grounding-DINO for object detection,
    with frozen parameters."
    """

    def __init__(self, cfg: GroundingConfig, use_mock: bool = True):
        self.cfg = cfg
        self.use_mock = use_mock
        if not use_mock:
            try:
                from groundingdino.util.inference import load_model
                # Real model: uncomment and provide paths
                # self.model = load_model("path/to/config", "path/to/weights")
                pass
            except ImportError:
                print("[GroundingDINO] groundingdino package not installed. "
                      "Install with: pip install groundingdino-py")
                self.use_mock = True

    def detect(self, image_np: np.ndarray, text_prompt: str) -> List[BoundingBox]:
        """
        Detect objects matching text_prompt in a single image frame.

        Parameters
        ----------
        image_np    : (H, W, 3) uint8 numpy array — one video frame
        text_prompt : object description from Ref-Thinker (T_f or T_s)

        Returns
        -------
        boxes : list of BoundingBox objects above the confidence thresholds
        """
        if self.use_mock:
            return self._mock_detect(image_np)

        # Real implementation (requires groundingdino package):
        # from groundingdino.util.inference import predict
        # import torchvision.transforms as T
        # transform = T.Compose([T.ToTensor(), T.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])])
        # image_tensor = transform(image_np)
        # boxes, logits, phrases = predict(
        #     model=self.model, image=image_tensor, caption=text_prompt,
        #     box_threshold=self.cfg.tau_bbox, text_threshold=self.cfg.tau_text
        # )
        # H, W = image_np.shape[:2]
        # return [BoundingBox(b[0]*W, b[1]*H, b[2]*W, b[3]*H, l.item())
        #         for b, l in zip(boxes.tolist(), logits.tolist())]
        return self._mock_detect(image_np)

    def _mock_detect(self, image_np: np.ndarray) -> List[BoundingBox]:
        """Synthetic bounding box for smoke testing."""
        H, W = image_np.shape[:2] if image_np.ndim == 3 else (256, 256)
        return [BoundingBox(
            x1=0.2 * W, y1=0.2 * H,
            x2=0.7 * W, y2=0.8 * H,
            score=0.85
        )]

    def detect_video(
        self, frames: List[np.ndarray], text_prompt: str
    ) -> List[Optional[BoundingBox]]:
        """
        Run detection on all N video frames for a single text prompt.

        Returns a list of N BoundingBox objects (or None if no object
        was detected in that frame above the threshold). A None box
        causes the Segment step to default to an all-background mask.
        """
        results = []
        for frame in frames:
            boxes = self.detect(frame, text_prompt)
            if boxes:
                best = max(boxes, key=lambda b: b.score)
                results.append(best if best.score >= self.cfg.tau_bbox else None)
            else:
                results.append(None)
        return results


# ─── SECTION 8: SAM2Wrapper — Frozen Mask Generation from Bounding Boxes ───────

class SAM2Wrapper:
    """
    Wrapper for the frozen SAM2 model used in TGS-Agent's Segment step.

    SAM2 receives bounding box prompts (one per frame) and produces
    binary segmentation masks. No SAM2 weights are modified during
    TGS-Agent training.

    Paper: "Notably, the prior SOTA method also employs SAM2 but requires
    fine-tuning its mask decoder. In contrast, our method leverages the
    frozen SAM2 and achieves superior performance."

    We use Hiera-Large-based SAM2 (paper default).

    Real usage:
        from sam2.build_sam import build_sam2_video_predictor
        self.predictor = build_sam2_video_predictor(config, checkpoint)

    This wrapper provides a mock implementation for smoke testing.
    """

    def __init__(self, use_mock: bool = True):
        self.use_mock = use_mock
        if not use_mock:
            try:
                from sam2.build_sam import build_sam2_video_predictor
                # self.predictor = build_sam2_video_predictor(config, checkpoint)
            except ImportError:
                print("[SAM2] sam2 package not found. Install from: "
                      "https://github.com/facebookresearch/segment-anything-2")
                self.use_mock = True

    def segment_video(
        self,
        frames: List[np.ndarray],
        boxes: List[Optional[BoundingBox]],
    ) -> List[np.ndarray]:
        """
        Generate binary segmentation masks for each video frame.

        For frames where no bounding box is provided (None), the mask
        defaults to all-background (all zeros).

        Parameters
        ----------
        frames : list of N (H, W, 3) uint8 numpy arrays
        boxes  : list of N BoundingBox objects or None

        Returns
        -------
        masks : list of N (H, W) bool numpy arrays
                True = foreground (referred object)
                False = background
        """
        if self.use_mock:
            return self._mock_segment(frames, boxes)

        # Real SAM2 implementation:
        # inference_state = self.predictor.init_state(video=frames)
        # for frame_idx, box in enumerate(boxes):
        #     if box is not None:
        #         self.predictor.add_new_points_or_box(
        #             inference_state, frame_idx, obj_id=1,
        #             box=np.array([box.x1, box.y1, box.x2, box.y2])
        #         )
        # video_segments = {}
        # for out_frame_idx, out_obj_ids, out_mask_logits in \
        #         self.predictor.propagate_in_video(inference_state):
        #     video_segments[out_frame_idx] = (out_mask_logits > 0).cpu().numpy()[0, 0]
        # return [video_segments.get(i, np.zeros(frames[0].shape[:2], bool))
        #         for i in range(len(frames))]
        return self._mock_segment(frames, boxes)

    def _mock_segment(
        self,
        frames: List[np.ndarray],
        boxes: List[Optional[BoundingBox]],
    ) -> List[np.ndarray]:
        """Synthetic binary masks for smoke testing."""
        masks = []
        for frame, box in zip(frames, boxes):
            H = frame.shape[0] if frame.ndim >= 2 else 256
            W = frame.shape[1] if frame.ndim >= 2 else 256
            mask = np.zeros((H, W), dtype=bool)
            if box is not None:
                x1 = int(max(0, box.x1)); y1 = int(max(0, box.y1))
                x2 = int(min(W, box.x2));  y2 = int(min(H, box.y2))
                mask[y1:y2, x1:x2] = True
            masks.append(mask)
        return masks


# ─── SECTION 9: TGSAgent — Complete Think-Ground-Segment Pipeline ───────────────

class TGSAgent(nn.Module):
    """
    TGS-Agent: An Object-aware Reasoning Agent for Referring Audio-Visual Segmentation.

    Full system implementing the Think → Ground → Segment pipeline:

        Think:   RefThinker(A, V, R, P) → T (reasoning chain + object description)
        Ground:  GroundingDINO(T_s or T_f, V) → B (bounding boxes per frame)
        Segment: SAM2(B, V) → M (binary masks per frame)

    The pipeline transforms:
        (audio stream + video frames + reference expression)
        → explicit object description
        → bounding boxes
        → pixel-level segmentation masks

    Key design choices:
    1. No pixel-level supervision required during inference
    2. SAM2 is used completely frozen (no mask decoder fine-tuning)
    3. Grounding-DINO is frozen (no detection head fine-tuning)
    4. Only Ref-Thinker is trained (Q-Formers, projectors, LoRA)
    5. Segmentation is interpretable at every step

    Paper: "Our TGS-Agent completes the transformation from audiovisual
    streams + reference text → object description → bbox → mask,
    demonstrating an object-aware, reliable, and explainable decision
    process for Ref-AVS task."
    """

    def __init__(
        self,
        ref_thinker_cfg: RefThinkerConfig,
        grounding_cfg: GroundingConfig,
        use_mock_tools: bool = True,
        use_fine_grained: bool = False,
    ):
        super().__init__()
        self.ref_thinker    = RefThinker(ref_thinker_cfg)
        self.parser         = RefThinkerOutputParser()
        self.detector       = GroundingDINOWrapper(grounding_cfg, use_mock=use_mock_tools)
        self.segmentor      = SAM2Wrapper(use_mock=use_mock_tools)
        self.use_fine_grained = use_fine_grained

    def forward(
        self,
        audio_feats: Tensor,
        visual_feats: Tensor,
        input_ids: Tensor,
        attention_mask: Optional[Tensor] = None,
        labels: Optional[Tensor] = None,
    ) -> Dict:
        """
        Training forward pass: computes cross-entropy loss on the
        think-answer generation task. Only Ref-Thinker is optimized.

        Parameters
        ----------
        audio_feats   : (B, T_a, audio_feat_dim)
        visual_feats  : (B, T_v, visual_feat_dim)
        input_ids     : (B, L) tokenized prompt + expected output
        attention_mask: (B, L)
        labels        : (B, L) — -100 for positions where loss is masked

        Returns
        -------
        dict from RefThinker.forward() with 'loss', 'logits', etc.
        """
        return self.ref_thinker(audio_feats, visual_feats,
                                input_ids, attention_mask, labels)

    @torch.no_grad()
    def inference(
        self,
        audio_feats: Tensor,
        visual_feats: Tensor,
        prompt_ids: Tensor,
        frames: List[np.ndarray],
        tokenizer_decode_fn=None,
    ) -> Dict:
        """
        Full inference pipeline: Think → Ground → Segment.

        Parameters
        ----------
        audio_feats  : (1, T_a, audio_feat_dim) — single video
        visual_feats : (1, T_v, visual_feat_dim)
        prompt_ids   : (1, L_p) tokenized user prompt
        frames       : list of N (H, W, 3) uint8 numpy arrays
        tokenizer_decode_fn : callable to convert token IDs to text
                              (use tokenizer.decode in real usage)

        Returns
        -------
        dict with:
            'think_chain' : str — the ... reasoning
            'f_object'    : str — fine-grained object description
            's_object'    : str — simplified category name
            'boxes'       : list of BoundingBox or None (N frames)
            'masks'       : list of (H,W) bool numpy arrays (N frames)
            'detection_prompt' : str used for Grounding-DINO
        """
        # ── Step 1: Think ─────────────────────────────────────────────────────
        generated_ids = self.ref_thinker.generate(
            audio_feats, visual_feats, prompt_ids
        )

        if tokenizer_decode_fn is not None:
            raw_text = tokenizer_decode_fn(generated_ids[0].tolist())
        else:
            raw_text = (
                " The video shows a woman playing a piano and a man "
                "playing a bassoon. The audio contains the sound of both instruments. "
                "The reference relates to the piano being played by the woman. "
                "a black grand piano"
                "piano"
            )

        parsed = self.parser.parse(raw_text)
        detection_prompt = self.parser.get_detection_prompt(
            parsed, use_fine_grained=self.use_fine_grained
        )

        # ── Step 2: Ground ─────────────────────────────────────────────────────
        if not detection_prompt:
            detection_prompt = parsed.get("f_object", "object")

        boxes = self.detector.detect_video(frames, detection_prompt)

        # ── Step 3: Segment ────────────────────────────────────────────────────
        masks = self.segmentor.segment_video(frames, boxes)

        return {
            "think_chain":       parsed["think"],
            "f_object":          parsed["f_object"],
            "s_object":          parsed["s_object"],
            "boxes":             boxes,
            "masks":             masks,
            "detection_prompt":  detection_prompt,
            "raw_output":        raw_text,
        }


# ─── SECTION 10: Loss Functions ────────────────────────────────────────────────

class RefThinkerLoss(nn.Module):
    """
    Autoregressive cross-entropy loss for Ref-Thinker training.

    Paper: "For both training phases, the autoregressive cross-entropy
    loss is used for optimization."

    The loss is applied only to the target token positions (labels != -100).
    The multimodal prefix tokens (audio and visual Q-Former outputs)
    are excluded from loss computation by setting their labels to -100.

    For phase 1 (Q-Former pretraining): the target is the caption text,
    and the LLM backbone is frozen.
    For phase 2 (instruction tuning): the target is the full think-answer
    chain, and LoRA parameters are updated via the same loss.
    """

    def __init__(self, label_smoothing: float = 0.0):
        super().__init__()
        self.label_smoothing = label_smoothing

    def forward(self, logits: Tensor, labels: Tensor,
                 ignore_index: int = -100) -> Tensor:
        """
        logits : (B, T, vocab_size) — model output logits (already shifted)
        labels : (B, T) — target token IDs (-100 where loss should be ignored)
        """
        return F.cross_entropy(
            logits.reshape(-1, logits.size(-1)),
            labels.reshape(-1),
            ignore_index=ignore_index,
            label_smoothing=self.label_smoothing,
        )


# ─── SECTION 11: Instruction Tuning Dataset ────────────────────────────────────

class RefAVSInstructionDataset(Dataset):
    """
    Instruction tuning dataset for Ref-Thinker Phase 2 training.

    Each sample provides:
    - Audio features extracted by BEATs encoder
    - Visual features extracted by CLIP-ViT-L/14 from video frames
    - Tokenized user prompt (with reference expression)
    - Tokenized expected think-answer chain (generated by Gemini-1.5-Pro)

    In real use, audio/visual features would be pre-extracted and cached.
    This dummy dataset generates random tensors for smoke testing.

    Paper: "We leverage Gemini-1.5-Pro with carefully designed prompts to
    analyze the video and generate an object-aware reasoning chain."
    """

    def __init__(
        self,
        n_samples: int,
        audio_feat_dim: int = 768,
        visual_feat_dim: int = 1024,
        audio_seq_len: int = 64,
        visual_seq_len: int = 128,
        max_text_len: int = 256,
        vocab_size: int = 32000,
    ):
        self.n = n_samples
        self.audio_feat_dim = audio_feat_dim
        self.visual_feat_dim = visual_feat_dim
        self.audio_seq_len = audio_seq_len
        self.visual_seq_len = visual_seq_len
        self.max_text_len = max_text_len
        self.vocab_size = vocab_size

    def __len__(self): return self.n

    def __getitem__(self, idx: int) -> Dict[str, Tensor]:
        L = self.max_text_len
        audio_feats  = torch.randn(self.audio_seq_len, self.audio_feat_dim)
        visual_feats = torch.randn(self.visual_seq_len, self.visual_feat_dim)
        input_ids    = torch.randint(0, self.vocab_size, (L,))
        attn_mask    = torch.ones(L, dtype=torch.long)
        # Labels: first 128 tokens are prompt (masked), rest is target
        labels = input_ids.clone()
        labels[:128] = -100
        return {
            "audio_feats":    audio_feats,
            "visual_feats":   visual_feats,
            "input_ids":      input_ids,
            "attention_mask": attn_mask,
            "labels":         labels,
        }


# ─── SECTION 12: Training Loop ─────────────────────────────────────────────────

def train_phase1_qformer(
    agent: TGSAgent,
    loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
    epoch: int,
) -> float:
    """
    Phase 1: Pretrain Q-Formers and projectors on captioning data.
    The LLM backbone is frozen throughout this phase.

    Paper: "During the pretraining phase, the LLM is frozen, and
    domain-specific caption datasets are independently used to train
    the audio and visual Q-Formers along with their corresponding
    projectors."
    """
    agent.ref_thinker.llm_backbone.requires_grad_(False)
    agent.ref_thinker.token_embedding.requires_grad_(False)

    agent.train()
    total_loss = 0.0
    for step, batch in enumerate(loader):
        audio_feats   = batch["audio_feats"].to(device)
        visual_feats  = batch["visual_feats"].to(device)
        input_ids     = batch["input_ids"].to(device)
        attn_mask     = batch["attention_mask"].to(device)
        labels        = batch["labels"].to(device)

        out  = agent(audio_feats, visual_feats, input_ids, attn_mask, labels)
        loss = out["loss"]
        optimizer.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(agent.parameters(), 1.0)
        optimizer.step()
        total_loss += loss.item()

        if step % 10 == 0:
            print(f"  [Phase1] Epoch {epoch} Step {step}/{len(loader)} "
                  f"loss={loss.item():.4f}")
    return total_loss / len(loader)


def train_phase2_instruction(
    agent: TGSAgent,
    loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    scheduler: Optional[torch.optim.lr_scheduler._LRScheduler],
    device: torch.device,
    epoch: int,
) -> float:
    """
    Phase 2: Instruction tuning with LoRA on think-answer reasoning chains.

    Paper: "We apply the LoRA technique for parameter-efficient tuning of
    the LLM. The batch size is 4, and the LLM is fine-tuned for 6 epochs.
    We use AdamW optimizer with an initial learning rate of 1e-4."
    """
    agent.train()
    total_loss = 0.0
    for step, batch in enumerate(loader):
        audio_feats  = batch["audio_feats"].to(device)
        visual_feats = batch["visual_feats"].to(device)
        input_ids    = batch["input_ids"].to(device)
        attn_mask    = batch["attention_mask"].to(device)
        labels       = batch["labels"].to(device)

        out  = agent(audio_feats, visual_feats, input_ids, attn_mask, labels)
        loss = out["loss"]
        optimizer.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(
            [p for p in agent.parameters() if p.requires_grad], 1.0
        )
        optimizer.step()
        if scheduler:
            scheduler.step()
        total_loss += loss.item()

        if step % 5 == 0:
            trainable = sum(p.numel() for p in agent.parameters() if p.requires_grad)
            print(f"  [Phase2] Epoch {epoch} Step {step}/{len(loader)} "
                  f"loss={loss.item():.4f} trainable_params={trainable:,}")
    return total_loss / len(loader)


# ─── SECTION 13: Inference & Evaluation Utilities ──────────────────────────────

def jaccard_index(pred_mask: np.ndarray, gt_mask: np.ndarray) -> float:
    """
    Compute Jaccard index (IoU) between predicted and ground-truth masks.

    Jaccard = |pred ∩ gt| / |pred ∪ gt|

    Used as the primary evaluation metric (J) in Ref-AVSBench and R²-AVSBench.
    """
    pred = pred_mask.astype(bool)
    gt   = gt_mask.astype(bool)
    inter = (pred & gt).sum()
    union = (pred | gt).sum()
    if union == 0:
        return 1.0 if inter == 0 else 0.0
    return float(inter) / float(union)


def f_score(pred_mask: np.ndarray, gt_mask: np.ndarray,
            beta: float = 1.0) -> float:
    """
    Compute F-score (F_β) between predicted and ground-truth masks.

    F_β = (1 + β²) · (precision · recall) / (β² · precision + recall)
    For β=1 (default), this is the harmonic mean of precision and recall.

    Used as the secondary evaluation metric (F) in Ref-AVSBench.
    """
    pred = pred_mask.astype(bool)
    gt   = gt_mask.astype(bool)
    tp = (pred & gt).sum()
    fp = (pred & ~gt).sum()
    fn = (~pred & gt).sum()
    if tp + fp == 0 or tp + fn == 0:
        return 0.0
    precision = tp / (tp + fp)
    recall    = tp / (tp + fn)
    if precision + recall == 0:
        return 0.0
    return (1 + beta**2) * precision * recall / (beta**2 * precision + recall)


def null_set_metric(pred_masks: List[np.ndarray],
                    background_areas: List[float]) -> float:
    """
    Compute the Null set metric S — the ratio of predicted mask area
    to background area, averaged across Null-set samples.

    In Null set videos, the referred object does not exist in the video.
    A lower S indicates fewer false-positive pixels. The ideal S is 0.0.

    Paper: "For the Null set, a metric S is used by computing the ratio
    between the predicted mask area and the background area."
    """
    ratios = []
    for mask, bg_area in zip(pred_masks, background_areas):
        pred_area = mask.astype(bool).sum()
        ratios.append(float(pred_area) / max(float(bg_area), 1))
    return float(np.mean(ratios)) if ratios else 0.0


def evaluate_tgs_agent(
    agent: TGSAgent,
    samples: List[Dict],
    device: torch.device,
) -> Dict[str, float]:
    """
    Full evaluation loop for TGS-Agent on Ref-AVSBench or R²-AVSBench.

    Each sample in `samples` should contain:
        audio_feats  : (1, T_a, D_a) tensor
        visual_feats : (1, T_v, D_v) tensor
        prompt_ids   : (1, L_p) tensor
        frames       : list of N (H, W, 3) numpy arrays
        gt_masks     : list of N (H, W) bool numpy arrays
        is_null      : bool — True if referred object is absent from video

    Returns
    -------
    dict with 'mean_J', 'mean_F', 'mean_JandF', 'null_S'
    """
    agent.eval()
    J_scores, F_scores, null_preds, null_bgs = [], [], [], []

    for sample in samples:
        result = agent.inference(
            audio_feats=sample["audio_feats"].to(device),
            visual_feats=sample["visual_feats"].to(device),
            prompt_ids=sample["prompt_ids"].to(device),
            frames=sample["frames"],
        )
        pred_masks = result["masks"]
        gt_masks   = sample["gt_masks"]

        if sample.get("is_null", False):
            for pm in pred_masks:
                total_pixels = pm.size
                null_preds.append(pm)
                null_bgs.append(total_pixels)
        else:
            for pm, gm in zip(pred_masks, gt_masks):
                J_scores.append(jaccard_index(pm, gm))
                F_scores.append(f_score(pm, gm))

    mean_J = float(np.mean(J_scores)) if J_scores else 0.0
    mean_F = float(np.mean(F_scores)) if F_scores else 0.0
    null_S = null_set_metric(null_preds, null_bgs) if null_preds else 0.0

    return {
        "mean_J":       round(mean_J * 100, 1),
        "mean_F":       round(mean_F * 100, 1),
        "mean_JandF":   round((mean_J + mean_F) / 2 * 100, 1),
        "null_S":       round(null_S, 3),
    }


# ─── SECTION 14: R²-AVSBench Reference Transformer ─────────────────────────────

class R2AVSBenchTransformer:
    """
    Reference transformation pipeline for creating R²-AVSBench.

    The R²-AVSBench is constructed by transforming references from
    Ref-AVSBench into more linguistically complex and reasoning-intensive
    equivalents using Gemini-1.5-Pro.

    Key transformations:
    1. Remove direct object names (e.g., 'man' → 'counterpart')
    2. Add relative pronouns for multi-step inference
    3. Require commonsense knowledge about object functions
    4. Require audio-functional reasoning
    5. Increase average reference length from 7.08 → 11.73 words

    Examples from the paper:
        "The couch sat by a woman"
        → "The item visually serving as a shared seating platform for the audio discourse"

        "The object making a sound by using a saw"
        → "The entity rhythmically wielding a tool known for forestry tasks in earlier eras"

    This class implements the transformation prompt template.
    In real use, this calls Gemini-1.5-Pro; here we return mock transformations.
    """

    TRANSFORMATION_PROMPT_TEMPLATE = """
    Given a referring expression for an audio-visual segmentation task,
    transform it into a more complex reference that:
    1. Avoids directly naming the target object
    2. Uses abstract or functional descriptions instead of object names
    3. Employs relative pronouns (whose, which, that)
    4. Requires reasoning about commonsense knowledge or context
    5. May incorporate audio-functional cues
    6. Ensures the new reference uniquely identifies the same target object
    7. Uses 10-15 words and varied sentence structure

    Original reference: {original_reference}
    Target object: {target_object}

    Output only the transformed reference. No explanation.
    """

    @classmethod
    def transform_reference(cls, reference: str, target_object: str,
                               use_mock: bool = True) -> str:
        """
        Transform a simple reference into a reasoning-intensive one.

        In production, this calls the Gemini-1.5-Pro API with the
        prompt template above. For smoke testing, returns a mock transformation.
        """
        if use_mock:
            return cls._mock_transform(reference, target_object)
        # Real: call Gemini API with TRANSFORMATION_PROMPT_TEMPLATE
        return reference

    @staticmethod
    def _mock_transform(reference: str, target_object: str) -> str:
        """Returns a mock complex reference for smoke testing."""
        return (f"The entity whose acoustic contribution is contextually "
                f"associated with the auditory scene in question, "
                f"specifically the {target_object}-related phenomenon.")


# ─── SECTION 15: Smoke Test ─────────────────────────────────────────────────────

def run_smoke_test():
    """
    End-to-end smoke test for TGS-Agent.

    Runs all three stages (Think, Ground, Segment) on synthetic data,
    verifies tensor shapes, computes loss, and exercises the evaluation
    utilities. No pretrained weights required.
    """
    print("=" * 65)
    print("  TGS-Agent Smoke Test — AAAI 2026")
    print("  Think-Ground-Segment for Ref-AVS")
    print("=" * 65)

    device = torch.device("cpu")
    rft_cfg = RefThinkerConfig(
        audio_feat_dim=768, visual_feat_dim=1024,
        llm_hidden_dim=512,  # reduced for smoke test
        num_query_tokens=8,
        qformer_layers=2, qformer_heads=4,
        lora_rank=4, lora_alpha=8,
        vocab_size=1000,
    )
    grnd_cfg = GroundingConfig(tau_bbox=0.1, tau_text=0.25)

    print("\n[1/6] Building TGS-Agent...")
    agent = TGSAgent(rft_cfg, grnd_cfg, use_mock_tools=True)
    total_params = sum(p.numel() for p in agent.parameters())
    trainable    = sum(p.numel() for p in agent.parameters() if p.requires_grad)
    print(f"  Total params:     {total_params:,}")
    print(f"  Trainable params: {trainable:,} ({100*trainable/total_params:.1f}%)")

    print("\n[2/6] Training forward pass (Phase 2 instruction tuning)...")
    B, T_a, T_v, L = 2, 32, 64, 64
    audio_feats  = torch.randn(B, T_a, rft_cfg.audio_feat_dim)
    visual_feats = torch.randn(B, T_v, rft_cfg.visual_feat_dim)
    input_ids    = torch.randint(0, rft_cfg.vocab_size, (B, L))
    attn_mask    = torch.ones(B, L)
    labels       = input_ids.clone()
    labels[:32] = -100

    out = agent(audio_feats, visual_feats, input_ids, attn_mask, labels)
    assert out["loss"] is not None and out["loss"].item() > 0
    print(f"  ✓ Loss: {out['loss'].item():.4f}")
    print(f"  ✓ Logits shape: {tuple(out['logits'].shape)}")

    print("\n[3/6] Inference pipeline (Think → Ground → Segment)...")
    prompt_ids = torch.randint(0, rft_cfg.vocab_size, (1, 20))
    N_frames, H, W = 5, 64, 64
    frames = [np.random.randint(0, 255, (H, W, 3), dtype=np.uint8)
              for _ in range(N_frames)]

    result = agent.inference(
        audio_feats=audio_feats[:1],
        visual_feats=visual_feats[:1],
        prompt_ids=prompt_ids,
        frames=frames,
    )
    assert len(result["masks"]) == N_frames
    assert len(result["boxes"]) == N_frames
    print(f"  ✓ Generated {N_frames} masks, shapes: {[m.shape for m in result['masks'][:2]]}...")
    print(f"  ✓ s_object: '{result['s_object']}'")
    print(f"  ✓ detection_prompt: '{result['detection_prompt']}'")

    print("\n[4/6] Evaluation metrics (J, F, Null-S)...")
    gt_mask = np.zeros((H, W), dtype=bool)
    gt_mask[15:50, 15:50] = True
    J = jaccard_index(result["masks"][0], gt_mask)
    F = f_score(result["masks"][0], gt_mask)
    print(f"  ✓ Jaccard (J): {J:.4f}  F-score (F): {F:.4f}  J&F: {(J+F)/2:.4f}")

    print("\n[5/6] Phase 1 training loop (Q-Former pretraining)...")
    ds = RefAVSInstructionDataset(
        n_samples=4, audio_feat_dim=rft_cfg.audio_feat_dim,
        visual_feat_dim=rft_cfg.visual_feat_dim,
        audio_seq_len=32, visual_seq_len=64,
        max_text_len=64, vocab_size=rft_cfg.vocab_size
    )
    loader = DataLoader(ds, batch_size=2, shuffle=True)
    opt = torch.optim.AdamW(
        [p for p in agent.parameters() if p.requires_grad], lr=1e-4
    )
    avg_loss = train_phase1_qformer(agent, loader, opt, device, epoch=1)
    print(f"  ✓ Phase 1 avg loss: {avg_loss:.4f}")

    print("\n[6/6] R²-AVSBench reference transformer...")
    orig_ref = "The object making a sound by using a saw"
    transformed = R2AVSBenchTransformer.transform_reference(
        orig_ref, target_object="man", use_mock=True
    )
    print(f"  Original : {orig_ref}")
    print(f"  Transformed: {transformed}")

    print(f"\n{'='*65}")
    print("  ✓ All checks passed. TGS-Agent is ready to use.")
    print("="*65)
    print("""
Next steps for full deployment:
  1. Replace RefThinker LLM backbone with LlamaForCausalLM:
       from transformers import LlamaForCausalLM
       self.llm = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf")

  2. Replace Q-Former encoders with pretrained models:
       CLIP visual: openai/clip-vit-large-patch14 (visual_feat_dim=1024)
       BEATs audio: microsoft/BEATs (audio_feat_dim=768)

  3. Install and configure Grounding-DINO:
       pip install groundingdino-py
       weights: groundingdino_swinT_ogc.pth (Swin-T, paper default)

  4. Install and configure SAM2:
       git clone https://github.com/facebookresearch/segment-anything-2
       weights: sam2_hiera_large.pt (paper default)

  5. Train Phase 1 on audio/visual caption data:
       AudioCaps (Kim et al. 2019), Video-LLaVA (Lin et al. 2023)

  6. Construct Phase 2 instruction set:
       Run Gemini-1.5-Pro on Ref-AVSBench training videos
       with the prompt template from Appendix A of the paper

  7. Train Phase 2 with LoRA on 4 × NVIDIA A100-40GB (bf16):
       batch_size=4, epochs=6, lr=1e-4, AdamW
       lora_rank=8, lora_alpha=16

  8. Evaluate on Ref-AVSBench and R²-AVSBench:
       python evaluate_tgs.py --split seen --benchmark ref_avs
       python evaluate_tgs.py --split unseen --benchmark r2_avs
""")


if __name__ == "__main__":
    torch.manual_seed(42)
    np.random.seed(42)
    run_smoke_test()