ModelTrain/module/vitg_encoder.py

"""
V-JEPA2 ViTG Encoder Wrapper for Tactile Image Processing

This module provides a frozen ViTG encoder that processes tactile images
and produces 1280-dimensional embeddings for the ACT policy.
"""

import torch
import torch.nn as nn
import torchvision.transforms as transforms
from typing import Optional

# Try to import V-JEPA2 model architecture (local copy for Python 3.8 compatibility)
try:
    # Use local vjepa2_compat directory (Python 3.8 compatible)
    from ModelTrain.vjepa2_compat.model_builder import create_vit_giant, create_vit_large, load_vjepa2_weights
    VJEPA_AVAILABLE = True
except ImportError as e:
    VJEPA_AVAILABLE = False
    print(f"WARNING: V-JEPA2 models not available: {e}")


class ViTGEncoder(nn.Module):
    """
    Wrapper for V-JEPA2 ViTG encoder to process tactile images.
    
    The encoder is frozen (all parameters have requires_grad=False) and produces
    1280-dimensional embeddings from input tactile images.
    """
    
    def __init__(self, ckpt_path: str, input_size: int = 224):
        """
        Initialize ViTG encoder from checkpoint.
        
        Args:
            ckpt_path: Path to the V-JEPA2 ViTG checkpoint (.pt file)
            input_size: Expected input image size (default: 224)
        """
        super().__init__()
        
        self.input_size = input_size
        self.embed_dim = 1280  # ViT-G standard embedding dimension
        
        # Load the checkpoint
        print(f"Loading ViTG checkpoint from: {ckpt_path}")
        checkpoint = torch.load(ckpt_path, map_location='cpu')
        
        # Extract the model from checkpoint
        # The checkpoint structure may vary, so we need to handle different formats
        if isinstance(checkpoint, dict):
            if 'model' in checkpoint:
                model_state = checkpoint['model']
            elif 'state_dict' in checkpoint:
                model_state = checkpoint['state_dict']
            elif 'encoder' in checkpoint:
                # V-JEPA may store encoder separately
                model_state = checkpoint['encoder']
            else:
                # Assume the checkpoint itself is the state dict
                model_state = checkpoint
        else:
            # If checkpoint is directly a model
            self.encoder = checkpoint
            model_state = None
        
        # If we have a state dict, we need to create the model architecture
        # For V-JEPA2 ViTG, we'll try to use the model directly if available
        if model_state is not None:
            # Try to infer model architecture from state dict keys
            # V-JEPA uses a vision transformer architecture
            try:
                # Attempt to create a compatible ViT-G architecture
                self.encoder = self._create_vitg_model()
                # Load the state dict, being permissive about mismatches
                missing_keys, unexpected_keys = self.encoder.load_state_dict(model_state, strict=False)
                if missing_keys:
                    print(f"Warning: Missing keys in checkpoint: {missing_keys[:5]}...")
                if unexpected_keys:
                    print(f"Warning: Unexpected keys in checkpoint: {unexpected_keys[:5]}...")
            except Exception as e:
                print(f"Error loading state dict: {e}")
                print("Attempting to use checkpoint directly as model...")
                self.encoder = checkpoint
        
        # Freeze all parameters
        self._freeze_encoder()
        
        # Set to eval mode
        self.encoder.eval()
        
        # Define preprocessing transforms
        # V-JEPA typically uses ImageNet normalization
        self.preprocess = transforms.Compose([
            transforms.Resize((self.input_size, self.input_size), antialias=True),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])
        
        print(f"ViTG encoder loaded successfully. Embedding dim: {self.embed_dim}")
    
    def _create_vitg_model(self):
        """
        Create a ViT-G model architecture.
        This is a placeholder - actual architecture depends on V-JEPA2 implementation.
        """
        # This would need to match the exact V-JEPA2 architecture
        # For now, we return a dummy module that will be replaced
        # In practice, you'd import the actual V-JEPA2 model class
        raise NotImplementedError(
            "Please ensure the checkpoint contains the full model, "
            "or import the V-JEPA2 model architecture explicitly."
        )
    
    def _freeze_encoder(self):
        """Freeze all encoder parameters."""
        for param in self.encoder.parameters():
            param.requires_grad = False
        print("ViTG encoder frozen (all parameters set to requires_grad=False)")
    
    def forward(self, x: torch.Tensor, return_cls_only: bool = True) -> torch.Tensor:
        """
        Forward pass through ViTG encoder.
        
        Args:
            x: Input tactile images, shape (B, C, H, W)
            return_cls_only: If True, return only CLS token embedding (B, 1280)
                           If False, return all patch embeddings (B, N, 1280)
        
        Returns:
            embeddings: Tensor of shape (B, 1280) if return_cls_only=True,
                       otherwise (B, N, 1280) where N is number of patches
        """
        # Preprocess images
        x = self.preprocess(x)
        
        # Pass through encoder (no gradient computation)
        with torch.no_grad():
            # V-JEPA encoders typically return a tuple or dict
            output = self.encoder(x)
            
            # Handle different output formats
            if isinstance(output, tuple):
                # Usually (features, intermediates) or similar
                features = output[0]
            elif isinstance(output, dict):
                # May have 'cls_token', 'patch_tokens', etc.
                if 'cls_token' in output:
                    features = output['cls_token']
                elif 'last_hidden_state' in output:
                    features = output['last_hidden_state']
                else:
                    # Take the first value
                    features = list(output.values())[0]
            else:
                features = output
            
            # Extract CLS token if needed
            if return_cls_only:
                if features.dim() == 3:  # (B, N, D)
                    # First token is typically CLS
                    features = features[:, 0, :]
                elif features.dim() == 2:  # (B, D)
                    # Already extracted
                    pass
                else:
                    raise ValueError(f"Unexpected feature shape: {features.shape}")
        
        return features
    
    def get_num_params(self):
        """Return the number of parameters in the encoder."""
        return sum(p.numel() for p in self.encoder.parameters())


