add vjepa2

app.py

@@ -1,7 +1,7 @@
 import os
 # Import spaces before torch to avoid CUDA initialization error
 try:
-    import spaces
+    import spaces # pyright: ignore[reportMissingImports]
 except ImportError:
     pass
 import gradio as gr
@@ -11,7 +11,10 @@ import torch
 from PIL import Image
 from pathlib import Path
 from sklearn.decomposition import PCA
-from transformers import VideoMAEImageProcessor, VideoMAEForPreTraining, VideoMAEModel
+from transformers import (
+    VideoMAEImageProcessor, VideoMAEForPreTraining, VideoMAEModel,
+    AutoVideoProcessor, AutoModel
+)
 from transformers.utils.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
 import matplotlib
 matplotlib.use('Agg')  # Use non-interactive backend
@@ -66,11 +69,9 @@ def load_video(video_path, num_frames=16, sample_rate=4):
     cap.release()
     return frames
 
-def load_model(model_name, model_type='pretraining'):
-    """
-    Load model and processor by name.
-    model_type: 'pretraining' for VideoMAEForPreTraining, 'base' for VideoMAEModel
-    """
+# ========== VideoMAE Functions ==========
+def load_videomae_model(model_name='MCG-NJU/videomae-base', model_type='pretraining'):
+    """Load VideoMAE model and processor"""
     processor = VideoMAEImageProcessor.from_pretrained(model_name)
     if model_type == 'base':
         model = VideoMAEModel.from_pretrained(model_name)
@@ -79,24 +80,8 @@ def load_model(model_name, model_type='pretraining'):
     model = model.to(device)
     return model, processor
 
-# Global model and processor
-model = None
-processor = None
-
-def initialize_model(model_name='MCG-NJU/videomae-base'):
-    """Initialize the model (call once at startup)"""
-    global model, processor
-    if model is None:
-        print(f"Loading model: {model_name}")
-        model, processor = load_model(model_name)
-        print("Model loaded successfully")
-    return model, processor
-
-def visualize_attention(video_frames, model, processor, layer_idx=-1):
-    """
-    Visualize attention maps from VideoMAE model.
-    Returns PIL Image for Gradio.
-    """
+def videomae_visualize_attention(video_frames, model, processor, layer_idx=-1):
+    """Visualize attention maps from VideoMAE model"""
     inputs = processor(video_frames, return_tensors="pt")
     pixel_values = inputs['pixel_values'].to(device)
     batch_size, time, num_channels, height, width = pixel_values.shape
@@ -105,13 +90,11 @@ def visualize_attention(video_frames, model, processor, layer_idx=-1):
     num_patches_per_frame = (height // patch_size) * (width // patch_size)
     num_temporal_patches = time // tubelet_size
     
-    # Use VideoMAEModel to get attention weights
     if hasattr(model, 'videomae'):
         encoder_model = model.videomae
     else:
         encoder_model = model
     
-    # Disable SDPA and use eager attention
     original_attn_impl = getattr(encoder_model.config, '_attn_implementation', None)
     encoder_model.config._attn_implementation = "eager"
     
@@ -128,7 +111,6 @@ def visualize_attention(video_frames, model, processor, layer_idx=-1):
     attention_weights = attentions[layer_idx][0]
     avg_attn = attention_weights.mean(dim=0)
     
-    # Unnormalize frames
     dtype = pixel_values.dtype
     mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device=device, dtype=dtype)[None, None, :, None, None]
     std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device=device, dtype=dtype)[None, None, :, None, None]
@@ -150,14 +132,12 @@ def visualize_attention(video_frames, model, processor, layer_idx=-1):
     avg_attn_received = avg_attn.mean(dim=0)
     attn_per_patch = avg_attn_received.reshape(num_temporal_patches, H_p, W_p)
     
-    # Create visualization for first frame
     frame_idx = 0
     frame_img = frames_unnorm[frame_idx * tubelet_size]
     attn_map = attn_per_patch[frame_idx].detach().cpu().numpy()
     attn_map = (attn_map - attn_map.min()) / (attn_map.max() - attn_map.min() + 1e-8)
     attn_map_upsampled = cv2.resize(attn_map, (width, height))
     
-    # Create overlay
     fig, ax = plt.subplots(1, 1, figsize=(10, 10))
     ax.imshow(frame_img)
     ax.imshow(attn_map_upsampled, alpha=0.5, cmap='jet')
@@ -166,29 +146,37 @@ def visualize_attention(video_frames, model, processor, layer_idx=-1):
     
     return fig_to_image(fig)
 
-def visualize_latent(video_frames, model, processor):
-    """
-    Visualize latent space representations from VideoMAE model.
-    Returns PIL Image for Gradio.
-    """
+def videomae_compute_attention_all_frames(video_frames, model, processor, layer_idx=-1):
+    """Compute attention maps for all frames - VideoMAE"""
     inputs = processor(video_frames, return_tensors="pt")
     pixel_values = inputs['pixel_values'].to(device)
