Rad_IR/radir/RadIR.py

import math
import copy
from contextlib import contextmanager
from functools import partial, wraps
from pathlib import Path

import torch
import torch.nn.functional as F
from torch import nn, einsum
from torch.utils.checkpoint import checkpoint
from einops import rearrange, repeat, reduce
from einops.layers.torch import Rearrange, Reduce


from transformers import BertTokenizer, BertModel
from radir import lorentz as L

from radir.data.hmclip_mimic_box import HyCoClipMimicBoxDataset_JPG, collate_hmclip_box
from radir.data.conditional_dataset_train import SimpleImageTextDataset, simple_collate



# NOTE: Import a Cross-Attn Module and PE Module
from positional_encodings.torch_encodings import PositionalEncoding3D, PositionalEncoding1D, PositionalEncoding2D
from radir.transformer_decoder import TransformerDecoder, TransformerDecoderLayer
# from transformer_decoder import TransformerDecoder, TransformerDecoderLayer

# NOTE

# helper functions

def identity(t, *args, **kwargs):
    return t

def exists(val):
    return val is not None

def default(val, d):
    return val if exists(val) else d

@contextmanager
def null_context():
    yield

def max_neg_value(dtype):
    return -torch.finfo(dtype).max

def cast_tuple(t):
    return t if isinstance(t, (tuple, list)) else (t,)

def masked_mean(t, mask, dim = 1, eps = 1e-6):
    t = t.masked_fill(~mask, 0.)
    numer = t.sum(dim = dim)
    denom = mask.sum(dim = dim).clamp(min = eps)
    return numer / denom

def log(t, eps = 1e-20):
    return torch.log(t + eps)

def l2norm(t):
    
    return F.normalize(t, dim = -1)

def matrix_diag(t):
    device = t.device
    i, j = t.shape[-2:]
    num_diag_el = min(i, j)
    i_range = torch.arange(i, device = device)
    j_range = torch.arange(j, device = device)
    diag_mask = rearrange(i_range, 'i -> i 1') == rearrange(j_range, 'j -> 1 j')
    diag_el = t.masked_select(diag_mask)
    return rearrange(diag_el, '(b d) -> b d', d = num_diag_el)

# checkpointing helper function

def make_checkpointable(fn):
    @wraps(fn)
    def inner(*args):
        input_needs_grad = any([isinstance(el, torch.Tensor) and el.requires_grad for el in args])

        if not input_needs_grad:
            return fn(*args)

        return checkpoint(fn, *args)

    return inner

# keyword argument helpers

def pick_and_pop(keys, d):
    values = list(map(lambda key: d.pop(key), keys))
    return dict(zip(keys, values))

def group_dict_by_key(cond, d):
    return_val = [dict(),dict()]
    for key in d.keys():
        match = bool(cond(key))
        ind = int(not match)
        return_val[ind][key] = d[key]
    return (*return_val,)

def string_begins_with(prefix, str):
    return str.startswith(prefix)

def group_by_key_prefix(prefix, d):
    return group_dict_by_key(partial(string_begins_with, prefix), d)

def groupby_prefix_and_trim(prefix, d):
    with_prefix, without_prefix = group_dict_by_key(partial(string_begins_with, prefix), d)
    trimmed = {k[len(prefix):]: v for k, v in with_prefix.items()}
    return trimmed, without_prefix


def safe_normalize(feat, max_norm=10.0, eps=1e-6):
    # 支持 (B,D) 或 (B,1,D)
    orig_shape = feat.shape
    feat = feat.reshape(-1, feat.size(-1))

    norm = feat.norm(dim=-1, keepdim=True)  # (B,1)
    scale = (max_norm / (norm + eps)).clamp(max=1.0)
    feat = feat * scale

    return feat.reshape(orig_shape)

# main clip class

