feature_extractor.py

import torch
import torch.nn.functional as F
from torchvision import transforms
import numpy as np
from scipy.stats import skew, kurtosis
import zlib

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

guesser_transform = transforms.Compose([
    transforms.CenterCrop(224),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                         std=[0.229, 0.224, 0.225])
])

def normalize_01_into_pm1(x):
    return x * 2 - 1

def extract_raw_log_probs(images, model, vae, is_var=False, condition=None, content_guesser=None, batch_size=16):
    """
    Extracts raw per-token log probabilities for a specific model.
    Returns:
        - RAR: log_p_cond, log_p_uncond each a list of 1 tensor: [(N, seq_len)]
        - VAR: log_p_cond, log_p_uncond each a list of num_scales tensors: [(N, scale_len_s)]
    """
    if not is_var:
        all_log_p_cond = []
        all_log_p_uncond = []
    else:
        all_log_p_cond = []
        all_log_p_uncond = []
        var_token_list = None

    with torch.no_grad():
        for i in range(0, len(images), batch_size):
            batch = images[i:i+batch_size].to(DEVICE)
            batch_cond = condition[i:i+batch_size].to(DEVICE) if condition is not None else None

            if condition is not None:
                predicted_labels = batch_cond
            else:
                assert content_guesser is not None, "Either condition or content_guesser must be provided"
                images_for_guesser = guesser_transform(batch)
                content_logits = content_guesser(images_for_guesser)
                predicted_labels = torch.argmax(content_logits, dim=1)

            if not is_var:  # RAR
                rar_tokens = vae.encode(batch)
                processed_condition = model.preprocess_condition(predicted_labels, cond_drop_prob=0.0)
                logits_cond, gt_labels = model(rar_tokens, processed_condition, return_labels=True)
                # logits_cond, gt_labels = model(rar_tokens, predicted_labels, return_labels=True)
                logits_uncond, _ = model(
                    rar_tokens, 
                    # torch.full_like(predicted_labels, fill_value=model.none_condition_id), 
                    model.get_none_condition(predicted_labels),
                    return_labels=True
                )

                loss_cond = F.cross_entropy(
                    logits_cond[:, :-1].reshape(-1, logits_cond.size(-1)),
                    gt_labels.view(-1), reduction='none'
                ).view(len(batch), -1)

                loss_uncond = F.cross_entropy(
                    logits_uncond[:, :-1].reshape(-1, logits_uncond.size(-1)),
                    gt_labels.view(-1), reduction='none'
                ).view(len(batch), -1)

                all_log_p_cond.append(-loss_cond.cpu())
                all_log_p_uncond.append(-loss_uncond.cpu())

            else:  # VAR
                batch = normalize_01_into_pm1(batch) # Normalize to [-1, 1] for VAR VAE
                
                var_token_list = vae.img_to_idxBl(batch)
                x_in = vae.quantize.idxBl_to_var_input(var_token_list)

                # Disable random condition dropping during inference
                original_drop_rate = model.cond_drop_rate
                model.cond_drop_rate = 0.0
                
                logits_cond = model(predicted_labels, x_in)
                logits_uncond = model(torch.full_like(predicted_labels, fill_value=model.num_classes), x_in)

                model.cond_drop_rate = original_drop_rate  # Restore original drop rate
                
                start_idx = 0
                for s, gt_tokens in enumerate(var_token_list):
                    scale_len = gt_tokens.shape[1]
                    end_idx = start_idx + scale_len

                    scale_logits_cond = logits_cond[:, start_idx:end_idx, :]
                    scale_logits_uncond = logits_uncond[:, start_idx:end_idx, :]

                    loss_cond = F.cross_entropy(
                        scale_logits_cond.reshape(-1, scale_logits_cond.size(-1)),
                        gt_tokens.view(-1), reduction='none'
                    ).view(len(batch), -1)

                    loss_uncond = F.cross_entropy(
                        scale_logits_uncond.reshape(-1, scale_logits_uncond.size(-1)),
                        gt_tokens.view(-1), reduction='none'
                    ).view(len(batch), -1)

                    all_log_p_cond.append(-loss_cond.cpu())
                    all_log_p_uncond.append(-loss_uncond.cpu())
                    start_idx = end_idx

            del batch, batch_cond
            torch.cuda.empty_cache()

    # Concatenate across batches
    if not is_var:
        # flat list of num_batches tensors, each (batch_size, seq_len)
        log_p_cond   = [torch.cat(all_log_p_cond,   dim=0)]  # list of 1 tensor: [(N, seq_len)]
        log_p_uncond = [torch.cat(all_log_p_uncond, dim=0)]  # list of 1 tensor: [(N, seq_len)]
    else:
        # flat list of (num_batches * num_scales) tensors, interleaved by scale
        num_scales   = len(var_token_list)  # 10, inferred from last batch
        log_p_cond   = [torch.cat(all_log_p_cond[s::num_scales],   dim=0) for s in range(num_scales)]
        log_p_uncond = [torch.cat(all_log_p_uncond[s::num_scales], dim=0) for s in range(num_scales)]

    return log_p_cond, log_p_uncond  # always a list of tensors: length 1 for RAR, length 10 for VAR