+    batch_size, time, num_channels, height, width = pixel_values.shape
+    tubelet_size = model.config.tubelet_size
+    patch_size = model.config.patch_size
+    num_patches_per_frame = (height // patch_size) * (width // patch_size)
+    num_temporal_patches = time // tubelet_size
     
     if hasattr(model, 'videomae'):
         encoder_model = model.videomae
     else:
         encoder_model = model
     
-    outputs = encoder_model(pixel_values, output_hidden_states=True)
-    hidden_states = outputs.last_hidden_state[0]
+    original_attn_impl = getattr(encoder_model.config, '_attn_implementation', None)
+    encoder_model.config._attn_implementation = "eager"
     
-    batch_size, time, num_channels, height, width = pixel_values.shape
-    tubelet_size = model.config.tubelet_size
-    patch_size = model.config.patch_size
-    num_patches_per_frame = (height // patch_size) * (width // patch_size)
-    num_temporal_patches = time // tubelet_size
+    try:
+        outputs = encoder_model(pixel_values, output_attentions=True)
+    finally:
+        if original_attn_impl is not None:
+            encoder_model.config._attn_implementation = original_attn_impl
+    
+    attentions = outputs.attentions
+    if layer_idx < 0:
+        layer_idx = len(attentions) + layer_idx
+    
+    attention_weights = attentions[layer_idx][0]
+    avg_attn = attention_weights.mean(dim=0)
     
-    # Unnormalize frames
     dtype = pixel_values.dtype
     mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device=device, dtype=dtype)[None, None, :, None, None]
     std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device=device, dtype=dtype)[None, None, :, None, None]
@@ -196,71 +184,31 @@ def visualize_latent(video_frames, model, processor):
     frames_unnorm = frames_unnorm[0].permute(0, 2, 3, 1).detach().cpu().numpy()
     frames_unnorm = np.clip(frames_unnorm, 0, 1)
     
-    seq_len = hidden_states.shape[0]
+    seq_len = avg_attn.shape[0]
+    H_p = height // patch_size
+    W_p = width // patch_size
     expected_seq_len = num_temporal_patches * num_patches_per_frame
     
     if seq_len != expected_seq_len:
         if seq_len % num_patches_per_frame == 0:
             num_temporal_patches = seq_len // num_patches_per_frame
         else:
