load_model.py

# Compatibility: transformers==4.30.2 calls hf_hub_download(use_auth_token=...),
# removed in huggingface_hub>=0.20.  Patch before importing transformers so the
# bound reference inside transformers.utils.hub picks up the fixed version.
import huggingface_hub as _hfhub
_hfhub_orig = _hfhub.hf_hub_download
def _hfhub_compat(*args, use_auth_token=None, token=None, **kwargs):
    return _hfhub_orig(*args, token=token or use_auth_token, **kwargs)
_hfhub.hf_hub_download = _hfhub_compat

import transformers

from torch.cuda.amp import autocast, GradScaler

from datasets import REFAVS
from configs import args
from torch.utils.data import DataLoader
from functools import partial
from models.llava import conversation as conversation_lib
# from  models.avs_model import VISAForCausalLM
from  models.avs_model import Simtoken_ForCausalLM
import torch
from torch.cuda import amp
from transformers import AutoConfig
from peft import LoraConfig, get_peft_model
from torch import optim
from torch.optim import AdamW
from transformers import get_cosine_schedule_with_warmup
from tqdm import tqdm

from utils import utility
import random
import numpy as np
import re
import time
import os
from PIL import Image


import warnings

from utils.metric.utility import mask_iou

warnings.filterwarnings("ignore")

from transformers import logging
logging.set_verbosity_error()


IGNORE_INDEX = -100
IMAGE_TOKEN_INDEX = -200
DEFAULT_IMAGE_TOKEN = "<image>"
DEFAULT_IMAGE_PATCH_TOKEN = "<im_patch>"
DEFAULT_IM_START_TOKEN = "<im_start>"
DEFAULT_IM_END_TOKEN = "<im_end>"
DEFAULT_VIDEO_TOKEN = "<video>"

AUDIO_TOKEN_INDEX = -300
DEFAULT_AUDIO_TOKEN = "<audio>"

def set_seed(seed=42):
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

def dict_to_cuda(input_dict):
    for k, v in input_dict.items():
        if isinstance(input_dict[k], torch.Tensor):
            input_dict[k] = v.cuda(non_blocking=True)
        elif (
                isinstance(input_dict[k], list)
                and len(input_dict[k]) > 0
                and isinstance(input_dict[k][0], torch.Tensor)
        ):
            input_dict[k] = [ele.cuda(non_blocking=True) for ele in v]
    return input_dict

def tokenizer_image_audio_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, audio_token_index=AUDIO_TOKEN_INDEX, num_frames=10, return_tensors=None):

    prompt_chunks = re.split(r'(<image>|<audio>|<video>)', prompt)

    prompt_chunks = [chunk for chunk in prompt_chunks if chunk]

    # divide prompt into two set
    text_chunks = []  # text
    token_types = []  # <image>/<audio>/<video>
    for chunk in prompt_chunks:
        if chunk == "<image>":
            token_types.append("image")
        elif chunk == "<audio>":
            token_types.append("audio")
        elif chunk == "<video>":
            token_types.append("video")
        else:
            text_chunks.append(chunk)

    # Tokenize the text
    tokenized_chunks = [tokenizer(chunk).input_ids for chunk in text_chunks]

    def insert_separators(text_chunks, tokenized_chunks, token_types, image_token_index, audio_token_index, num_frames):
        input_ids = []
        offset = 0
        if (
                len(tokenized_chunks) > 0
                and len(tokenized_chunks[0]) > 0
                and tokenized_chunks[0][0] == tokenizer.bos_token_id
        ):
            offset = 1
            input_ids.append(tokenized_chunks[0][0])

        min_length = min(len(text_chunks), len(token_types))
        for i in range(min_length):

            input_ids.extend(tokenized_chunks[i][offset:])

            if token_types[i] == "image":
                input_ids.append(image_token_index)
            elif token_types[i] == "audio":
                input_ids.append(audio_token_index)
            elif token_types[i] == "video":
                input_ids.extend([image_token_index] * num_frames)


        if len(text_chunks) > min_length:
            input_ids.extend(tokenized_chunks[min_length][offset:])

        return input_ids

    input_ids = insert_separators(text_chunks, tokenized_chunks, token_types, image_token_index, audio_token_index, num_frames)

    if return_tensors is not None:
        if return_tensors == "pt":
            return torch.tensor(input_ids, dtype=torch.long)
        raise ValueError(f"Unsupported tensor type: {return_tensors}")
    return input_ids

def collate_fn(batch, tokenizer=None):
    vids = []
    images = []
    image_clips = []
    masks = []
    conversations = []
    audio_feats = []
    image_feats = []
    resizes = []
    orgsizes = []
    first_refs = []

    refs = []
    first_refs = []
    refs_num = []
    fids = []


    for data in batch:
        vids.append(data['vid'])
        images.append(data['image'])
        image_clips.append(data['img_clip'])
        masks.append(data['mask'])
        conversations.append(data['conversation'])
        audio_feats.append(data['feat_aud'])
        resizes.append(data['resize'])
        orgsizes.append(data['orgsize'])
        image_feats.append(data['feat_sam'])
        refs_num.append(len(data['ref']))
        fids.append(data['fids'])

        refs.append(data['ref'])
        first_refs.append(data['ref'][0])

    input_ids = [tokenizer_image_audio_token(conv, tokenizer, return_tensors="pt") for conv in conversations]  # list
    input_ids = torch.nn.utils.rnn.pad_sequence(input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)
    attention_masks = input_ids.ne(tokenizer.pad_token_id)

    ref_ids = [tokenizer_image_audio_token(ref, tokenizer, return_tensors="pt") for ref in first_refs]

    conv = conversation_lib.default_conversation.copy()
    labels = input_ids.clone()

    sep = 'Sure, it is [SEG]'

    for conversation, target in zip(conversations, labels):
        parts = conversation.split(sep)
        cur_len = 1
        target[:cur_len] = IGNORE_INDEX

        sep_len = len(tokenizer_image_audio_token(sep, tokenizer)) - 1

        for i in range(len(parts)-1):
            part_len = len(tokenizer_image_audio_token(parts[i], tokenizer)) - 2
            target[cur_len: cur_len + part_len] = IGNORE_INDEX
            cur_len += part_len + sep_len

        target[cur_len:] = IGNORE_INDEX


    return {"vids": vids,
            "images": images,  # list[B]:[T, 3, 1024, 1024]
            "images_clip": image_clips,  # list[B]:[T, 3, 224, 224]
            "masks": masks,  # list[B]:[num_ref, T, H, W]
            "convs": conversations,  # list[B]: str
            "input_ids": input_ids,  # list[B]:[max_len]
            "attention_masks": attention_masks,  # list[B]:[max_len]
            "labels": labels,  # list[B]:[max_len]
            "audio_feats": audio_feats,  # list[B]:[10, 128]
            "resizes": resizes,  # list[B]
            "orgsizes": orgsizes,  # list[B]
            "image_feats": image_feats,
            "ref_ids": ref_ids,  # list[B]: [ref_id_len]
            "refs_num": refs_num,
            "fids": fids,
            "refs": refs,
    }