def extract_cfg_gap_mink_features(images, content_guesser, rar_models, var_models, rar_vae, var_vae, k_values=[10, 20, 30, 40, 50]):
    """
    Extracts CFG-Gap features for both RAR and VAR models:
    - For each model: NLL with condition vs NLL without condition, then take the gap.
    Returns: np.array of shape (batch, num_models)
    """
    features = []

    # --- RAR ---
    rar_names = ["rarb", "rarl", "rarxl", "rarxxl"]
    with torch.no_grad():
        images_for_guesser = guesser_transform(images.to(DEVICE))
        content_logits = content_guesser(images_for_guesser)
        predicted_labels = torch.argmax(content_logits, dim=1)
        rar_tokens = rar_vae.encode(images.to(DEVICE))
        for name in rar_names:
            model = rar_models[name]
            # With Condition
            processed_condition = model.preprocess_condition(predicted_labels, cond_drop_prob=0.0)
            logits_cond, gt_labels = model(rar_tokens, processed_condition, return_labels=True)
            logits_cond = logits_cond[:, :-1]
            # loss_cond = F.cross_entropy(
            #     logits_cond.reshape(-1, logits_cond.size(-1)),
            #     gt_labels.view(-1),
            #     reduction='none'
            # ).view(logits_cond.size(0), -1).mean(dim=1).cpu().numpy()
            loss_cond = F.cross_entropy(
                logits_cond.reshape(-1, logits_cond.size(-1)),
                gt_labels.view(-1),
                reduction='none'
            ).view(logits_cond.size(0), -1)
            
            # Without Condition
            # dummy_labels = torch.full_like(predicted_labels, fill_value=model.none_condition_id)
            dummy_labels = model.get_none_condition(predicted_labels)
            logits_uncond, _ = model(rar_tokens, dummy_labels, return_labels=True)
            logits_uncond = logits_uncond[:, :-1]
            # loss_uncond = F.cross_entropy(
            #     logits_uncond.reshape(-1, logits_uncond.size(-1)),
            #     gt_labels.view(-1),
            #     reduction='none'
            # ).view(logits_uncond.size(0), -1).mean(dim=1).cpu().numpy() 
            loss_uncond = F.cross_entropy(
                logits_uncond.reshape(-1, logits_uncond.size(-1)),
                gt_labels.view(-1),
                reduction='none'
            ).view(logits_uncond.size(0), -1)
            
            # # Convert NLL to raw probabilities p(x)
            # probs_cond = torch.exp(-loss_cond)
            # probs_uncond = torch.exp(-loss_uncond)
            # # Calculate CFG-Gap on probabilities: p(x|c) - p(x|c_null)
            # # The paper uses this difference *per-token* as input to MIAs. 
            # # Here, we take the mean across the sequence to get a single scalar feature per image.
            # cfg_gap = (probs_cond - probs_uncond) # (Batch, Seq_Len)
            # sorted_gap, _ = torch.sort(cfg_gap, dim=1, descending=True)
            # seq_len = sorted_gap.size(1)
            # selected_cfg_gap = []
            # for k in k_values:
            #     top_k_idx = max(1, int(seq_len * k / 100.0))  # Ensure at least 1 token is selected
            #     mink_gap_mean = sorted_gap[:, :top_k_idx].mean(dim=1).cpu().numpy()  # Mean of the top-k% smallest gaps
            #     selected_cfg_gap.append(mink_gap_mean)
            # features.append(np.stack(selected_cfg_gap, axis=1)) # (Batch, len(k_values))
            # # features.append(loss_uncond - loss_cond) # CFG-Gap
            
            cfg_gap = loss_uncond - loss_cond
            sorted_gap, _ = torch.sort(cfg_gap, dim=1, descending=True)
            seq_len = sorted_gap.size(1)
            selected_cfg_gap = []
            for k in k_values:
                top_k_idx = max(1, int(seq_len * k / 100.0))  # Ensure at least 1 token is selected
                mink_gap_mean = sorted_gap[:, :top_k_idx].mean(dim=1).cpu().numpy()  # Mean of the top-k% smallest gaps
                selected_cfg_gap.append(mink_gap_mean)
            features.append(np.stack(selected_cfg_gap, axis=1)) # (Batch, len(k_values))

    # --- VAR ---
    var_names = ["var16", "var20", "var24", "var30"]
    with torch.no_grad():
        images = normalize_01_into_pm1(images) # Normalize to [-1, 1] for VAR VAE
        var_tokens_list = var_vae.img_to_idxBl(images.to(DEVICE))
        x_in = var_vae.quantize.idxBl_to_var_input(var_tokens_list)
        var_gt = torch.cat(var_tokens_list, dim=1)
        for name in var_names:
            model = var_models[name]
            
            original_drop_rate = model.cond_drop_rate
            model.cond_drop_rate = 0.0  # Disable random condition dropping during inference
            
            # With Condition
            logits_cond = model(predicted_labels, x_in)
            # loss_cond = F.cross_entropy(
            #     logits_cond.view(-1, logits_cond.size(-1)),
            #     var_gt.view(-1),
            #     reduction='none'
            # ).view(logits_cond.size(0), -1).mean(dim=1).cpu().numpy()
            loss_cond = F.cross_entropy(
                logits_cond.view(-1, logits_cond.size(-1)),
                var_gt.view(-1),
                reduction='none'
            ).view(logits_cond.size(0), -1)
            
            # Without Condition
            dummy_labels = torch.full_like(predicted_labels, fill_value=model.num_classes)  # = 1000
            logits_uncond = model(dummy_labels, x_in)
            # loss_uncond = F.cross_entropy(
            #     logits_uncond.view(-1, logits_uncond.size(-1)),
            #     var_gt.view(-1),
            #     reduction='none'
            # ).view(logits_uncond.size(0), -1).mean(dim=1).cpu().numpy()
            loss_uncond = F.cross_entropy(
                logits_uncond.view(-1, logits_uncond.size(-1)),
                var_gt.view(-1),
                reduction='none'
            ).view(logits_uncond.size(0), -1)
            
            model.cond_drop_rate = original_drop_rate  # Restore original drop rate
            
            # # Convert NLL to raw probabilities p(x)
            # probs_cond = torch.exp(-loss_cond)
            # probs_uncond = torch.exp(-loss_uncond)
            # # Calculate CFG-Gap on probabilities: p(x|c) - p(x|c_null)
            # cfg_gap = (probs_cond - probs_uncond) # (Batch, Seq_Len)
            # sorted_gap, _ = torch.sort(cfg_gap, dim=1, descending=True)
            # seq_len = sorted_gap.size(1)
            # selected_cfg_gap = []
            # for k in k_values:
            #     top_k_idx = max(1, int(seq_len * k / 100.0))  # Ensure at least 1 token is selected
            #     mink_gap_mean = sorted_gap[:, :top_k_idx].mean(dim=1).cpu().numpy()  # Mean of the top-k% smallest gaps
            #     selected_cfg_gap.append(mink_gap_mean)
            # features.append(np.stack(selected_cfg_gap, axis=1)) # (Batch, len(k_values))
            # # features.append(loss_uncond - loss_cond) # CFG-Gap
            
            cfg_gap = loss_uncond - loss_cond
            sorted_gap, _ = torch.sort(cfg_gap, dim=1, descending=True)
            seq_len = sorted_gap.size(1)
            selected_cfg_gap = []
            for k in k_values:
                top_k_idx = max(1, int(seq_len * k / 100.0))  # Ensure at least 1 token is selected
                mink_gap_mean = sorted_gap[:, :top_k_idx].mean(dim=1).cpu().numpy()  # Mean of the top-k% smallest gaps
                selected_cfg_gap.append(mink_gap_mean)
            features.append(np.stack(selected_cfg_gap, axis=1)) # (Batch, len(k_values))
    
    # Concatenate features for all models: (batch, num_models=8)
    # return np.stack(features, axis=1)
    return np.concatenate(features, axis=1) # (Batch, 8*len(k_values))