class ViTGEncoderSimple(nn.Module):
    """
    V-JEPA2 ViT encoder wrapper for tactile image processing.
    Loads V-JEPA2 checkpoint and creates frozen encoder.
    Supports both ViT-Giant (1408-dim) and ViT-Large (1024-dim).
    """
    
    def __init__(self, ckpt_path: str, embed_dim: int = None, input_size: int = 224, model_type: str = 'vitg'):
        super().__init__()
        
        self.model_type = model_type
        self.input_size = input_size
        
        # Set embed_dim based on model_type if not explicitly provided
        if embed_dim is None:
            if model_type == 'vitg':
                self.embed_dim = 1408
            elif model_type == 'vitl':
                self.embed_dim = 1024
            else:
                raise ValueError(f"Unknown model_type: {model_type}. Choose 'vitg' or 'vitl'")
        else:
            self.embed_dim = embed_dim
        
        print(f"Loading ViT-{model_type.upper()} checkpoint from: {ckpt_path}")
        
        # Load checkpoint directly to GPU to save RAM
        checkpoint = torch.load(ckpt_path, map_location='cuda')
        
        # Handle different checkpoint formats
        if hasattr(checkpoint, 'eval'):
            # Checkpoint is already a model
            self.encoder = checkpoint
            print("Loaded full model from checkpoint")
        elif isinstance(checkpoint, dict):
            # Checkpoint is a dictionary with state_dicts
            # print(f"Checkpoint keys: {list(checkpoint.keys())}")
            
            # V-JEPA2 checkpoints contain state_dicts, need to instantiate model
            if not VJEPA_AVAILABLE:
                raise ImportError(
                    "V-JEPA2 model architecture not available.\n"
                    "The local V-JEPA2 model files should be in ModelTrain/vjepa2_compat/\n"
                    "Check that backbones.py and vision_transformer.py exist there."
                )
            
            # Create V-JEPA2 ViT model based on model_type
            # Note: Must match checkpoint architecture (tubelet_size=2 for video models)
            print(f"Creating V-JEPA2 ViT-{model_type.upper()} model (img_size={input_size}, tubelet_size=2)")
            if model_type == 'vitg':
                self.encoder = create_vit_giant(
                    img_size=input_size,
                    patch_size=16,
                    num_frames=2,  # Match checkpoint (will duplicate frames for static images)
                    tubelet_size=2,  # Match checkpoint architecture
                )
            elif model_type == 'vitl':
                self.encoder = create_vit_large(
                    img_size=input_size,
                    patch_size=16,
                    num_frames=2,  # Match checkpoint (will duplicate frames for static images)
                    tubelet_size=2,  # Match checkpoint architecture
                )
            else:
                raise ValueError(f"Unknown model_type: {model_type}. Choose 'vitg' or 'vitl'")
            
            # Load weights using helper function
            use_target = 'target_encoder' in checkpoint
            self.encoder = load_vjepa2_weights(self.encoder, ckpt_path, use_target_encoder=use_target)
        else:
            raise ValueError(f"Unsupported checkpoint format: {type(checkpoint)}")
        
        # Move to GPU and freeze encoder
        self.encoder.cuda()
        for param in self.encoder.parameters():
            param.requires_grad = False
        
        self.encoder.eval()
        
        print(f"ViT-{model_type.upper()} encoder loaded and frozen. Embed dim: {self.embed_dim}")
    
    def forward(self, x: torch.Tensor, return_all_tokens: bool = False) -> torch.Tensor:
        """
        Forward pass returning CLS token embeddings or all patch tokens.
        
        Args:
            x: Input images (B, C, H, W), assumed to be already resized and normalized
            return_all_tokens: If True, return all patch tokens (B, num_patches, embed_dim)
                             If False, return only CLS token (B, embed_dim)
        
        Returns:
            CLS embeddings (B, embed_dim) if return_all_tokens=False
            All patch tokens (B, num_patches, embed_dim) if return_all_tokens=True
        """
        # V-JEPA2 expects images to be already normalized (done in dataset)
        # No additional preprocessing needed here
        
        # Forward pass without gradients
        with torch.no_grad():
            # V-JEPA2 models expect (B, C, num_frames, H, W) for videos
            # For static images, duplicate the frame to match num_frames=2
            if x.dim() == 4:  # (B, C, H, W)
                x = x.unsqueeze(2)  # (B, C, 1, H, W)
                # Duplicate frame to match model's expected num_frames=2
                x = x.repeat(1, 1, 2, 1, 1)  # (B, C, 2, H, W)
            
            # Forward through V-JEPA2 encoder
            features = self.encoder(x)
            
            # V-JEPA2 outputs patch tokens: (B, num_patches, embed_dim)
            # Extract CLS token or return all tokens based on flag
            if isinstance(features, (tuple, list)):
                features = features[0]
            
            if features.dim() == 3:  # (B, num_patches, embed_dim)
                if return_all_tokens:
                    # Return all patch tokens
                    return features  # (B, num_patches, embed_dim)
                else:
                    # Use CLS token (first token)
                    features = features[:, 0, :]  # Take CLS token
            elif features.dim() == 2:  # (B, embed_dim)
                # Already extracted (shouldn't happen with return_all_tokens=True)
                pass
        
        return features