-            raise ValueError(f"Cannot reshape hidden states: seq_len={seq_len}, expected={expected_seq_len}")
-    
-    hidden_states_reshaped = hidden_states.reshape(num_temporal_patches, num_patches_per_frame, -1)
-    hidden_size = hidden_states_reshaped.shape[-1]
-    hidden_states_flat = hidden_states_reshaped.reshape(-1, hidden_size).detach().cpu().numpy()
-    
-    pca = PCA(n_components=3)
-    pca_components = pca.fit_transform(hidden_states_flat)
-    pca_reshaped = pca_components.reshape(num_temporal_patches, num_patches_per_frame, 3)
-    
-    H_p = int(np.sqrt(num_patches_per_frame))
-    W_p = H_p
-    
-    if H_p * W_p == num_patches_per_frame:
-        pca_spatial = pca_reshaped.reshape(num_temporal_patches, H_p, W_p, 3)
-    else:
-        factors = []
-        for i in range(1, int(np.sqrt(num_patches_per_frame)) + 1):
-            if num_patches_per_frame % i == 0:
-                factors.append((i, num_patches_per_frame // i))
-        if factors:
-            H_p, W_p = factors[-1]
-            pca_spatial = pca_reshaped.reshape(num_temporal_patches, H_p, W_p, 3)
-        else:
-            raise ValueError(f"Cannot reshape {num_patches_per_frame} patches into a 2D grid")
-    
-    # Normalize components
-    for t in range(num_temporal_patches):
-        for c in range(3):
-            comp = pca_spatial[t, :, :, c]
-            comp_min, comp_max = comp.min(), comp.max()
-            if comp_max > comp_min:
-                pca_spatial[t, :, :, c] = (comp - comp_min) / (comp_max - comp_min)
-            else:
-                pca_spatial[t, :, :, c] = 0.5
+            raise ValueError(f"Cannot reshape attention: seq_len={seq_len}, expected={expected_seq_len}")
     
-    # Show first frame
-    frame_idx = 0
-    frame_img = frames_unnorm[frame_idx * tubelet_size]
-    rgb_image = pca_spatial[frame_idx]
-    upscale_factor = 8
-    rgb_image_upscaled = cv2.resize(rgb_image, (W_p * upscale_factor, H_p * upscale_factor), interpolation=cv2.INTER_NEAREST)
+    avg_attn_received = avg_attn.mean(dim=0)
+    attn_per_patch = avg_attn_received.reshape(num_temporal_patches, H_p, W_p)
     
-    fig = plt.figure(figsize=(6,6))
-    ax = fig.add_subplot(1, 1, 1)
-    ax.imshow(rgb_image_upscaled)
-    ax.set_title(f"PCA Components (RGB = PC1, PC2, PC3)")
-    ax.axis('off')
-    plt.suptitle(f"Explained Variance: {pca.explained_variance_ratio_.sum():.2%}", fontsize=12)
-    plt.tight_layout()
+    attention_maps = []
+    for frame_idx in range(num_temporal_patches):
+        attn_map = attn_per_patch[frame_idx].detach().cpu().numpy()
+        attn_map = (attn_map - attn_map.min()) / (attn_map.max() - attn_map.min() + 1e-8)
+        attn_map_upsampled = cv2.resize(attn_map, (width, height))
+        attention_maps.append(attn_map_upsampled)
     
-    return fig_to_image(fig)
+    return frames_unnorm, attention_maps
 
-def compute_reconstruction_all_frames(video_frames, model, processor):
-    """
-    Compute reconstruction for all frames and return as numpy arrays.
-    Returns: (original_frames, reconstructed_frames) as numpy arrays
-    """
+def videomae_compute_reconstruction_all_frames(video_frames, model, processor):
+    """Compute reconstruction for all frames - VideoMAE"""
     inputs = processor(video_frames, return_tensors="pt")
     T, C, H, W = inputs['pixel_values'][0].shape
     tubelet_size = model.config.tubelet_size
@@ -332,117 +280,113 @@ def compute_reconstruction_all_frames(video_frames, model, processor):
     
     return original_frames, reconstructed_frames_np
 
-def visualize_reconstruction(video_frames, model, processor):
-    """
-    Visualize reconstruction from VideoMAE model.
-    Returns PIL Image for Gradio.
-    """
+def videomae_compute_latent_all_frames(video_frames, model, processor):
+    """Compute PCA latent visualizations for all frames - VideoMAE"""
     inputs = processor(video_frames, return_tensors="pt")
-    T, C, H, W = inputs['pixel_values'][0].shape
-    tubelet_size = model.config.tubelet_size
-    patch_size = model.config.patch_size
-    T = T//tubelet_size
+    pixel_values = inputs['pixel_values'].to(device)
     
-    num_patches = (model.config.image_size // model.config.patch_size) ** 2
-    num_masked = int(0.9 * num_patches * (model.config.num_frames // model.config.tubelet_size))
-    total_patches = (model.config.num_frames // model.config.tubelet_size) * num_patches
-    batch_size = inputs['pixel_values'].shape[0]
-    bool_masked_pos = torch.zeros((batch_size, total_patches), dtype=torch.bool)
-
-    for b in range(batch_size):
-        mask_indices = np.random.choice(total_patches, num_masked, replace=False)
-        bool_masked_pos[b, mask_indices] = True
-
-    inputs['bool_masked_pos'] = bool_masked_pos.to(device)
-    inputs['pixel_values'] = inputs['pixel_values'].to(device)
+    if hasattr(model, 'videomae'):
+        encoder_model = model.videomae
+    else:
+        encoder_model = model
     
-    outputs = model(**inputs)
-    logits = outputs.logits 
+    outputs = encoder_model(pixel_values, output_hidden_states=True)
+    hidden_states = outputs.last_hidden_state[0]
     
-    pixel_values = inputs['pixel_values']
     batch_size, time, num_channels, height, width = pixel_values.shape
     tubelet_size = model.config.tubelet_size
     patch_size = model.config.patch_size
     num_patches_per_frame = (height // patch_size) * (width // patch_size)
     num_temporal_patches = time // tubelet_size
-    total_patches = num_temporal_patches * num_patches_per_frame
     
     dtype = pixel_values.dtype
     mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device=device, dtype=dtype)[None, None, :, None, None]
     std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device=device, dtype=dtype)[None, None, :, None, None]
     frames_unnorm = pixel_values * std + mean
+    frames_unnorm = frames_unnorm[0].permute(0, 2, 3, 1).detach().cpu().numpy()
+    frames_unnorm = np.clip(frames_unnorm, 0, 1)
     
-    frames_patched = frames_unnorm.view(
-        batch_size, time // tubelet_size, tubelet_size, num_channels,
-        height // patch_size, patch_size, width // patch_size, patch_size,
-    )
-    frames_patched = frames_patched.permute(0, 1, 4, 6, 2, 5, 7, 3).contiguous()
-    videos_patch = frames_patched.view(
-        batch_size, total_patches, tubelet_size * patch_size * patch_size * num_channels,
-    )
+    seq_len = hidden_states.shape[0]
+    expected_seq_len = num_temporal_patches * num_patches_per_frame
     
-    if model.config.norm_pix_loss:
-        patch_mean = videos_patch.mean(dim=-2, keepdim=True)
-        patch_std = (videos_patch.var(dim=-2, unbiased=True, keepdim=True).sqrt() + 1e-6)
-        logits_denorm = logits * patch_std + patch_mean
-    else:
-        logits_denorm = torch.clamp(logits, 0.0, 1.0)
+    if seq_len != expected_seq_len:
+        if seq_len % num_patches_per_frame == 0:
+            num_temporal_patches = seq_len // num_patches_per_frame
+        else:
+            raise ValueError(f"Cannot reshape hidden states: seq_len={seq_len}, expected={expected_seq_len}")
     
-    reconstructed_patches = videos_patch.clone()
-    reconstructed_patches[bool_masked_pos] = logits_denorm.reshape(-1, tubelet_size * patch_size * patch_size * num_channels)
+    hidden_states_reshaped = hidden_states.reshape(num_temporal_patches, num_patches_per_frame, -1)
+    hidden_size = hidden_states_reshaped.shape[-1]
+    hidden_states_flat = hidden_states_reshaped.reshape(-1, hidden_size).detach().cpu().numpy()
     
-    reconstructed_patches_reshaped = reconstructed_patches.view(
-        batch_size, time // tubelet_size, height // patch_size, width // patch_size,
-        tubelet_size, patch_size, patch_size, num_channels,
-    )
-    reconstructed_patches_reshaped = reconstructed_patches_reshaped.permute(0, 1, 4, 7, 2, 5, 3, 6).contiguous()
-    reconstructed_frames = reconstructed_patches_reshaped.view(
-        batch_size, time, num_channels, height, width,
-    )
+    pca = PCA(n_components=3)
+    pca_components = pca.fit_transform(hidden_states_flat)
+    pca_reshaped = pca_components.reshape(num_temporal_patches, num_patches_per_frame, 3)
     
-    original_frames = frames_unnorm[0].permute(0, 2, 3, 1).detach().cpu().numpy()
-    reconstructed_frames_np = reconstructed_frames[0].permute(0, 2, 3, 1).detach().cpu().numpy()
+    H_p = int(np.sqrt(num_patches_per_frame))
+    W_p = H_p
     
-    original_frames = np.clip(original_frames, 0, 1)
-    reconstructed_frames_np = np.clip(reconstructed_frames_np, 0, 1)
+    if H_p * W_p == num_patches_per_frame:
+        pca_spatial = pca_reshaped.reshape(num_temporal_patches, H_p, W_p, 3)
+    else:
+        factors = []
+        for i in range(1, int(np.sqrt(num_patches_per_frame)) + 1):
+            if num_patches_per_frame % i == 0:
+                factors.append((i, num_patches_per_frame // i))
+        if factors:
+            H_p, W_p = factors[-1]
+            pca_spatial = pca_reshaped.reshape(num_temporal_patches, H_p, W_p, 3)
+        else:
+            raise ValueError(f"Cannot reshape {num_patches_per_frame} patches into a 2D grid")
     
-    # Show first frame
-    frame_idx = 0
-    fig = plt.figure(figsize=(6,6))
-    ax = plt.subplot(111)
-
-    ax.imshow(reconstructed_frames_np[frame_idx * tubelet_size])
-    ax.set_title(f"Reconstructed Frame: {frame_idx * tubelet_size}")
-    ax.axis('off')
+    for t in range(num_temporal_patches):
+        for c in range(3):
+            comp = pca_spatial[t, :, :, c]
+            comp_min, comp_max = comp.min(), comp.max()
+            if comp_max > comp_min:
+                pca_spatial[t, :, :, c] = (comp - comp_min) / (comp_max - comp_min)
+            else:
+                pca_spatial[t, :, :, c] = 0.5
     
-    return fig_to_image(fig)
+    upscale_factor = 8
+    pca_images = []
+    for t_idx in range(num_temporal_patches):
+        rgb_image = pca_spatial[t_idx]
+        rgb_image_upscaled = cv2.resize(rgb_image, (W_p * upscale_factor, H_p * upscale_factor), interpolation=cv2.INTER_NEAREST)
+        pca_images.append(rgb_image_upscaled)
+    
+    return frames_unnorm, pca_images
 
-def compute_attention_all_frames(video_frames, model, processor, layer_idx=-1):
-    """
-    Compute attention maps for all frames.
-    Returns: (original_frames, attention_maps) as numpy arrays
-    """
-    inputs = processor(video_frames, return_tensors="pt")
-    pixel_values = inputs['pixel_values'].to(device)
-    batch_size, time, num_channels, height, width = pixel_values.shape
+# ========== V-JEPA 2 Functions ==========
+def load_vjepa2_model(model_name='facebook/vjepa2-vitl-fpc64-256'):
+    """Load V-JEPA 2 model and processor"""
+    processor = AutoVideoProcessor.from_pretrained(model_name)
+    model = AutoModel.from_pretrained(model_name)
+    model = model.to(device)
+    model.eval()
+    return model, processor
+
+def vjepa2_compute_attention_all_frames(video_frames, model, processor, layer_idx=-1):
+    """Compute attention maps for all frames - V-JEPA 2"""
+    video_array = np.stack([np.array(frame) for frame in video_frames])
+    
+    inputs = processor(video_array, return_tensors="pt")
+    pixel_values_videos = inputs['pixel_values_videos'].to(device)
+    batch_size, num_frames, num_channels, height, width = pixel_values_videos.shape
+    
     tubelet_size = model.config.tubelet_size
     patch_size = model.config.patch_size
     num_patches_per_frame = (height // patch_size) * (width // patch_size)
-    num_temporal_patches = time // tubelet_size
+    num_temporal_patches = num_frames // tubelet_size
     
-    if hasattr(model, 'videomae'):
-        encoder_model = model.videomae
-    else:
-        encoder_model = model
-    
-    original_attn_impl = getattr(encoder_model.config, '_attn_implementation', None)
-    encoder_model.config._attn_implementation = "eager"
+    original_attn_impl = getattr(model.config, '_attn_implementation', None)
+    model.config._attn_implementation = "eager"
     
     try:
-        outputs = encoder_model(pixel_values, output_attentions=True)
+        outputs = model(pixel_values_videos=pixel_values_videos, output_attentions=True, skip_predictor=True)
     finally:
         if original_attn_impl is not None:
-            encoder_model.config._attn_implementation = original_attn_impl
+            model.config._attn_implementation = original_attn_impl
     
     attentions = outputs.attentions
     if layer_idx < 0:
@@ -451,10 +395,10 @@ def compute_attention_all_frames(video_frames, model, processor, layer_idx=-1):
     attention_weights = attentions[layer_idx][0]
     avg_attn = attention_weights.mean(dim=0)
     
-    dtype = pixel_values.dtype
+    dtype = pixel_values_videos.dtype
     mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device=device, dtype=dtype)[None, None, :, None, None]
     std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device=device, dtype=dtype)[None, None, :, None, None]
-    frames_unnorm = pixel_values * std + mean
+    frames_unnorm = pixel_values_videos * std + mean
     frames_unnorm = frames_unnorm[0].permute(0, 2, 3, 1).detach().cpu().numpy()
     frames_unnorm = np.clip(frames_unnorm, 0, 1)
     
@@ -472,7 +416,6 @@ def compute_attention_all_frames(video_frames, model, processor, layer_idx=-1):
     avg_attn_received = avg_attn.mean(dim=0)
     attn_per_patch = avg_attn_received.reshape(num_temporal_patches, H_p, W_p)
     
-    # Create attention maps for all temporal patches
     attention_maps = []
     for frame_idx in range(num_temporal_patches):
         attn_map = attn_per_patch[frame_idx].detach().cpu().numpy()
@@ -482,32 +425,75 @@ def compute_attention_all_frames(video_frames, model, processor, layer_idx=-1):
     
     return frames_unnorm, attention_maps
 
-def compute_latent_all_frames(video_frames, model, processor):
-    """
-    Compute PCA latent visualizations for all frames.
-    Returns: (original_frames, pca_images) as numpy arrays
-    """
-    inputs = processor(video_frames, return_tensors="pt")
-    pixel_values = inputs['pixel_values'].to(device)
+def vjepa2_compute_reconstruction_all_frames(video_frames, model, processor):
+    """Compute feature visualizations for all frames - V-JEPA 2"""
+    video_array = np.stack([np.array(frame) for frame in video_frames])
     
-    if hasattr(model, 'videomae'):
-        encoder_model = model.videomae
-    else:
-        encoder_model = model
+    inputs = processor(video_array, return_tensors="pt")
+    pixel_values_videos = inputs['pixel_values_videos'].to(device)
+    batch_size, num_frames, num_channels, height, width = pixel_values_videos.shape
     
-    outputs = encoder_model(pixel_values, output_hidden_states=True)
+    with torch.no_grad():
+        outputs = model(pixel_values_videos=pixel_values_videos, skip_predictor=False)
+    
+    encoder_outputs = outputs.last_hidden_state[0]
+    
+    tubelet_size = model.config.tubelet_size
+    patch_size = model.config.patch_size
+    num_patches_per_frame = (height // patch_size) * (width // patch_size)
+    num_temporal_patches = num_frames // tubelet_size
+    
+    dtype = pixel_values_videos.dtype
+    mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device=device, dtype=dtype)[None, None, :, None, None]
+    std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device=device, dtype=dtype)[None, None, :, None, None]
+    frames_unnorm = pixel_values_videos * std + mean
+    frames_unnorm = frames_unnorm[0].permute(0, 2, 3, 1).detach().cpu().numpy()
+    frames_unnorm = np.clip(frames_unnorm, 0, 1)
+    
+    seq_len = encoder_outputs.shape[0]
+    if seq_len != num_temporal_patches * num_patches_per_frame:
+        if seq_len % num_patches_per_frame == 0:
+            num_temporal_patches = seq_len // num_patches_per_frame
+        else:
+            raise ValueError(f"Cannot reshape encoder outputs: seq_len={seq_len}")
+    
+    encoder_reshaped = encoder_outputs.reshape(num_temporal_patches, num_patches_per_frame, -1)
+    feature_maps = encoder_reshaped[:, :, 0].detach().cpu().numpy()
+    
+    H_p = height // patch_size
+    W_p = width // patch_size
+    feature_maps_spatial = feature_maps.reshape(num_temporal_patches, H_p, W_p)
+    
+    feature_visualizations = []
+    for t_idx in range(num_temporal_patches):
+        feat_map = feature_maps_spatial[t_idx]
+        feat_map = (feat_map - feat_map.min()) / (feat_map.max() - feat_map.min() + 1e-8)
+        feat_map_upsampled = cv2.resize(feat_map, (width, height))
+        feat_rgb = np.stack([feat_map_upsampled] * 3, axis=-1)
+        feature_visualizations.append(feat_rgb)
+    
+    return frames_unnorm, feature_visualizations
+
+def vjepa2_compute_latent_all_frames(video_frames, model, processor):
+    """Compute PCA latent visualizations for all frames - V-JEPA 2"""
+    video_array = np.stack([np.array(frame) for frame in video_frames])
+    
+    inputs = processor(video_array, return_tensors="pt")
+    pixel_values_videos = inputs['pixel_values_videos'].to(device)
+    
+    outputs = model(pixel_values_videos=pixel_values_videos, output_hidden_states=True, skip_predictor=True)
     hidden_states = outputs.last_hidden_state[0]
     
-    batch_size, time, num_channels, height, width = pixel_values.shape
+    batch_size, num_frames, num_channels, height, width = pixel_values_videos.shape
     tubelet_size = model.config.tubelet_size
     patch_size = model.config.patch_size
     num_patches_per_frame = (height // patch_size) * (width // patch_size)
-    num_temporal_patches = time // tubelet_size
+    num_temporal_patches = num_frames // tubelet_size
     
-    dtype = pixel_values.dtype
+    dtype = pixel_values_videos.dtype
     mean = torch.as_tensor(IMAGENET_DEFAULT_MEAN).to(device=device, dtype=dtype)[None, None, :, None, None]
     std = torch.as_tensor(IMAGENET_DEFAULT_STD).to(device=device, dtype=dtype)[None, None, :, None, None]
-    frames_unnorm = pixel_values * std + mean
+    frames_unnorm = pixel_values_videos * std + mean
     frames_unnorm = frames_unnorm[0].permute(0, 2, 3, 1).detach().cpu().numpy()
     frames_unnorm = np.clip(frames_unnorm, 0, 1)
     
@@ -544,7 +530,6 @@ def compute_latent_all_frames(video_frames, model, processor):
         else:
             raise ValueError(f"Cannot reshape {num_patches_per_frame} patches into a 2D grid")
     
-    # Normalize components
     for t in range(num_temporal_patches):
         for c in range(3):
             comp = pca_spatial[t, :, :, c]
@@ -554,7 +539,6 @@ def compute_latent_all_frames(video_frames, model, processor):
             else:
                 pca_spatial[t, :, :, c] = 0.5
     
-    # Create upscaled images for all frames
     upscale_factor = 8
     pca_images = []
     for t_idx in range(num_temporal_patches):
@@ -564,81 +548,108 @@ def compute_latent_all_frames(video_frames, model, processor):
     
     return frames_unnorm, pca_images
 
-# Dummy function for backward compatibility
-def process_video(video_path):
-    cap = cv2.VideoCapture(video_path)
-    frames = []
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        frames.append(frame)
-    cap.release()
-    visualizations = [cv2.applyColorMap((f * 0.5).astype(np.uint8), cv2.COLORMAP_JET) for f in frames]
-    return frames, visualizations
+# ========== Unified Model Management ==========
+# Global model and processor
+videomae_model = None
+videomae_processor = None
+vjepa2_model = None
+vjepa2_processor = None
+current_model_type = None
+
+def initialize_model(model_type='videomae', model_name=None):
+    """Initialize the selected model"""
+    global videomae_model, videomae_processor, vjepa2_model, vjepa2_processor, current_model_type
+    
+    if model_type == 'videomae':
+        if videomae_model is None:
+            if model_name is None:
+                model_name = 'MCG-NJU/videomae-base'
+            print(f"Loading VideoMAE model: {model_name}")
+            videomae_model, videomae_processor = load_videomae_model(model_name)
+            print("VideoMAE model loaded successfully")
+        current_model_type = 'videomae'
+        return videomae_model, videomae_processor
+    elif model_type == 'vjepa2':
+        if vjepa2_model is None:
+            if model_name is None:
+                model_name = 'facebook/vjepa2-vitl-fpc64-256'
+            print(f"Loading V-JEPA 2 model: {model_name}")
+            vjepa2_model, vjepa2_processor = load_vjepa2_model(model_name)
+            print("V-JEPA 2 model loaded successfully")
+        current_model_type = 'vjepa2'
+        return vjepa2_model, vjepa2_processor
 
 # Global state to store frames after upload
 stored_frames = []
-stored_viz = []
-
-# Cache for visualization results: {video_path: {mode: {frame_idx: image}}}
 visualization_cache = {}
 current_video_path = None
 
-def on_upload(video_path, mode):
-    global stored_frames, stored_viz, model, processor, visualization_cache, current_video_path
+def on_upload(video_path, mode, model_type):
+    global stored_frames, visualization_cache, current_video_path
     if video_path is None:
         return gr.update(maximum=0), None, None
     
     # Initialize model if needed
-    if model is None:
-        model, processor = initialize_model()
+    model, processor = initialize_model(model_type)
     
-    # Check if we need to recompute (new video or mode not cached)
+    # Check if we need to recompute (new video, mode, or model type)
     video_path_str = str(video_path)
+    cache_key = f"{video_path_str}_{model_type}"
     need_to_load_video = (video_path_str != current_video_path)
-    need_to_compute_mode = (video_path_str not in visualization_cache or mode not in visualization_cache[video_path_str])
+    need_to_compute_mode = (cache_key not in visualization_cache or mode not in visualization_cache[cache_key])
     
     if need_to_load_video:
-        # Load video frames
+        # Adjust frame count based on model type
+        num_frames = 64 if model_type == 'vjepa2' else 16
+        sample_rate = 1 if model_type == 'vjepa2' else 4
         print(f"Loading video: {video_path_str}")
-        video_frames = load_video(video_path)
+        video_frames = load_video(video_path, num_frames=num_frames, sample_rate=sample_rate)
         stored_frames = video_frames
         current_video_path = video_path_str
     else:
-        # Reuse already loaded frames
         video_frames = stored_frames
     
-    # Initialize cache for this video
-    if video_path_str not in visualization_cache:
-        visualization_cache[video_path_str] = {}
+    # Initialize cache for this video+model combination
+    if cache_key not in visualization_cache:
+        visualization_cache[cache_key] = {}
     
     if need_to_compute_mode:
-        # Compute all visualizations and cache them
-        print(f"Computing {mode} visualization for all frames...")
+        print(f"Computing {mode} visualization for all frames using {model_type}...")
         num_frames = len(stored_frames)
         tubelet_size = model.config.tubelet_size
         
         if mode == "reconstruction":
-            original_frames, reconstructed_frames = compute_reconstruction_all_frames(video_frames, model, processor)
-            # Cache as images per frame - map model frames to stored frames
-            visualization_cache[video_path_str][mode] = {}
-            for i in range(num_frames):
-                # Map stored frame index to model frame index
-                model_frame_idx = min(i, len(reconstructed_frames) - 1)
-                fig = plt.figure(figsize=(6, 6))
-                ax = plt.subplot(111)
-                ax.imshow(reconstructed_frames[model_frame_idx])
-                ax.set_title(f"Reconstructed Frame: {i}")
-                ax.axis('off')
-                visualization_cache[video_path_str][mode][i] = fig_to_image(fig)
+            if model_type == 'videomae':
+                original_frames, reconstructed_frames = videomae_compute_reconstruction_all_frames(video_frames, model, processor)
+                visualization_cache[cache_key][mode] = {}
+                for i in range(num_frames):
+                    model_frame_idx = min(i, len(reconstructed_frames) - 1)
+                    fig = plt.figure(figsize=(6, 6))
+                    ax = plt.subplot(111)
+                    ax.imshow(reconstructed_frames[model_frame_idx])
+                    ax.set_title(f"Reconstructed Frame: {i}")
+                    ax.axis('off')
+                    visualization_cache[cache_key][mode][i] = fig_to_image(fig)
+            else:  # vjepa2
+                original_frames, feature_visualizations = vjepa2_compute_reconstruction_all_frames(video_frames, model, processor)
+                visualization_cache[cache_key][mode] = {}
+                for i in range(num_frames):
+                    temporal_patch_idx = min(i // tubelet_size, len(feature_visualizations) - 1)
+                    fig = plt.figure(figsize=(6, 6))
+                    ax = plt.subplot(111)
+                    ax.imshow(feature_visualizations[temporal_patch_idx])
+                    ax.set_title(f"Encoder Features - Frame {i}")
+                    ax.axis('off')
+                    visualization_cache[cache_key][mode][i] = fig_to_image(fig)
         
         elif mode == "attention":
-            original_frames, attention_maps = compute_attention_all_frames(video_frames, model, processor)
-            visualization_cache[video_path_str][mode] = {}
+            if model_type == 'videomae':
+                original_frames, attention_maps = videomae_compute_attention_all_frames(video_frames, model, processor)
+            else:  # vjepa2
+                original_frames, attention_maps = vjepa2_compute_attention_all_frames(video_frames, model, processor)
+            
+            visualization_cache[cache_key][mode] = {}
             for i in range(num_frames):
-                # Map stored frame to temporal patch
                 temporal_patch_idx = min(i // tubelet_size, len(attention_maps) - 1)
                 model_frame_idx = min(i, len(original_frames) - 1)
                 if temporal_patch_idx < len(attention_maps):
@@ -648,13 +659,16 @@ def on_upload(video_path, mode):
                     ax.imshow(attention_maps[temporal_patch_idx], alpha=0.5, cmap='jet')
                     ax.set_title(f"Attention Map - Frame {i}")
                     ax.axis('off')
-                    visualization_cache[video_path_str][mode][i] = fig_to_image(fig)
+                    visualization_cache[cache_key][mode][i] = fig_to_image(fig)
         
         elif mode == "latent":
-            original_frames, pca_images = compute_latent_all_frames(video_frames, model, processor)
-            visualization_cache[video_path_str][mode] = {}
+            if model_type == 'videomae':
+                original_frames, pca_images = videomae_compute_latent_all_frames(video_frames, model, processor)
+            else:  # vjepa2
+                original_frames, pca_images = vjepa2_compute_latent_all_frames(video_frames, model, processor)
+            
+            visualization_cache[cache_key][mode] = {}
             for i in range(num_frames):
-                # Map stored frame to temporal patch
                 temporal_patch_idx = min(i // tubelet_size, len(pca_images) - 1)
                 if temporal_patch_idx < len(pca_images):
                     fig = plt.figure(figsize=(6, 6))
@@ -662,34 +676,30 @@ def on_upload(video_path, mode):
                     ax.imshow(pca_images[temporal_patch_idx])
                     ax.set_title(f"PCA Components - Frame {i}")
                     ax.axis('off')
-                    visualization_cache[video_path_str][mode][i] = fig_to_image(fig)
+                    visualization_cache[cache_key][mode][i] = fig_to_image(fig)
         
-        print(f"Caching complete for {mode} mode")
+        print(f"Caching complete for {mode} mode using {model_type}")
     
     # Load from cache
     max_idx = len(stored_frames) - 1
     frame_idx = 0
     
-    # Get original frame
     if isinstance(stored_frames[0], Image.Image):
         first_frame = np.array(stored_frames[0])
     else:
         first_frame = stored_frames[0]
     
-    # Get visualization from cache
-    if video_path_str in visualization_cache and mode in visualization_cache[video_path_str]:
-        if frame_idx in visualization_cache[video_path_str][mode]:
-            viz_img = visualization_cache[video_path_str][mode][frame_idx]
+    if cache_key in visualization_cache and mode in visualization_cache[cache_key]:
+        if frame_idx in visualization_cache[cache_key][mode]:
+            viz_img = visualization_cache[cache_key][mode][frame_idx]
         else:
-            # Fallback if frame not in cache
             viz_img = Image.fromarray(first_frame)
     else:
-        # Fallback if not cached
         viz_img = Image.fromarray(first_frame)
     
     return gr.update(maximum=max_idx, value=0), first_frame, viz_img
 
-def update_frame(idx, mode):
+def update_frame(idx, mode, model_type):
     global stored_frames, visualization_cache, current_video_path
     if not stored_frames:
         return None, None
@@ -698,49 +708,57 @@ def update_frame(idx, mode):
     if frame_idx >= len(stored_frames):
         frame_idx = len(stored_frames) - 1
     
-    # Get frame
     if isinstance(stored_frames[frame_idx], Image.Image):
         frame = np.array(stored_frames[frame_idx])
     else:
         frame = stored_frames[frame_idx]
     
-    # Load visualization from cache (fast!)
     video_path_str = current_video_path
-    if video_path_str and video_path_str in visualization_cache:
-        if mode in visualization_cache[video_path_str]:
-            if frame_idx in visualization_cache[video_path_str][mode]:
-                viz_img = visualization_cache[video_path_str][mode][frame_idx]
+    cache_key = f"{video_path_str}_{model_type}" if video_path_str else None
+    if cache_key and cache_key in visualization_cache:
+        if mode in visualization_cache[cache_key]:
+            if frame_idx in visualization_cache[cache_key][mode]:
+                viz_img = visualization_cache[cache_key][mode][frame_idx]
             else:
-                # Fallback if frame not in cache
                 viz_img = Image.fromarray(frame)
         else:
-            # Mode not cached, return frame
             viz_img = Image.fromarray(frame)
     else:
-        # Not cached, return frame
         viz_img = Image.fromarray(frame)
     
     return frame, viz_img
 
-
 def load_example_video(video_file):
-    def _load_example_video(  mode):
+    def _load_example_video(mode, model_type):
         """Load the predefined example video"""
         example_path = f"examples/{video_file}"
-        return on_upload(example_path, mode)
+        return on_upload(example_path, mode, model_type)
     
     return _load_example_video
 
+def on_model_change(model_type, video_path, mode):
+    """Handle model type change"""
+    if video_path is None:
+        return gr.update(maximum=0), None, None
+    return on_upload(video_path, mode, model_type)
+
 # --- Gradio UI Layout ---
-with gr.Blocks(title="VideoMAE Representation Explorer") as demo:
-    gr.Markdown("## 🎥 VideoMAE Frame-by-Frame Representation Explorer")
-    
-    mode_radio = gr.Radio(
-        choices=["reconstruction", "attention", "latent"],
-        value="reconstruction",
-        label="Visualization Mode",
-        info="Choose the type of visualization"
-    )
+with gr.Blocks(title="Video Representation Explorer") as demo:
+    gr.Markdown("## 🎥 Video Representation Explorer - VideoMAE & V-JEPA 2")
+    
+    with gr.Row():
+        model_radio = gr.Radio(
+            choices=["videomae", "vjepa2"],
+            value="videomae",
+            label="Model Type",
+            info="Choose between VideoMAE and V-JEPA 2"
+        )
+        mode_radio = gr.Radio(
+            choices=["reconstruction", "attention", "latent"],
+            value="reconstruction",
+            label="Visualization Mode",
+            info="Choose the type of visualization"
+        )
     
     with gr.Row():
         with gr.Column():
@@ -752,47 +770,49 @@ with gr.Blocks(title="VideoMAE Representation Explorer") as demo:
 
     # Event Listeners
     video_lists = os.listdir("examples") if os.path.exists("examples") else []
-    video_lists =  os.listdir("app/examples") if os.path.exists("app/examples") else video_lists
+    video_lists = os.listdir("app/examples") if os.path.exists("app/examples") else video_lists
     with gr.Row():
         video_input = gr.Video(label="Upload Video")
         with gr.Column():
             for video_file in video_lists:
                 load_example_btn = gr.Button(f"Load Example Video ({video_file})", variant="secondary")
-                load_example_btn.click(load_example_video(video_file), inputs=mode_radio, outputs=[frame_slider, orig_output, viz_output])
+                load_example_btn.click(
+                    load_example_video(video_file), 
+                    inputs=[mode_radio, model_radio], 
+                    outputs=[frame_slider, orig_output, viz_output]
+                )
             
-        # load_example_btn = gr.Button("Load Example Video (dog.mp4)", variant="secondary")
-    video_input.change(on_upload, inputs=[video_input, mode_radio], outputs=[frame_slider, orig_output, viz_output])
+    video_input.change(
+        on_upload, 
+        inputs=[video_input, mode_radio, model_radio], 
+        outputs=[frame_slider, orig_output, viz_output]
+    )
+    
+    frame_slider.change(
+        update_frame, 
+        inputs=[frame_slider, mode_radio, model_radio], 
+        outputs=[orig_output, viz_output]
+    )
     
-    frame_slider.change(update_frame, inputs=[frame_slider, mode_radio], outputs=[orig_output, viz_output])
-    def on_mode_change(video_path, mode):
-        """Handle mode change - compute if not cached, otherwise use cache"""
-        global stored_frames, model, processor, visualization_cache, current_video_path
+    def on_mode_change(video_path, mode, model_type):
+        """Handle mode change"""
         if video_path is None:
             return gr.update(maximum=0), None, None
-        
-        video_path_str = str(video_path)
-        
-        # If video is already loaded and mode is cached, just return cached result
-        if video_path_str == current_video_path and video_path_str in visualization_cache:
-            if mode in visualization_cache[video_path_str]:
-                max_idx = len(stored_frames) - 1
-                frame_idx = 0
-                if isinstance(stored_frames[0], Image.Image):
-                    first_frame = np.array(stored_frames[0])
-                else:
-                    first_frame = stored_frames[0]
-                if frame_idx in visualization_cache[video_path_str][mode]:
-                    viz_img = visualization_cache[video_path_str][mode][frame_idx]
-                else:
-                    viz_img = Image.fromarray(first_frame)
-                return gr.update(maximum=max_idx, value=0), first_frame, viz_img
-        
-        # Otherwise, compute (will use cached video frames if available)
-        return on_upload(video_path, mode)
+        return on_upload(video_path, mode, model_type)
+    
+    mode_radio.change(
+        on_mode_change, 
+        inputs=[video_input, mode_radio, model_radio], 
+        outputs=[frame_slider, orig_output, viz_output]
+    )
     
-    mode_radio.change(on_mode_change, inputs=[video_input, mode_radio], outputs=[frame_slider, orig_output, viz_output])
+    model_radio.change(
+        on_model_change, 
+        inputs=[model_radio, video_input, mode_radio], 
+        outputs=[frame_slider, orig_output, viz_output]
+    )
 
 if __name__ == "__main__":
-    # Initialize model at startup
-    initialize_model()
-    demo.launch()
+    # Initialize default model at startup
+    initialize_model('videomae')
+    demo.launch(share=True)

examples/easy_motion_tube2.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f38d7badf4d90059fa50c905daec70a107e9bfd225d7fba4e89416f20927656
+size 49261