def extract_nll_features(images, content_guesser, rar_models, var_models, rar_vae, var_vae, advanced=False):
    """
    Extracts the 8-dimensional NLL feature vector for a batch of images.
    Input: Images (N, 3, 256, 256)
    Output: Numpy Array (N, 8) -> [rarb, rarl, rarxl, rarxxl, var16, var20, var24, var30]
    """
    rar_features = []
    var_features = []
    
    # Ensure the feature vector always follows this EXACT order
    rar_names = ["rarb", "rarl", "rarxl", "rarxxl"]
    var_names = ["var16", "var20", "var24", "var30"]
    
    content_logits = content_guesser(images.to(DEVICE))
    predicted_labels = torch.argmax(content_logits, dim=1)
    # label_tensor = torch.tensor(labels, dtype=torch.long, device=DEVICE)
    
    with torch.no_grad():
        # Encode once for all 4 RAR models
        rar_tokens = rar_vae.encode(images.to(DEVICE))

        for name in rar_names:
            model = rar_models[name]
            
            processed_condition = model.preprocess_condition(predicted_labels, cond_drop_prob=0.0)
            
            # condition = torch.zeros(len(images), dtype=torch.long, device=DEVICE) # Dummy condition (class 0)
            logits, gt_labels = model(rar_tokens, processed_condition, return_labels=True)
            logits = logits[:, :-1] # Remove last prediction to match GT
            # Calculate Mean NLL per image
            loss = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)), 
                gt_labels.view(-1), 
                reduction='none'
            )
            # Reshape to (Batch, Seq_Len)
            loss_per_image = loss.view(len(images), -1).cpu().numpy()
            
            if advanced:
                stats = []
                for img_loss in loss_per_image:
                    stats.append(
                        [
                            img_loss.mean(),
                            img_loss.std(),
                            np.min(img_loss),
                            np.max(img_loss),
                            np.median(img_loss),
                            skew(img_loss),
                            kurtosis(img_loss)
                        ]
                    )
                rar_features.append(np.array(stats)) # (Batch, 7)
            else:
                # mean over sequence -> (Batch,)
                rar_features.append(loss_per_image.mean(axis=1)) # (Batch,)

        # Tokenize once for all 4 VAR models
        images = normalize_01_into_pm1(images) # Ensure images are in the expected range for tokenization
        
        var_tokens_list = var_vae.img_to_idxBl(images.to(DEVICE))
        x_in = var_vae.quantize.idxBl_to_var_input(var_tokens_list)
        var_gt = torch.cat(var_tokens_list, dim=1)
        
        for name in var_names:
            model = var_models[name]
            
            model.original_drop_rate = model.cond_drop_rate
            model.cond_drop_rate = 0.0  # Disable random condition dropping during inference
            
            # label_B = torch.zeros(len(images), dtype=torch.long, device=DEVICE)
            logits = model(predicted_labels, x_in) 
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)), 
                var_gt.view(-1), 
                reduction='none'
            )
            loss_per_image = loss.view(len(images), -1).cpu().numpy() # (Batch, Seq_Len)
            
            model.cond_drop_rate = model.original_drop_rate  # Restore original drop rate
            
            if advanced:
                stats = []
                for img_loss in loss_per_image:
                    stats.append(
                        [
                            img_loss.mean(),
                            img_loss.std(),
                            np.min(img_loss),
                            np.max(img_loss),
                            np.median(img_loss),
                            skew(img_loss),
                            kurtosis(img_loss)
                        ]
                    )
                var_features.append(np.array(stats)) # (Batch, 7)
            else:
                var_features.append(loss_per_image.mean(axis=1)) # (Batch,)
    
    # Normalization
    if advanced:
        rar_features = np.stack(rar_features, axis=1).reshape(len(images), -1) # (Batch, 28)
        var_features = np.stack(var_features, axis=1).reshape(len(images), -1) # (Batch, 28)
        return np.hstack([rar_features, var_features]) # (Batch, 56)
    else:
        rar_features = np.stack(rar_features, axis=1) # (Batch, 4)
        var_features = np.stack(var_features, axis=1) # (Batch, 4)
        return np.hstack([rar_features, var_features]) # (Batch, 8)
    # rar_features = (rar_features - rar_features.mean(axis=0)) / (rar_features.std(axis=0) + 1e-8)
    # var_features = (var_features - var_features.mean(axis=0)) / (var_features.std(axis=0) + 1e-8)
    
    # Stack a list of 8 arrays, each with shape (Batch_Size,), into (Batch_Size, 8), so a matrix where columns = [rarb, rarl, ..., var30]
    # Each value is the mean negative log-likelihood (NLL) of one image under one model:
    # return np.hstack([rar_features, var_features]) # (Batch, 8)


