model/rayzer_official_v3.py

# Copyright (c) 2025 Haian Jin. Created for the LVSM project (ICLR 2025).
import copy
from datetime import datetime
import os
import math
import numpy as np
import torch
import torch.nn as nn
from easydict import EasyDict as edict
from skimage.metrics import structural_similarity
from einops.layers.torch import Rearrange
from einops import rearrange, repeat
import traceback
import torch.distributed as dist
from transformers import AutoImageProcessor, AutoModel
import cv2
from utils import camera_utils, data_utils
from .transformer import (
    QK_Norm_TransformerBlock,
    init_weights,
    _init_weights,
    _init_weights_layerwise,
)
from .loss import LossComputer_official
import torch.nn.functional as F
from utils.pe_utils_official import get_1d_sincos_pos_emb_from_grid, get_2d_sincos_pos_embed
from utils.pose_utils import rot6d2mat, quat2mat
from utils.metric_utils import compute_psnr, compute_lpips, compute_ssim
from utils.data_utils import create_video_from_frames
from utils.training_utils import bilinear_resize
from PIL import Image


def _gn_groups(out_ch: int, prefer: int = 32) -> int:
    g = math.gcd(out_ch, prefer)
    return g if g > 0 else 1


class Up2x(nn.Module):
    def __init__(self, in_ch, out_ch, gn_groups=32):
        super().__init__()
        g = _gn_groups(out_ch, gn_groups)

        # 2x upsample (bilinear) then change channels in_ch -> out_ch
        self.conv_up = nn.Conv2d(in_ch, out_ch, kernel_size=1)

        # refinement at target channels
        self.refine = nn.Sequential(
            nn.GroupNorm(g, out_ch),
            nn.GELU(),
            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
            nn.GroupNorm(g, out_ch),
            nn.GELU(),
        )

    def forward(self, x):
        # single upsample
        x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=False)
        # project channels to out_ch
        x = self.conv_up(x)
        # refine
        x = self.refine(x)
        return x