class RADIR(nn.Module):
    def __init__(
            self,
            image_encoder = None,
            text_encoder = None,
            tokenizer = None,
            fusion_module = None,   # NOTE: Pass in a Transformer Decoder Module
            dim_text = 768,
            dim_image = 768,
            dim_latent = 768,
            text_seq_len = 256,
            text_has_cls_token = False,
            text_pad_id = 0,
            text_causal_mask = False,
            text_eos_id = None,
            text_encode_without_mask = False,
            visual_image_size = 224,
            visual_has_cls_token = False,
            channels = 3,
            use_all_token_embeds = False,
            decoupled_contrastive_learning = False,
            use_mlm = False,
            text_ssl_loss_weight = 0.05,
            use_visual_ssl = False,
            visual_ssl = None,
            embed_dim = 512,
            image_ssl_loss_weight = 0.05,
            multiview_loss_weight = 0.1,
            curv_init: float = 1.0,
            learn_curv: bool = True,
            entail_weight: float = 0.3,
            local_aperture_thresh: float = 1.2,
            global_aperture_thresh: float = 0.7,
            use_repo_ce_loss: bool = True
    ):
        super().__init__()
        



        #assert use_all_token_embeds or (visual_has_cls_token or text_has_cls_token), 'CLS token must be included on both vision and text transformers if you are not using fine-grained contrastive learning loss'

        # store some parameters for access

        self.dim_text = dim_text
        self.dim_image = dim_image
        self.dim_latent = dim_latent

        self.image_size = visual_image_size

        # instantiate text transformer

        self.text_pad_id = text_pad_id
        self.text_has_cls_token = text_has_cls_token
        self.text_seq_len = text_seq_len

        self.text_encode_without_mask = text_encode_without_mask # whether to pass in text mask to text encoder

        self.text_causal_mask = text_causal_mask
        self.text_eos_id = text_eos_id

        assert not (text_causal_mask and not exists(text_eos_id)), 'text EOS token id must be given if using causal mask in text transformer'


        self.text_transformer = text_encoder


        # instantiate image transformer

        self.visual_has_cls_token = visual_has_cls_token

        self.visual_transformer = image_encoder

        # NOTE: Pass in a Transformer Decoder Module and a Positional-Encoding Module
        if exists(fusion_module):
            self.fusion_module = fusion_module
        else:
            decoder_layer = TransformerDecoderLayer(d_model=dim_latent, nhead=8, normalize_before=True)
            decoder_norm = nn.LayerNorm(dim_latent)
            self.fusion_module = TransformerDecoder(decoder_layer=decoder_layer, num_layers=4, norm=decoder_norm)

        self.pos_embedding = PositionalEncoding2D(dim_latent)(torch.zeros(1, 24, 24, dim_latent)) # 1 H W Dim
        # 修改形状为1 (HW) Dim
        self.pos_embedding = rearrange(self.pos_embedding, 'b h w d -> b (h w) d')
        
        # NOTE

        # text ssl

        self.use_mlm = use_mlm
        self.text_ssl_loss_weight = text_ssl_loss_weight if use_mlm else 0

        # image ssl

        self.use_visual_ssl = use_visual_ssl or exists(visual_ssl)
        self.image_ssl_loss_weight = image_ssl_loss_weight if use_visual_ssl else 0


        self.to_text_latent = nn.Linear(dim_text, dim_latent, bias = False)


        self.to_visual_latent = nn.Linear(dim_image, dim_latent, bias = False)
        self.to_fused_latent = nn.Linear(dim_latent, dim_latent, bias = False)

        # Initialize curvature parameter. Hyperboloid curvature will be `-curv`.
        self.curv = nn.Parameter(
            torch.tensor(curv_init).log(), requires_grad=learn_curv
        )
        
        # When learning the curvature parameter, restrict it in this interval to
        # prevent training instability.
        self._curv_minmax = {
            "max": math.log(curv_init * 10),
            "min": math.log(curv_init / 10),
        }

        # temperature

        self.temperature = nn.Parameter(torch.tensor(1.))
        self.use_repo_ce_loss = use_repo_ce_loss
        self.logit_scale = nn.Parameter(torch.tensor(1 / 0.07).log())  # 与HycoClip一致的可学习温度
        self.entail_weight = entail_weight
        self.global_aperture_thresh = global_aperture_thresh
        self.local_aperture_thresh = local_aperture_thresh
        self.visual_alpha = nn.Parameter(torch.tensor(embed_dim**-0.5).log())
        self.textual_alpha = nn.Parameter(torch.tensor(embed_dim**-0.5).log())


        # from https://arxiv.org/abs/2111.07783 (FILIP paper)
        self.use_all_token_embeds = use_all_token_embeds

        # proposed in https://arxiv.org/abs/2110.06848 (DCL) and https://arxiv.org/abs/2110.11316 (CLOOB)
        self.decoupled_contrastive_learning = decoupled_contrastive_learning

        self.multiview_loss_weight = multiview_loss_weight

        if tokenizer != None:
            self.tokenizer=tokenizer
        else:
            raise ValueError('Tokenzier is not defined.')
        

    def load(self, path,allow_partial_load = True):
        print(f'** MODEL ** Load Checkpoint from {path}')
        path = Path(path)
        assert path.exists()
        pkg = torch.load(path, map_location='cpu', weights_only=False)
        model_dict = self.state_dict()
        if 'model' in pkg:
            pkg_state_dict = pkg['model']  # 假设 pkg['model'] 是模型的状态字典
        else:
            pkg_state_dict = pkg
        if allow_partial_load:
            # 检查差异
            unexpected_state_dict = [k for k in pkg_state_dict.keys() if k not in model_dict.keys()]
            missing_state_dict = [k for k in model_dict.keys() if k not in pkg_state_dict.keys()]
            unmatchd_state_dict = [k for k, v in pkg_state_dict.items() if k in model_dict.keys() and v.shape != model_dict[k].shape]
            # 加载部分参数
            state_dict = {k: v for k, v in pkg_state_dict.items() if k in model_dict.keys() and v.shape == model_dict[k].shape}
            model_dict.update(state_dict)
            self.load_state_dict(model_dict)
            # 可以选择输出具体的参数情况
            # print('** MODEL ** The following parameters are UNEXPECTED in checkpoint:\n')
            # print(unexpected_state_dict)
            # print('** MODEL ** The following parameters are MISSING in checkpoint:\n')
            # print(missing_state_dict)
            # print('** MODEL ** The following parameters have DIFFERENT SHAPES in checkpoint:\n')
            # print(unmatchd_state_dict)
            # ('** MODEL ** The following parameters are LOADED in:\n')
            # print(state_dict.keys())
        else:
            self.load_state_dict(pkg_state_dict)
  
    def tokenize(self, prompt):
        text_tokens=self.tokenizer(prompt, return_tensors="pt", padding="max_length", truncation=True, max_length=512).to('cuda')
        return text_tokens
    
    def token_embedding(self,input_ids):
        input_shape = input_ids.size()
        batch_size, seq_length = input_shape
        if hasattr(self.text_transformer.embeddings, "token_type_ids"):
            print("hahatrue")

        buffered_token_type_ids = self.text_transformer.embeddings.token_type_ids[:, :seq_length]
        buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
        token_type_ids = buffered_token_type_ids_expanded
        text_embeddings = self.text_transformer.embeddings(input_ids = input_ids, token_type_ids = token_type_ids)
        return text_embeddings

    def freeze_text_encoder(self):
        for param in self.text_transformer.parameters():
            param.requires_grad = False
        for param in self.to_text_latent.parameters():
            param.requires_grad = False
            
    def open_text_encoder(self):
        for param in self.text_transformer.parameters():
            param.requires_grad = True
        for param in self.to_text_latent.parameters():
            param.requires_grad = True
    
    def freeze_vision_encoder(self):
        for param in self.visual_transformer.parameters():
            param.requires_grad = False
        for param in self.to_visual_latent.parameters():
            param.requires_grad = False
    
    def open_vision_encoder(self):
        for param in self.visual_transformer.parameters():
            param.requires_grad = True
        for param in self.to_visual_latent.parameters():
            param.requires_grad = True
    
    def open_fusion_module(self):
        for param in self.fusion_module.parameters():
            param.requires_grad = True
        for param in self.to_fused_latent.parameters():
            param.requires_grad = True
            
    def freeze_fusion_module(self):
        for param in self.fusion_module.parameters():
            param.requires_grad = False
        for param in self.to_fused_latent.parameters():
            param.requires_grad = False

    def get_log_curv(self):
        return torch.clamp(
            self.curv,
            min=self._curv_minmax["min"],
            max=self._curv_minmax["max"],
        )

    def get_curv(self):
        with torch.cuda.amp.autocast(enabled=False):
            curv = torch.exp(self.get_log_curv().float())
            return torch.clamp(curv, min=1e-4)

    def project_img(self, image_feats):
        image_feats = image_feats * self.visual_alpha.exp()
        with torch.cuda.amp.autocast(enabled=False):
            curv = self.get_curv()
            image_feats = L.exp_map0(image_feats.float(), curv)
        return image_feats

    def project_txt(self, text_feats):
        text_feats = text_feats * self.textual_alpha.exp()
        with torch.cuda.amp.autocast(enabled=False):
            curv = self.get_curv()
            text_feats = L.exp_map0(text_feats.float(), curv)
        return text_feats



    def mask_average_pool(self, x, mask):
        """
        对 token 特征做 masked average pooling

        Args:
            x:    [B, T, D]          (例如 word_feats: [8,128,512])
            mask: [B, Tm] 或 [B, K, Tm]
                目前的情况通常是 [B,111]

        Returns:
            pooled: [B, K, D]
        """
        import torch
        import torch.nn.functional as F

        # x: [B, T, D]
        T = x.size(1)

        m = mask
        # [B, Tm] -> [B, 1, Tm]
        if m.dim() == 2:
            m = m.unsqueeze(1)

        # 对齐最后一维到 T（补 0 或截断）
        Tm = m.size(-1)
        if Tm < T:
            m = F.pad(m, (0, T - Tm), value=0)   # 右侧补0
        elif Tm > T:
            m = m[..., :T]

        # einsum 用 float dtype
        m = m.to(dtype=x.dtype)

        # 避免全0 mask 导致除0
        den = m.sum(dim=-1, keepdim=True).clamp_min(1.0)

        # [B, K, T] x [B, T, D] -> [B, K, D]
        return torch.einsum('bkt,btd->bkd', m, x) / den

    def _unpack_text_batch(self, text):
        if isinstance(text, dict):
            input_ids = text.get('caption_ids', text.get('input_ids'))
            token_type_ids = text.get('token_type_ids', None)
            attention_mask = text.get('attention_mask', None)
        else:
            input_ids = text.input_ids
            token_type_ids = getattr(text, 'token_type_ids', None)
            attention_mask = getattr(text, 'attention_mask', None)

        return input_ids, token_type_ids, attention_mask


    def forward(
            self,
            text,   # B
            image,  # B C H W D
            device,
            return_loss = True,
            is_condition = False,
            modal_embedding = False,
            phrases_bboxes_mask=None,
            phrase_text_mask=None,
            modal_indexs = None,
            return_latents=True,  # 🟢 显式声明
            **kwargs
    ):
        if is_condition:

            
            B, device = text.input_ids.shape[0], device

            # derive text mask

            text_mask =text.attention_mask

            # get encoded text

            text_args = (text.input_ids,text.attention_mask)

            if not self.text_encode_without_mask:
                text_args = (*text_args, text_mask)

            text_embeddings = self.text_transformer(text.input_ids, attention_mask = text.attention_mask )
            enc_text = text_embeddings[0]   # B 512(token_len) 768
            

            if image.ndim == 3:                 # [H, W, ?] 的情况先 unsqueeze
                image = image.unsqueeze(0)
            if image.ndim == 4:                 # B×C×H×W
                if image.shape[1] != 1:
                    image = image.mean(dim=1, keepdim=True)
            elif image.ndim == 5:               # B×C×D×H×W
                if image.shape[1] != 1:
                    image = image.mean(dim=1, keepdim=True)
            else:
                raise ValueError(...)

            


    
            # NOTE
            modal_indexs = modal_indexs.view(-1)  # (B * local_B) 重复之后，condition就能和uncon的情况一样了
            enc_image = self.visual_transformer(image, return_encoded_tokens=True,modal_embedding=modal_embedding,modal_indexs=modal_indexs,is_condition=True)   # B*N 24 24 24 512(vis_dim)


            image_embeds = rearrange(enc_image, 'b t d -> t b d')  
            

            text_embeds = enc_text[:,0,:] # B 768

            text_latents = self.to_text_latent(text_embeds).unsqueeze(0) # 1 B*N 512
            image_latents = image_embeds
            t = image_latents.shape[0]
            temp_pos_embedding = self.pos_embedding[:,:t, :]  # 1 t 512

            pos = repeat(temp_pos_embedding, 'one d dim -> d (one bn) dim', bn=image_latents.shape[1]).to(device)  # 24*24, B*N, 512
            
            fused_latents, _ = self.fusion_module(text_latents, image_latents, pos=pos) # 1 B*N 512


            fused_latents = fused_latents[0]
            fused_latents = self.to_fused_latent(fused_latents) # B*N 512
            
            norms = torch.norm(fused_latents, dim=-1)
            # print('fused_latents范数:', '最小值:', norms.min(), '最大值:', norms.max(), '存在零:', (norms == 0).any())
            if (norms < 1e-6).any():
                print('警告：检测到极小范数')
                # 可选：添加微小噪声避免零范数问题
                fused_latents = fused_latents + torch.randn_like(fused_latents) * 1e-8
            if torch.isnan(fused_latents).any() or torch.isinf(fused_latents).any():
                print('NaN/Inf in vec **before** l2norm')
            
            if fused_latents.isnan().any() or fused_latents.isinf().any():
                print('NaN/Inf in vec **after** l2norm')
            if (fused_latents == 0).any():
                print('Zero in vec **after** l2norm')

            # NOTE
            
            if return_latents:
                return text_latents, image_latents, fused_latents, self.temperature.exp()

            # ——保持与 else 分支一致 ——
            
            # 全局文本/图像向量
            text_latents_global = l2norm(self.to_text_latent(enc_text[:, 0, :]))     # B x D_latent
            image_global       = l2norm(fused_latents.mean(dim=1))                   # B x D_latent

            # 双曲投影与全局蕴含
            #_curv       = self.curv.exp()
            with torch.cuda.amp.autocast(enabled=False):
                _curv = self.get_curv()
            text_feats  = self.project_txt(text_latents_global)
            image_feats = self.project_img(image_global)
            _angle      = L.oxy_angle(text_feats, image_feats, _curv)
            _aperture   = L.half_aperture(text_feats, _curv)
            text_image_entailment_loss = torch.clamp(_angle - self.global_aperture_thresh * _aperture, min=0).mean()

            # 温度缩放（与仓库版相同）
            
            #改curv1
            #self.logit_scale.data = torch.clamp(self.logit_scale.data, max=math.log(100.0))
            with torch.no_grad():
                self.logit_scale.clamp_(max=math.log(100.0))
            #改curv2
            
            scale = self.logit_scale.exp()

            # B×B logits（-双曲距离），对称 CE
            targets     = torch.arange(image_global.size(0), device=image_global.device)
            image_logits = -L.pairwise_dist(image_feats, text_feats, _curv)  # [B, B]
            text_logits  = -L.pairwise_dist(text_feats, image_feats, _curv)  # [B, B]
            # ABLATION: 去掉 image→text CE，只保留 text→image CE
            i2t_ce_loss = F.cross_entropy(scale * image_logits, targets)
            t2i_ce_loss = F.cross_entropy(scale * text_logits, targets)
            ce_loss = t2i_ce_loss
            # ce_loss = 0.5 * (i2t_ce_loss + t2i_ce_loss)

            loss = ce_loss + self.entail_weight * text_image_entailment_loss

            entailment_loss = text_image_entailment_loss

            return loss, ce_loss, entailment_loss, text_image_entailment_loss
        else:
            #text_mask：文本的注意力掩码（布尔矩阵），形状是 B x seq_len。
            #text_args：传入文本编码器的集合，包括输入的 token ids 和注意力掩码。
            #text_embeddings：文本编码器的输出，包含每个 token 的向量表示。形状是 B x seq_len x dim_text。
            #enc_text：文本编码器的最终输出，通常是取 [CLS] token 的向量表示。形状是 B x seq_len x dim_text。
            #text_latents：把整句输出文本压成一个固定长度的向量，代表这句话的整体语义。形状是 B x 512。
            #text_embeds: 文本的全局向量表示，取自文本编码器的 [CLS] token 输出。形状是 B x dim_text。

            #local_B： 单个图像的局部块数量（local_B）。
            #modal_indexs：表示每个图像块所属的模态索引，形状是 B x local_B。
            #enc_image：图像编码器的输出，包含每个图像块的向量表示。形状是 (B * local_B) x dim_image。
            #image_embeds：把所有图像块的向量表示重新排列成 (HWD) x (B * local_B) x dim_image 的形状。
            #image_latents：图像块向量，这些局部向量按顺序排成一串，形状是 (HWD) x (B * local_B) x 512。
            #box_image_feats_eu：池化后与文本短语对应的图像块特征，形状是 B x K x D_latent。
            #box_text_feats_eu：池化后与图像区域对应的文本短语特征，形状是 B x K x D_latent。
            

            input_ids, token_type_ids, attention_mask = self._unpack_text_batch(text)
            B = input_ids.size(0)


            with torch.cuda.amp.autocast(enabled=False):
                _curv = self.get_curv()

            if image.ndim == 3:                 # [H, W, ?] 的情况先 unsqueeze
                image = image.unsqueeze(0)
            if image.ndim == 4:                 # B×C×H×W
                if image.shape[1] != 1:
                    image = image.mean(dim=1, keepdim=True)
            elif image.ndim == 5:               # B×C×D×H×W
                if image.shape[1] != 1:
                    image = image.mean(dim=1, keepdim=True)
            else:
                raise ValueError(...)



            # ====== 文本编码 ======
            text_embeddings = self.text_transformer(
                input_ids=input_ids,
                token_type_ids=token_type_ids,
                attention_mask=attention_mask
            )
            enc_text = text_embeddings[0]  # B×T_txt×dim_text
            
            #text_latents_global = self.to_text_latent(enc_text[:, 0, :])
            text_latents_global = enc_text[:, 0, :]      #  不再过 Linear
            
            modality_dict = {'CT': 0, 'CXR': 1}
            modal_indexs = torch.tensor([modality_dict['CXR']], dtype=torch.long, device=device)

            modal_indexs = modal_indexs.view(-1)  # (B * local_B) 重复之后，condition就能和uncon的情况一样了

            # ====== 图像编码（双路）======
            # Patch 路：box / cross entailment（保留空间 token 网格）
            enc_image_patch = self.visual_transformer(
                image,
                return_encoded_tokens=True,
                modal_embedding=modal_embedding,
                modal_indexs=None,
                is_condition=True,
            )  # [B, T_img, D]

            # Global 路：global CE + global entailment（enc_cls_transformer 整图 CLS）
            enc_image_global = self.visual_transformer(
                image,
                return_encoded_tokens=True,
                modal_embedding=modal_embedding,
                modal_indexs=None,
                is_condition=False,
            )  # [B, D]

            if self.use_all_token_embeds:
                text_embeds = enc_text[:, 1:] if self.text_has_cls_token else enc_text
            else:
                text_embeds = enc_text[:, :] if enc_text.ndim == 3 else enc_text

            text_embeds = text_embeds[:, 0, :]
            text_latents = l2norm(text_embeds)

            if enc_image_global.dim() == 2:
                image_latents = l2norm(enc_image_global)
            elif enc_image_global.dim() == 3:
                image_latents = l2norm(enc_image_global[:, 0, :])
            else:
                raise ValueError(
                    f"Unexpected global visual feature shape: {enc_image_global.shape}"
                )

            image_global = image_latents

            word_feats = enc_text
            patch_feats = enc_image_patch

            if phrase_text_mask is not None:
                if phrase_text_mask.dim() == 2:  # B×T_txt
                    phrase_text_mask = phrase_text_mask.unsqueeze(1)  # B×1×T_txt
                box_text_feats_eu = self.mask_average_pool(word_feats, phrase_text_mask)  # B×K×dim_latent

            #patch_feats = rearrange(patch_feats_bn, '(b n) t d -> b n t d', n=local_B).mean(dim=1)  # B×T_img×dim_latent
            if phrases_bboxes_mask is not None:
                if phrases_bboxes_mask.dim() == 2:
                    phrases_bboxes_mask = phrases_bboxes_mask.unsqueeze(1)  # B×1×T_img
                box_image_feats_eu = self.mask_average_pool(patch_feats, phrases_bboxes_mask)

            # 双曲投影 + 蕴含损失计算（仅当提供 box mask 时）
            if phrase_text_mask is not None and phrases_bboxes_mask is not None:
                text_feats = self.project_txt(text_latents_global) # B×D_hyp
                #image_global = patch_feats.mean(dim=1)          # B×D_latent
                image_feats = self.project_img(image_global)                # B×D_hyp
                box_text_feats = self.project_txt(box_text_feats_eu)        # B×K×D_hyp
                box_image_feats = self.project_img(box_image_feats_eu)      # B×K×D_hyp
                # 后续四项蕴含损失逻辑保持不变



            if not return_latents:
                return text_latents, image_latents, self.temperature.exp()


            #celoss用，防止被normalize拉爆
            image_logits = -L.pairwise_dist(image_feats, text_feats, _curv) # [B, B]
            text_logits = -L.pairwise_dist(text_feats, image_feats, _curv) # [B, B]

            

            #四项蕴含损失
            _angle = L.oxy_angle(text_feats, image_feats, _curv)

            _aperture = L.half_aperture(text_feats, _curv)

            _box_angle = L.oxy_angle(box_text_feats, box_image_feats, _curv)

            _box_aperture = L.half_aperture(box_text_feats, _curv)

            _cross_image_angle = L.oxy_angle(box_image_feats, image_feats.unsqueeze(1), _curv)

            _box_image_aperture = L.half_aperture(box_image_feats, _curv)

            _cross_text_angle = L.oxy_angle(box_text_feats, text_feats.unsqueeze(1), _curv)

            _box_text_aperture = L.half_aperture(box_text_feats, _curv)




            #全文-全图蕴含
            text_image_entailment_loss = torch.clamp(_angle - self.global_aperture_thresh * _aperture, min=0).mean()

            box_text_image_entailment_loss = torch.clamp(_box_angle - self.global_aperture_thresh * _box_aperture, min=0).mean()

            cross_image_entailment_loss = torch.clamp(_cross_image_angle - self.local_aperture_thresh * _box_image_aperture, min=0).mean()
    
            cross_text_entailment_loss = torch.clamp(_cross_text_angle - self.local_aperture_thresh * _box_text_aperture, min=0).mean()

            entailment_loss = 0.5 * (
                text_image_entailment_loss
                + box_text_image_entailment_loss
                + cross_image_entailment_loss
                + cross_text_entailment_loss
            )



            
            # 3) 温度缩放
            
            self.logit_scale.data = torch.clamp(self.logit_scale.data, max=4.6052)


            targets = torch.arange(image_global.size(0), device=image_global.device)
            scale = self.logit_scale.exp()

            # ABLATION: 去掉 image→text CE，只保留 text→image CE
            i2t_ce_loss = F.cross_entropy(scale * image_logits, targets)
            t2i_ce_loss = F.cross_entropy(scale * text_logits, targets)
            ce_loss = t2i_ce_loss
            # ce_loss = 0.5 * (i2t_ce_loss + t2i_ce_loss)

            loss = ce_loss + self.entail_weight * entailment_loss

            return loss, ce_loss, entailment_loss, text_image_entailment_loss
                   