def extract_mia_features(images, content_guesser, rar_models, var_models, rar_vae, var_vae, use_cfg=True):
    """
    Aligned with `privacy_attacks_against_iars-main/src/attacks/` feature extraction.
    Computes per-model: 
    mean_loss; min_k_loss (k=10, 20, 30%); min_kplus_loss; zlib; above_mean; slope; hinge.
    Returns: np.array of shape (batch, num_features * num_models)
    """
    features = []
    rar_names = ["rarb", "rarl", "rarxl", "rarxxl"]
    var_names = ["var16", "var20", "var24", "var30"]

    def compute_all_features(logits, tokens, k_values=[10, 20, 30], max_entropies=[5.0, 7.0]):
        """
        Replicates `compute_all()` in `llm_mia.py`
        logits: (B, seq_len, vocab) — either raw or CFG-gap logits
        tokens: (B, seq_len) — ground truth token indices
        Returns: (B, num_features)
        """
        B, N, V = logits.shape
        

        # Token-level losses (positive values)
        token_losses = F.cross_entropy(
            logits.reshape(-1, V),
            tokens.reshape(-1).long(),
            reduction='none'
        ).reshape(B, N)  # (B, N)

        # Token-level log probs (negative values)
        log_probs = F.log_softmax(logits, dim=-1)
        token_logprobas = torch.gather(
            log_probs, 2, tokens.unsqueeze(2).long()
        )[..., 0]  # (B, N)

        # Probabilities for Min-K%++
        probas = F.softmax(logits, dim=-1)  # (B, N, V)

        # Mean loss
        mean_loss = token_losses.mean(dim=1)  # (B,)

        # Min-K% log probability (LOWEST k% of log probs = least confident)
        min_k_features = []
        for k in k_values:
            k_count = int(k * N / 100)
            if k_count > 0:
                # Take k% SMALLEST log probs (most negative = least confident)
                min_k_logprob = torch.topk(token_logprobas, k_count, dim=1, largest=False).values.mean(dim=1)
                min_k_features.append(-min_k_logprob)  # Negate to make it a "loss-like" feature

        # Min-K%++ (normalized by local mean/std of log probs)
        mu = (probas * log_probs).sum(dim=2)       # (B, N) expected log prob
        sigma = (probas * log_probs.pow(2)).sum(dim=2) - mu.pow(2)  # (B, N) variance
        normalized = (token_logprobas - mu) / (sigma + 1e-6)  # (B, N)
        minkplus_features = []
        for k in k_values:
            k_count = int(k * N / 100)
            if k_count > 0:
                # Most negative normalized scores = most surprising tokens
                minkplus = torch.topk(normalized, k_count, dim=1, largest=False).values.mean(dim=1)
                minkplus_features.append(-minkplus) # Negate for consistency

        # # Zlib ratio: mean_loss / zlib_entropy_per_token
        # # Approximated as loss relative to a flat baseline
        # zlib_entropy = torch.log(torch.tensor(V, dtype=torch.float32, device=logits.device))
        # zlib_ratio = mean_loss / zlib_entropy  # (B,)
        
        # Zlib ratio
        zlib_entropies = torch.tensor([
            len(zlib.compress(bytes(str(t.cpu().tolist()), "utf-8")))
            for t in tokens
        ], device=tokens.device, dtype=torch.float32)
        zlib_ratio = -token_logprobas.mean(dim=1) / zlib_entropies  # (B,)

        # # Above-mean loss
        # mean_per_img = token_losses.mean(dim=1, keepdim=True)
        # above_mean = (token_losses * (token_losses > mean_per_img).float()).sum(dim=1) / \
        #              (token_losses > mean_per_img).float().sum(dim=1).clamp(min=1)  # (B,)

        # # Slope (linear trend across sequence positions)
        # positions = torch.arange(N, device=logits.device, dtype=torch.float32)
        # positions = positions - positions.mean()
        # slope = (token_losses * positions).sum(dim=1) / (positions.pow(2).sum() + 1e-6)  # (B,)

        # Hinge loss
        token_logits = torch.gather(logits, 2, tokens.unsqueeze(2).long())[..., 0]  # (B, N)
        logits_copy = logits.clone()
        logits_copy.scatter_(2, tokens.unsqueeze(2).long(),
                             torch.full_like(token_logits.unsqueeze(2), float('-inf')))
        second_best = logits_copy.max(dim=2).values  # (B, N)
        hinge = -(token_logits - second_best).mean(dim=1)  # (B,)
        
        # === CAMIA Features (5 sub-features) ===
        # 1. Below mean log probability
        below_mean = (token_logprobas < token_logprobas.mean(dim=1, keepdim=True)).float().mean(dim=1)
        
        # 2. Below mean min_k_plus
        below_mean_minkplus = (normalized < normalized.mean(dim=1, keepdim=True)).float().mean(dim=1)
        
        # 3. Below previous mean (cumulative mean)
        cumsum = token_logprobas.cumsum(dim=1)
        indices = torch.arange(1, N + 1, device=token_logprobas.device).unsqueeze(0)
        prev_mean = cumsum / indices
        below_prev_mean = -(token_logprobas < prev_mean).float().mean(dim=1)
        
        # 4. Below previous mean min_k_plus
        cumsum_minkplus = normalized.cumsum(dim=1)
        prev_mean_minkplus = cumsum_minkplus / indices
        below_prev_mean_minkplus = -(normalized < prev_mean_minkplus).float().mean(dim=1)
        
        # 5. Slope
        x = torch.arange(N, dtype=torch.float32, device=token_logprobas.device)
        x_mean = x.mean()
        token_logprobas_mean = token_logprobas.mean(dim=1, keepdim=True)
        slope_num = ((x - x_mean) * (token_logprobas - token_logprobas_mean)).sum(dim=1)
        slope_den = ((x - x_mean) ** 2).sum()
        slope = slope_num / slope_den
        
        # === SURP Features (surprise-based) ===
        surp_features = []
        for k_pct in k_values:
            k_count = int(k_pct * N / 100)
            if k_count > 0:
                bound_k = torch.topk(-token_logprobas, k_count, dim=1).values.max(dim=1, keepdim=True).values
                
                for max_entropy in max_entropies:
                    # Entropy-based surp
                    mask_entropy = (-mu < max_entropy) & (token_logprobas < bound_k)
                    surp_k_entropy = (token_logprobas * mask_entropy).sum(dim=1) / mask_entropy.sum(dim=1).clamp(min=1e-6)
                    surp_features.append(-surp_k_entropy)
                    
        for k_pct in k_values:
            k_count = int(k_pct * N / 100)
            if k_count > 0:
                bound_k = torch.topk(-token_logprobas, k_count, dim=1).values.max(dim=1, keepdim=True).values
                
                for max_entropy in max_entropies:
                    # Counter-based surp
                    mask_counter = (-mu < max_entropy) & (token_logprobas < bound_k)
                    surp_counter_k = -mask_counter.float().mean(dim=1)
                    surp_features.append(surp_counter_k)

        # Stack all: (B, 3 + 3 + 3 + 5 + 3 * 2 * 2) = (B, 26)
        all_features = torch.stack(
            [mean_loss, zlib_ratio, hinge] + 
            min_k_features + 
            minkplus_features +
            [below_mean, below_mean_minkplus, below_prev_mean, below_prev_mean_minkplus, slope] +
            surp_features,
            dim=1
        )
        return all_features  # (B, 26)

    with torch.no_grad():
        images_for_guesser = guesser_transform(images.to(DEVICE))
        content_logits = content_guesser(images_for_guesser)
        predicted_labels = torch.argmax(content_logits, dim=1)

        # --- RAR ---
        rar_tokens = rar_vae.encode(images.to(DEVICE))
        for name in rar_names:
            model = rar_models[name]
            processed_condition = model.preprocess_condition(predicted_labels, cond_drop_prob=0.0)
            logits_cond, gt_labels = model(rar_tokens, processed_condition, return_labels=True)            
            # logits_cond, gt_labels = model(rar_tokens, predicted_labels, return_labels=True)
            logits_cond = logits_cond[:, :-1]  # (B, seq_len, vocab)
            # gt_labels: (B, seq_len)

            if use_cfg: # CFG gap in logit space
                logits_uncond, _ = model(
                    rar_tokens,
                    model.get_none_condition(predicted_labels), # This is correct
                    # torch.full_like(predicted_labels, fill_value=model.none_condition_id), # This is correct
                    # torch.full_like(predicted_labels, fill_value=model.num_classes), # This is a common alternative (using an out-of-vocab class as "unconditional")
                    return_labels=True
                )
                logits_uncond = logits_uncond[:, :-1]
                logits_input = logits_cond - logits_uncond
            else: # raw conditional only
                logits_input = logits_cond

            feats = compute_all_features(logits_input, gt_labels)
            features.append(feats.cpu().numpy())  # (B, 11)

        # --- VAR ---
        images = normalize_01_into_pm1(images) # Ensure images are in the expected range for tokenization
        
        var_tokens_list = var_vae.img_to_idxBl(images.to(DEVICE))
        x_in = var_vae.quantize.idxBl_to_var_input(var_tokens_list)
        var_gt = torch.cat(var_tokens_list, dim=1)  # (B, total_tokens)

        for name in var_names:
            model = var_models[name]
            
            # Disable random condition dropping during inference
            original_drop_rate = model.cond_drop_rate
            model.cond_drop_rate = 0.0
            
            logits_cond = model(predicted_labels, x_in)  # (B, total_tokens, vocab)

            if use_cfg:
                logits_uncond = model(
                    torch.full_like(predicted_labels, fill_value=model.num_classes),
                    x_in
                )
                logits_input = logits_cond - logits_uncond  # CFG gap
            else:
                logits_input = logits_cond
            
            model.cond_drop_rate = original_drop_rate  # Restore original drop rate

            feats = compute_all_features(logits_input, var_gt)
            features.append(feats.cpu().numpy())  # (B, 26)

    # (B, 26 * 8) = (B, 208)
    return np.concatenate(features, axis=1)