class MotionMaskPredictor(nn.Module):
    """
    Three-modality motion mask predictor: DINOv3 + Image Tokens + Camera Poses.

    Architecture:
        1. Feature Fusion: Combines semantic (DINO), appearance (tokens), and geometric (pose) cues
        2. Motion Reasoning Transformer: Processes fused features to understand motion patterns
        3. Improved DPT Decoder: Multi-scale upsampling with skip connections for crisp masks

    This is the M-step of the EM training loop: learns to predict motion masks from GT images.
    """

    def __init__(self, config, dinov3_backbone, dinov3_processor):
        super().__init__()
        self.config = config
        # Store DINOv3 as non-persistent reference (moves to device but not saved in state_dict)
        # We use object.__setattr__ to bypass nn.Module's __setattr__ which would save it
        object.__setattr__(self, "_dinov3_backbone_ref", dinov3_backbone)
        object.__setattr__(self, "_dinov3_processor_ref", dinov3_processor)
        self.use_grad_checkpoint = config.training.get("grad_checkpoint_every", 1) > 0

        # Manually register DINOv3's device placement hook (moves to device without saving)
        # This ensures it follows the parent model's device
        self._register_dinov3_device_hook()

        # Dimensions
        d_main = config.model.transformer.d  # Main transformer dim (e.g., 768)
        d_dino = dinov3_backbone.config.hidden_size  # DINOv3 dim (e.g., 1024 for ViT-L)
        d_fused = d_main  # Output fusion dimension

        # Store DINOv3 config
        self.patch_size_dino = getattr(dinov3_backbone.config, "patch_size", 16)
        self.num_register_tokens = getattr(dinov3_backbone.config, "num_register_tokens", 4)

        # Normalization buffers for DINOv3 (from processor)
        mean = torch.tensor(dinov3_processor.image_mean, dtype=torch.float32).view(1, 3, 1, 1)
        std = torch.tensor(dinov3_processor.image_std, dtype=torch.float32).view(1, 3, 1, 1)
        self.register_buffer("dino_norm_mean", mean)
        self.register_buffer("dino_norm_std", std)

        # 1. Feature Fusion Module
        # Project each modality to common dimension, then fuse
        self.dino_proj = nn.Sequential(
            nn.LayerNorm(d_dino, bias=False),
            nn.Linear(d_dino, d_fused, bias=True),
        )
        self.token_proj = nn.Sequential(
            nn.LayerNorm(d_main, bias=False),
            nn.Linear(d_main, d_fused, bias=True),
        )
        self.pose_proj = nn.Sequential(
            nn.LayerNorm(d_main, bias=False),  # Plücker embeddings already in d_main
            nn.Linear(d_main, d_fused, bias=True),
        )

        # Fusion via concatenation + MLP (3*d_fused -> d_fused)
        self.fusion_mlp = nn.Sequential(
            nn.LayerNorm(d_fused * 3, bias=False),
            nn.Linear(d_fused * 3, d_fused, bias=True),
            nn.GELU(),
            nn.Linear(d_fused, d_fused, bias=True),
        )

        use_qk_norm = config.model.transformer.get("use_qk_norm", False)
        d_head = d_fused // 12  # 12 heads for 768-dim
        self.motion_transformer = nn.ModuleList(
            [QK_Norm_TransformerBlock(d_fused, d_head, use_qk_norm=use_qk_norm) for _ in range(4)]
        )

        # 3. Improved DPT Head with Multi-Scale Features
        # Following original DPT: extract features at multiple scales and fuse
        self.dpt_readout = nn.Sequential(
            nn.LayerNorm(d_fused, bias=False),
            nn.Linear(d_fused, d_fused, bias=True),
        )

        # Multi-scale feature extraction (3 scales)
        # Scale 1: 16x16 (patch level)
        self.scale1_refine = nn.Sequential(
            nn.Conv2d(d_fused, d_fused // 2, kernel_size=3, padding=1),
            nn.GroupNorm(num_groups=32, num_channels=d_fused // 2),
            nn.GELU(),
            nn.Conv2d(d_fused // 2, d_fused // 2, kernel_size=3, padding=1),
        )
        self.up2 = Up2x(d_fused // 2, d_fused // 4)  # 16 -> 32
        self.up3 = Up2x(d_fused // 4, d_fused // 8)  # 32 -> 64
        self.up4 = Up2x(d_fused // 8, d_fused // 16)  # 64 -> 128
        self.up5 = Up2x(d_fused // 16, d_fused // 32)  # 128 -> 256

        self.scale2_refine = nn.Sequential(
            nn.Conv2d(d_fused // 4, d_fused // 4, kernel_size=3, padding=1),
            nn.GroupNorm(num_groups=32, num_channels=d_fused // 4),
            nn.GELU(),
        )
        self.scale3_refine = nn.Sequential(
            nn.Conv2d(d_fused // 8, d_fused // 8, kernel_size=3, padding=1),
            nn.GroupNorm(num_groups=32, num_channels=d_fused // 8),
            nn.GELU(),
        )
        self.scale4_refine = nn.Sequential(
            nn.Conv2d(d_fused // 16, d_fused // 16, kernel_size=3, padding=1),
            nn.GroupNorm(num_groups=16, num_channels=d_fused // 16),
            nn.GELU(),
        )
        self.scale5_refine = nn.Sequential(
            nn.Conv2d(d_fused // 32, d_fused // 32, kernel_size=3, padding=1),
            nn.GroupNorm(num_groups=8, num_channels=d_fused // 32),
            nn.GELU(),
        )

        # Final prediction head
        self.head = nn.Sequential(
            nn.Conv2d(d_fused // 32, d_fused // 32, kernel_size=3, padding=1),
            nn.GELU(),
            nn.Conv2d(d_fused // 32, 1, kernel_size=1, bias=True),
        )

        # Initialize weights
        for m in [
            self.dino_proj,
            self.token_proj,
            self.pose_proj,
            self.fusion_mlp,
            self.motion_transformer,
            self.scale1_refine,
            self.scale2_refine,
            self.scale3_refine,
            self.scale4_refine,
            self.scale5_refine,
            self.head,
        ]:
            m.apply(_init_weights)
        self.up2.apply(_init_weights)
        self.up3.apply(_init_weights)
        self.up4.apply(_init_weights)
        self.up5.apply(_init_weights)

    def _register_dinov3_device_hook(self):
        """Register a hook to move DINOv3 backbone when parent model moves to a device."""

        def move_dinov3_to_device(module, inputs):
            if hasattr(module, "_dinov3_backbone_ref"):
                backbone = object.__getattribute__(module, "_dinov3_backbone_ref")
                if backbone is not None:
                    device = next(module.parameters()).device
                    backbone.to(device)

        # Register the hook
        self.register_forward_pre_hook(move_dinov3_to_device)

    @torch.no_grad()
    def _extract_dino_features(self, images_224):
        """
        Args:
            images_224: [B, 3, 224, 224] in [0,1], already resized to 224
        Returns:
            [B, 16*16, D_dino] patch tokens
        """
        img_norm = (images_224 - self.dino_norm_mean) / self.dino_norm_std
        outputs = self._dinov3_backbone_ref(img_norm, output_hidden_states=False)
        skip = 1 + self.num_register_tokens
        return outputs.last_hidden_state[:, skip:, :]  # [B, 256, D_dino] => 16×16

    def forward(self, target_images, target_image_tokens, target_plucker_emb):
        """
        Predict motion masks from three modalities.

        Args:
            target_images: [B*V, 3, H, W] GT target images in [0, 1]
            target_image_tokens: [B*V, N_main, D_main] encoded image tokens from main transformer
            target_plucker_emb: [B*V, N_main, D_main] tokenized Plücker embeddings (camera poses)

        Returns:
            mask_logits: [B*V, 1, H_out, W_out] motion mask logits (H_out=256, W_out=256)
        """
        BV, _, H, W = target_images.shape
        N_main = target_image_tokens.shape[1]
        Hm = Wm = int(N_main**0.5)  # e.g., 16x16
        assert Hm * Wm == N_main, f"Hm * Wm ({Hm}*{Wm}) != N_main ({N_main})"

        # Use full resolution (256x256) for DINOv3 to match main model's grid
        dino_in = target_images  # Already 256x256, no resize needed
        dino_features = self._extract_dino_features(dino_in)  # [BV, 256, D_dino]

        Nd = dino_features.shape[1]
        Hd = Wd = int(Nd**0.5)

        # 2. Align spatial dimensions (interpolate DINO to main grid size)
        if (Hd != Hm) or (Wd != Wm):
            dino_spatial = rearrange(dino_features, "bv (h w) d -> bv d h w", h=Hd, w=Wd)
            dino_spatial = bilinear_resize(dino_spatial, size=(Hm, Wm))
            dino_features = rearrange(dino_spatial, "bv d h w -> bv (h w) d")

        dino_proj = self.dino_proj(dino_features)
        token_proj = self.token_proj(target_image_tokens)
        pose_proj = self.pose_proj(target_plucker_emb)
        fused = torch.cat([dino_proj, token_proj, pose_proj], dim=-1)
        fused = self.fusion_mlp(fused)

        # 5. Motion reasoning via transformer (with gradient checkpointing)
        x = fused
        if self.use_grad_checkpoint and self.training:
            for layer in self.motion_transformer:
                x = torch.utils.checkpoint.checkpoint(layer, x, use_reentrant=False)
        else:
            for layer in self.motion_transformer:
                x = layer(x)

        # readout -> spatial 16×16
        x = self.dpt_readout(x)  # [BV, 256, D_fused]
        x = rearrange(x, "bv (h w) d -> bv d h w", h=Hm, w=Wm)  # [BV, D_fused, 16, 16]

        # decoder 16→32→64→128→256
        x1 = self.scale1_refine(x)  # 16×16, C/2
        x2 = self.up2(x1)
        x2 = self.scale2_refine(x2)
        x3 = self.up3(x2)
        x3 = self.scale3_refine(x3)
        x4 = self.up4(x3)
        x4 = self.scale4_refine(x4)
        x5 = self.up5(x4)
        x5 = self.scale5_refine(x5)

        mask_logits = self.head(x5)
        return mask_logits


class DinoV3UncertaintyPseudoLabelMaker(nn.Module):
    """
    DINOv3-based pseudo-label generator for motion/uncertainty masks.
    Uses min(DINO cosine dissimilarity, SSIM dissimilarity) with percentile-based soft masking.

    This is the improved version from our Jupyter notebook analysis:
    - Combines DINOv3 features + SSIM for robustness
    - Generates soft masks (continuous probabilities) not binary
    - Uses percentile-based normalization (P75, P85, P90)
    - NO gradient computation (frozen DINOv3)
    - DDP-safe: device placement handled automatically by parent module
    """

    def __init__(
        self,
        dinov3_backbone,
        dinov3_processor,
        percentile_threshold=75,
        use_ssim=True,
        target_size=256,
        use_coseg_binary=False,
        coseg_k=64,
        coseg_consistency_min_frames=4,
        coseg_frame_saliency_quantile=0.75,
        coseg_cluster_top_percent=0.05,
        coseg_morph_kernel=3,
        coseg_morph_iters=1,
        coseg_min_component_area_ratio=0.0025,
        coseg_grabcut_kernel=7,
        w_dino=0.5,
        w_ssim=0.5,
    ):
        super().__init__()
        self.percentile_threshold = percentile_threshold
        self.use_ssim = use_ssim
        self.target_size = target_size
        self.use_coseg_binary = use_coseg_binary
        # Online co-seg params (lightweight)
        self.coseg_k = coseg_k
        self.coseg_consistency_min_frames = coseg_consistency_min_frames
        self.coseg_frame_saliency_quantile = coseg_frame_saliency_quantile
        self.coseg_cluster_top_percent = coseg_cluster_top_percent
        self.coseg_morph_kernel = coseg_morph_kernel
        self.coseg_morph_iters = coseg_morph_iters
        self.coseg_min_component_area_ratio = coseg_min_component_area_ratio
        self.coseg_grabcut_kernel = coseg_grabcut_kernel
        self.w_dino = w_dino
        self.w_ssim = w_ssim

        # Store DINOv3 backbone as non-persistent reference (moves to device but not saved in state_dict)
        object.__setattr__(self, "backbone", dinov3_backbone)
        object.__setattr__(self, "processor", dinov3_processor)

        # Manually register DINOv3's device placement hook
        self._register_dinov3_device_hook()

        # DINOv3 config - read actual patch_size from model
        self.patch_size = getattr(dinov3_backbone.config, "patch_size", 16)
        self.num_register_tokens = getattr(dinov3_backbone.config, "num_register_tokens", 4)
        self.embed_dim = dinov3_backbone.config.hidden_size

        # Normalization parameters (as buffers for automatic device placement)
        mean = torch.tensor(dinov3_processor.image_mean, dtype=torch.float32).view(1, 3, 1, 1)
        std = torch.tensor(dinov3_processor.image_std, dtype=torch.float32).view(1, 3, 1, 1)
        self.register_buffer("img_norm_mean", mean)
        self.register_buffer("img_norm_std", std)

    @property
    def device(self):
        """Get current device from registered buffers (DDP-safe)."""
        return self.img_norm_mean.device

    def _register_dinov3_device_hook(self):
        """Register a hook to move DINOv3 backbone when parent model moves to a device."""

        def move_dinov3_to_device(module, inputs):
            if hasattr(module, "backbone"):
                backbone = object.__getattribute__(module, "backbone")
                if backbone is not None:
                    device = module.device
                    backbone.to(device)

        # Register the hook
        self.register_forward_pre_hook(move_dinov3_to_device)

    @torch.no_grad()
    def _extract_patch_features(self, img_tensor):
        """
        Extract patch-level features from DINOv3.

        Args:
            img_tensor: [B, 3, H, W] in range [0, 1]

        Returns:
            features: [B, D, H_patches, W_patches] L2-normalized patch features
        """
        # Get original image dimensions
        B, _, H, W = img_tensor.shape
        
        # Ensure backbone is on the same device as input
        device = img_tensor.device
        if hasattr(self, "backbone"):
            backbone = object.__getattribute__(self, "backbone")
            if backbone is not None and next(backbone.parameters()).device != device:
                backbone.to(device)
        
        # Normalize using DINOv3's preprocessing
        img_normalized = (
            img_tensor - self.img_norm_mean.to(img_tensor.device)
        ) / self.img_norm_std.to(img_tensor.device)

        # Forward pass
        outputs = self.backbone(img_normalized, output_hidden_states=False)

        # Extract patch tokens (skip CLS + register tokens)
        skip_tokens = 1 + self.num_register_tokens  # CLS + 4 register tokens
        patch_tokens = outputs.last_hidden_state[:, skip_tokens:, :]  # [B, N_patches, D]

        B, N, D = patch_tokens.shape
        
        # Infer patch grid dimensions from N
        # The model might internally resize images, so we can't rely on input dimensions
        
        # First, try assuming square output
        sqrt_N = int(np.sqrt(N))
        if sqrt_N * sqrt_N == N:
            # Perfect square - use it
            H_patches = W_patches = sqrt_N
        else:
            # Not square - need to factorize N into H_patches × W_patches
            # Try to match the input image aspect ratio
            aspect_ratio = W / H
            
            # Find factors of N that are closest to the aspect ratio
            best_h = best_w = None
            min_diff = float('inf')
            
            for h in range(1, int(np.sqrt(N)) + 1):
                if N % h == 0:
                    w = N // h
                    ratio_diff = abs(w / h - aspect_ratio)
                    if ratio_diff < min_diff:
                        min_diff = ratio_diff
                        best_h, best_w = h, w
            
            if best_h is None:
                raise ValueError(f"Cannot factorize {N} patches (N must be composite or square)")
            
            H_patches, W_patches = best_h, best_w

        # Reshape to spatial grid and L2-normalize
        features = patch_tokens.transpose(1, 2).reshape(B, D, H_patches, W_patches)
        features = F.normalize(features, dim=1)  # L2 normalize along channel dim

        return features

    def _refine_with_grabcut(self, gt_img_rgb: np.ndarray, init_mask: np.ndarray) -> np.ndarray:
        """
        Refine binary mask using GrabCut (matches offline script functionality).

        Args:
            gt_img_rgb: [H, W, 3] uint8 RGB image
            init_mask: [H, W] uint8 {0, 1} initial mask

        Returns:
            refined_mask: [H, W] uint8 {0, 1} GrabCut-refined mask
        """
        if self.coseg_grabcut_kernel <= 0:
            return init_mask

        # Erode to get confident FG/BG seeds
        kernel = np.ones((self.coseg_grabcut_kernel, self.coseg_grabcut_kernel), np.uint8)
        fg_seed = cv2.erode(np.uint8(init_mask), kernel)
        bg_seed = cv2.erode(np.uint8(1 - init_mask), kernel)

        # Build GrabCut mask
        full_mask = np.ones((gt_img_rgb.shape[0], gt_img_rgb.shape[1]), np.uint8) * cv2.GC_PR_FGD
        full_mask[bg_seed == 1] = cv2.GC_BGD
        full_mask[fg_seed == 1] = cv2.GC_FGD

        # Run GrabCut
        bgdModel = np.zeros((1, 65), np.float64)
        fgdModel = np.zeros((1, 65), np.float64)
        try:
            cv2.grabCut(gt_img_rgb, full_mask, None, bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_MASK)
            grabcut_mask = np.where((full_mask == 2) | (full_mask == 0), 0, 1).astype("uint8")
        except Exception:
            # Fallback to init mask if GrabCut fails
            grabcut_mask = init_mask.astype("uint8")

        return grabcut_mask

    @torch.no_grad()
    def generate_binary_coseg_masks(
        self, gt_img, pred_img, save_dir=None, frame_indices=None, return_intermediates=False
    ):
        """
        Online co-segmentation: KMeans over GT patch features + fused saliency voting.
        GPU-optimized with vectorized PyTorch ops.
        Returns [B,1,target_size,target_size] uint8 {0,1}.

        Args:
            gt_img: [B, 3, H, W] ground truth images
            pred_img: [B, 3, H, W] predicted images
            save_dir: Optional directory to save masks as PNG files
            frame_indices: Optional list of frame indices for filename
            return_intermediates: If True, return dict with intermediate results for visualization

        Returns:
            If return_intermediates=False: torch.Tensor [B,1,H,W] final masks
            If return_intermediates=True: dict with keys:
                - 'masks': final masks [B,1,H,W]
                - 'clusters': K-means cluster assignments [B,H,W]
                - 'fg_binary': binary mask after K-means selection, before refinement [B,1,H,W]
        """
        gt_features = self._extract_patch_features(gt_img)  # [B, D, Hp, Wp]
        pred_features = self._extract_patch_features(pred_img)
        B, D, Hp, Wp = gt_features.shape
        device = gt_img.device

        gt_norm = F.normalize(gt_features, dim=1)
        pred_norm = F.normalize(pred_features, dim=1)
        cos_sim = (gt_norm * pred_norm).sum(dim=1, keepdim=True)  # [B,1,Hp,Wp]
        cos_dissim = 1.0 - cos_sim

        if self.use_ssim:
            ssim_dissim = self._compute_ssim_dissimilarity(gt_img, pred_img)  # [B,1,T,T]
            ssim_down = F.interpolate(ssim_dissim, size=(Hp, Wp), mode="area")

            def z(x):
                mu = x.mean(dim=(2, 3), keepdim=True)
                sd = x.std(dim=(2, 3), keepdim=True) + 1e-8
                return (x - mu) / sd

            fused_patch = self.w_dino * z(cos_dissim) + self.w_ssim * z(ssim_down)
        else:
            fused_patch = cos_dissim

        # KMeans on GT descriptors (GPU-accelerated, deterministic)
        gt_desc = gt_features.permute(0, 2, 3, 1).reshape(B, Hp * Wp, D)  # [B, Hp*Wp, D]
        X_all = gt_desc.reshape(-1, D)  # [B*Hp*Wp, D]

        # GPU KMeans with k-means++ initialization (matches sklearn behavior)
        N = X_all.shape[0]
        K = self.coseg_k

        # k-means++ initialization (deterministic with fixed seed)
        torch.manual_seed(42)
        centroids = []
        # Pick first centroid randomly
        first_idx = torch.randint(0, N, (1,), device=device)
        centroids.append(X_all[first_idx])

        for _ in range(K - 1):
            # Compute distance to nearest centroid
            centroids_stack = torch.cat(centroids, dim=0)  # [num_centroids, D]
            dists = torch.cdist(X_all, centroids_stack)  # [N, num_centroids]
            min_dists = dists.min(dim=1)[0]  # [N]
            # Sample proportional to distance squared (k-means++)
            probs = min_dists**2
            probs = probs / (probs.sum() + 1e-10)
            next_idx = torch.multinomial(probs, 1)
            centroids.append(X_all[next_idx])

        centroids = torch.cat(centroids, dim=0)  # [K, D]

        # Lloyd's algorithm (20 iterations for convergence)
        for iteration in range(20):
            # Assignment step
            dists = torch.cdist(X_all, centroids)  # [N, K]
            labels_flat = dists.argmin(dim=1)  # [N]

            # Update step
            centroids_new = centroids.clone()
            for k in range(K):
                mask_k = labels_flat == k
                if mask_k.sum() > 0:
                    centroids_new[k] = X_all[mask_k].mean(dim=0)

            # Check convergence
            if iteration > 0 and torch.allclose(centroids, centroids_new, atol=1e-4):
                break
            centroids = centroids_new

        labels_per_image = labels_flat.reshape(B, Hp * Wp)  # [B, Hp*Wp]

        # Vectorized cluster scoring (all on GPU)
        fused_flat = fused_patch.squeeze(1).reshape(B, Hp * Wp)  # [B, Hp*Wp]
        Kc = self.coseg_k

        # Per-cluster global average saliency: [K]
        cluster_avg_sals = torch.zeros(Kc, device=device)
        for k in range(Kc):
            mask_k = labels_per_image == k  # [B, Hp*Wp]
            vals_k = fused_flat[mask_k]
            cluster_avg_sals[k] = vals_k.mean() if vals_k.numel() > 0 else 0.0

        # Per-cluster per-frame median saliency: [K, B]
        cluster_per_frame_medians = torch.zeros(Kc, B, device=device)
        for k in range(Kc):
            for f in range(B):
                mask_kf = labels_per_image[f] == k  # [Hp*Wp]
                vals_kf = fused_flat[f, mask_kf]
                cluster_per_frame_medians[k, f] = vals_kf.median() if vals_kf.numel() > 0 else 0.0

        # Thresholds
        rank_thr = torch.quantile(cluster_avg_sals, 1.0 - self.coseg_cluster_top_percent)
        per_frame_thr = torch.quantile(fused_flat, self.coseg_frame_saliency_quantile)

        # Select foreground clusters (vectorized)
        high_sal_mask = cluster_avg_sals >= rank_thr  # [K]
        consistent_mask = (cluster_per_frame_medians >= per_frame_thr).sum(
            dim=1
        ) >= self.coseg_consistency_min_frames  # [K]
        fg_cluster_mask = high_sal_mask & consistent_mask  # [K]
        fg_cluster_indices = torch.nonzero(fg_cluster_mask, as_tuple=True)[0]  # list of K indices

        # Build masks (vectorized per-frame)
        masks = []
        fg_binary_before_refine = []  # Store binary mask before GrabCut for visualization
        labels_grid = labels_per_image.reshape(B, Hp, Wp)  # [B, Hp, Wp]
        
        for f in range(B):
            # Binary mask: pixel is FG if its cluster is in fg_cluster_indices
            fg_mask = torch.zeros(Hp, Wp, dtype=torch.bool, device=device)
            for k_idx in fg_cluster_indices:
                fg_mask |= labels_grid[f] == k_idx

            fg_uint8 = fg_mask.to(torch.uint8)  # [Hp, Wp]
            # Upsample to target_size
            up = F.interpolate(
                fg_uint8[None, None].float(),
                size=(self.target_size, self.target_size),
                mode="nearest",
            )
            up = up.squeeze().cpu().numpy().astype(np.uint8)

            # Morphological ops + connected components (still on CPU with cv2)
            if self.coseg_morph_kernel > 0 and self.coseg_morph_iters > 0:
                k = cv2.getStructuringElement(
                    cv2.MORPH_ELLIPSE, (self.coseg_morph_kernel, self.coseg_morph_kernel)
                )
                up = cv2.morphologyEx(up, cv2.MORPH_CLOSE, k, iterations=self.coseg_morph_iters)
                up = cv2.morphologyEx(up, cv2.MORPH_OPEN, k, iterations=self.coseg_morph_iters)

            num_labels, labels_cc, stats, _ = cv2.connectedComponentsWithStats(up, connectivity=8)
            min_area = int(
                self.coseg_min_component_area_ratio * self.target_size * self.target_size
            )
            clean = np.zeros_like(up, dtype=np.uint8)
            for lab in range(1, num_labels):
                if stats[lab, cv2.CC_STAT_AREA] >= min_area:
                    clean[labels_cc == lab] = 1

            # Store binary mask before GrabCut refinement
            fg_binary_before_refine.append(clean)

            # Apply GrabCut refinement (matches offline script)
            gt_img_rgb = (gt_img[f].permute(1, 2, 0).cpu().numpy() * 255.0).astype(
                np.uint8
            )  # [H, W, 3]
            grabcut_mask = self._refine_with_grabcut(gt_img_rgb, clean)
            masks.append(grabcut_mask)

        masks_np = np.stack(masks, axis=0)[:, None, :, :]  # [B,1,H,W]
        masks_tensor = torch.from_numpy(masks_np).to(device)

        # Optionally save masks to disk
        if save_dir is not None:
            import os

            os.makedirs(save_dir, exist_ok=True)
            for b_idx in range(B):
                mask_uint8 = masks_np[b_idx, 0]  # [H, W] uint8 {0, 1}
                mask_uint8 = (mask_uint8 * 255).astype(np.uint8)  # Scale to {0, 255}

                # Generate filename
                if frame_indices is not None and b_idx < len(frame_indices):
                    frame_idx = frame_indices[b_idx]
                    filename = f"coseg_mask_frame_{frame_idx:06d}.png"
                else:
                    filename = f"coseg_mask_{b_idx:02d}.png"

                Image.fromarray(mask_uint8, mode="L").save(os.path.join(save_dir, filename))

        # Return intermediate results if requested
        if return_intermediates:
            # Upsample cluster labels to target_size for visualization
            clusters_upsampled = F.interpolate(
                labels_grid.unsqueeze(1).float(),
                size=(self.target_size, self.target_size),
                mode="nearest",
            )  # [B, 1, T, T]
            
            fg_binary_np = np.stack(fg_binary_before_refine, axis=0)[:, None, :, :]  # [B,1,H,W]
            fg_binary_tensor = torch.from_numpy(fg_binary_np).to(device).float()
            
            return {
                "masks": masks_tensor,
                "clusters": clusters_upsampled.squeeze(1),  # [B, H, W]
                "fg_binary": fg_binary_tensor,  # [B, 1, H, W]
            }
        
        return masks_tensor

    def _compute_cosine_dissimilarity(self, gt_img, pred_img):
        """
        Compute cosine dissimilarity (1 - cosine_similarity) between GT and predicted images.

        Args:
            gt_img: [B, 3, H, W] in range [0, 1]
            pred_img: [B, 3, H, W] in range [0, 1]

        Returns:
            cos_dissim: [B, 1, target_size, target_size] dissimilarity map
        """
        # Extract features
        gt_features = self._extract_patch_features(gt_img)  # [B, D, H_p, W_p]
        pred_features = self._extract_patch_features(pred_img)  # [B, D, H_p, W_p]

        # Cosine similarity (features already L2-normalized)
        gt_norm = F.normalize(gt_features, dim=1)
        pred_norm = F.normalize(pred_features, dim=1)
        cos_sim = (gt_norm * pred_norm).sum(dim=1, keepdim=True)  # [B, 1, H_p, W_p]

        # Dissimilarity = 1 - similarity
        cos_dissim = 1.0 - cos_sim

        # Upsample to target size
        cos_dissim_resized = F.interpolate(
            cos_dissim,
            size=(self.target_size, self.target_size),
            mode="bilinear",
            align_corners=False,
        )

        return cos_dissim_resized

    def _compute_ssim_dissimilarity(self, gt_img, pred_img):
        """
        Compute SSIM dissimilarity (1 - SSIM) between GT and predicted images.

        Args:
            gt_img: [B, 3, H, W] in range [0, 1]
            pred_img: [B, 3, H, W] in range [0, 1]

        Returns:
            ssim_dissim: [B, 1, target_size, target_size] dissimilarity map
        """

        B = gt_img.shape[0]

        # Resize to target size first
        gt_resized = F.interpolate(
            gt_img, size=(self.target_size, self.target_size), mode="bilinear", align_corners=False
        )
        pred_resized = F.interpolate(
            pred_img,
            size=(self.target_size, self.target_size),
            mode="bilinear",
            align_corners=False,
        )

        # Compute SSIM for each image in the batch
        ssim_dissim_list = []
        for b in range(B):
            # Convert to numpy uint8 [0, 255]
            gt_np = (gt_resized[b].permute(1, 2, 0).cpu().numpy() * 255.0).astype(np.uint8)
            pred_np = (pred_resized[b].permute(1, 2, 0).cpu().numpy() * 255.0).astype(np.uint8)

            # Compute SSIM map
            ssim_map = structural_similarity(
                gt_np, pred_np, win_size=11, channel_axis=2, data_range=255, full=True
            )[
                1
            ]  # Get the full SSIM map

            # Average across channels if needed
            ssim_map_gray = ssim_map.mean(axis=2) if ssim_map.ndim == 3 else ssim_map

            # Dissimilarity = 1 - SSIM
            ssim_dissim = 1.0 - ssim_map_gray
            ssim_dissim_list.append(ssim_dissim)

        # Stack and convert to tensor
        ssim_dissim = np.stack(ssim_dissim_list, axis=0)  # [B, H, W]
        ssim_dissim = torch.from_numpy(ssim_dissim).unsqueeze(1).to(gt_img.device)  # [B, 1, H, W]

        return ssim_dissim

    def _create_soft_mask(self, dissimilarity_map):
        """
        Create soft mask using percentile-based normalization.

        Args:
            dissimilarity_map: [B, 1, H, W] combined dissimilarity

        Returns:
            soft_mask: [B, 1, H, W] soft mask in range [0, 1]
        """
        B = dissimilarity_map.shape[0]
        soft_masks = []

        for b in range(B):
            dissim = dissimilarity_map[b, 0]  # [H, W]

            # Percentile-based normalization
            p_thresh = np.percentile(dissim.cpu().numpy(), self.percentile_threshold)
            p_max = dissim.max().item()

            # Normalize: [p_thresh, p_max] -> [0, 1]
            soft_mask = torch.clamp((dissim - p_thresh) / (p_max - p_thresh + 1e-8), 0.0, 1.0)
            soft_masks.append(soft_mask)

        soft_masks = torch.stack(soft_masks, dim=0).unsqueeze(1)  # [B, 1, H, W]
        return soft_masks

    @torch.no_grad()
    def forward(self, gt_img, pred_img, debug_save_path=None, use_comprehensive_viz=False):
        """
        Generate motion/uncertainty pseudo-labels (soft or binary based on use_coseg_binary flag).

        Args:
            gt_img: [B, 3, H, W] ground truth images in range [0, 1]
            pred_img: [B, 3, H, W] predicted images in range [0, 1]
            debug_save_path: Optional path to save debug visualization
            use_comprehensive_viz: If True, save comprehensive visualization showing full pipeline
                                  (GT, Rendering, MSE, SSIM, DINO, min, K-means, FG, Final)

        Returns:
            motion_mask: [B, 1, target_size, target_size] mask in range [0, 1]
                        If use_coseg_binary=True: binary {0,1} via KMeans+GrabCut
                        If use_coseg_binary=False: soft [0,1] via percentile normalization
        """
        # Compute DINOv3 cosine dissimilarity
        dino_dissim = self._compute_cosine_dissimilarity(gt_img, pred_img)

        # Optionally combine with SSIM
        if self.use_ssim:
            ssim_dissim = self._compute_ssim_dissimilarity(gt_img, pred_img)
            # Conservative: min dissimilarity (both must agree it's different)
            combined_dissim = torch.minimum(dino_dissim, ssim_dissim)
        else:
            ssim_dissim = None
            combined_dissim = dino_dissim
        
        # Initialize vars for comprehensive visualization
        kmeans_clusters = None
        fg_mask_binary = None
        mse_dissim = None
        
        if self.use_coseg_binary:
            # Produce binary motion masks (1=motion, 0=static) at target_size
            if use_comprehensive_viz and debug_save_path is not None:
                # Request intermediate results for visualization
                result = self.generate_binary_coseg_masks(gt_img, pred_img, return_intermediates=True)
                motion_mask = result["masks"].to(combined_dissim.device).to(torch.float32)
                kmeans_clusters = result["clusters"]
                fg_mask_binary = result["fg_binary"]
            else:
                bin_motion = (
                    self.generate_binary_coseg_masks(gt_img, pred_img)
                    .to(combined_dissim.device)
                    .to(torch.float32)
                )
                motion_mask = bin_motion
        else:
            # Create soft mask using percentile normalization (uncertainty-style)
            motion_mask = self._create_soft_mask(combined_dissim)

        # Debug visualization
        if debug_save_path is not None:
            if use_comprehensive_viz:
                # Compute MSE for comprehensive visualization
                mse_dissim = self._compute_mse_dissimilarity(gt_img, pred_img)
                
                # Save comprehensive visualization
                self._save_comprehensive_visualization(
                    gt_img,
                    pred_img,
                    mse_dissim,
                    ssim_dissim,
                    dino_dissim,
                    combined_dissim,
                    kmeans_clusters,
                    fg_mask_binary,
                    motion_mask,
                    debug_save_path,
                )
            else:
                # Save basic debug visualization
                self._save_debug_visualization(
                    gt_img,
                    pred_img,
                    dino_dissim,
                    ssim_dissim,
                    combined_dissim,
                    motion_mask,
                    debug_save_path,
                )

        return motion_mask

    def _compute_mse_dissimilarity(self, gt_img, pred_img):
        """
        Compute MSE dissimilarity between GT and predicted images.

        Args:
            gt_img: [B, 3, H, W] in range [0, 1]
            pred_img: [B, 3, H, W] in range [0, 1]

        Returns:
            mse_dissim: [B, 1, target_size, target_size] dissimilarity map
        """
        # Resize to target size
        gt_resized = F.interpolate(
            gt_img, size=(self.target_size, self.target_size), mode="bilinear", align_corners=False
        )
        pred_resized = F.interpolate(
            pred_img,
            size=(self.target_size, self.target_size),
            mode="bilinear",
            align_corners=False,
        )

        # Compute MSE per pixel (average over channels)
        mse = ((gt_resized - pred_resized) ** 2).mean(dim=1, keepdim=True)  # [B, 1, H, W]

        return mse

    def _save_comprehensive_visualization(
        self,
        gt_img,
        pred_img,
        mse_dissim,
        ssim_dissim,
        dino_dissim,
        combined_dissim,
        kmeans_clusters,
        fg_mask_binary,
        final_mask,
        save_path,
    ):
        """
        Save comprehensive visualization showing the full pipeline:
        GT, Rendering, MSE, SSIM, DINO, min(SSIM,DINO), K-means, FG selection, Final mask.

        Args:
            gt_img: [B, 3, H, W]
            pred_img: [B, 3, H, W]
            mse_dissim: [B, 1, H, W] MSE dissimilarity
            ssim_dissim: [B, 1, H, W] SSIM dissimilarity
            dino_dissim: [B, 1, H, W] DINO dissimilarity
            combined_dissim: [B, 1, H, W] min(SSIM, DINO)
            kmeans_clusters: [B, H, W] cluster assignments (for visualization)
            fg_mask_binary: [B, 1, H, W] binary mask after K-means selection
            final_mask: [B, 1, H, W] final mask after all post-processing
            save_path: Path to save figure
        """
        import matplotlib.pyplot as plt
        from pathlib import Path
        import matplotlib.colors as mcolors

        Path(save_path).parent.mkdir(parents=True, exist_ok=True)

        # Take first sample from batch (convert bfloat16 to float32 for numpy)
        gt = gt_img[0].float().cpu().permute(1, 2, 0).numpy()  # [H, W, 3]
        pred = pred_img[0].float().cpu().permute(1, 2, 0).numpy()  # [H, W, 3]
        mse = mse_dissim[0, 0].float().cpu().numpy()  # [H, W]
        ssim = ssim_dissim[0, 0].float().cpu().numpy() if ssim_dissim is not None else None
        dino = dino_dissim[0, 0].float().cpu().numpy()  # [H, W]
        combined = combined_dissim[0, 0].float().cpu().numpy()  # [H, W]
        clusters = kmeans_clusters[0].float().cpu().numpy() if kmeans_clusters is not None else None
        fg_binary = (
            fg_mask_binary[0, 0].float().cpu().numpy() if fg_mask_binary is not None else None
        )
        final = final_mask[0, 0].float().cpu().numpy()  # [H, W]

        # Create figure with 2 rows x 5 columns
        fig, axes = plt.subplots(2, 5, figsize=(25, 10))

        # Row 1: Full pipeline
        # 1. GT Image
        axes[0, 0].imshow(gt)
        axes[0, 0].set_title("GT Image", fontsize=12, fontweight="bold")
        axes[0, 0].axis("off")

        # 2. Rendered Image
        axes[0, 1].imshow(pred)
        axes[0, 1].set_title("Rendered Image", fontsize=12, fontweight="bold")
        axes[0, 1].axis("off")

        # 3. MSE Dissimilarity
        im_mse = axes[0, 2].imshow(mse, cmap="hot")
        axes[0, 2].set_title(
            f"MSE Dissim\n(mean={mse.mean():.3f}, max={mse.max():.3f})", fontsize=12, fontweight="bold"
        )
        axes[0, 2].axis("off")
        plt.colorbar(im_mse, ax=axes[0, 2], fraction=0.046)

        # 4. DINO Dissimilarity
        im_dino = axes[0, 3].imshow(dino, cmap="hot", vmin=0, vmax=1)
        axes[0, 3].set_title(
            f"DINO Cosine Dissim\n(mean={dino.mean():.3f}, max={dino.max():.3f})",
            fontsize=12,
            fontweight="bold",
        )
        axes[0, 3].axis("off")
        plt.colorbar(im_dino, ax=axes[0, 3], fraction=0.046)

        # 5. SSIM Dissimilarity
        if ssim is not None:
            im_ssim = axes[0, 4].imshow(ssim, cmap="hot", vmin=0, vmax=1)
            axes[0, 4].set_title(
                f"SSIM Dissim\n(mean={ssim.mean():.3f}, max={ssim.max():.3f})",
                fontsize=12,
                fontweight="bold",
            )
            axes[0, 4].axis("off")
            plt.colorbar(im_ssim, ax=axes[0, 4], fraction=0.046)
        else:
            axes[0, 4].axis("off")

        # Row 2: Clustering and final results
        # 1. min(SSIM, DINO)
        im_combined = axes[1, 0].imshow(combined, cmap="hot", vmin=0, vmax=1)
        axes[1, 0].set_title(
            f"min(DINO, SSIM)\n(mean={combined.mean():.3f}, max={combined.max():.3f})",
            fontsize=12,
            fontweight="bold",
        )
        axes[1, 0].axis("off")
        plt.colorbar(im_combined, ax=axes[1, 0], fraction=0.046)

        # 2. K-means Clusters (if available)
        if clusters is not None:
            # Create a colorful visualization of clusters
            n_clusters = int(clusters.max()) + 1
            im_clusters = axes[1, 1].imshow(clusters, cmap="tab20", interpolation="nearest")
            axes[1, 1].set_title(
                f"K-means Clusters\n(K={n_clusters})", fontsize=12, fontweight="bold"
            )
            axes[1, 1].axis("off")
        else:
            axes[1, 1].axis("off")

        # 3. Binary FG mask (after K-means selection, before refinement)
        if fg_binary is not None:
            im_fg = axes[1, 2].imshow(fg_binary, cmap="hot", vmin=0, vmax=1)
            axes[1, 2].set_title(
                f"FG Selection (binary)\n(coverage={fg_binary.mean():.3f})",
                fontsize=12,
                fontweight="bold",
            )
            axes[1, 2].axis("off")
            plt.colorbar(im_fg, ax=axes[1, 2], fraction=0.046)
        else:
            axes[1, 2].axis("off")

        # 4. Final mask (after all post-processing)
        im_final = axes[1, 3].imshow(final, cmap="hot", vmin=0, vmax=1)
        axes[1, 3].set_title(
            f"Final Mask\n(coverage={final.mean():.3f})", fontsize=12, fontweight="bold"
        )
        axes[1, 3].axis("off")
        plt.colorbar(im_final, ax=axes[1, 3], fraction=0.046)

        # 5. Final mask overlay on GT
        axes[1, 4].imshow(gt)
        axes[1, 4].imshow(final, cmap="jet", alpha=0.5, vmin=0, vmax=1)
        axes[1, 4].set_title("Final Mask on GT", fontsize=12, fontweight="bold")
        axes[1, 4].axis("off")

        # Title
        mode = "Binary Co-Seg" if self.use_coseg_binary else "Soft Percentile"
        param_str = f"mode={mode}, w_dino={self.w_dino:.2f}, w_ssim={self.w_ssim:.2f}"
        if self.use_coseg_binary:
            param_str += f", K={self.coseg_k}"
        else:
            param_str += f", P={self.percentile_threshold}"

        plt.suptitle(f"DINOv3 Motion Mask Pipeline - {param_str}", fontsize=14, fontweight="bold")
        plt.tight_layout()
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
        plt.close()

    def _save_debug_visualization(
        self, gt_img, pred_img, dino_dissim, ssim_dissim, combined_dissim, motion_mask, save_path
    ):
        """Save debug visualization showing all intermediate steps."""
        import matplotlib.pyplot as plt
        from pathlib import Path

        Path(save_path).parent.mkdir(parents=True, exist_ok=True)

        # Take first sample from batch (convert bfloat16 to float32 for numpy)
        gt = gt_img[0].float().cpu().permute(1, 2, 0).numpy()  # [H, W, 3]
        pred = pred_img[0].float().cpu().permute(1, 2, 0).numpy()  # [H, W, 3]
        dino = dino_dissim[0, 0].float().cpu().numpy()  # [H, W]
        ssim = (
            ssim_dissim[0, 0].float().cpu().numpy() if ssim_dissim is not None else None
        )  # [H, W]
        combined = combined_dissim[0, 0].float().cpu().numpy()  # [H, W]
        mask = motion_mask[0, 0].float().cpu().numpy()  # [H, W]

        # Create figure
        n_cols = 6 if ssim is not None else 5
        fig, axes = plt.subplots(2, n_cols, figsize=(n_cols * 4, 8))

        # Row 1: Images and dissimilarity maps
        axes[0, 0].imshow(gt)
        axes[0, 0].set_title("GT Image", fontsize=14, fontweight="bold")
        axes[0, 0].axis("off")

        axes[0, 1].imshow(pred)
        axes[0, 1].set_title("Rendered Image", fontsize=14, fontweight="bold")
        axes[0, 1].axis("off")

        im2 = axes[0, 2].imshow(dino, cmap="hot", vmin=0, vmax=1)
        axes[0, 2].set_title(
            f"DINO Dissim\n(mean={dino.mean():.3f}, max={dino.max():.3f})",
            fontsize=14,
            fontweight="bold",
        )
        axes[0, 2].axis("off")
        plt.colorbar(im2, ax=axes[0, 2], fraction=0.046)

        if ssim is not None:
            im3 = axes[0, 3].imshow(ssim, cmap="hot", vmin=0, vmax=1)
            axes[0, 3].set_title(
                f"SSIM Dissim\n(mean={ssim.mean():.3f}, max={ssim.max():.3f})",
                fontsize=14,
                fontweight="bold",
            )
            axes[0, 3].axis("off")
            plt.colorbar(im3, ax=axes[0, 3], fraction=0.046)

            im4 = axes[0, 4].imshow(combined, cmap="hot", vmin=0, vmax=1)
            axes[0, 4].set_title(
                f"Combined (min)\n(mean={combined.mean():.3f}, max={combined.max():.3f})",
                fontsize=14,
                fontweight="bold",
            )
            axes[0, 4].axis("off")
            plt.colorbar(im4, ax=axes[0, 4], fraction=0.046)

            im5 = axes[0, 5].imshow(mask, cmap="hot", vmin=0, vmax=1)
            axes[0, 5].set_title(
                f"Final Mask (P{self.percentile_threshold})\n(mean={mask.mean():.3f}, max={mask.max():.3f})",
                fontsize=14,
                fontweight="bold",
            )
            axes[0, 5].axis("off")
            plt.colorbar(im5, ax=axes[0, 5], fraction=0.046)
        else:
            im4 = axes[0, 3].imshow(combined, cmap="hot", vmin=0, vmax=1)
            axes[0, 3].set_title(
                f"Combined\n(mean={combined.mean():.3f}, max={combined.max():.3f})",
                fontsize=14,
                fontweight="bold",
            )
            axes[0, 3].axis("off")
            plt.colorbar(im4, ax=axes[0, 3], fraction=0.046)

            im5 = axes[0, 4].imshow(mask, cmap="hot", vmin=0, vmax=1)
            axes[0, 4].set_title(
                f"Final Mask (P{self.percentile_threshold})\n(mean={mask.mean():.3f}, max={mask.max():.3f})",
                fontsize=14,
                fontweight="bold",
            )
            axes[0, 4].axis("off")
            plt.colorbar(im5, ax=axes[0, 4], fraction=0.046)

        # Row 2: Overlays on GT image
        axes[1, 0].imshow(gt)
        axes[1, 0].set_title("GT (reference)", fontsize=14)
        axes[1, 0].axis("off")

        axes[1, 1].imshow(pred)
        axes[1, 1].set_title("Rendered (reference)", fontsize=14)
        axes[1, 1].axis("off")

        # DINO overlay
        axes[1, 2].imshow(gt)
        axes[1, 2].imshow(dino, cmap="hot", alpha=0.5, vmin=0, vmax=1)
        axes[1, 2].set_title("GT + DINO overlay", fontsize=14)
        axes[1, 2].axis("off")

        if ssim is not None:
            # SSIM overlay
            axes[1, 3].imshow(gt)
            axes[1, 3].imshow(ssim, cmap="hot", alpha=0.5, vmin=0, vmax=1)
            axes[1, 3].set_title("GT + SSIM overlay", fontsize=14)
            axes[1, 3].axis("off")

            # Combined overlay
            axes[1, 4].imshow(gt)
            axes[1, 4].imshow(combined, cmap="hot", alpha=0.5, vmin=0, vmax=1)
            axes[1, 4].set_title("GT + Combined overlay", fontsize=14)
            axes[1, 4].axis("off")

            # Final mask overlay
            axes[1, 5].imshow(gt)
            axes[1, 5].imshow(mask, cmap="hot", alpha=0.5, vmin=0, vmax=1)
            axes[1, 5].set_title("GT + Final Mask overlay", fontsize=14)
            axes[1, 5].axis("off")
        else:
            # Combined overlay
            axes[1, 3].imshow(gt)
            axes[1, 3].imshow(combined, cmap="hot", alpha=0.5, vmin=0, vmax=1)
            axes[1, 3].set_title("GT + Combined overlay", fontsize=14)
            axes[1, 3].axis("off")

            # Final mask overlay
            axes[1, 4].imshow(gt)
            axes[1, 4].imshow(mask, cmap="hot", alpha=0.5, vmin=0, vmax=1)
            axes[1, 4].set_title("GT + Final Mask overlay", fontsize=14)
            axes[1, 4].axis("off")

        plt.suptitle(
            f"DINOv3 Pseudo-Label Debug (w_dino={self.w_dino:.2f}, w_ssim={self.w_ssim:.2f})",
            fontsize=16,
            fontweight="bold",
        )
        plt.tight_layout()
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
        plt.close()


class Images2Latent4D(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config

        self.process_data = data_utils.ProcessData(config)
        # self.transform_input = TransformInput(config)

        # image tokenizer
        if not config.model.image_tokenizer.get("use_off_the_shelf", False):
            self.image_tokenizer = nn.Sequential(
                Rearrange(
                    "b v c (hh ph) (ww pw) -> (b v) (hh ww) (ph pw c)",
                    ph=self.config.model.image_tokenizer.patch_size,
                    pw=self.config.model.image_tokenizer.patch_size,
                ),
                nn.Linear(
                    config.model.image_tokenizer.in_channels
                    * (config.model.image_tokenizer.patch_size**2),
                    config.model.transformer.d,
                    bias=False,
                ),
            )
            self.image_tokenizer.apply(_init_weights)
        else:
            # self.image_tokenizer = CustomImageTokenizer(config)
            pass

        # image positional embedding embedder
        self.use_pe_embedding_layer = config.model.get("input_with_pe", True)
        self.pe_embedder = (
            nn.Sequential(
                nn.Linear(
                    config.model.transformer.d,
                    config.model.transformer.d,
                ),
                nn.SiLU(),
                nn.Linear(
                    config.model.transformer.d,
                    config.model.transformer.d,
                ),
            )
            if self.use_pe_embedding_layer
            else nn.Identity()
        )
        self.pe_embedder.apply(_init_weights)

        # latent scene representation
        self.scene_code = nn.Parameter(
            torch.randn(
                config.model.scene_latent.length,
                config.model.transformer.d,
            )
        )
        nn.init.trunc_normal_(self.scene_code, std=0.02)

        # pose tokens
        self.cam_code = nn.Parameter(
            torch.randn(
                self.config.model.pose_latent.get("length", 1),
                config.model.transformer.d,
            )
        )
        nn.init.trunc_normal_(self.cam_code, std=0.02)

        # pose pe temporal embedder
        self.temporal_pe_embedder = (
            nn.Sequential(
                nn.Linear(
                    config.model.transformer.d,
                    config.model.transformer.d,
                ),
                nn.SiLU(),
                nn.Linear(
                    config.model.transformer.d,
                    config.model.transformer.d,
                ),
            )
            if self.use_pe_embedding_layer
            else nn.Identity()
        )
        self.temporal_pe_embedder.apply(_init_weights)

        # qk norm settings
        use_qk_norm = config.model.transformer.get("use_qk_norm", False)

        # transformer encoder and init
        self.transformer_encoder = [
            QK_Norm_TransformerBlock(
                config.model.transformer.d,
                config.model.transformer.d_head,
                use_qk_norm=use_qk_norm,
            )
            for _ in range(config.model.transformer.encoder_n_layer)
        ]
        if config.model.transformer.get("special_init", False):
            if config.model.transformer.get("depth_init", False):
                for idx in range(len(self.transformer_encoder)):
                    weight_init_std = 0.02 / (2 * (idx + 1)) ** 0.5
                    self.transformer_encoder[idx].apply(
                        lambda module: _init_weights_layerwise(module, weight_init_std)
                    )
            else:
                for idx in range(len(self.transformer_encoder)):
                    weight_init_std = 0.02 / (2 * config.model.transformer.encoder_n_layer) ** 0.5
                    self.transformer_encoder[idx].apply(
                        lambda module: _init_weights_layerwise(module, weight_init_std)
                    )
            self.transformer_encoder = nn.ModuleList(self.transformer_encoder)
        else:
            self.transformer_encoder = nn.ModuleList(self.transformer_encoder)
            self.transformer_encoder.apply(_init_weights)

        # transformer encoder2 and init
        self.transformer_encoder_geom = [
            QK_Norm_TransformerBlock(
                config.model.transformer.d,
                config.model.transformer.d_head,
                use_qk_norm=use_qk_norm,
            )
            for _ in range(config.model.transformer.encoder_geom_n_layer)
        ]
        if config.model.transformer.get("special_init", False):
            if config.model.transformer.get("depth_init", False):
                for idx in range(len(self.transformer_encoder_geom)):
                    weight_init_std = 0.02 / (2 * (idx + 1)) ** 0.5
                    self.transformer_encoder_geom[idx].apply(
                        lambda module: _init_weights_layerwise(module, weight_init_std)
                    )
            else:
                for idx in range(len(self.transformer_encoder_geom)):
                    weight_init_std = (
                        0.02 / (2 * config.model.transformer.encoder_geom_n_layer) ** 0.5
                    )
                    self.transformer_encoder_geom[idx].apply(
                        lambda module: _init_weights_layerwise(module, weight_init_std)
                    )
            self.transformer_encoder_geom = nn.ModuleList(self.transformer_encoder_geom)
        else:
            self.transformer_encoder_geom = nn.ModuleList(self.transformer_encoder_geom)
            self.transformer_encoder_geom.apply(_init_weights)

        # ln before decoder
        self.decoder_ln = nn.LayerNorm(config.model.transformer.d, bias=False)

        # transformer decoder and init
        self.transformer_decoder = [
            QK_Norm_TransformerBlock(
                config.model.transformer.d,
                config.model.transformer.d_head,
                use_qk_norm=use_qk_norm,
            )
            for _ in range(config.model.transformer.decoder_n_layer)
        ]
        if config.model.transformer.get("special_init", False):
            if config.model.transformer.depth_init:
                for idx in range(len(self.transformer_decoder)):
                    weight_init_std = 0.02 / (2 * (idx + 1)) ** 0.5
                    self.transformer_decoder[idx].apply(
                        lambda module: _init_weights_layerwise(module, weight_init_std)
                    )
            else:
                for idx in range(len(self.transformer_decoder)):
                    weight_init_std = 0.02 / (2 * config.model.transformer.decoder_n_layer) ** 0.5
                    self.transformer_decoder[idx].apply(
                        lambda module: _init_weights_layerwise(module, weight_init_std)
                    )
            self.transformer_decoder = nn.ModuleList(self.transformer_decoder)
        else:
            self.transformer_decoder = nn.ModuleList(self.transformer_decoder)
            self.transformer_decoder.apply(_init_weights)

        # pose predictor
        self.pose_predictor = PoseEstimator(config)

        # target pose tokenizer
        self.target_latent_h = (
            config.model.target_image.height // config.model.target_image.patch_size
        )
        self.target_latent_w = (
            config.model.target_image.width // config.model.target_image.patch_size
        )
        self.target_pose_tokenizer = nn.Sequential(
            Rearrange(
                "b v c (hh ph) (ww pw) -> (b v) (hh ww) (ph pw c)",
                ph=self.config.model.target_pose_tokenizer.patch_size,
                pw=self.config.model.target_pose_tokenizer.patch_size,
            ),
            nn.Linear(
                config.model.target_pose_tokenizer.in_channels
                * (config.model.target_pose_tokenizer.patch_size**2),
                config.model.transformer.d,
                bias=False,
            ),
        )
        self.target_pose_tokenizer.apply(_init_weights)

        self.target_pose_tokenizer2 = nn.Sequential(
            Rearrange(
                "b v c (hh ph) (ww pw) -> (b v) (hh ww) (ph pw c)",
                ph=self.config.model.target_pose_tokenizer.patch_size,
                pw=self.config.model.target_pose_tokenizer.patch_size,
            ),
            nn.Linear(
                config.model.target_pose_tokenizer.in_channels
                * (config.model.target_pose_tokenizer.patch_size**2),
                config.model.transformer.d,
                bias=False,
            ),
        )
        self.target_pose_tokenizer2.apply(_init_weights)

        # fuse mlp
        self.mlp_fuse = nn.Sequential(
            nn.LayerNorm(config.model.transformer.d * 2, bias=False),
            nn.Linear(
                config.model.transformer.d * 2,
                config.model.transformer.d,
                bias=True,
            ),
            nn.SiLU(),
            nn.Linear(
                config.model.transformer.d,
                config.model.transformer.d,
                bias=True,
            ),
        )
        self.mlp_fuse.apply(_init_weights)

        # output regresser
        self.image_token_decoder = nn.Sequential(
            nn.LayerNorm(config.model.transformer.d, bias=False),
            nn.Linear(
                config.model.transformer.d,
                (config.model.target_image.patch_size**2) * 3,
                bias=False,
            ),
            nn.Sigmoid(),
        )
        self.image_token_decoder.apply(_init_weights)

        # Motion mask system (two components):
        # 1. DINOv3 pseudo-label maker: generates GT masks from (rendered, GT) pairs
        # 2. MotionMaskPredictor: learns to predict masks from encoded tokens
        self.use_motion_mask = config.model.get("use_motion_mask", False)
        self.use_dinov3_pseudolabel = config.model.get("use_dinov3_pseudolabel", False)

        # Component 1: DINOv3 pseudo-label generator (frozen, for supervision)
        if self.use_dinov3_pseudolabel:
            dinov3_cfg = config.model.get("dinov3_pseudolabel", {})
            model_name = dinov3_cfg.get("model_name", "facebook/dinov3-vit7b16-pretrain-lvd1689m")

            # Load shared DINOv3 backbone
            dinov3_backbone = AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16)
            dinov3_processor = AutoImageProcessor.from_pretrained(model_name)
            dinov3_backbone.eval()

            # Move to GPU (will be handled by .to(device) in __init__, but ensure dtype is set first)
            # Note: Don't call .cuda() here, let the parent model's .to(device) handle it

            # Freeze all DINOv3 parameters
            for p in dinov3_backbone.parameters():
                p.requires_grad = False

            self.dinov3_pseudolabel_maker = DinoV3UncertaintyPseudoLabelMaker(
                dinov3_backbone=dinov3_backbone,
                dinov3_processor=dinov3_processor,
                percentile_threshold=dinov3_cfg.get("percentile_threshold", 75),
                use_ssim=dinov3_cfg.get("use_ssim", True),
                target_size=dinov3_cfg.get("target_size", 256),
                use_coseg_binary=dinov3_cfg.get("use_coseg_binary", False),
                coseg_k=dinov3_cfg.get("coseg_k", 64),
                coseg_consistency_min_frames=dinov3_cfg.get("coseg_consistency_min_frames", 4),
                coseg_frame_saliency_quantile=dinov3_cfg.get("coseg_frame_saliency_quantile", 0.75),
                coseg_cluster_top_percent=dinov3_cfg.get("coseg_cluster_top_percent", 0.05),
                coseg_morph_kernel=dinov3_cfg.get("coseg_morph_kernel", 3),
                coseg_morph_iters=dinov3_cfg.get("coseg_morph_iters", 1),
                coseg_min_component_area_ratio=dinov3_cfg.get(
                    "coseg_min_component_area_ratio", 0.0025
                ),
                coseg_grabcut_kernel=dinov3_cfg.get("coseg_grabcut_kernel", 7),
                w_dino=dinov3_cfg.get("w_dino", 0.5),
                w_ssim=dinov3_cfg.get("w_ssim", 0.5),
            )

            percentile = dinov3_cfg.get("percentile_threshold", 75)
            use_ssim = dinov3_cfg.get("use_ssim", True)
            print(f"✓ DINOv3 pseudo-label maker initialized: P{percentile}, SSIM={use_ssim}")
        else:
            self.dinov3_pseudolabel_maker = None
            dinov3_backbone = None
            dinov3_processor = None

        # Component 2: Learned motion mask predictor (trainable)
        if self.use_motion_mask:
            if self.use_dinov3_pseudolabel and self.dinov3_pseudolabel_maker is not None:
                # Share the DINOv3 backbone from pseudo-label maker
                self.motion_mask_predictor = MotionMaskPredictor(
                    config,
                    self.dinov3_pseudolabel_maker.backbone,
                    self.dinov3_pseudolabel_maker.processor,
                )
                print("✓ Motion mask predictor initialized (sharing DINOv3 backbone)")
            else:
                # Load separate DINOv3 if pseudo-label maker not enabled
                model_name = config.model.get("dinov3_pseudolabel", {}).get(
                    "model_name", "facebook/dinov3-vit7b16-pretrain-lvd1689m"
                )
                dinov3_backbone = AutoModel.from_pretrained(model_name, torch_dtype=torch.bfloat16)
                dinov3_processor = AutoImageProcessor.from_pretrained(model_name)
                dinov3_backbone.eval()
                for p in dinov3_backbone.parameters():
                    p.requires_grad = False

                self.motion_mask_predictor = MotionMaskPredictor(
                    config,
                    dinov3_backbone,
                    dinov3_processor,
                )
                print("✓ Motion mask predictor initialized (separate DINOv3 backbone)")
        else:
            self.motion_mask_predictor = None

        # Opportunistic caching for DINOv3 pseudo-labels (saves compute on frozen renderings)
        self.enable_mask_cache = config.model.dinov3_pseudolabel.get("enable_mask_cache", False)
        if self.enable_mask_cache:
            self.mask_cache_dir = config.model.dinov3_pseudolabel.get(
                "mask_cache_dir", "./motion_mask_cache"
            )
            print(f"✓ Opportunistic mask caching enabled: {self.mask_cache_dir}")
        else:
            self.mask_cache_dir = None

        # loss
        self.loss_computer = LossComputer_official(config)

        # config backup
        self.config_bk = copy.deepcopy(config)
        self.render_interpolate = config.training.get("render_interpolate", False)

        copy_paste_cfg = (
            getattr(config.training, "copy_paste", {}) if hasattr(config, "training") else {}
        )

        inference_cfg = getattr(config, "inference", None)
        if hasattr(inference_cfg, "get"):
            is_inference_mode = inference_cfg.get("if_inference", False)
        else:
            is_inference_mode = bool(inference_cfg)

        self.copy_paste_training_enabled = not is_inference_mode and copy_paste_cfg.get(
            "enabled", False
        )

        if config.model.transformer.get("fix_decoder", False):
            self.freeze_weights()
            if self.copy_paste_training_enabled:
                self._unfreeze_renderer_weights()

        # training settings
        if config.inference or config.get("evaluation", False):
            if config.training.get("random_inputs", False):
                self.random_index = True
            else:
                self.random_index = False
        else:
            self.random_index = config.training.get("random_split", False)

        print("Use random index:", self.random_index)

        self.motion_mask_threshold = config.model.get("motion_mask_threshold", 0.1)
        if config.model.get("use_mae_masking", False):
            self.noise_token = nn.Parameter(torch.randn(1, 1, self.config.model.transformer.d))
            nn.init.trunc_normal_(self.noise_token, std=0.02)

        # Stage 3: Freeze motion mask predictor if motion_mask_only_training is False
        # (This means we're training the rendering model, not the motion mask)
        motion_mask_only = config.model.get("motion_mask_only_training", False)
        if self.use_motion_mask and not motion_mask_only:
            if self.copy_paste_training_enabled:
                print("✓ Stage 3 + copy-paste detected: keeping motion mask predictor trainable")
            else:
                print(
                    "✓ Stage 3 detected: Freezing motion mask predictor (training rendering only)"
                )
                self._freeze_motion_mask_predictor()

    def freeze_weights(self):
        for param in self.target_pose_tokenizer.parameters():
            param.requires_grad = False

        for param in self.image_token_decoder.parameters():
            param.requires_grad = False

        for param in self.transformer_decoder.parameters():
            param.requires_grad = False

    def _unfreeze_renderer_weights(self):
        for module in [self.transformer_decoder, self.image_token_decoder]:
            for param in module.parameters():
                param.requires_grad = True
        print("✓ Copy-paste enabled: renderer weights unfrozen")

    def _freeze_motion_mask_predictor(self):
        """Freeze motion mask predictor for Stage 3 training (rendering only)."""
        if self.motion_mask_predictor is not None:
            for param in self.motion_mask_predictor.parameters():
                param.requires_grad = False
            print(
                f"  Frozen {sum(p.numel() for p in self.motion_mask_predictor.parameters())} motion mask parameters"
            )

    def train(self, mode=True):
        # override the train method to keep the fronzon modules in eval mode
        super().train(mode)
        self.loss_computer.eval()

    def get_overview(self):
        count_train_params = lambda model: sum(
            p.numel() for p in model.parameters() if p.requires_grad
        )
        overview = edict(
            image_tokenizer=count_train_params(self.image_tokenizer),
            scene_latent=self.scene_code.data.numel(),
            transformer=count_train_params(self.transformer_decoder)
            + count_train_params(self.transformer_encoder),
            image_token_decoder=count_train_params(self.image_token_decoder),
        )
        return overview

    def _try_load_mask_from_cache(self, uid):
        """Try to load cached mask for a given scene UID."""
        if not self.enable_mask_cache or self.mask_cache_dir is None:
            return None

        cache_file = os.path.join(self.mask_cache_dir, f"mask_{uid:08d}.pt")
        if os.path.exists(cache_file):
            try:
                mask = torch.load(cache_file, map_location="cpu")
                return mask
            except Exception as e:
                print(f"Warning: Failed to load cache {cache_file}: {e}")
                return None
        return None

    def _save_mask_to_cache(self, uid, mask):
        """Save computed mask to cache for future use."""
        if not self.enable_mask_cache or self.mask_cache_dir is None:
            return

        os.makedirs(self.mask_cache_dir, exist_ok=True)
        cache_file = os.path.join(self.mask_cache_dir, f"mask_{uid:08d}.pt")
        try:
            torch.save(mask.cpu(), cache_file)
        except Exception as e:
            print(f"Warning: Failed to save cache {cache_file}: {e}")

    def _apply_token_dropping(
        self,
        img_tokens_fused,  # [b, v_input*n, d] - AFTER fusion!
        predicted_masks,  # [b*v_input, 1, 256, 256] - motion mask logits OR binary probs
        b,
        v_input,
        n,
        d,
        dataset_sources=None,
        using_gt_masks=False,
        copy_paste_mask_input=None,
    ):
        """
        Apply token dropping to FUSED tokens using predicted motion masks.
        Zeros out dynamic tokens instead of replacing with learnable noise.

        For dynamic scenes (DRE10K): Use predicted motion masks
        For static scenes (RE10K): Apply random dropout (20% of tokens)

        Args:
            img_tokens_fused: [b, v_input*n, d] fused tokens (img + plucker)
            predicted_masks: [b*v_input, 1, 256, 256] motion masks (logits OR binary probs)
            b, v_input, n, d: batch size, num input views, num tokens, token dim
            dataset_sources: List[str] of length b, indicating dataset source per batch element
            using_gt_masks: bool, if True predicted_masks are already binary probs {0,1}

        Returns:
            img_tokens_dropped: [b, v_input*n, d] tokens with dynamic positions zeroed
            dynamic_mask: [b*v_input, 1, H_patch, W_patch] binary mask for visualization
        """
        if copy_paste_mask_input is None and (predicted_masks is None or not self.use_motion_mask):
            # No dropping
            return img_tokens_fused, None

        # Downsample masks from 256×256 to token grid (16×16)
        H_patch = W_patch = int(n**0.5)

        # Determine which masks to use for token dropping
        if copy_paste_mask_input is not None:
            mask_tokens_probs = copy_paste_mask_input.to(img_tokens_fused.device).float()
            # mask_tokens = F.adaptive_avg_pool2d(mask_tokens_probs, (H_patch, W_patch))
            # dynamic_mask = (mask_tokens > 0.5).float()
            use_copy_paste_masks = True
        else:
            use_copy_paste_masks = False
            if using_gt_masks:
                mask_tokens_probs = predicted_masks  # Already binary probabilities {0, 1}
            else:
                mask_tokens_probs = torch.sigmoid(predicted_masks)  # Convert logits to probs
        mask_tokens = F.adaptive_avg_pool2d(mask_tokens_probs, (H_patch, W_patch))
        dynamic_mask = (mask_tokens > self.motion_mask_threshold).float()

        # Apply random dropout based on dataset source
        if dataset_sources is not None and self.training and not use_copy_paste_masks:
            for batch_idx in range(b):
                dataset_source = dataset_sources[batch_idx]

                if dataset_source == "re10k" and self.config.model.get(
                    "re10k_random_dropout", False
                ):
                    # RE10K (static): Replace predicted masks with semantic box masks
                    dropout_ratio = self.config.model.get("re10k_dropout_ratio", 0.2)  # Default 20%
                    box_masks = torch.zeros(
                        v_input, 1, H_patch, W_patch, device=dynamic_mask.device
                    )

                    for v_idx in range(v_input):
                        # Number of boxes per view (1-3)
                        num_boxes = torch.randint(1, 4, (1,)).item()
                        target_area_per_box = (dropout_ratio * H_patch * W_patch) / num_boxes

                        for _ in range(num_boxes):
                            # Random box size (aspect ratio between 0.5 and 2.0)
                            aspect_ratio = 0.5 + 1.5 * torch.rand(1).item()  # [0.5, 2.0]
                            box_h = int((target_area_per_box * aspect_ratio) ** 0.5)
                            box_w = int(target_area_per_box / box_h)

                            # Clamp to valid range
                            box_h = min(max(box_h, 2), H_patch)
                            box_w = min(max(box_w, 2), W_patch)

                            # Random position
                            y = torch.randint(0, H_patch - box_h + 1, (1,)).item()
                            x = torch.randint(0, W_patch - box_w + 1, (1,)).item()

                            # Fill box region
                            box_masks[v_idx, 0, y : y + box_h, x : x + box_w] = 1.0

                    # Replace predicted masks with box masks for RE10K
                    start_idx = batch_idx * v_input
                    end_idx = start_idx + v_input
                    dynamic_mask[start_idx:end_idx] = box_masks

                # DRE10K: Add random dropout ON TOP of motion mask (not elif!)
                if dataset_source == "dre10k" and self.config.model.get(
                    "dre10k_random_dropout", False
                ):
                    # DRE10K (dynamic): Add additional random dropout on top of motion mask
                    dropout_ratio = self.config.model.get(
                        "dre10k_dropout_ratio", 0.05
                    )  # Default 5%

                    for v_idx in range(v_input):
                        # Generate small random box masks (1-2 boxes)
                        num_boxes = torch.randint(1, 3, (1,)).item()
                        target_area_per_box = (dropout_ratio * H_patch * W_patch) / num_boxes

                        for _ in range(num_boxes):
                            # Random box size (aspect ratio between 0.5 and 2.0)
                            aspect_ratio = 0.5 + 1.5 * torch.rand(1).item()  # [0.5, 2.0]
                            box_h = int((target_area_per_box * aspect_ratio) ** 0.5)
                            box_w = int(target_area_per_box / box_h)

                            # Clamp to valid range
                            box_h = min(max(box_h, 2), H_patch)
                            box_w = min(max(box_w, 2), W_patch)

                            # Random position
                            y = torch.randint(0, H_patch - box_h + 1, (1,)).item()
                            x = torch.randint(0, W_patch - box_w + 1, (1,)).item()

                            # Fill box region (add to existing motion mask using OR operation)
                            view_global_idx = batch_idx * v_input + v_idx
                            dynamic_mask[view_global_idx, 0, y : y + box_h, x : x + box_w] = 1.0

        # Reshape to match token sequence
        dynamic_mask_flat = rearrange(
            dynamic_mask, "(b v) 1 h w -> b (v h w) 1", b=b, v=v_input
        )  # [b, v_input*n, 1]

        # Token dropping: zero out dynamic tokens (no learnable noise!)
        img_tokens_dropped = img_tokens_fused * (1.0 - dynamic_mask_flat)

        return img_tokens_dropped, dynamic_mask

    def _apply_mae_masking(
        self,
        img_tokens_input,
        predicted_masks,
        b,
        v_input,
        n,
        d,
        dataset_sources=None,
        using_gt_masks=False,
        copy_paste_mask_input=None,
    ):
        """
        Apply MAE-style token masking to input tokens using predicted motion masks.
        """
        if copy_paste_mask_input is None and (predicted_masks is None or not self.use_motion_mask):
            return img_tokens_input

        # Downsample masks from 256×256 to token grid (16×16)
        H_patch = W_patch = int(n**0.5)

        if copy_paste_mask_input is not None:
            mask_tokens_probs = copy_paste_mask_input.to(img_tokens_input.device).float()
            # mask_tokens = F.adaptive_avg_pool2d(mask_tokens_probs, (H_patch, W_patch))
            # dynamic_mask = (mask_tokens > 0.5).float()
        else:
            # Binarize: sigmoid + threshold
            if using_gt_masks:
                mask_tokens_probs = predicted_masks  # Already binary probabilities {0, 1}
            else:
                mask_tokens_probs = torch.sigmoid(predicted_masks)
        mask_tokens = F.adaptive_avg_pool2d(mask_tokens_probs, (H_patch, W_patch))
        dynamic_mask = (mask_tokens > self.motion_mask_threshold).float()  # [b*v_input, 1, 16, 16]

        # Apply random dropout based on dataset source
        if dataset_sources is not None and self.training and copy_paste_mask_input is None:
            for batch_idx in range(b):
                dataset_source = dataset_sources[batch_idx]

                if dataset_source == "re10k" and self.config.model.get(
                    "re10k_random_dropout", False
                ):
                    # RE10K (static): Replace predicted masks with semantic box masks
                    dropout_ratio = self.config.model.get("re10k_dropout_ratio", 0.2)  # Default 20%
                    box_masks = torch.zeros(
                        v_input, 1, H_patch, W_patch, device=dynamic_mask.device
                    )

                    for v_idx in range(v_input):
                        # Number of boxes per view (1-3)
                        num_boxes = torch.randint(1, 4, (1,)).item()
                        target_area_per_box = (dropout_ratio * H_patch * W_patch) / num_boxes

                        for _ in range(num_boxes):
                            # Random box size (aspect ratio between 0.5 and 2.0)
                            aspect_ratio = 0.5 + 1.5 * torch.rand(1).item()  # [0.5, 2.0]
                            box_h = int((target_area_per_box * aspect_ratio) ** 0.5)
                            box_w = int(target_area_per_box / box_h)

                            # Clamp to valid range
                            box_h = min(max(box_h, 2), H_patch)
                            box_w = min(max(box_w, 2), W_patch)

                            # Random position
                            y = torch.randint(0, H_patch - box_h + 1, (1,)).item()
                            x = torch.randint(0, W_patch - box_w + 1, (1,)).item()

                            # Fill box region
                            box_masks[v_idx, 0, y : y + box_h, x : x + box_w] = 1.0

                    # Replace predicted masks with box masks for RE10K
                    start_idx = batch_idx * v_input
                    end_idx = start_idx + v_input
                    dynamic_mask[start_idx:end_idx] = box_masks

                # DRE10K: Add random dropout ON TOP of motion mask (not elif!)
                if dataset_source == "dre10k" and self.config.model.get(
                    "dre10k_random_dropout", False
                ):
                    # DRE10K (dynamic): Add additional random dropout on top of motion mask
                    dropout_ratio = self.config.model.get(
                        "dre10k_dropout_ratio", 0.05
                    )  # Default 5%

                    for v_idx in range(v_input):
                        # Generate small random box masks (1-2 boxes)
                        num_boxes = torch.randint(1, 3, (1,)).item()
                        target_area_per_box = (dropout_ratio * H_patch * W_patch) / num_boxes

                        for _ in range(num_boxes):
                            # Random box size (aspect ratio between 0.5 and 2.0)
                            aspect_ratio = 0.5 + 1.5 * torch.rand(1).item()  # [0.5, 2.0]
                            box_h = int((target_area_per_box * aspect_ratio) ** 0.5)
                            box_w = int(target_area_per_box / box_h)

                            # Clamp to valid range
                            box_h = min(max(box_h, 2), H_patch)
                            box_w = min(max(box_w, 2), W_patch)

                            # Random position
                            y = torch.randint(0, H_patch - box_h + 1, (1,)).item()
                            x = torch.randint(0, W_patch - box_w + 1, (1,)).item()

                            # Fill box region (add to existing motion mask using OR operation)
                            view_global_idx = batch_idx * v_input + v_idx
                            dynamic_mask[view_global_idx, 0, y : y + box_h, x : x + box_w] = 1.0

        # Reshape to match token sequence
        dynamic_mask_flat = rearrange(
            dynamic_mask, "(b v) 1 h w -> b (v h w) 1", b=b, v=v_input
        )  # [b, v_input*n, 1]

        # MAE-style replacement: keep static tokens, replace dynamic with learnable noise
        noise_tokens = self.noise_token.expand(b, v_input * n, d)
        img_tokens_masked = (
            img_tokens_input * (1.0 - dynamic_mask_flat)  # Keep static
            + noise_tokens * dynamic_mask_flat  # Replace dynamic with noise
        )
        return img_tokens_masked, dynamic_mask

    @torch.no_grad()
    def reconstruct_images(self, data_batch):
        """
        Minimal API used by train.py for visualization.
        Delegates to forward with create_visual=True to produce renderings and visuals.
        Returns an edict with at least: input, target, render, and optional render_video.
        """
        return self.forward(data_batch, create_visual=True, render_video=False)

    def forward(self, data, create_visual=False, render_video=False, iter=0):
        """
        Split all images into two sets, use one set to get scene representation,
        use the other to render & train
        """
        input, target = self.process_data(
            data, has_target_image=True, target_has_input=True, compute_rays=False
        )
        image = input.image * 2.0 - 1.0  # [b, v, c, h, w], range (0,1) to (-1,1)
        b, v_input, c, h, w = image.shape
        image_all = data["image"] * 2.0 - 1.0  # [b, v_all, c, h, w], range (0,1) to (-1,1)
        v_all = image_all.shape[1]
        v_target = (
            target.image.shape[1] if hasattr(target, "image") and target.image is not None else 0
        )
        device = image.device
        batch_idx = torch.arange(b).unsqueeze(1).to(device)

        """se3 pose prediction for all views"""
        # tokenize images, add spatial-temporal p.e.
        img_tokens = self.image_tokenizer(image_all)  # [b * v, n, d]
        _, n, d = img_tokens.shape
        if self.use_pe_embedding_layer:
            img_tokens = self.add_sptial_temporal_pe(img_tokens, b, v_all, h, w)
        img_tokens = rearrange(img_tokens, "(b v) n d -> b (v n) d", b=b, v=v_all)  # [b, v * n, d]

        # get camera tokens, add temporal p.e.
        cam_tokens = self.get_camera_tokens(b, v_all)  # [b, v_all * n_cam, d]
        n_cam = cam_tokens.shape[1] // v_all
        assert n_cam == 1
        cam_tokens = rearrange(cam_tokens, "b (v n) d -> b v n d", v=v_all)  # [b, v_all, n_cam, d]
        cam_tokens = rearrange(cam_tokens, "b v n d -> b (v n) d")  # [b, v_all * n_cam, d]

        # pose estimation for all views
        all_tokens = torch.cat([cam_tokens, img_tokens], dim=1)
        all_tokens = self.run_encoder(all_tokens)
        cam_tokens, _ = all_tokens.split([v_all * n_cam, v_all * n], dim=1)

        # get se3 poses and intrinsics
        cam_tokens = rearrange(cam_tokens, "b (v n) d -> (b v) n d", b=b, v=v_all, n=n_cam)[
            :, 0
        ]  # [b * v_all, d]
        cam_info = self.pose_predictor(
            cam_tokens, v_all
        )  # [b * v_all, num_pose_element+3+4], rot, 3d trans, 4d fxfycxcy
        c2w, fxfycxcy = get_cam_se3(cam_info)  # [b * v_all,4,4], [b * v_all,4]
        normalized = True

        # get plucker ray and embeddings
        plucker_rays = cam_info_to_plucker(
            c2w, fxfycxcy, self.config.model.target_image, normalized=normalized
        )
        plucker_rays = rearrange(plucker_rays, "(b v) c h w -> b v c h w", b=b, v=v_all)
        # Split plucker rays into input and target
        plucker_input = plucker_rays[:, :v_input]  # [b, v_input, 6, h, w]
        plucker_target = (
            plucker_rays[:, v_input : v_input + v_target] if v_target > 0 else None
        )  # [b, v_target, 6, h, w]

        plucker_emb_input_tokens = self.target_pose_tokenizer(plucker_input)  # [b * v_input, n, d]
        plucker_emb_input_render_tokens = self.target_pose_tokenizer2(
            plucker_input
        )  # [b * v_input, n, d]
        plucker_emb_target_tokens = (
            self.target_pose_tokenizer2(plucker_target) if plucker_target is not None else None
        )  # [b * v_target, n, d]

        # ===== COPY-PASTE AUGMENTATION MASKS (GROUND TRUTH TRANSIENTS) =====
        copy_paste_mask = data.get("copy_paste_mask", None)
        copy_paste_mask_input = None
        copy_paste_mask_target = None
        copy_paste_mask_input_flat = None
        copy_paste_mask_target_flat = None

        if copy_paste_mask is not None:
            if not isinstance(copy_paste_mask, torch.Tensor):
                copy_paste_mask = torch.tensor(copy_paste_mask)
            copy_paste_mask = copy_paste_mask.to(device=device, dtype=torch.float32)

            if copy_paste_mask.dim() == 4:
                # Assume shape (views, 1, H, W) → add batch dim
                copy_paste_mask = copy_paste_mask.unsqueeze(0)

            # Ensure shape [b, v_all, 1, H, W]
            shape_is_valid = (
                copy_paste_mask.dim() == 5
                and copy_paste_mask.shape[0] == b
                and copy_paste_mask.shape[1] == v_all
            )

            if not shape_is_valid:
                print(
                    "[Model] Warning: copy_paste_mask shape mismatch. "
                    f"Got {copy_paste_mask.shape}, expected ({b}, {v_all}, 1, H, W). "
                    "Ignoring copy-paste masks for this batch."
                )
                copy_paste_mask = None
            else:
                if v_input > 0:
                    copy_paste_mask_input = copy_paste_mask[:, :v_input]
                    copy_paste_mask_input_flat = rearrange(
                        copy_paste_mask_input, "b v c h w -> (b v) c h w"
                    )

                if v_target > 0:
                    copy_paste_mask_target = copy_paste_mask[:, v_input : v_input + v_target]
                    copy_paste_mask_target_flat = rearrange(
                        copy_paste_mask_target, "b v c h w -> (b v) c h w"
                    )

        # ===== PREDICT MOTION MASKS BEFORE RESHAPING =====
        # Use per-view Plücker embeddings [b*v, n, d] for motion mask predictor
        predicted_masks_input = None
        predicted_masks_target = None
        dinov3_masks = None
        if self.use_motion_mask and self.motion_mask_predictor is not None:

            input_image_tokens = rearrange(img_tokens, "b (v n) d -> (b v) n d", v=v_all)[
                : b * v_input
            ]  # [b*v_input, n, d]
            input_image_for_motion = input.image.reshape(b * v_input, c, h, w)
            # Extract target image tokens from all tokens
            target_image_tokens = rearrange(img_tokens, "b (v n) d -> (b v) n d", v=v_all)[
                b * v_input :
            ]  # [b*v_target, n, d]

            # Prepare target images for motion mask prediction (in range [0, 1])
            target_images_for_motion = target.image.reshape(
                b * v_target, c, h, w
            )  # [b*v_target, 3, h, w]

            # Predict motion masks BEFORE reshaping plucker_emb_input
            predicted_masks_input = self.motion_mask_predictor(
                input_image_for_motion, input_image_tokens, plucker_emb_input_tokens
            )  # [b*v_input, 1, 256, 256]

            if v_target > 0:
                predicted_masks_target = self.motion_mask_predictor(
                    target_images_for_motion, target_image_tokens, plucker_emb_target_tokens
                )  # [b*v_target, 1, 256, 256]

            # Override with GT motion masks if provided (inference only)
            # Store flag to skip sigmoid later (GT masks are already binary probabilities)
            using_gt_masks_input = False
            using_gt_masks_target = False
            # Safely check inference config (might be bool or dict)
            inference_cfg = self.config.inference
            if hasattr(inference_cfg, "get"):
                use_gt_motion_masks = inference_cfg.get("use_gt_motion_masks", False)
            else:
                use_gt_motion_masks = False
            if not self.training and use_gt_motion_masks:
                if hasattr(input, "gt_motion_masks") and input.gt_motion_masks is not None:
                    gt_masks = input.gt_motion_masks.to(predicted_masks_input.device).float()
                    predicted_masks_input = gt_masks.view(-1, 1, 256, 256)
                    using_gt_masks_input = True
                    print(
                        "[Model] Using GT motion masks for INPUT views "
                        f"(shape: {predicted_masks_input.shape})"
                    )

                if (
                    v_target > 0
                    and hasattr(target, "gt_motion_masks")
                    and target.gt_motion_masks is not None
                ):
                    gt_masks = target.gt_motion_masks.to(predicted_masks_target.device).float()
                    predicted_masks_target = gt_masks.view(-1, 1, 256, 256)
                    using_gt_masks_target = True
                    print(
                        "[Model] Using GT motion masks for TARGET views "
                        f"(shape: {predicted_masks_target.shape})"
                    )

        # ===== NOW RESHAPE FOR RENDERING PIPELINE =====
        plucker_emb_input_flat = rearrange(
            plucker_emb_input_tokens, "(b v) n d -> b (v n) d", b=b, v=v_input
        )  # [b, v_input * n, d]

        # breakpoint()
        # ## ===== TEMPORARY VISUALIZATION CODE =====
        # import matplotlib.pyplot as plt
        # import numpy as np
        # from pathlib import Path

        # vis_dir = Path("./experiments/temp_vis")
        # vis_dir.mkdir(parents=True, exist_ok=True)

        # # Get first sample from batch
        # sample_idx = 0

        # # Get input images [b, v_input, 3, 256, 256]
        # input_imgs = data["image"][sample_idx, :v_input].cpu()  # [v_input, 3, 256, 256]

        # # Get target images [b, v_target, 3, 256, 256]
        # target_imgs = data["image"][sample_idx, v_input:].cpu()  # [v_target, 3, 256, 256]

        # # Get masks for this sample (detach to avoid gradient issues)
        # input_masks_sample = predicted_masks_input[sample_idx * v_input:(sample_idx + 1) * v_input].detach().cpu()  # [v_input, 1, 256, 256]
        # target_masks_sample = predicted_masks_target[sample_idx * v_target:(sample_idx + 1) * v_target].detach().cpu()  # [v_target, 1, 256, 256]

        # # Convert to numpy and denormalize images
        # def to_numpy_img(img):
        #     # img: [3, 256, 256] in range [0, 1]
        #     img = img.permute(1, 2, 0).numpy()  # [256, 256, 3]
        #     img = np.clip(img, 0, 1)
        #     return img

        # def mask_to_numpy(mask):
        #     # mask: [1, 256, 256] logits
        #     mask_prob = torch.sigmoid(mask).squeeze(0).float().numpy()  # [256, 256]
        #     return mask_prob

        # # Create figure: 2 rows (input, target), columns = max(v_input, v_target)
        # n_cols = max(v_input, v_target)
        # fig, axes = plt.subplots(2, n_cols, figsize=(3 * n_cols, 6))
        # if n_cols == 1:
        #     axes = axes.reshape(2, 1)

        # # Row 0: Input views with masks
        # for i in range(n_cols):
        #     ax = axes[0, i]
        #     if i < v_input:
        #         img = to_numpy_img(input_imgs[i])
        #         mask = mask_to_numpy(input_masks_sample[i])

        #         ax.imshow(img)
        #         ax.imshow(mask, cmap='jet', alpha=0.5, vmin=0, vmax=1)
        #         ax.set_title(f'Input {i}\n(mask overlay)', fontsize=10)
        #     else:
        #         ax.axis('off')
        #     ax.set_xticks([])
        #     ax.set_yticks([])

        # # Row 1: Target views with masks
        # for i in range(n_cols):
        #     ax = axes[1, i]
        #     if i < v_target:
        #         img = to_numpy_img(target_imgs[i])
        #         mask = mask_to_numpy(target_masks_sample[i])

        #         ax.imshow(img)
        #         ax.imshow(mask, cmap='jet', alpha=0.5, vmin=0, vmax=1)
        #         ax.set_title(f'Target {i}\n(mask overlay)', fontsize=10)
        #     else:
        #         ax.axis('off')
        #     ax.set_xticks([])
        #     ax.set_yticks([])

        # plt.suptitle(f'Predicted Motion Masks (Iter {iter}, Sample {sample_idx})', fontsize=14, fontweight='bold')
        # plt.tight_layout()

        # save_path = vis_dir / f'masks_iter_{iter:08d}_{datetime.now().strftime("%Y%m%d_%H%M%S")}.png'
        # plt.savefig(save_path, dpi=150, bbox_inches='tight')
        # plt.close(fig)

        # print(f"[DEBUG] Saved mask visualization to {save_path}")
        # # ===== END TEMPORARY VISUALIZATION CODE =====
        # breakpoint()

        # Generate DINOv3 pseudo-labels (only after rendering, in loss computation)
        # This will be done in the loss computer to ensure rendered images are available

        """predict scene representation using (posed) input views"""
        # get posed image representation
        img_tokens_input = rearrange(img_tokens, "b (v n) d -> b v n d", v=v_all)[
            :, :v_input
        ]  # [b, v_input, n, d]
        img_tokens_input = rearrange(img_tokens_input, "b v n d -> b (v n) d")
        img_tokens_input = torch.cat(
            [img_tokens_input, plucker_emb_input_flat], dim=-1
        )  # [b, v_input * n, 2d]
        img_tokens_input = self.mlp_fuse(img_tokens_input)  # [b, v_input * n, d]

        # ===== TOKEN MASKING OR DROPPING (APPLY TO INPUT TOKENS) =====
        use_token_dropping = self.config.model.get("use_token_dropping", True)

        if use_token_dropping:
            img_tokens_input_masked_fused, input_patch_mask = self._apply_token_dropping(
                img_tokens_input,
                predicted_masks_input,
                b,
                v_input,
                n,
                d,
                dataset_sources=data.get("dataset_source", [None] * b),
                using_gt_masks=using_gt_masks_input,
                copy_paste_mask_input=copy_paste_mask_input_flat,
            )
        else:
            img_tokens_input_masked_fused, input_patch_mask = self._apply_mae_masking(
                img_tokens_input,
                predicted_masks_input,
                b,
                v_input,
                n,
                d,
                dataset_sources=data.get("dataset_source", [None] * b),
                using_gt_masks=using_gt_masks_input,
                copy_paste_mask_input=copy_paste_mask_input_flat,
            )

        ## TEMPORARY VISUALIZATION CODE ###
        ## ===== VISUALIZATION =====
        # breakpoint()
        # from pathlib import Path
        # sample_idx = 0
        # vis_dir = Path("./temp_vis")
        # vis_dir.mkdir(parents=True, exist_ok=True)

        # # Input images and their patch-masks (these masks include any RE10K random-dropout logic)
        # input_imgs = data["image"][sample_idx, :v_input].cpu()  # [v_input, 3, 256, 256]
        # input_masks = input_patch_mask[sample_idx * v_input:(sample_idx + 1) * v_input].detach().cpu()  # [(v_input),1,Hpatch,Wpatch]

        # # Optional: target masks (if you also predicted them)
        # # target_imgs = data["image"][sample_idx, v_input:].cpu()      # [v_target, 3, 256, 256]
        # # target_masks = target_patch_mask[sample_idx*v_target:(sample_idx+1)*v_target].detach().cpu()

        # # Save a grid for inputs
        # n_cols = v_input
        # fig, axes = plt.subplots(1, n_cols, figsize=(3*n_cols, 3))
        # if n_cols == 1:
        #     axes = [axes]

        # for i in range(v_input):
        #     overlay = self.red_patch_overlay(input_imgs[i], input_masks[i])   # PIL
        #     axes[i].imshow(overlay)
        #     axes[i].set_title(f"Input {i}")
        #     axes[i].axis("off")

        # fig.tight_layout()
        # ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        # save_path = vis_dir / f"inputs_overlay_iter_{iter:08d}_{ts}.png"
        # plt.savefig(save_path, dpi=150, bbox_inches="tight")
        # plt.close(fig)
        # print(f"[viz] saved {save_path}")
        # breakpoint()
        ### END TEMPORARY VISUALIZATION CODE ###

        # replicate scene tokens
        scene_tokens = self.scene_code.expand(b, -1, -1)  # [b, n_scene, d]
        n_scene = scene_tokens.shape[1]
        all_tokens = torch.cat(
            [scene_tokens, img_tokens_input_masked_fused], dim=1
        )  # [b, n_scene + v_input * n, d]

        # encoder layers, update scene representation
        all_tokens = self.run_encoder_geom(all_tokens)
        scene_tokens, _ = all_tokens.split(
            [n_scene, v_input * n], dim=1
        )  # [b, n_scene, d], [b, v_input * n, d]

        """render with scene representation from input views and pose of target views"""
        # render target images
        render_input_images = None
        # Safely check inference config (might be bool or dict)
        inference_cfg = self.config.inference
        render_input_preds = False
        if hasattr(inference_cfg, "get"):
            render_input_preds = inference_cfg.get("render_input_predictions", False)

        if v_input > 0 and not self.training and render_input_preds:
            render_input_results = self.render_images(scene_tokens, plucker_emb_input_render_tokens)
            render_input_images = render_input_results.rendered_images

        if plucker_emb_target_tokens is not None and v_target > 0:
            render_results = self.render_images(scene_tokens, plucker_emb_target_tokens)  # dict

            # Generate DINOv3 pseudo-labels after rendering (for motion mask supervision)
            if (
                self.use_dinov3_pseudolabel
                and self.dinov3_pseudolabel_maker is not None
                and v_target > 0
            ):
                # Get dataset sources - PyTorch's default collate converts "dataset_source" (per sample)
                # into "dataset_source" (list) NOT "dataset_sources" (plural)
                dataset_sources = data.get("dataset_source", [None] * b)
                if not isinstance(dataset_sources, list):
                    # Fallback: if it's a single string, replicate it
                    dataset_sources = [dataset_sources] * b
                dinov3_masks_list = []
                cache_hits = 0
                cache_misses = 0

                for batch_idx_i in range(b):
                    is_re10k = dataset_sources[batch_idx_i] == "re10k"
                    uid = int(target.index[batch_idx_i, 0, -1].item())

                    sample_copy_mask = None
                    if copy_paste_mask_target is not None:
                        sample_copy_mask = copy_paste_mask_target[batch_idx_i]
                        if sample_copy_mask.numel() == 0:
                            sample_copy_mask = None

                    if sample_copy_mask is not None:
                        mask_per_sample = sample_copy_mask
                    elif is_re10k:
                        # Static RE10K - skip DINOv3 computation, just use all-zeros
                        # This teaches the model: "static scenes → zero motion mask"
                        mask_per_sample = torch.zeros(v_target, 1, 256, 256, device=device)
                    else:
                        # Dynamic DRE10K - compute or load from cache
                        # Try loading from cache if enabled
                        # SAFE during motion mask training: rendering is frozen, so masks are deterministic!
                        cached_mask = None
                        if self.enable_mask_cache:
                            cached_mask = self._try_load_mask_from_cache(uid)
                            if cached_mask is not None:
                                mask_per_sample = cached_mask.to(device)
                                cache_hits += 1

                        # Compute if cache miss
                        if cached_mask is None:
                            cache_misses += 1

                            # Dynamic DRE10K - generate actual pseudo-labels with DINOv3
                            # Prepare GT and rendered images for this batch element
                            gt_sample = target.image[batch_idx_i]  # [v_target, 3, h, w]
                            pred_sample = render_results.rendered_images[
                                batch_idx_i
                            ]  # [v_target, 3, h, w]

                            # Generate pseudo-labels using DINOv3 with debug visualization
                            debug_path = None
                            if create_visual and batch_idx_i == 0:
                                # Save debug visualization for first batch only
                                if not dist.is_initialized() or dist.get_rank() == 0:
                                    debug_path = (
                                        f"./debug_dinov3/step_{iter:06d}_batch{batch_idx_i}.png"
                                    )

                            mask_per_sample = self.dinov3_pseudolabel_maker(
                                gt_sample, pred_sample, debug_save_path=debug_path
                            )  # [v_target, 1, 256, 256]

                            # Save to cache for future use
                            if self.enable_mask_cache:
                                self._save_mask_to_cache(uid, mask_per_sample)

                    dinov3_masks_list.append(mask_per_sample)

                # Log cache stats every iteration (only for DRE10K samples that were actually computed)
                if self.enable_mask_cache and cache_misses > 0:
                    total_dre10k = cache_hits + cache_misses
                    hit_rate = cache_hits / total_dre10k * 100
                    print(
                        f"[Iter {iter}] DINOv3 mask cache (DRE10K only): {cache_hits}/{total_dre10k} hits ({hit_rate:.1f}% hit rate)"
                    )

                dinov3_masks = torch.stack(dinov3_masks_list, dim=0)  # [b, v_target, 1, 256, 256]
                dinov3_masks = rearrange(
                    dinov3_masks, "b v c h w -> (b v) c h w"
                )  # [b*v_target, 1, 256, 256]

        if dinov3_masks is not None:
            # NOTE: Copy-paste masks are already merged per-sample in dinov3_masks_list
            # at line 2051-2052. No need to replace here.

            # compute loss with motion mask distillation
            # Pass dataset_sources to enable per-dataset masking logic (RE10K vs DRE10K)
            dataset_sources = data.get("dataset_source", None)
            target_images_for_loss = (
                target.image_clean if hasattr(target, "image_clean") else target.image
            )
            loss_metrics = self.loss_computer(
                render_results.rendered_images,
                target_images_for_loss,
                input,
                predicted_target_masks=predicted_masks_target,
                dinov3_target_masks=dinov3_masks,
                predicted_input_masks=predicted_masks_input,
                copy_paste_input_masks=copy_paste_mask_input_flat,
                create_visual=create_visual,
                dataset_sources=dataset_sources,
            )
        else:
            # No pseudo-labels (e.g. inference with use_dinov3_pseudolabel=false).
            # Keep the real render_results from above; only produce a dummy loss.
            device = scene_tokens.device
            loss_metrics = edict(
                loss=torch.tensor(0.0, device=device, requires_grad=True),
                l2_loss=torch.tensor(0.0, device=device),
                psnr=torch.tensor(0.0, device=device),
                lpips_loss=torch.tensor(0.0, device=device),
                perceptual_loss=torch.tensor(0.0, device=device),
                ssim_loss=torch.tensor(0.0, device=device),
                visual=None,
            )
            if v_target == 0 or plucker_emb_target_tokens is None:
                render_results = edict(rendered_images=None)

        vis_only_results = None
        if (
            create_visual
            and v_target > 0
            and not (self.training and self.copy_paste_training_enabled)
        ):
            with torch.no_grad():
                c2w_all_reshaped = rearrange(c2w, "(b v) c d -> b v c d", v=v_all)
                fxfycxcy_all_reshaped = rearrange(fxfycxcy, "(b v) c -> b v c", v=v_all)

                c2w_target = c2w_all_reshaped[
                    :, v_input : v_input + v_target
                ]  # [b, v_target, 4, 4]
                fxfycxcy_target = fxfycxcy_all_reshaped[
                    :, v_input : v_input + v_target
                ]  # [b, v_target, 4]
                c2w_target = rearrange(c2w_target, "b v c d -> (b v) c d")
                fxfycxcy_target = rearrange(fxfycxcy_target, "b v c -> (b v) c")
                vis_only_results = self.render_images_video(
                    scene_tokens, c2w_target, fxfycxcy_target, normalized=normalized
                )

        # Reshape masks for visualization (from flattened to per-batch)
        if dinov3_masks is not None:
            dinov3_masks_viz = rearrange(dinov3_masks, "(b v) c h w -> b v c h w", b=b, v=v_target)
        else:
            dinov3_masks_viz = None

        if predicted_masks_target is not None:
            predicted_masks_viz = rearrange(
                predicted_masks_target, "(b v) c h w -> b v c h w", b=b, v=v_target
            )
        else:
            predicted_masks_viz = None

        # return results
        if hasattr(target, "image_clean"):
            target.image_aug = target.image
            target.image = target.image_clean

        result = edict(
            input=input,
            target=target,
            loss_metrics=loss_metrics,
            render=render_results.rendered_images,
            c2w=rearrange(c2w, "(b v) c d -> b v c d", b=b, v=v_all),
            fxfycxcy=fxfycxcy,
        )

        if render_input_images is not None:
            result.render_input = render_input_images

        # Add motion mask visualization keys if available
        if dinov3_masks_viz is not None:
            result.dinov3_masks = dinov3_masks_viz.detach()
        if predicted_masks_viz is not None:
            result.predicted_target_masks = predicted_masks_viz.detach()

        if input_patch_mask is not None:
            result.input_patch_mask = input_patch_mask.detach()  # [b*v_input, 1, H_patch, W_patch]
        if predicted_masks_input is not None:
            result.predicted_input_masks = (
                predicted_masks_input.detach()
            )  # [b*v_input, 1, H_patch, W_patch]
        if predicted_masks_target is not None:
            result.predicted_target_masks = (
                predicted_masks_target.detach()
            )  # [b*v_target, 1, H_patch, W_patch]

        if vis_only_results is not None:
            result.video_rendering = vis_only_results.rendered_images_video.detach()

        return result

    def render_images(self, scene_tokens, target_tokens):
        """
        Render target views based on the scene representation, target view pose tokens and target view tokens
        Args:
            scene_tokens: [b, n_scene, d]
            target_tokens: plucker embedding tokens of input images [b*v, n_target, d]
        Return:
            rendered_images: rendered target views in [b,v,c,h,w]
        """
        b, _, d = scene_tokens.shape
        bv = target_tokens.shape[0]
        v = bv // b

        # repeat scene tokens
        scene_tokens = scene_tokens.unsqueeze(1).repeat(1, v, 1, 1)
        scene_tokens = rearrange(scene_tokens, "b v n d -> (b v) n d")  # [b*v, n_scene, d]

        # get all tokens
        n_target, n_scene = target_tokens.shape[1], scene_tokens.shape[1]
        all_tokens = torch.cat(
            [target_tokens, scene_tokens], dim=1
        )  # [b*v, n_target+n_pose+n_scene, d]

        # render
        rendered_images_all = self.render(all_tokens, n_target, n_scene, v)  # [b,v,c,h,w]

        render_results = edict(rendered_images=rendered_images_all)
        return render_results

    def render(self, all_tokens, n_target, n_scene, v):
        """
        Run decoder layers and output mlp to render target views
        Args:
            all_tokens: [b*v, n_target+n_pose+n_scene, d]
        Return:
            rendered views: [b, v, c, h, w]
        """
        all_tokens = self.decoder_ln(all_tokens)
        all_tokens = self.run_decoder(all_tokens)

        # split tokens
        target_tokens, _ = all_tokens.split([n_target, n_scene], dim=1)  # [b*v, n_target, d]

        # regress image
        rendered_images_all = self.image_token_decoder(target_tokens)  # [b*v, n_target, p*p*3]
        patch_size = self.config.model.target_image.patch_size
        rendered_images_all = rearrange(
            rendered_images_all,
            "(b v) (h w) (p1 p2 c) -> b v c (h p1) (w p2)",
            v=v,
            h=self.target_latent_h,
            w=self.target_latent_w,
            p1=patch_size,
            p2=patch_size,
            c=3,
        )
        return rendered_images_all

    def get_camera_tokens(self, b, v):
        n, d = self.cam_code.shape[-2:]

        cam_tokens = rearrange(self.cam_code, "n d -> 1 1 n d")  # [1, 1, n_cam, d]
        cam_tokens = repeat(cam_tokens, "1 1 n d -> 1 v n d", v=v)  # [1, v, n_cam, d]
        cam_tokens = rearrange(cam_tokens, "1 v n d -> 1 (v n) d")  # [1, v*n_cam, d]

        # get temporal pe
        img_indices = torch.arange(v).repeat_interleave(n)  # [v*n_cam]
        img_indices = img_indices.to(cam_tokens.device)
        temporal_pe = get_1d_sincos_pos_emb_from_grid(
            embed_dim=d, pos=img_indices, device=cam_tokens.device
        ).to(
            cam_tokens.dtype
        )  # [v*n_cam, d]
        temporal_pe = temporal_pe.reshape(1, v * n, d)  # [1, v*n_cam, d]
        temporal_pe = self.temporal_pe_embedder(temporal_pe)  # [1, v*n_cam, d]

        return (cam_tokens + temporal_pe).repeat(b, 1, 1)  # [b, v*n_cam, d]

    def add_sptial_temporal_pe(self, img_tokens, b, v, h_origin, w_origin):
        """
        Adding spatial-temporal pe to input image tokens
        Args:
            img_tokens: shape [b*v, n, d]
        Return:
            image tokens with positional embedding
        """
        patch_size = self.config.model.image_tokenizer.patch_size
        num_h_tokens = h_origin // patch_size
        num_w_tokens = w_origin // patch_size
        assert (num_h_tokens * num_w_tokens) == img_tokens.shape[1]
        bv, n, d = img_tokens.shape

        # get temporal pe
        img_indices = torch.arange(v).repeat_interleave(n)  # [v*n]
        img_indices = img_indices.unsqueeze(0).repeat(b, 1).reshape(-1)  # [b*v*n]
        img_indices = img_indices.to(img_tokens.device)
        temporal_pe = get_1d_sincos_pos_emb_from_grid(
            embed_dim=d // 2, pos=img_indices, device=img_tokens.device
        ).to(
            img_tokens.dtype
        )  # [b*v*n, d2]
        temporal_pe = temporal_pe.reshape(b, v, n, d // 2)  # [b,v,n,d2]

        # get spatial pe
        spatial_pe = get_2d_sincos_pos_embed(
            embed_dim=d // 2,
            grid_size=(num_h_tokens, num_w_tokens),
            device=img_tokens.device,
        ).to(
            img_tokens.dtype
        )  # [n, d3]
        spatial_pe = spatial_pe.reshape(1, 1, n, d // 2).repeat(b, v, 1, 1)  # [b,v,n,d3]

        # embed pe
        pe = self.pe_embedder(
            torch.cat([spatial_pe, temporal_pe], dim=-1).reshape(bv, n, d)
        )  # [b*v,n,d]

        return img_tokens + pe

    def run_layers_encoder(self, start, end):
        def custom_forward(concat_nerf_img_tokens):
            for i in range(start, min(end, len(self.transformer_encoder))):
                concat_nerf_img_tokens = self.transformer_encoder[i](concat_nerf_img_tokens)
            return concat_nerf_img_tokens

        return custom_forward

    def run_layers_encoder_geom(self, start, end):
        def custom_forward(concat_nerf_img_tokens):
            for i in range(start, min(end, len(self.transformer_encoder_geom))):
                concat_nerf_img_tokens = self.transformer_encoder_geom[i](concat_nerf_img_tokens)
            return concat_nerf_img_tokens

        return custom_forward

    def run_layers_decoder(self, start, end):
        def custom_forward(concat_nerf_img_tokens):
            for i in range(start, min(end, len(self.transformer_decoder))):
                concat_nerf_img_tokens = self.transformer_decoder[i](concat_nerf_img_tokens)
            return concat_nerf_img_tokens

        return custom_forward

    def run_encoder(self, all_tokens_encoder):
        checkpoint_every = self.config.training.grad_checkpoint_every
        for i in range(0, len(self.transformer_encoder), checkpoint_every):
            all_tokens_encoder = torch.utils.checkpoint.checkpoint(
                self.run_layers_encoder(i, i + 1),
                all_tokens_encoder,
                use_reentrant=False,
            )
            if checkpoint_every > 1:
                all_tokens_encoder = self.run_layers_encoder(i + 1, i + checkpoint_every)(
                    all_tokens_encoder
                )
        return all_tokens_encoder

    def run_encoder_geom(self, all_tokens_encoder):
        checkpoint_every = self.config.training.grad_checkpoint_every
        for i in range(0, len(self.transformer_encoder_geom), checkpoint_every):
            all_tokens_encoder = torch.utils.checkpoint.checkpoint(
                self.run_layers_encoder_geom(i, i + 1),
                all_tokens_encoder,
                use_reentrant=False,
            )
            if checkpoint_every > 1:
                all_tokens_encoder = self.run_layers_encoder_geom(i + 1, i + checkpoint_every)(
                    all_tokens_encoder
                )
        return all_tokens_encoder

    def run_decoder(self, all_tokens_encoder):
        checkpoint_every = self.config.training.grad_checkpoint_every
        for i in range(0, len(self.transformer_decoder), checkpoint_every):
            all_tokens_encoder = torch.utils.checkpoint.checkpoint(
                self.run_layers_decoder(i, i + 1),
                all_tokens_encoder,
                use_reentrant=False,
            )
            if checkpoint_every > 1:
                all_tokens_encoder = self.run_layers_decoder(i + 1, i + checkpoint_every)(
                    all_tokens_encoder
                )
        return all_tokens_encoder

    def red_patch_overlay(self, img_3chw, patch_mask_1hw, alpha=0.5):
        """
        img_3chw: torch.Tensor [3, H, W], values in [0,1]
        patch_mask_1hw: torch.Tensor [1, Hpatch, Wpatch] with {0,1} per token
        returns: PIL.Image (H, W, 3) with red patches overlayed
        """
        H, W = img_3chw.shape[-2:]
        # nearest upsample so each token becomes a solid square
        mask_hw = F.interpolate(patch_mask_1hw.unsqueeze(0), size=(H, W), mode="nearest").squeeze(
            0
        )  # [1,H,W]
        img = img_3chw.permute(1, 2, 0).cpu().numpy()  # [H,W,3], 0..1
        mask = mask_hw.squeeze(0).cpu().numpy()[..., None]  # [H,W,1], {0,1}

        red = np.zeros_like(img)
        red[..., 0] = 1.0  # pure red
        vis = img * (1.0 - alpha * mask) + red * (alpha * mask)  # simple alpha blend
        vis = (np.clip(vis, 0, 1) * 255).astype(np.uint8)
        return Image.fromarray(vis)

    def render_images_video(self, scene_tokens_all, c2w_all, fxfycxcy_all, normalized=False):
        """
        scene_tokens_all: [b, n_scene, d]
        c2w_all: [b*v, 4, 4]
        fxfycxcy_all: [b*v, 4]
        """
        with torch.no_grad():
            scene_tokens_all = scene_tokens_all.detach()
            c2w_all = c2w_all.detach()
            fxfycxcy_all = fxfycxcy_all.detach()

            b, _, d = scene_tokens_all.shape
            bv = c2w_all.shape[0]
            v = bv // b
            c2w_all = rearrange(c2w_all, "(b v) x y -> b v x y", v=v)
            fxfycxcy_all = rearrange(fxfycxcy_all, "(b v) x -> b v x", v=v)
            device = scene_tokens_all.device

            all_renderings = []
            num_frames = 30
            traj_type = "interpolate"
            order_poses = False

            for i in range(b):
                scene_tokens = scene_tokens_all[i]  # [n_scene, d]
                c2ws = c2w_all[i]  # [v, 4, 4]
                fxfycxcy = fxfycxcy_all[i]  # [v, 4]
                if traj_type == "interpolate":
                    # build Ks from fxfycxcy
                    Ks = torch.zeros((c2ws.shape[0], 3, 3), device=device)
                    Ks[:, 0, 0] = fxfycxcy[:, 0]
                    Ks[:, 1, 1] = fxfycxcy[:, 1]
                    Ks[:, 0, 2] = fxfycxcy[:, 2]
                    Ks[:, 1, 2] = fxfycxcy[:, 3]
                    c2ws, Ks = camera_utils.get_interpolated_poses_many(
                        c2ws[:, :3, :4], Ks, num_frames, order_poses=order_poses
                    )
                    frame_c2ws = torch.cat(
                        [
                            c2ws.to(device),
                            torch.tensor([[[0, 0, 0, 1]]], device=device).repeat(
                                c2ws.shape[0], 1, 1
                            ),
                        ],
                        dim=1,
                    )  # [v',4,4]
                    frame_fxfycxcy = torch.zeros((c2ws.shape[0], 4), device=device)
                    frame_fxfycxcy[:, 0] = Ks[:, 0, 0]  # [v',4]
                    frame_fxfycxcy[:, 1] = Ks[:, 1, 1]
                    frame_fxfycxcy[:, 2] = Ks[:, 0, 2]
                    frame_fxfycxcy[:, 3] = Ks[:, 1, 2]
                elif traj_type == "same":
                    frame_c2ws = c2ws.clone()
                    frame_fxfycxcy = fxfycxcy.clone()
                else:
                    raise NotImplementedError

                plucker_rays = cam_info_to_plucker(
                    frame_c2ws,
                    frame_fxfycxcy,
                    self.config.model.target_image,
                    normalized=normalized,
                )  # [v',6,h,w]
                plucker_embeddings = self.target_pose_tokenizer2(
                    plucker_rays.unsqueeze(0)
                )  # [v',n_target,d]

                v_render = plucker_embeddings.shape[0]

                scene_tokens = scene_tokens.unsqueeze(0).repeat(v_render, 1, 1)  # [v',n_scene,d]
                all_tokens = torch.cat(
                    [plucker_embeddings, scene_tokens], dim=1
                )  # [v', n_target+n_scene, d]

                # render
                n_target, n_scene = plucker_embeddings.shape[1], scene_tokens.shape[1]
                rendered_images_all = self.render(all_tokens, n_target, n_scene, v_render).squeeze(
                    0
                )  # [v',c,h,w]
                all_renderings.append(rendered_images_all)

            all_renderings = torch.stack(all_renderings)  # [b,v',c,h,w]

        render_results = edict(rendered_images_video=all_renderings)
        return render_results

    @torch.no_grad()
    def visualize_dinov3_pseudolabels(self, data_batch, save_dir, render_only=False):
        """
        Visualize complete DINOv3 pseudo-mask generation pipeline.

        Shows the full pipeline: GT, Rendered, SSIM dissim, DINO dissim,
        soft mask (percentile-based), and binary co-seg mask (KMeans+GrabCut).

        Args:
            data_batch: Input batch from dataloader
            save_dir: Directory to save visualizations
            render_only: If True, skip pseudo-label generation (just show GT vs rendered)

        Saves:
            - 4-column grid: GT | Rendered | Soft Mask on GT | Binary Mask on GT
            - Per-view metrics (PSNR, LPIPS, SSIM)
            - Debug visualization with SSIM/DINO/Combined dissimilarity
        """
        import matplotlib.pyplot as plt
        from matplotlib import cm

        os.makedirs(save_dir, exist_ok=True)

        # Run forward pass to get renderings
        result = self.forward(data_batch, create_visual=False, render_video=False)

        input, target = result.input, result.target
        rendered_images = result.render  # [b, v_target, 3, h, w]

        b = input.image.shape[0]
        v_target = (
            target.image.shape[1] if hasattr(target, "image") and target.image is not None else 0
        )
        h, w = target.image.shape[-2:]
        device = target.image.device

        # Generate complete pseudo-labels (soft + binary)
        soft_masks = None
        binary_masks = None
        debug_info = None

        if not render_only and self.use_dinov3_pseudolabel and self.dinov3_pseudolabel_maker:
            dataset_sources = data_batch.get("dataset_sources", [None] * b)

            soft_masks_list = []
            binary_masks_list = []
            debug_info_list = []

            for batch_idx_i in range(b):
                is_re10k = dataset_sources[batch_idx_i] == "re10k"

                if is_re10k:
                    # Static RE10K - use all-zeros
                    soft_mask = torch.zeros(v_target, 1, h, w, device=device)
                    binary_mask = torch.zeros(v_target, 1, h, w, device=device)
                    debug = None
                else:
                    # Dynamic RE10K - compute full pseudo-mask pipeline
                    gt_sample = target.image[batch_idx_i]  # [v_target, 3, h, w]
                    pred_sample = rendered_images[batch_idx_i]  # [v_target, 3, h, w]

                    # Generate soft mask (DINO + SSIM percentile-based)
                    self.dinov3_pseudolabel_maker.use_coseg_binary = False
                    soft_mask = self.dinov3_pseudolabel_maker(gt_sample, pred_sample)

                    # Generate binary co-seg mask (KMeans + GrabCut)
                    self.dinov3_pseudolabel_maker.use_coseg_binary = True
                    binary_mask = self.dinov3_pseudolabel_maker(gt_sample, pred_sample)

                    # Extract debug info (SSIM, DINO dissimilarity)
                    dino_dissim = self.dinov3_pseudolabel_maker._compute_cosine_dissimilarity(
                        gt_sample, pred_sample
                    )
                    ssim_dissim = self.dinov3_pseudolabel_maker._compute_ssim_dissimilarity(
                        gt_sample, pred_sample
                    )
                    combined_dissim = torch.minimum(dino_dissim, ssim_dissim)

                    debug = {
                        "dino_dissim": dino_dissim.cpu(),
                        "ssim_dissim": ssim_dissim.cpu(),
                        "combined_dissim": combined_dissim.cpu(),
                    }

                soft_masks_list.append(soft_mask)
                binary_masks_list.append(binary_mask)
                debug_info_list.append(debug)

            soft_masks = torch.stack(soft_masks_list, dim=0)  # [b, v_target, 1, h, w]
            binary_masks = torch.stack(binary_masks_list, dim=0)  # [b, v_target, 1, h, w]
            debug_info = debug_info_list

        # Compute metrics
        psnr_per_view = []
        lpips_per_view = []
        ssim_per_view = []

        for batch_idx_i in range(b):
            psnr = compute_psnr(target.image[batch_idx_i], rendered_images[batch_idx_i])
            lpips = compute_lpips(target.image[batch_idx_i], rendered_images[batch_idx_i])
            ssim = compute_ssim(
                target.image[batch_idx_i].to(torch.float32),
                rendered_images[batch_idx_i].to(torch.float32),
            )
            psnr_per_view.append(psnr)
            lpips_per_view.append(lpips)
            ssim_per_view.append(ssim)

        # Save visualizations
        for batch_idx_i in range(b):
            scene_idx = int(target.index[batch_idx_i, 0, -1].item())
            dataset_source = data_batch.get("dataset_sources", ["unknown"] * b)[batch_idx_i]

            # Get scene name and original filenames for saving masks
            scene_name = data_batch.get("scene_name", [f"scene_{scene_idx:06d}"] * b)[batch_idx_i]
            original_filenames = data_batch.get("original_filenames", [None] * b)[batch_idx_i]

            scene_dir = os.path.join(save_dir, f"scene_{scene_idx:06d}_{dataset_source}")
            os.makedirs(scene_dir, exist_ok=True)

            # Create masks directory for .npy files
            masks_dir = os.path.join(scene_dir, "masks_npy")
            os.makedirs(masks_dir, exist_ok=True)

            gt_imgs = target.image[batch_idx_i].cpu()  # [v_target, 3, h, w]
            rendered_imgs = rendered_images[batch_idx_i].cpu()  # [v_target, 3, h, w]

            v_target_curr = gt_imgs.shape[0]

            # Save binary masks as .npy files with original image names
            if binary_masks is not None and original_filenames is not None:
                binary_masks_curr = binary_masks[batch_idx_i].cpu().numpy()  # [v_target, 1, h, w]
                soft_masks_curr_np = soft_masks[batch_idx_i].cpu().numpy()  # [v_target, 1, h, w]

                # Get number of input views to correctly index target filenames
                v_input = input.image.shape[1]

                for v_idx in range(v_target_curr):
                    # Get original filename for this target view
                    # original_filenames contains all views (input + target)
                    # Target views start after input views
                    original_path = original_filenames[v_input + v_idx]
                    # Extract filename without extension (e.g., "00042" from "images/scene/00042.png")
                    from pathlib import Path

                    image_name = Path(original_path).stem  # e.g., "00042"

                    # Save binary mask as .npy
                    binary_mask_np = binary_masks_curr[v_idx, 0]  # [h, w]
                    binary_mask_path = os.path.join(masks_dir, f"{image_name}_binary.npy")
                    np.save(binary_mask_path, binary_mask_np.astype(np.float32))

                    # Also save soft mask for comparison
                    soft_mask_np = soft_masks_curr_np[v_idx, 0]  # [h, w]
                    soft_mask_path = os.path.join(masks_dir, f"{image_name}_soft.npy")
                    np.save(soft_mask_path, soft_mask_np.astype(np.float32))

                # Save a metadata file for this scene
                metadata = {
                    "scene_name": scene_name,
                    "scene_idx": scene_idx,
                    "dataset_source": dataset_source,
                    "num_views": v_target_curr,
                    "mask_files": [
                        f"{Path(original_filenames[v_input + v_idx]).stem}_binary.npy"
                        for v_idx in range(v_target_curr)
                    ],
                    "original_filenames": [
                        original_filenames[v_input + v_idx] for v_idx in range(v_target_curr)
                    ],
                }
                metadata_path = os.path.join(masks_dir, "metadata.json")
                import json

                with open(metadata_path, "w") as f:
                    json.dump(metadata, f, indent=2)

                print(f"  ✓ Saved {v_target_curr} mask pairs (.npy) for scene {scene_name}")

            # Main visualization: GT | Rendered | Soft Mask on GT | Binary Mask on GT
            if soft_masks is not None and binary_masks is not None:
                fig, axes = plt.subplots(v_target_curr, 4, figsize=(20, 5 * v_target_curr))
                if v_target_curr == 1:
                    axes = axes.reshape(1, -1)

                soft_masks_curr = soft_masks[batch_idx_i].cpu()  # [v_target, 1, h, w]
                binary_masks_curr = binary_masks[batch_idx_i].cpu()  # [v_target, 1, h, w]
                cmap = cm.get_cmap("hot")

                for v_idx in range(v_target_curr):
                    gt_np = gt_imgs[v_idx].permute(1, 2, 0).float().numpy()
                    rendered_np = rendered_imgs[v_idx].permute(1, 2, 0).float().numpy()
                    soft_mask_np = soft_masks_curr[v_idx, 0].float().numpy()
                    binary_mask_np = binary_masks_curr[v_idx, 0].float().numpy()

                    psnr_val = psnr_per_view[batch_idx_i][v_idx].item()
                    lpips_val = lpips_per_view[batch_idx_i][v_idx].item()
                    ssim_val = ssim_per_view[batch_idx_i][v_idx].item()

                    # Column 1: GT
                    axes[v_idx, 0].imshow(gt_np)
                    title = f"GT View {v_idx}"
                    if v_idx == 0:
                        title += f"\n({dataset_source.upper()})"
                    axes[v_idx, 0].set_title(title, fontsize=12, fontweight="bold")
                    axes[v_idx, 0].axis("off")

                    # Column 2: Rendered
                    axes[v_idx, 1].imshow(rendered_np)
                    axes[v_idx, 1].set_title(
                        f"Rendered\nPSNR:{psnr_val:.2f} SSIM:{ssim_val:.3f}",
                        fontsize=12,
                        fontweight="bold",
                    )
                    axes[v_idx, 1].axis("off")

                    # Column 3: Soft Mask on GT
                    axes[v_idx, 2].imshow(gt_np)
                    mask_colored = cmap(soft_mask_np)[:, :, :3]
                    axes[v_idx, 2].imshow(mask_colored, alpha=0.6)
                    axes[v_idx, 2].set_title(
                        f"Soft Mask (Percentile)\n[0,1] continuous", fontsize=12, fontweight="bold"
                    )
                    axes[v_idx, 2].axis("off")

                    # Column 4: Binary Mask on GT
                    axes[v_idx, 3].imshow(gt_np)
                    binary_colored = cmap(binary_mask_np)[:, :, :3]
                    axes[v_idx, 3].imshow(binary_colored, alpha=0.6)
                    axes[v_idx, 3].set_title(
                        f"Binary Mask (Co-Seg)\n{{0,1}} KMeans+GrabCut",
                        fontsize=12,
                        fontweight="bold",
                    )
                    axes[v_idx, 3].axis("off")

                plt.tight_layout()
                plt.savefig(
                    os.path.join(scene_dir, "pseudolabels.png"), dpi=150, bbox_inches="tight"
                )
                plt.close()

                # Debug visualization: SSIM | DINO | Combined dissimilarity
                if debug_info[batch_idx_i] is not None:
                    debug = debug_info[batch_idx_i]
                    fig, axes = plt.subplots(v_target_curr, 3, figsize=(15, 5 * v_target_curr))
                    if v_target_curr == 1:
                        axes = axes.reshape(1, -1)

                    for v_idx in range(v_target_curr):
                        dino = debug["dino_dissim"][v_idx, 0].float().numpy()
                        ssim_d = debug["ssim_dissim"][v_idx, 0].float().numpy()
                        combined = debug["combined_dissim"][v_idx, 0].float().numpy()

                        # SSIM dissimilarity
                        im = axes[v_idx, 0].imshow(ssim_d, cmap="hot", vmin=0, vmax=1)
                        axes[v_idx, 0].set_title(
                            f"SSIM Dissimilarity\nView {v_idx}", fontsize=12, fontweight="bold"
                        )
                        axes[v_idx, 0].axis("off")
                        plt.colorbar(im, ax=axes[v_idx, 0], fraction=0.046)

                        # DINO dissimilarity
                        im = axes[v_idx, 1].imshow(dino, cmap="hot", vmin=0, vmax=1)
                        axes[v_idx, 1].set_title(
                            f"DINO Dissimilarity\nView {v_idx}", fontsize=12, fontweight="bold"
                        )
                        axes[v_idx, 1].axis("off")
                        plt.colorbar(im, ax=axes[v_idx, 1], fraction=0.046)

                        # Combined (min)
                        im = axes[v_idx, 2].imshow(combined, cmap="hot", vmin=0, vmax=1)
                        axes[v_idx, 2].set_title(
                            f"Combined (min)\nView {v_idx}", fontsize=12, fontweight="bold"
                        )
                        axes[v_idx, 2].axis("off")
                        plt.colorbar(im, ax=axes[v_idx, 2], fraction=0.046)

                    plt.tight_layout()
                    plt.savefig(
                        os.path.join(scene_dir, "debug_dissimilarity.png"),
                        dpi=150,
                        bbox_inches="tight",
                    )
                    plt.close()

            else:
                # Fallback: GT vs Rendered only
                fig, axes = plt.subplots(v_target_curr, 2, figsize=(10, 5 * v_target_curr))
                if v_target_curr == 1:
                    axes = axes.reshape(1, -1)

                for v_idx in range(v_target_curr):
                    gt_np = gt_imgs[v_idx].permute(1, 2, 0).float().numpy()
                    rendered_np = rendered_imgs[v_idx].permute(1, 2, 0).float().numpy()

                    psnr_val = psnr_per_view[batch_idx_i][v_idx].item()
                    ssim_val = ssim_per_view[batch_idx_i][v_idx].item()

                    axes[v_idx, 0].imshow(gt_np)
                    axes[v_idx, 0].set_title(f"GT View {v_idx}", fontsize=12)
                    axes[v_idx, 0].axis("off")

                    axes[v_idx, 1].imshow(rendered_np)
                    axes[v_idx, 1].set_title(
                        f"Rendered\nPSNR:{psnr_val:.2f} SSIM:{ssim_val:.3f}", fontsize=12
                    )
                    axes[v_idx, 1].axis("off")

                plt.tight_layout()
                plt.savefig(os.path.join(scene_dir, "rendering.png"), dpi=150, bbox_inches="tight")
                plt.close()

            # Save metrics
            with open(os.path.join(scene_dir, "metrics.txt"), "w") as f:
                f.write(f"Scene: {scene_idx:06d}\n")
                f.write(f"Dataset: {dataset_source}\n")
                f.write(f"Num views: {v_target_curr}\n\n")
                f.write("Per-view metrics:\n")
                for v_idx in range(v_target_curr):
                    f.write(f"  View {v_idx}: ")
                    f.write(f"PSNR={psnr_per_view[batch_idx_i][v_idx].item():.2f}, ")
                    f.write(f"LPIPS={lpips_per_view[batch_idx_i][v_idx].item():.4f}, ")
                    f.write(f"SSIM={ssim_per_view[batch_idx_i][v_idx].item():.4f}\n")
                f.write(f"\nAverage:\n")
                f.write(f"  PSNR: {psnr_per_view[batch_idx_i].mean().item():.2f}\n")
                f.write(f"  LPIPS: {lpips_per_view[batch_idx_i].mean().item():.4f}\n")
                f.write(f"  SSIM: {ssim_per_view[batch_idx_i].mean().item():.4f}\n")

            print(
                f"✓ Saved scene {scene_idx:06d} ({dataset_source}): "
                f"PSNR={psnr_per_view[batch_idx_i].mean().item():.2f}"
            )

        print(f"\n✓ Visualization complete! Saved to: {save_dir}")

    @torch.no_grad()
    def save_visuals(self, out_dir, result, batch, save_all=False):
        os.makedirs(out_dir, exist_ok=True)

        input, target = result.input, result.target

        # save comparison
        if result.loss_metrics.visual is not None:
            uids = [target.index[b, 0, -1].item() for b in range(target.index.size(0))]

            uid_firstlast = f"{uids[0]:08}_{uids[-1]:08}"
            Image.fromarray(result.loss_metrics.visual).save(
                os.path.join(out_dir, f"supervision_{uid_firstlast}.jpg")
            )
            with open(os.path.join(out_dir, f"uids.txt"), "w") as f:
                uids = "_".join([f"{uid:08}" for uid in uids])
                f.write(uids)

        # save input, rendered interpolated imaged, video
        for b in range(input.image.size(0)):
            uid = input.index[b, 0, -1].item()

            # vis input
            vis_image = rearrange(input.image[b], "v c h w -> h (v w) c")
            vis_image = (vis_image.cpu().numpy() * 255.0).clip(0.0, 255.0).astype(np.uint8)
            Image.fromarray(vis_image).save(os.path.join(out_dir, f"input_{uid}.png"))

            if ("render_video" in result.keys()) and torch.is_tensor(result.render_video):
                render_video = result.render_video[b].detach().cpu()  # [v, c, h, w]
                render_video = np.ascontiguousarray(np.array(render_video.to(torch.float32)))
                render_video = rearrange(render_video, "v c h w -> v h w c")
                create_video_from_frames(
                    render_video, f"{out_dir}/render_video_{uid}.mp4", framerate=30
                )

            if not save_all:
                break

    @torch.no_grad()
    def save_evaluations(self, out_dir, result, batch, dataset):
        os.makedirs(out_dir, exist_ok=True)
        input, target = result.input, result.target
        for b in range(input.image.size(0)):
            uid = input.index[b, 0, -1].item()
            curr_out_dir = os.path.join(out_dir, "%06d" % (uid))
            os.makedirs(curr_out_dir, exist_ok=True)
            curr_out_color_dir = curr_out_dir
            os.makedirs(curr_out_color_dir, exist_ok=True)

            target_idx = target.index[b, :, 0].cpu().numpy()
            _, v, _, h, w = input.image.size()

            # vis gt and predicted target views
            im = torch.cat((result.target.image[b], result.render[b]), dim=2).detach().cpu()
            im = rearrange(im, "v c h w -> h (v w) c")
            im = (im.cpu().numpy() * 255.0).clip(0.0, 255.0).astype(np.uint8)
            Image.fromarray(im).save(os.path.join(curr_out_color_dir, f"gt_vs_pred.png"))

            # vis input views
            input_im = result.input.image[b]  # V C H W
            input_im = rearrange(input_im, "v c h w -> h (v w) c")
            input_im = (input_im.cpu().numpy() * 255.0).clip(0.0, 255.0).astype(np.uint8)
            Image.fromarray(input_im).save(os.path.join(curr_out_color_dir, f"input.png"))

            # # vis interpolated views
            # if torch.is_tensor(result.render_interpolate):
            #     interpolate_im = result.render_interpolate[b].detach().cpu().to(torch.float32)
            #     interpolate_im = rearrange(interpolate_im, 'v c h w -> h (v w) c')
            #     interpolate_im = (interpolate_im.cpu().numpy() * 255.0).clip(0.0, 255.0).astype(np.uint8)
            #     Image.fromarray(interpolate_im).save(
            #         os.path.join(curr_out_color_dir, f"pred_interpolate.png"))

            # get metrics
            psnr = compute_psnr(target.image[b], result.render[b])  # .mean()
            lpips = compute_lpips(target.image[b], result.render[b])  # .mean()
            ssim = compute_ssim(
                target.image[b].to(torch.float32), result.render[b].to(torch.float32)
            )  # .mean()

            with open(os.path.join(curr_out_dir, "per_view_metrics.txt"), "w") as f:
                for i in range(psnr.shape[0]):
                    f.write(f"view: {target_idx[i]}, psnr: {psnr[i]},")
                    f.write(f"lpips: {lpips[i]},")
                    f.write(f"ssim: {ssim[i]}\n")

            with open(os.path.join(curr_out_dir, "metrics.txt"), "w") as f:
                metrics = f"psnr: {psnr.mean()}\nlpips: {lpips.mean()}\nssim: {ssim.mean()}\n"
                f.write(metrics)

            # save rendered video
            if torch.is_tensor(result.get("render_video", None)):
                all_frames = result.render_video[b].detach().cpu()  # [v h w c]
                all_frames = np.ascontiguousarray(np.array(all_frames.to(torch.float32)))
                all_frames = rearrange(all_frames, "v c h w -> v h w c")
                create_video_from_frames(all_frames, f"{curr_out_dir}/input_traj.mp4", framerate=30)

            # save rendered video
            if torch.is_tensor(result.get("render_video", None)):
                all_frames = result.render_video[b].detach().cpu()  # [v h w c]
                all_frames = np.ascontiguousarray(np.array(all_frames.to(torch.float32)))
                all_frames = rearrange(all_frames, "v c h w -> v h w c")
                create_video_from_frames(all_frames, f"{curr_out_dir}/input_traj.mp4", framerate=30)

            # visualize cameras (placeholder - implement if needed)
            # c2w = result.c2w[b].detach().cpu()
            # input_idx_cur = input_idx[b].detach().cpu()
            # c2w_input = c2w[input_idx_cur]
            # TODO: Implement camera trajectory visualization if needed
            # pass


def get_cam_se3(cam_info):
    """
    cam_info: [b,num_pose_element+3+4], rot, 3d trans, 4d fxfycxcy
    """
    b, n = cam_info.shape

    if n == 13:
        rot_6d = cam_info[:, :6]
        R = rot6d2mat(rot_6d)  # [b,3,3]
        t = cam_info[:, 6:9].unsqueeze(-1)  # [b,3,1]
        fxfycxcy = cam_info[:, 9:]  # normalized by resolution / shift from average, [b,4]
    elif n == 11:
        rot_quat = cam_info[:, :4]
        R = quat2mat(rot_quat)
        t = cam_info[:, 4:7].unsqueeze(-1)  # [b,3,1]
        fxfycxcy = cam_info[:, 7:]  # normalized by resolution / shift from average, [b,4]
    else:
        raise NotImplementedError

    Rt = torch.cat([R, t], dim=2)  # [b,3,4]
    c2w = torch.cat(
        [
            Rt,
            torch.tensor([0, 0, 0, 1], dtype=R.dtype, device=R.device)
            .view(1, 1, 4)
            .repeat(b, 1, 1),
        ],
        dim=1,
    )  # [b,4,4]
    return c2w, fxfycxcy


def cam_info_to_plucker(c2w, fxfycxcy, target_imgs_info, normalized=True):
    """
    c2w: [b,4,4]
    fxfycxcy: [b,4]
    """
    b = c2w.shape[0]
    device = c2w.device
    h, w = target_imgs_info.height, target_imgs_info.width

    fxfycxcy = fxfycxcy.clone()
    if normalized:
        fxfycxcy[:, 0] *= w
        fxfycxcy[:, 1] *= h
        fxfycxcy[:, 2] *= w
        fxfycxcy[:, 3] *= h

    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
    y, x = y.to(c2w), x.to(c2w)
    x = x[None, :, :].expand(b, -1, -1).reshape(b, -1)
    y = y[None, :, :].expand(b, -1, -1).reshape(b, -1)
    x = (x + 0.5 - fxfycxcy[:, 2:3]) / fxfycxcy[:, 0:1]
    y = (y + 0.5 - fxfycxcy[:, 3:4]) / fxfycxcy[:, 1:2]
    z = torch.ones_like(x)
    ray_d = torch.stack([x, y, z], dim=2)  # [b, h*w, 3]
    ray_d = torch.bmm(ray_d, c2w[:, :3, :3].transpose(1, 2))  # [b, h*w, 3]
    ray_d = ray_d / torch.norm(ray_d, dim=2, keepdim=True)  # [b, h*w, 3]
    ray_o = c2w[:, :3, 3][:, None, :].expand_as(ray_d)  # [b, h*w, 3]

    ray_o = ray_o.reshape(b, h, w, 3).permute(0, 3, 1, 2)  # [b,3,h,w]
    ray_d = ray_d.reshape(b, h, w, 3).permute(0, 3, 1, 2)

    plucker = torch.cat(
        [
            torch.cross(ray_o, ray_d, dim=1),
            ray_d,
        ],
        dim=1,
    )
    return plucker  # [b,c=6,h,w]


class PoseEstimator(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.config = config

        self.pose_rep = self.config.model.pose_latent.get("representation", "6d")
        print("Pose representation:", self.pose_rep)
        if self.pose_rep == "6d":
            self.num_pose_element = 6
        elif self.pose_rep == "quat":
            self.num_pose_element = 4
        else:
            raise NotImplementedError

        self.rel_head = nn.Sequential(
            nn.Linear(
                config.model.transformer.d * 2,
                config.model.transformer.d,
                bias=True,
            ),
            nn.SiLU(),
            nn.Linear(
                config.model.transformer.d,
                self.num_pose_element + 3,
                bias=True,
            ),
        )
        self.rel_head.apply(_init_weights)

        self.canonical_k_head = nn.Sequential(
            nn.Linear(
                config.model.transformer.d,
                config.model.transformer.d,
                bias=True,
            ),
            nn.SiLU(),
            nn.Linear(
                config.model.transformer.d,
                1,
                bias=True,
            ),
        )
        self.canonical_k_head.apply(_init_weights)

        self.f_bias = 1.25

    def forward(self, x, v):
        """
        x: [b*v, d]
        """
        canonical = self.config.model.pose_latent.get("canonical", "first")
        x = rearrange(x, "(b v) d -> b v d", v=v)
        b = x.shape[0]
        if canonical == "first":
            x_canonical = x[:, 0:1]  # [b,1,d]
            x_rel = x[:, 1:]  # [b,v-1,d]
        elif canonical == "middle":
            cano_idx = v // 2
            rel_indices = torch.cat([torch.arange(cano_idx), torch.arange(cano_idx + 1, v)])
            x_canonical = x[:, cano_idx : cano_idx + 1]  # [b,1,d]
            x_rel = x[:, rel_indices]  # [b,v-1,d]
        else:
            raise NotImplementedError

        fxfy_canonical = self.canonical_k_head(x_canonical[:, 0]) + self.f_bias  # [b,1]
        fxfy_canonical = fxfy_canonical.unsqueeze(1).repeat(1, 1, 2)  # [b,1,2]

        if self.pose_rep == "6d":
            rt_canonical = (
                torch.tensor([1, 0, 0, 0, 1, 0, 0, 0, 0])
                .reshape(1, 1, 9)
                .to(fxfy_canonical)
                .repeat(b, 1, 1)
            )  # [b,1,9]
        elif self.pose_rep == "quat":
            rt_canonical = (
                torch.tensor([1, 0, 0, 0, 0, 0, 0])
                .reshape(1, 1, 7)
                .to(fxfy_canonical)
                .repeat(b, 1, 1)
            )  # [b,1,7]
        info_canonical = torch.cat([rt_canonical, fxfy_canonical], dim=-1)  # [b,1,11]

        feat_rel = torch.cat([x_canonical.repeat(1, v - 1, 1), x_rel], dim=-1)  # [b,v-1,2*d]
        info_rel = self.rel_head(feat_rel)  # [b,v-1,num_pose_element+3]
        info_all = info_canonical.repeat(1, v, 1)  # [b,v,num_pose_element+3+2]

        if canonical == "first":
            info_all[:, 1:, : self.num_pose_element + 3] += info_rel
        elif canonical == "middle":
            info_all[:, rel_indices, : self.num_pose_element + 3] += info_rel
        else:
            raise NotImplementedError

        cxcy_all = torch.tensor([0.5, 0.5]).reshape(1, 1, 2).repeat(b, v, 1).to(info_all)
        info_all = torch.cat([info_all, cxcy_all], dim=-1)
        return rearrange(info_all, "b v d -> (b v) d")