if __name__ == '__main__':
    import os
    import argparse
    import torch

    from transformer_maskgit.ctvit import CTViT
    from transformers import AutoTokenizer, AutoModel, AutoConfig

    # --------------------------------------------------------
    # Args
    # --------------------------------------------------------
    parser = argparse.ArgumentParser(description='Debug RADIR with batch')

    parser.add_argument('--batch', type=int, default=8)
    parser.add_argument('--channels', type=int, default=1)
    parser.add_argument('--height', type=int, default=224)
    parser.add_argument('--width', type=int, default=224)
    parser.add_argument('--modal_embedding', action='store_true')

    # Dataset paths
    parser.add_argument(
        '--dataset_root',
        type=str,
        default=r"/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset"
    )
    parser.add_argument(
        '--matched_dir',
        type=str,
        default=r"/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/processed/matched"
    )
    parser.add_argument(
        '--caption_file',
        type=str,
        default="/mnt/data-3/users/zhangzeyu/HMClip/RadIR/dataset/all.jsonl"   # 相对 dataset_root
    )

    # Local HF model path
    parser.add_argument(
        '--model_path',
        type=str,
        default=r"/mnt/data-3/users/zhangzeyu/HMClip/RadIR/hf_models/microsoft/BiomedVLP-CXR-BERT-specialized"
    )

    args = parser.parse_args()

    # --------------------------------------------------------
    # Device
    # --------------------------------------------------------
    if not torch.cuda.is_available():
        raise RuntimeError('需要 GPU 环境（CUDA 可用）。')

    device = torch.device('cuda')

    # --------------------------------------------------------
    # Debug HF loading
    # --------------------------------------------------------
    from transformers.utils import cached_file
    try:
        path = cached_file(args.model_path, "tokenizer_config.json", local_files_only=True)
        print(">>> HF tokenizer_config.json loaded from:", path)
    except Exception as e:
        print(">>> Failed to load tokenizer_config.json:", e)

    print(">>> Checking config loading...")
    cfg = AutoConfig.from_pretrained(
        args.model_path,
        trust_remote_code=True,
        local_files_only=True
    )
    print(">>> model_type:", cfg.model_type)
    print(">>> config path:", cfg._name_or_path)

    # --------------------------------------------------------
    # Load text encoder & tokenizer
    # --------------------------------------------------------
    tokenizer = AutoTokenizer.from_pretrained(
        args.model_path,
        trust_remote_code=True,
        local_files_only=True
    )

    text_model = AutoModel.from_pretrained(
        args.model_path,
        trust_remote_code=True,
        local_files_only=True
    ).to(device)

    # --------------------------------------------------------
    # Load image encoder (CTViT)
    # --------------------------------------------------------
    image_encoder = CTViT(
        dim=512,
        codebook_size=8192,
        image_size=args.height,
        patch_size=16,
        temporal_patch_size=10,
        spatial_depth=8,
        temporal_depth=6,
        cls_depth=4,
        dim_head=32,
        heads=8,
        channels=args.channels
    ).to(device)

    # --------------------------------------------------------
    # Build RADIR
    # --------------------------------------------------------
    from radir.RadIR import RADIR

    clip = RADIR(
        tokenizer=tokenizer,
        image_encoder=image_encoder,
        text_encoder=text_model,
        dim_text=768,
        dim_image=512,
        dim_latent=512,
        use_mlm=False,
        use_all_token_embeds=False
    ).to(device)

    # --------------------------------------------------------
    # Build dataset
    # --------------------------------------------------------
    from radir.data.hmclip_mimic_box import (
        HyCoClipMimicBoxDataset_JPG,
        collate_hmclip_box
    )
    from torch.utils.data import DataLoader

    dataset = HyCoClipMimicBoxDataset_JPG(
        split="train",
        tokenizer=tokenizer,
        dataset_root=args.dataset_root,
        cig_matched_dir=args.matched_dir,
        caption_file=args.caption_file,
        transform=None,     # 使用 Dataset 内置 2D + MONAI
        max_words=128,
        imsize=224
    )

    loader = DataLoader(
        dataset,
        batch_size=args.batch,
        shuffle=True,
        num_workers=0,
        collate_fn=collate_hmclip_box,
    )

    batch = next(iter(loader))

    imgs = batch["imgs"].to(device)  # [B, 1, 224, 224]

    enc_text = {
        "input_ids": batch["caption_ids"].to(device),
        "token_type_ids": batch["token_type_ids"].to(device),
        "attention_mask": batch["attention_mask"].to(device),
    }

    phrase_text_mask = batch["phrase_text_mask"].to(device)
    phrases_bboxes_mask = batch["phrases_bboxes_mask"].to(device)

    print(f">>> Loaded batch:")
    print(f"    imgs: {imgs.shape}")
    print(f"    input_ids: {enc_text['input_ids'].shape}")
    print(f"    phrase_text_mask: {phrase_text_mask.shape}")
    print(f"    phrases_bboxes_mask: {phrases_bboxes_mask.shape}")

    # --------------------------------------------------------
    # Estimate visual token length T_img
    # --------------------------------------------------------
    with torch.no_grad():
        enc_image_tokens = clip.visual_transformer(
            imgs,
            return_encoded_tokens=True,
            modal_embedding=args.modal_embedding,
            modal_indexs=None,
            is_condition=False
        )
        T_img = enc_image_tokens.shape[1]
        print(">>> T_img =", T_img)

    # --------------------------------------------------------
    # Adjust masks
    # --------------------------------------------------------
    phrase_text_mask = phrase_text_mask.unsqueeze(1)         # [B,1,T_txt]
    phrases_bboxes_mask = phrases_bboxes_mask.unsqueeze(1)   # [B,1,T_img?]

    cur_T_img = phrases_bboxes_mask.shape[-1]
    if cur_T_img != T_img:
        if cur_T_img > T_img:
            phrases_bboxes_mask = phrases_bboxes_mask[..., :T_img]
        else:
            pad = torch.ones(
                imgs.size(0), 1, T_img - cur_T_img,
                device=device
            )
            phrases_bboxes_mask = torch.cat(
                [phrases_bboxes_mask, pad],
                dim=-1
            )
        print(f"[DBG] adjusted phrases_bboxes_mask: {cur_T_img} -> {T_img}")

    # --------------------------------------------------------
    # Modality index (CXR = 1)
    # --------------------------------------------------------
    modality_dict = {'CT': 0, 'CXR': 1}
    modal_indexs = torch.tensor(
        [modality_dict['CXR']],
        dtype=torch.long,
        device=device
    )

    # --------------------------------------------------------
    # Forward
    # --------------------------------------------------------
    print(">>> Forward (return_loss=True)")

    losses = clip(
        text=enc_text,
        image=imgs,
        device=device,
        is_condition=False,
        return_latents=True,
        return_loss=True,
        modal_embedding=args.modal_embedding,
        phrase_text_mask=phrase_text_mask,
        phrases_bboxes_mask=phrases_bboxes_mask,
        modal_indexs=modal_indexs
    )

    loss, ce_loss, entailment_loss, text_image_entailment_loss = losses

    print(f"loss_total: {float(loss):.6f}")
    print(f"loss_ce(image<->text): {float(ce_loss):.6f}")
    print(f"loss_entail_global: {float(entailment_loss):.6f}")
    print(f"loss_entail_image_to_text: {float(text_image_entailment_loss):.6f}")