import torch.multiprocessing as mp
if __name__ == "__main__":
    mp.set_start_method("spawn", force=True)
    set_seed(42)
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        args.mllm,
        cache_dir=None,
        model_max_length=2048,  # 2048
        padding_side="right",
        use_fast=False,
    )

    tokenizer.pad_token = tokenizer.unk_token
    num_added_tokens = tokenizer.add_tokens("[SEG]")
    seg_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids[0]  # 32000
    print("seg_token_idx: ", seg_token_idx)


    _split = args.eval_split
    _dataset = REFAVS(_split, args, tokenizer, input_type='refer')
    _dataloader = DataLoader(_dataset, batch_size=1, shuffle=False, num_workers=0, collate_fn=partial(collate_fn, tokenizer=tokenizer))



    model_args = {
        "train_mask_decoder": True,
        "out_dim": 256,  # 256
        "ce_loss_weight": 1.0,
        "dice_loss_weight": 0.5,
        "bce_loss_weight": 2.0,
        "seg_token_idx": seg_token_idx,
        "vision_pretrained": args.vision_pretrained,  # sam_vit_h_xxx.pth
        "vision_tower": args.vision_tower,
        "use_im_start_end": False,
        "compress": args.compress,
        "start": args.start,
    }


    # model = Simtoken_ForCausalLM.from_pretrained(args.mllm, torch_dtype=torch.float32, low_cpu_mem_usage=True, **model_args)
    model = Simtoken_ForCausalLM.from_pretrained(args.mllm, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True,
                                                 **model_args)

    print("\nmodel loaded")

    model.config.eos_token_id = tokenizer.eos_token_id
    model.config.bos_token_id = tokenizer.bos_token_id
    model.config.pad_token_id = tokenizer.pad_token_id

    model.enable_input_require_grads()
    model.gradient_checkpointing_enable()

    model.get_model().initialize_vision_modules(model.get_model().config)
    vision_tower = model.get_model().get_vision_tower()
    vision_tower.to(dtype=torch.float32, device="cuda")

    model_args_from_pt = AutoConfig.from_pretrained(args.mllm)
    model_args_from_pt.use_cluster = True
    model_args_from_pt.freeze = False
    model_args_from_pt.mm_tune = True
    model_args_from_pt.spatial_cluster_rate0 = 64
    model_args_from_pt.spatial_cluster_rate1 = 32
    model_args_from_pt.spatial_cluster_rate2 = 16
    model_args_from_pt.temporal_cluster_rate = 0.0625
    model_args_from_pt.use_cluster = True
    model_args_from_pt.vision_tune = False
    model.get_model().initialize_cluster_modules(model_args_from_pt)

    model.get_model().initialize_lisa_modules(model.get_model().config)

    for p in vision_tower.parameters():
        p.requires_grad = False
    for p in model.get_model().mm_projector.parameters():
        p.requires_grad = False

    lora_r = 8
    target_modules = "q_proj,v_proj"
    if lora_r > 0:
        def find_linear_layers(model, lora_target_modules):
            cls = torch.nn.Linear
            lora_module_names = set()

            for name, module in model.named_modules():
                if (
                        isinstance(module, cls)
                        and all(
                    [
                        x not in name
                        for x in [
                        "visual_model",
                        "vision_tower",
                        "mm_projector",
                        "text_hidden_fcs",
                        "audio_feature_layer",
                    ]
                    ]
                )
                        and any([x in name for x in lora_target_modules])
                ):
                    lora_module_names.add(name)
            return sorted(list(lora_module_names))


        lora_alpha = 16
        lora_dropout = 0.05

        lora_target_modules = find_linear_layers(
            model, target_modules.split(",")
        )
        lora_config = LoraConfig(
            r=lora_r,
            lora_alpha=lora_alpha,
            target_modules=lora_target_modules,
            lora_dropout=lora_dropout,
            bias="none",
            task_type="CAUSAL_LM",
        )

        model = get_peft_model(model, lora_config)
        print("\nLora deployed")

        model.print_trainable_parameters()

    model = model.to("cuda")
    model.resize_token_embeddings(len(tokenizer))

    model.load_state_dict(torch.load(args.saved_model), strict=False)
    print("saved model loaded")


    save_root = args.visualization_root

    def visualization(model, dataloader, save_root, name):
        save_root = os.path.join(save_root, name)
        os.makedirs(save_root, exist_ok=True)
        print(f"save_root: {save_root}")
        model.eval()
        for batch in tqdm(dataloader, desc=f"Visualization on {name} "):
            input_dict = dict_to_cuda(batch)
            with torch.no_grad():
                output_dict = model.forward(images=input_dict["images"],
                                            images_clip=input_dict["images_clip"],
                                            audio_features=input_dict["audio_feats"],
                                            image_features=input_dict["image_feats"],
                                            input_ids=input_dict["input_ids"],
                                            labels=input_dict["labels"],
                                            attention_masks=input_dict["attention_masks"],
                                            masks_list=input_dict["masks"],
                                            resize_list=input_dict["resizes"],
                                            orgsize_list=input_dict["orgsizes"],
                                            conversation_list=input_dict["convs"],
                                            refs_num=input_dict["refs_num"],
                                            fids=input_dict["fids"],
                                            vids=input_dict["vids"],
                                            contrast=args.ct_weight,
                                            ref_ids=input_dict["ref_ids"],
                                            inference=True)
            pred_masks = output_dict["pred_masks"]  # list[B]:[num_seg, T, H, W]
            gt_masks = output_dict["gt_masks"]  # list[B]:[num_seg, T, H, W]

            for b in range(len(pred_masks)):
                sample = torch.sigmoid(pred_masks[b])  # [num_seg, T, H, W]
                vid = input_dict["vids"][b]
                vid_root = os.path.join(save_root, vid)
                os.makedirs(vid_root, exist_ok=True)
                # print("vid_root:", vid_root)

                binary_sample = (sample > 0.4).to(torch.uint8)
                num_seg, T, H, W = sample.shape

                for seg_idx in range(num_seg):
                    ref = input_dict["refs"][b][seg_idx]
                    ref_root = os.path.join(vid_root, ref)
                    os.makedirs(ref_root, exist_ok=True)
                    # print("ref_root:", ref_root)

                    for t in range(T):
                        mask_np = binary_sample[seg_idx, t].cpu().numpy() * 255
                        mask_img = Image.fromarray(mask_np.astype(np.uint8))

                        save_path = os.path.join(ref_root, f"frame{t}.png")
                        mask_img.save(save_path)
                        # print(f"image saved as {save_path}")
        print("visualization finished")


    def valuate(model, dataloader, name, max_rows=-1):
        model.eval()

        total_iou = 0
        total_fscore = 0
        count = 0

        _total = min(max_rows, len(dataloader)) if max_rows > 0 else len(dataloader)
        for i, batch in enumerate(tqdm(dataloader, desc=f"Evaluating on {name}", total=_total)):
            if 0 < max_rows <= i:
                break
            input_dict = dict_to_cuda(batch)

            with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
                with torch.no_grad():
                    output_dict = model.forward(images=input_dict["images"],
                                                images_clip=input_dict["images_clip"],
                                                audio_features=input_dict["audio_feats"],
                                                image_features=input_dict["image_feats"],
                                                input_ids=input_dict["input_ids"],
                                                labels=input_dict["labels"],
                                                attention_masks=input_dict["attention_masks"],
                                                masks_list=input_dict["masks"],
                                                resize_list=input_dict["resizes"],
                                                orgsize_list=input_dict["orgsizes"],
                                                conversation_list=input_dict["convs"],
                                                refs_num=input_dict["refs_num"],
                                                fids=input_dict["fids"],
                                                vids=input_dict["vids"],
                                                contrast=args.ct_weight,
                                                ref_ids=input_dict["ref_ids"],
                                                inference=True)
            pred_masks = output_dict["pred_masks"]  # list[B]:[num_seg, T, H, W]
            gt_masks = output_dict["gt_masks"]  # list[B]:[num_seg, T, H, W]
            for i in range(len(pred_masks)):
                num_seg = pred_masks[i].shape[0]
                T = pred_masks[i].shape[1]
                iou = utility.mask_iou(pred_masks[i], gt_masks[i])
                fscore = utility.Eval_Fmeasure(pred_masks[i], gt_masks[i], None)

                total_iou += iou * num_seg * T
                total_fscore += fscore * num_seg * T
                count += num_seg * T

        print(f"\n  valuate on {name}:  miou: {total_iou/count}  fscore: {total_fscore/count}")


    def valuate_Null(model, dataloader, max_rows=-1):
        model.eval()

        total_metric = 0
        count = 0

        _total = min(max_rows, len(dataloader)) if max_rows > 0 else len(dataloader)
        for i, batch in enumerate(tqdm(dataloader, desc=f"Evaluating on Null", total=_total)):
            if 0 < max_rows <= i:
                break
            input_dict = dict_to_cuda(batch)
            with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
                with torch.no_grad():
                    output_dict = model.forward(images=input_dict["images"],
                                                images_clip=input_dict["images_clip"],
                                                audio_features=input_dict["audio_feats"],
                                                image_features=input_dict["image_feats"],
                                                input_ids=input_dict["input_ids"],
                                                labels=input_dict["labels"],
                                                attention_masks=input_dict["attention_masks"],
                                                masks_list=input_dict["masks"],
                                                resize_list=input_dict["resizes"],
                                                orgsize_list=input_dict["orgsizes"],
                                                conversation_list=input_dict["convs"],
                                                refs_num=input_dict["refs_num"],
                                                fids=input_dict["fids"],
                                                vids=input_dict["vids"],
                                                contrast=args.ct_weight,
                                                ref_ids=input_dict["ref_ids"],
                                                inference=True)
            pred_masks = output_dict["pred_masks"]  # list[B]:[num_seg, T, H, W]
            gt_masks = output_dict["gt_masks"]  # list[B]:[num_seg, T, H, W]
            for i in range(len(pred_masks)):
                num_seg = pred_masks[i].shape[0]
                T = pred_masks[i].shape[1]
                null_metric = utility.metric_s_for_null(pred_masks[i])

                total_metric += null_metric * num_seg * T
                count += num_seg * T

        print(f"\n  valuate on test_n_refer, metric: {total_metric / count}")




    from seg_ltpo import (
        LTPOConfig, ltpo_optimize, best_of_2_optimize, decode_full_video,
        get_sam_model, get_anchor_indices,
        QLTPOConfig, q_ltpo_autograd, check_grad_connectivity,
        reset_q_ltpo_stats, get_q_ltpo_stats,
        q_ltpo_frame_adaptive, decode_full_video_adaptive,
        _compute_avt_proxy_reward,
    )

    def print_q_ltpo_stats(name: str) -> None:
        stats = get_q_ltpo_stats()
        if not stats:
            return
        n = len(stats)
        acc_rate       = sum(s["accepted"]         for s in stats) / n
        mean_gain      = sum(s["reward_gain"]       for s in stats) / n
        mean_drift     = sum(s["drift"]             for s in stats) / n
        clip_rate      = sum(s["hit_clip"]          for s in stats) / n
        mean_iou_init  = sum(s["R_iou_pred_init"]   for s in stats) / n
        mean_iou_best  = sum(s["R_iou_pred_best"]   for s in stats) / n
        mean_area_init = sum(s["area_hard_init"]    for s in stats) / n
        mean_area_best = sum(s["area_hard_best"]    for s in stats) / n
        # Null safety: reward improved but predicted area grew >20 %
        null_risk = sum(
            1 for s in stats
            if s["reward_gain"] > 0 and s["area_hard_best"] > s["area_hard_init"] * 1.2
        ) / n
        gains = sorted(s["reward_gain"] for s in stats)
        def _pct(v, p): return v[max(0, int(len(v) * p / 100) - 1)]
        mean_e0 = sum(s["e0"] for s in stats) / n
        mean_mask_iou      = sum(s.get("mask_soft_iou",      0.0) for s in stats) / n
        mean_iou_contrib   = sum(s.get("R_iou_contrib_gain", 0.0) for s in stats) / n
        mean_soft_area_init = sum(s.get("r_area_soft_init",  0.0) for s in stats) / n
        mean_soft_area_best = sum(s.get("r_area_soft_best",  0.0) for s in stats) / n
        # B1 activation diagnostics
        b1_excesses    = sorted(s.get("b1_peak_excess", 0.0) for s in stats)
        b1_act_rate    = sum(1 for v in b1_excesses if v > 1e-8) / n
        b1_mean_excess = sum(b1_excesses) / n
        print(f"\n  [q-LTPO stats | {name} | n={n}]")
        print(f"    acceptance rate      : {acc_rate:.3f}")
        print(f"    mean e0 (exist prior): {mean_e0:.4f}  ← should differ Null vs Seen")
        print(f"    mean reward gain     : {mean_gain:+.4f}")
        print(f"    reward_gain p10/50/90: {_pct(gains,10):+.4f} / {_pct(gains,50):+.4f} / {_pct(gains,90):+.4f}")
        print(f"    mean drift ‖q−q₀‖   : {mean_drift:.4f}")
        print(f"    hit-clip ratio       : {clip_rate:.3f}")
        print(f"    R_iou_pred init→best : {mean_iou_init:.4f} → {mean_iou_best:.4f}")
        print(f"    R_iou_contrib_gain   : {mean_iou_contrib:+.4f}  ← λ_iou·e0·Δiou")
        print(f"    mask soft-IoU(init,best): {mean_mask_iou:.4f}  ← 1.0=mask不变")
        print(f"    area (hard) init→best: {mean_area_init:.4f} → {mean_area_best:.4f}")
        print(f"    soft area init→best  : {mean_soft_area_init:.4f} → {mean_soft_area_best:.4f}")
        print(f"    B1 activation rate   : {b1_act_rate:.3f}  ← frac(peak_area > e0)")
        print(f"    B1 mean excess       : {b1_mean_excess:.5f}  ← mean ReLU(peak_area - e0)")
        print(f"    B1 excess p10/50/90  : {_pct(b1_excesses,10):.5f} / {_pct(b1_excesses,50):.5f} / {_pct(b1_excesses,90):.5f}")
        print(f"    reward↑ & area+20%↑  : {null_risk:.3f}  ← Null safety indicator")
        # Direction II: frame-adaptive delta diagnostics
        delta_norms = [s.get("delta_norm", 0.0) for s in stats]
        if any(v > 0 for v in delta_norms):
            print(f"    mean delta ‖Δ‖       : {sum(delta_norms)/n:.4f}  ← per-anchor residual norm")

    def valuate_ltpo(model, dataloader, name, ltpo_cfg, optimize_fn=None,
                     max_rows=-1, multimask=False, use_edge=False):
        if optimize_fn is None:
            optimize_fn = ltpo_optimize
        """
        Evaluate with SEG-LTPO test-time optimisation + optional boundary refinement.

        decode_mode:
          multimask=False, use_edge=False : original single-mask decode (default)
          multimask=True,  use_edge=False : 3 candidates, SAM iou_pred selection (step 1a)
          multimask=True,  use_edge=True  : 3 candidates, boundary-edge score (step 1b)
        """
        model.eval()
        sam_model    = get_sam_model(model)
        model_dtype  = torch.bfloat16
        num_frames   = 10
        anchor_indices = get_anchor_indices(num_frames, ltpo_cfg.num_anchors)

        total_iou    = 0
        total_fscore = 0
        count        = 0

        _total = min(max_rows, len(dataloader)) if max_rows > 0 else len(dataloader)
        for i, batch in enumerate(tqdm(dataloader, desc=f"LTPO Evaluating on {name}", total=_total)):
            if 0 < max_rows <= i:
                break
            input_dict = dict_to_cuda(batch)

            # ── Step 1: standard forward pass (LLM + SAM decode) ──────────
            with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
                with torch.no_grad():
                    output_dict = model.forward(
                        images=input_dict["images"],
                        images_clip=input_dict["images_clip"],
                        audio_features=input_dict["audio_feats"],
                        image_features=input_dict["image_feats"],
                        input_ids=input_dict["input_ids"],
                        labels=input_dict["labels"],
                        attention_masks=input_dict["attention_masks"],
                        masks_list=input_dict["masks"],
                        resize_list=input_dict["resizes"],
                        orgsize_list=input_dict["orgsizes"],
                        conversation_list=input_dict["convs"],
                        refs_num=input_dict["refs_num"],
                        fids=input_dict["fids"],
                        vids=input_dict["vids"],
                        contrast=args.ct_weight,
                        ref_ids=input_dict["ref_ids"],
                        inference=True,
                    )

            gt_masks          = output_dict["gt_masks"]         # list[B]:[num_seg, T, H, W]
            seg_emb_list      = output_dict["seg_embeddings"]   # list[B]:[num_seg, 256]

            for b in range(len(input_dict["images"])):
                image_embeds_b = input_dict["image_feats"][b]   # [T, 256, 64, 64]
                resize_b       = input_dict["resizes"][b]
                orgsize_b      = input_dict["orgsizes"][b]
                rgb_b = input_dict["images"][b] if use_edge else None  # [T,3,H,W]

                # Convert initial Fseg to float32 for stable optimisation.
                # seg_emb_list[b]: [num_seg, 256] in bfloat16
                F_init_b = seg_emb_list[b].detach().float()     # [num_seg, 256]

                pred_masks_ltpo = []
                for seg_idx in range(F_init_b.shape[0]):
                    fseg_init = F_init_b[seg_idx : seg_idx + 1]  # [1, 256]

                    # ── Step 2: optimisation (float32, outside autocast) ──────
                    best_fseg = optimize_fn(
                        fseg_init, image_embeds_b, anchor_indices,
                        sam_model, model_dtype, ltpo_cfg,
                    )  # [1, 256] float32

                    # ── Step 3: decode full video with best Fseg ──────────────
                    pred_mask = decode_full_video(
                        best_fseg, image_embeds_b, sam_model,
                        resize_b, orgsize_b, model_dtype,
                        rgb_frames=rgb_b, multimask=multimask,
                    )  # [T, H, W]
                    pred_masks_ltpo.append(pred_mask)

                pred_masks_b = torch.stack(pred_masks_ltpo, dim=0)  # [num_seg, T, H, W]

                num_seg = pred_masks_b.shape[0]
                T_      = pred_masks_b.shape[1]
                iou     = utility.mask_iou(pred_masks_b, gt_masks[b])
                fscore  = utility.Eval_Fmeasure(pred_masks_b, gt_masks[b], None)

                total_iou    += iou    * num_seg * T_
                total_fscore += fscore * num_seg * T_
                count        += num_seg * T_

        print(f"\n  LTPO valuate on {name}:  miou: {total_iou/count:.4f}  fscore: {total_fscore/count:.4f}")


    def valuate_ltpo_null(model, dataloader, ltpo_cfg, optimize_fn=None, max_rows=-1):
        if optimize_fn is None:
            optimize_fn = ltpo_optimize
        """LTPO evaluation for Null split: measures S metric (lower = fewer false-positive masks)."""
        model.eval()
        sam_model      = get_sam_model(model)
        model_dtype    = torch.bfloat16
        num_frames     = 10
        anchor_indices = get_anchor_indices(num_frames, ltpo_cfg.num_anchors)

        total_metric = 0
        count        = 0

        _total = min(max_rows, len(dataloader)) if max_rows > 0 else len(dataloader)
        for i, batch in enumerate(tqdm(dataloader, desc="LTPO Evaluating on Null", total=_total)):
            if 0 < max_rows <= i:
                break
            input_dict = dict_to_cuda(batch)

            with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
                with torch.no_grad():
                    output_dict = model.forward(
                        images=input_dict["images"],
                        images_clip=input_dict["images_clip"],
                        audio_features=input_dict["audio_feats"],
                        image_features=input_dict["image_feats"],
                        input_ids=input_dict["input_ids"],
                        labels=input_dict["labels"],
                        attention_masks=input_dict["attention_masks"],
                        masks_list=input_dict["masks"],
                        resize_list=input_dict["resizes"],
                        orgsize_list=input_dict["orgsizes"],
                        conversation_list=input_dict["convs"],
                        refs_num=input_dict["refs_num"],
                        fids=input_dict["fids"],
                        vids=input_dict["vids"],
                        contrast=args.ct_weight,
                        ref_ids=input_dict["ref_ids"],
                        inference=True,
                    )

            seg_emb_list = output_dict["seg_embeddings"]  # list[B]:[num_seg, 256]

            for b in range(len(input_dict["images"])):
                image_embeds_b = input_dict["image_feats"][b]
                resize_b       = input_dict["resizes"][b]
                orgsize_b      = input_dict["orgsizes"][b]
                F_init_b       = seg_emb_list[b].detach().float()

                pred_masks_ltpo = []
                for seg_idx in range(F_init_b.shape[0]):
                    fseg_init = F_init_b[seg_idx : seg_idx + 1]
                    best_fseg = optimize_fn(
                        fseg_init, image_embeds_b, anchor_indices,
                        sam_model, model_dtype, ltpo_cfg,
                    )
                    pred_mask = decode_full_video(
                        best_fseg, image_embeds_b, sam_model,
                        resize_b, orgsize_b, model_dtype,
                    )
                    pred_masks_ltpo.append(pred_mask)

                pred_masks_b = torch.stack(pred_masks_ltpo, dim=0)  # [num_seg, T, H, W]
                num_seg      = pred_masks_b.shape[0]
                T_           = pred_masks_b.shape[1]
                null_metric  = utility.metric_s_for_null(pred_masks_b)

                total_metric += null_metric * num_seg * T_
                count        += num_seg * T_

        print(f"\n  LTPO valuate on Null:  S metric: {total_metric/count:.4f}")


    def valuate_ltpo_adaptive(model, dataloader, name, ltpo_cfg, max_rows=-1):
        """Evaluate with Direction II frame-adaptive token optimization."""
        model.eval()
        sam_model      = get_sam_model(model)
        model_dtype    = torch.bfloat16
        num_frames     = 10
        anchor_indices = get_anchor_indices(num_frames, ltpo_cfg.num_anchors)

        total_iou    = 0
        total_fscore = 0
        count        = 0

        _total = min(max_rows, len(dataloader)) if max_rows > 0 else len(dataloader)
        for i, batch in enumerate(tqdm(dataloader, desc=f"FA-LTPO Evaluating on {name}", total=_total)):
            if 0 < max_rows <= i:
                break
            input_dict = dict_to_cuda(batch)

            with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
                with torch.no_grad():
                    output_dict = model.forward(
                        images=input_dict["images"],
                        images_clip=input_dict["images_clip"],
                        audio_features=input_dict["audio_feats"],
                        image_features=input_dict["image_feats"],
                        input_ids=input_dict["input_ids"],
                        labels=input_dict["labels"],
                        attention_masks=input_dict["attention_masks"],
                        masks_list=input_dict["masks"],
                        resize_list=input_dict["resizes"],
                        orgsize_list=input_dict["orgsizes"],
                        conversation_list=input_dict["convs"],
                        refs_num=input_dict["refs_num"],
                        fids=input_dict["fids"],
                        vids=input_dict["vids"],
                        contrast=args.ct_weight,
                        ref_ids=input_dict["ref_ids"],
                        inference=True,
                    )

            gt_masks     = output_dict["gt_masks"]       # list[B]:[num_seg, T, H, W]
            seg_emb_list = output_dict["seg_embeddings"] # list[B]:[num_seg, 256]

            for b in range(len(input_dict["images"])):
                image_embeds_b = input_dict["image_feats"][b]
                resize_b       = input_dict["resizes"][b]
                orgsize_b      = input_dict["orgsizes"][b]
                F_init_b       = seg_emb_list[b].detach().float()

                pred_masks_ltpo = []
                for seg_idx in range(F_init_b.shape[0]):
                    fseg_init = F_init_b[seg_idx : seg_idx + 1]

                    q_global, delta = q_ltpo_frame_adaptive(
                        fseg_init, image_embeds_b, anchor_indices,
                        sam_model, model_dtype, ltpo_cfg,
                    )

                    pred_mask = decode_full_video_adaptive(
                        q_global, delta, anchor_indices,
                        image_embeds_b, sam_model,
                        resize_b, orgsize_b, model_dtype,
                    )
                    pred_masks_ltpo.append(pred_mask)

                pred_masks_b = torch.stack(pred_masks_ltpo, dim=0)
                num_seg = pred_masks_b.shape[0]
                T_      = pred_masks_b.shape[1]
                iou     = utility.mask_iou(pred_masks_b, gt_masks[b])
                fscore  = utility.Eval_Fmeasure(pred_masks_b, gt_masks[b], None)

                total_iou    += iou    * num_seg * T_
                total_fscore += fscore * num_seg * T_
                count        += num_seg * T_

        print(f"\n  FA-LTPO valuate on {name}:  miou: {total_iou/count:.4f}  fscore: {total_fscore/count:.4f}")

    # ── Step A0: reward–metric correlation study ─────────────────────────

    def _print_correlation_report(per_sample: list) -> None:
        import numpy as np
        n = len(per_sample)
        if n == 0:
            return

        r_iou   = np.array([s["reward_gain"]   for s in per_sample], dtype=float)
        r_avt   = np.array([s["r_avt_gain"]    for s in per_sample], dtype=float)
        r_avt_c = np.array([s["r_avt_c_gain"]  for s in per_sample], dtype=float)
        dm      = np.array([s["delta_miou"]     for s in per_sample], dtype=float)
        df      = np.array([s["delta_f"]        for s in per_sample], dtype=float)

        def pearson(x, y):
            x = x - x.mean(); y = y - y.mean()
            denom = np.sqrt((x ** 2).sum() * (y ** 2).sum())
            return float((x * y).sum() / (denom + 1e-12))

        def wrong_frac(gains, deltas):
            return sum(1 for g, d in zip(gains, deltas) if g > 0 and d < 0) / n

        print(f"\n  [Step A0: Reward–Metric Correlation | n={n}]")
        print(f"    mean ΔmIoU : {dm.mean():+.4f}  (std {dm.std():.4f})")
        print(f"    mean ΔF    : {df.mean():+.4f}  (std {df.std():.4f})")
        print(f"\n    Pearson r  with ΔmIoU :")
        print(f"      R_iou_pred_gain   : {pearson(r_iou,   dm):+.3f}   ← current proxy")
        print(f"      R_avt_gain        : {pearson(r_avt,   dm):+.3f}   ← cos(z_in, q_init)")
        print(f"      R_avt_c_gain      : {pearson(r_avt_c, dm):+.3f}   ← cos(z_in,q)-β·cos(z_out,q)")
        print(f"\n    Pearson r  with ΔF :")
        print(f"      R_iou_pred_gain   : {pearson(r_iou,   df):+.3f}")
        print(f"      R_avt_gain        : {pearson(r_avt,   df):+.3f}")
        print(f"      R_avt_c_gain      : {pearson(r_avt_c, df):+.3f}")
        print(f"\n    Wrong direction  (gain>0  but  Δ<0):")
        print(f"      R_iou / ΔmIoU : {wrong_frac(r_iou,   dm):.3f}")
        print(f"      R_avt / ΔmIoU : {wrong_frac(r_avt,   dm):.3f}")
        print(f"      R_iou / ΔF    : {wrong_frac(r_iou,   df):.3f}")
        print(f"      R_avt / ΔF    : {wrong_frac(r_avt,   df):.3f}")

    def valuate_ltpo_correlation_study(model, dataloader, ltpo_cfg, max_rows=-1):
        """Step A0: per-sample reward–metric correlation study.

        For each (video, segment) sample runs:
          1. Baseline decode  (q_init → mask → IoU/F)
          2. q-LTPO s1        (q_best → mask → IoU/F)
        Records reward signals and ΔmIoU / ΔF per sample, then prints
        Pearson correlation table to identify which reward best predicts
        actual metric improvement.
        """
        model.eval()
        sam_model      = get_sam_model(model)
        model_dtype    = torch.bfloat16
        anchor_indices = get_anchor_indices(10, ltpo_cfg.num_anchors)

        per_sample = []

        _total = min(max_rows, len(dataloader)) if max_rows > 0 else len(dataloader)
        for i, batch in enumerate(
            tqdm(dataloader, desc="Correlation study (s1)", total=_total)
        ):
            if 0 < max_rows <= i:
                break
            input_dict = dict_to_cuda(batch)

            with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
                with torch.no_grad():
                    output_dict = model.forward(
                        images=input_dict["images"],
                        images_clip=input_dict["images_clip"],
                        audio_features=input_dict["audio_feats"],
                        image_features=input_dict["image_feats"],
                        input_ids=input_dict["input_ids"],
                        labels=input_dict["labels"],
                        attention_masks=input_dict["attention_masks"],
                        masks_list=input_dict["masks"],
                        resize_list=input_dict["resizes"],
                        orgsize_list=input_dict["orgsizes"],
                        conversation_list=input_dict["convs"],
                        refs_num=input_dict["refs_num"],
                        fids=input_dict["fids"],
                        vids=input_dict["vids"],
                        contrast=args.ct_weight,
                        ref_ids=input_dict["ref_ids"],
                        inference=True,
                    )

            gt_masks     = output_dict["gt_masks"]       # list[B]:[num_seg, T, H, W]
            seg_emb_list = output_dict["seg_embeddings"] # list[B]:[num_seg, 256]

            for b in range(len(input_dict["images"])):
                image_embeds_b = input_dict["image_feats"][b]
                resize_b       = input_dict["resizes"][b]
                orgsize_b      = input_dict["orgsizes"][b]
                F_init_b       = seg_emb_list[b].detach().float()

                for seg_idx in range(F_init_b.shape[0]):
                    q_init  = F_init_b[seg_idx : seg_idx + 1]          # [1, 256]
                    gt_seg  = gt_masks[b][seg_idx : seg_idx + 1]        # [1, T, H, W]

                    # Baseline decode (q_init, no LTPO)
                    with torch.no_grad():
                        pred_base = decode_full_video(
                            q_init, image_embeds_b, sam_model,
                            resize_b, orgsize_b, model_dtype,
                        ).unsqueeze(0)                                  # [1, T, H, W]
                    iou_base = utility.mask_iou(pred_base, gt_seg)
                    f_base   = utility.Eval_Fmeasure(pred_base, gt_seg, None)

                    # LTPO (s1) — also computes r_avt inside q_ltpo_autograd
                    reset_q_ltpo_stats()
                    q_best = q_ltpo_autograd(
                        q_init, image_embeds_b, anchor_indices,
                        sam_model, model_dtype, ltpo_cfg,
                    )
                    stat = get_q_ltpo_stats()[0]

                    with torch.no_grad():
                        pred_ltpo = decode_full_video(
                            q_best, image_embeds_b, sam_model,
                            resize_b, orgsize_b, model_dtype,
                        ).unsqueeze(0)
                    iou_ltpo = utility.mask_iou(pred_ltpo, gt_seg)
                    f_ltpo   = utility.Eval_Fmeasure(pred_ltpo, gt_seg, None)

                    per_sample.append({
                        "reward_gain":  stat["reward_gain"],
                        "r_avt_gain":   stat.get("r_avt_gain",   0.0),
                        "r_avt_c_gain": stat.get("r_avt_c_gain", 0.0),
                        "e0":           stat["e0"],
                        "accepted":     stat["accepted"],
                        "delta_miou":   float(iou_ltpo  - iou_base),
                        "delta_f":      float(f_ltpo    - f_base),
                    })

        _print_correlation_report(per_sample)

    # ── Stage 0: gradient connectivity check ─────────────────────────────
    # Loads one image_embed directly from disk — no dataloader, no gt_mask,
    # no media frames required.  F_init is a unit-scale random vector that
    # mimics the distribution of Fseg (SAM prompt embeddings are in ℝ^256
    # with per-dim std ≈ 0.05–0.3; we use std=0.1 as a neutral initialisation).
    def run_stage0_check():
        import glob
        sam_model   = get_sam_model(model)
        model_dtype = torch.bfloat16

        embed_files = sorted(glob.glob(os.path.join(args.data_dir, "image_embed", "*.pt")))
        if not embed_files:
            print("[Stage 0] ERROR: no .pt files found in data/image_embed/")
            return False

        img_embs = torch.load(embed_files[0], map_location="cuda")  # [T, 256, 64, 64]
        if img_embs.dim() == 3:                                      # [256,64,64] → [1,256,64,64]
            img_embs = img_embs.unsqueeze(0)

        torch.manual_seed(42)
        F_init = torch.randn(1, 256, device="cuda") * 0.1           # [1, 256] float32

        anchors = get_anchor_indices(img_embs.shape[0], 4)
        diag = check_grad_connectivity(F_init, img_embs, anchors, sam_model, model_dtype)
        print("\n[Stage 0] Gradient connectivity check:")
        print(f"  file used          : {os.path.basename(embed_files[0])}")
        print(f"  gradient_connected : {diag['gradient_connected']}")
        print(f"  grad_norm (step 0) : {diag['grad_norm_step0']:.6f}")
        print(f"  reward trajectory  : {[f'{r:.4f}' for r in diag['reward_trajectory']]}")
        return diag["gradient_connected"]

    # ── Bypass equivalence test ───────────────────────────────────────────
    # Three controlled tests to verify that fseg.unsqueeze(1) (bypass) is
    # numerically equivalent to prompt_encoder(text_embeds=fseg.unsqueeze(1)):
    #   Test 1 — dense_emb dtype: dense_A.to(bfloat16) vs dense_emb_bf16 (exact 0?)
    #   Test 2 — matched-prec anchor decode: same decoder, same inputs, both bfloat16
    #   Test 3 — full-video (all T frames) matched-prec decode
    # If all pass, delta_bypass_init = 0 and the +4.22% is purely from optimization.
    def run_bypass_test():
        from seg_ltpo import _precompute_dense_emb

        sam_model   = get_sam_model(model)
        pe          = sam_model.prompt_encoder
        mask_dec    = sam_model.mask_decoder
        model_dtype = torch.bfloat16

        # Get one real Fseg via a standard forward pass on the first batch
        batch = next(iter(_dataloader))
        input_dict = dict_to_cuda(batch)
        with torch.cuda.amp.autocast(dtype=torch.bfloat16, enabled=True):
            with torch.no_grad():
                output_dict = model.forward(
                    images=input_dict["images"],
                    images_clip=input_dict["images_clip"],
                    audio_features=input_dict["audio_feats"],
                    image_features=input_dict["image_feats"],
                    input_ids=input_dict["input_ids"],
                    labels=input_dict["labels"],
                    attention_masks=input_dict["attention_masks"],
                    masks_list=input_dict["masks"],
                    resize_list=input_dict["resizes"],
                    orgsize_list=input_dict["orgsizes"],
                    conversation_list=input_dict["convs"],
                    refs_num=input_dict["refs_num"],
                    fids=input_dict["fids"],
                    vids=input_dict["vids"],
                    contrast=args.ct_weight,
                    ref_ids=input_dict["ref_ids"],
                    inference=True,
                )

        fseg         = output_dict["seg_embeddings"][0][0:1].detach()  # [1,256] bfloat16
        image_embeds = input_dict["image_feats"][0]                    # [T,256,64,64] float32
        device       = fseg.device

        anchor_indices = get_anchor_indices(image_embeds.shape[0], 4)
        img_anc        = image_embeds[anchor_indices]                          # [A,256,64,64] float32
        dense_emb_bf16 = _precompute_dense_emb(sam_model, model_dtype, device)  # [1,256,64,64] bfloat16
        dense_pe       = pe.get_dense_pe().to(device)                          # float32

        def _decode(img, sparse_emb, dense_emb):
            return mask_dec(
                image_embeddings=img,
                image_pe=dense_pe,
                sparse_prompt_embeddings=sparse_emb,
                dense_prompt_embeddings=dense_emb,
                multimask_output=False,
            )

        def _check(label, tensor_a, tensor_b, exact=False):
            err = (tensor_a.float() - tensor_b.float()).abs().max().item()
            tol = 0.0 if exact else 1e-4
            status = "PASS" if err <= tol else "FAIL"
            print(f"  [{status}]  {label:50s}  max|A-B| = {err:.2e}")
            return err <= tol

        print(f"\n[Bypass Test]  fseg dtype={fseg.dtype}  norm={fseg.float().norm().item():.4f}")

        with torch.no_grad():
            # Get prompt_encoder outputs (called outside autocast → float32)
            sparse_A, dense_A = pe(points=None, boxes=None, masks=None,
                                   text_embeds=fseg.unsqueeze(1))
            sparse_B = fseg.unsqueeze(1)   # bypass sparse: identical tensor

            # ── Test 1: dense_emb dtype artifact ────────────────────────────────
            # Hypothesis: dense_A (float32) and dense_emb_bf16 differ only because
            # no_mask_embed.weight is float32; casting to bfloat16 should give exact 0.
            print("\n  [Test 1] dense_emb dtype artifact (expected: exact 0)")
            t1 = _check("dense_A.to(bfloat16) vs dense_emb_bf16",
                        dense_A.to(torch.bfloat16), dense_emb_bf16, exact=True)

            # ── Test 2: matched-precision decode on anchors ──────────────────────
            # Both paths use bfloat16 sparse + bfloat16 dense.
            # If sparse_emb is identical and dense_emb is identical (per Test 1),
            # masks and iou_preds must be identical (same decoder, same inputs).
            print("\n  [Test 2] matched-precision anchor decode (expected: exact 0)")
            dense_A_bf16 = dense_A.to(model_dtype)
            masks_A, iou_A = _decode(img_anc, sparse_A,  dense_A_bf16)
            masks_B, iou_B = _decode(img_anc, sparse_B,  dense_emb_bf16)
            _check("sparse_emb",                    sparse_A, sparse_B,  exact=True)
            t2m = _check("masks   (anchors, matched prec)", masks_A,  masks_B,  exact=True)
            t2i = _check("iou_preds (anchors, matched prec)", iou_A,  iou_B,    exact=True)
            t2  = t2m and t2i

            # ── Test 3: full-video bypass-init baseline (all T frames) ──────────
            # Extend Test 2 to all T frames; quantifies delta_bypass_init over
            # the complete video rather than just the 4 anchor frames.
            print(f"\n  [Test 3] full-video matched-precision decode (T={image_embeds.shape[0]} frames)")
            masks_full_A, _ = _decode(image_embeds, sparse_A,  dense_A_bf16)
            masks_full_B, _ = _decode(image_embeds, sparse_B,  dense_emb_bf16)
            t3 = _check("masks (all frames, matched prec)", masks_full_A, masks_full_B, exact=True)

        print("\n  ── Verdict ──────────────────────────────────────────────────────")
        if t1 and t2 and t3:
            print("  ALL PASS — bypass is algebraically and numerically equivalent to")
            print("  prompt_encoder path under matched precision.  delta_bypass_init = 0.")
            print("  The +4.22% mIoU improvement is purely from q-LTPO optimization.")
        else:
            failures = []
            if not t1: failures.append("Test 1 (dense dtype)")
            if not t2: failures.append("Test 2 (anchor decode)")
            if not t3: failures.append("Test 3 (full-video decode)")
            print(f"  FAIL in: {', '.join(failures)}")
            print("  delta_bypass_init ≠ 0; need per-sample mIoU comparison to quantify.")

    # ── Run evaluation ────────────────────────────────────────────────────

    ltpo_cfg          = LTPOConfig()
    q_ltpo_cfg_s1     = QLTPOConfig(stage=1)
    q_ltpo_cfg_s2     = QLTPOConfig(stage=2)
    q_ltpo_cfg_s21    = QLTPOConfig(stage=21)   # P1a: tether probe
    q_ltpo_cfg_s22    = QLTPOConfig(stage=22)   # P1b: faithful ext-ref

    # ── Direction B: boundary precision probes ──────────────────────────────
    q_ltpo_cfg_b1_w03 = QLTPOConfig(stage=1, lambda_area_inc=0.3, area_inc_tau=0.0)
    q_ltpo_cfg_b1_w10 = QLTPOConfig(stage=1, lambda_area_inc=1.0, area_inc_tau=0.0)

    # ── Direction II: Frame-adaptive token optimization ─────────────────────
    # fa_c03: delta clipped at 0.3×‖q_init‖ — moderate constraint.
    #   First probe to answer: "does constrained frame-adaptive beat shared q?"
    #   If yes → ablate tighter/looser constraints and smoothness in follow-up.
    q_ltpo_cfg_fa_c03  = QLTPOConfig(stage=1, lambda_residual=0.001, lambda_smooth_temp=0.0, max_delta_drift_scale=0.3)

    max_rows          = args.max_eval_rows       # -1 = all rows

    # --max_eval_rows 0  → Stage 0 + bypass equivalence check, then exit
    if max_rows == 0:
        run_stage0_check()
        run_bypass_test()
    elif _split == 'test_n':
        # Null safety check: baseline + Stage 1 + frame-adaptive
        valuate_Null(model, _dataloader, max_rows=max_rows)
        for cfg_name, cfg in [("s1", q_ltpo_cfg_s1)]:
            reset_q_ltpo_stats()
            valuate_ltpo_null(model, _dataloader, cfg,
                              optimize_fn=q_ltpo_autograd, max_rows=max_rows)
            print_q_ltpo_stats(f"null_q_ltpo_{cfg_name}")
        reset_q_ltpo_stats()
        valuate_ltpo_adaptive(model, _dataloader, "null_fa_c03",
                              q_ltpo_cfg_fa_c03, max_rows=max_rows)
        print_q_ltpo_stats("null_fa_c03")
    else:
        valuate(model, _dataloader, _split, max_rows=max_rows)
        # Step A0: reward–metric correlation study (s1 + AVT proxy signals)
        valuate_ltpo_correlation_study(
            model, _dataloader, q_ltpo_cfg_s1, max_rows=max_rows
        )