revert

app.py

@@ -1,7 +1,7 @@
 import os
 # Import spaces before torch to avoid CUDA initialization error
 try:
-    import spaces # pyright: ignore[reportMissingImports]
+    import spaces
 except ImportError:
     pass
 import gradio as gr
@@ -11,10 +11,7 @@ import torch
 from PIL import Image
 from pathlib import Path
 from sklearn.decomposition import PCA
-from transformers import (
-    VideoMAEImageProcessor, VideoMAEForPreTraining, VideoMAEModel,
-    AutoVideoProcessor, AutoModel
-)
+from transformers import VideoMAEImageProcessor, VideoMAEForPreTraining, VideoMAEModel
 from transformers.utils.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
 import matplotlib
 matplotlib.use('Agg')  # Use non-interactive backend
@@ -69,9 +66,11 @@ def load_video(video_path, num_frames=16, sample_rate=4):
     cap.release()
     return frames
 
-# ========== VideoMAE Functions ==========
-def load_videomae_model(model_name='MCG-NJU/videomae-base', model_type='pretraining'):
-    """Load VideoMAE model and processor"""
+def load_model(model_name, model_type='pretraining'):
+    """
+    Load model and processor by name.
+    model_type: 'pretraining' for VideoMAEForPreTraining, 'base' for VideoMAEModel
+    """
     processor = VideoMAEImageProcessor.from_pretrained(model_name)
     if model_type == 'base':
         model = VideoMAEModel.from_pretrained(model_name)
@@ -80,8 +79,24 @@ def load_videomae_model(model_name='MCG-NJU/videomae-base', model_type='pretrain
     model = model.to(device)
     return model, processor
 
-def videomae_visualize_attention(video_frames, model, processor, layer_idx=-1):
-    """Visualize attention maps from VideoMAE model"""
+# 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.
+    """
     inputs = processor(video_frames, return_tensors="pt")
     pixel_values = inputs['pixel_values'].to(device)
     batch_size, time, num_channels, height, width = pixel_values.shape
@@ -90,11 +105,13 @@ def videomae_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"
     
@@ -111,6 +128,7 @@ def videomae_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]
@@ -132,12 +150,14 @@ def videomae_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')
@@ -146,37 +166,29 @@ def videomae_visualize_attention(video_frames, model, processor, layer_idx=-1):
     
     return fig_to_image(fig)
 
-def videomae_compute_attention_all_frames(video_frames, model, processor, layer_idx=-1):
-    """Compute attention maps for all frames - VideoMAE"""
+def visualize_latent(video_frames, model, processor):
+    """
+    Visualize latent space representations from VideoMAE model.
+    Returns PIL Image for Gradio.
+    """
     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
     
-    original_attn_impl = getattr(encoder_model.config, '_attn_implementation', None)
-    encoder_model.config._attn_implementation = "eager"
-    
-    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
+    outputs = encoder_model(pixel_values, output_hidden_states=True)
+    hidden_states = outputs.last_hidden_state[0]
     
-    attention_weights = attentions[layer_idx][0]
-    avg_attn = attention_weights.mean(dim=0)
+    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
     
+    # 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]
@@ -184,31 +196,71 @@ def videomae_compute_attention_all_frames(video_frames, model, processor, layer_
     frames_unnorm = frames_unnorm[0].permute(0, 2, 3, 1).detach().cpu().numpy()
     frames_unnorm = np.clip(frames_unnorm, 0, 1)
     
-    seq_len = avg_attn.shape[0]
-    H_p = height // patch_size
-    W_p = width // patch_size
+    seq_len = hidden_states.shape[0]
     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 attention: seq_len={seq_len}, expected={expected_seq_len}")
+            raise ValueError(f"Cannot reshape hidden states: seq_len={seq_len}, expected={expected_seq_len}")
     
-    avg_attn_received = avg_attn.mean(dim=0)
-    attn_per_patch = avg_attn_received.reshape(num_temporal_patches, H_p, W_p)
+    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()
     
-    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)
+    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)
     
-    return frames_unnorm, attention_maps
+    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
+    
+    # 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)
+    
+    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()
+    
+    return fig_to_image(fig)
 
-def videomae_compute_reconstruction_all_frames(video_frames, model, processor):
-    """Compute reconstruction for all frames - VideoMAE"""
+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
+    """
     inputs = processor(video_frames, return_tensors="pt")
     T, C, H, W = inputs['pixel_values'][0].shape
     tubelet_size = model.config.tubelet_size
@@ -280,113 +332,117 @@ def videomae_compute_reconstruction_all_frames(video_frames, model, processor):
     
     return original_frames, reconstructed_frames_np
 
-def videomae_compute_latent_all_frames(video_frames, model, processor):
-    """Compute PCA latent visualizations for all frames - VideoMAE"""
+def visualize_reconstruction(video_frames, model, processor):
+    """
+    Visualize reconstruction from VideoMAE model.
+    Returns PIL Image for Gradio.
+    """
     inputs = processor(video_frames, return_tensors="pt")
-    pixel_values = inputs['pixel_values'].to(device)
+    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
     
-    if hasattr(model, 'videomae'):
-        encoder_model = model.videomae
-    else:
-        encoder_model = model
+    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)
     
-    outputs = encoder_model(pixel_values, output_hidden_states=True)
-    hidden_states = outputs.last_hidden_state[0]
+    outputs = model(**inputs)
+    logits = outputs.logits 
     
+    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)
-    
-    seq_len = hidden_states.shape[0]
-    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()
+    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,
+    )
     
-    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)
+    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)
     
-    H_p = int(np.sqrt(num_patches_per_frame))
-    W_p = H_p
+    reconstructed_patches = videos_patch.clone()
+    reconstructed_patches[bool_masked_pos] = logits_denorm.reshape(-1, tubelet_size * patch_size * patch_size * num_channels)
     
-    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")
+    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,
+    )
     
-    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
+    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()
     
-    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)
+    original_frames = np.clip(original_frames, 0, 1)
+    reconstructed_frames_np = np.clip(reconstructed_frames_np, 0, 1)
     
-    return frames_unnorm, pca_images
-
-# ========== 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
+    # Show first frame
+    frame_idx = 0
+    fig = plt.figure(figsize=(6,6))
+    ax = plt.subplot(111)
 
-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
+    ax.imshow(reconstructed_frames_np[frame_idx * tubelet_size])
+    ax.set_title(f"Reconstructed Frame: {frame_idx * tubelet_size}")
+    ax.axis('off')
     
+    return fig_to_image(fig)
+
+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
     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
+    num_temporal_patches = time // tubelet_size
     
-    original_attn_impl = getattr(model.config, '_attn_implementation', None)
-    model.config._attn_implementation = "eager"
+    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"
     
     try:
-        outputs = model(pixel_values_videos=pixel_values_videos, output_attentions=True, skip_predictor=True)
+        outputs = encoder_model(pixel_values, output_attentions=True)
     finally:
         if original_attn_impl is not None:
-            model.config._attn_implementation = original_attn_impl
+            encoder_model.config._attn_implementation = original_attn_impl
     
     attentions = outputs.attentions
     if layer_idx < 0:
@@ -395,10 +451,10 @@ def vjepa2_compute_attention_all_frames(video_frames, model, processor, layer_id
     attention_weights = attentions[layer_idx][0]
     avg_attn = attention_weights.mean(dim=0)
     
-    dtype = pixel_values_videos.dtype
+    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_videos * std + mean
+    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)
     
@@ -416,6 +472,7 @@ def vjepa2_compute_attention_all_frames(video_frames, model, processor, layer_id
     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()
@@ -425,75 +482,32 @@ def vjepa2_compute_attention_all_frames(video_frames, model, processor, layer_id
     
     return frames_unnorm, attention_maps
 
-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])
-    
-    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
-    
-    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])
+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)
     
-    inputs = processor(video_array, return_tensors="pt")
-    pixel_values_videos = inputs['pixel_values_videos'].to(device)
+    if hasattr(model, 'videomae'):
+        encoder_model = model.videomae
+    else:
+        encoder_model = model
     
-    outputs = model(pixel_values_videos=pixel_values_videos, output_hidden_states=True, skip_predictor=True)
+    outputs = encoder_model(pixel_values, output_hidden_states=True)
     hidden_states = outputs.last_hidden_state[0]
     
-    batch_size, num_frames, num_channels, height, width = pixel_values_videos.shape
+    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 = num_frames // tubelet_size
+    num_temporal_patches = time // tubelet_size
     
-    dtype = pixel_values_videos.dtype
+    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_videos * std + mean
+    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)
     
@@ -530,6 +544,7 @@ def vjepa2_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]
@@ -539,6 +554,7 @@ def vjepa2_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):
@@ -548,108 +564,81 @@ def vjepa2_compute_latent_all_frames(video_frames, model, processor):
     
     return frames_unnorm, pca_images
 
-# ========== 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
+# 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
 
 # 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, model_type):
-    global stored_frames, visualization_cache, current_video_path
+def on_upload(video_path, mode):
+    global stored_frames, stored_viz, model, processor, visualization_cache, current_video_path
     if video_path is None:
         return gr.update(maximum=0), None, None
     
     # Initialize model if needed
-    model, processor = initialize_model(model_type)
+    if model is None:
+        model, processor = initialize_model()
     
-    # Check if we need to recompute (new video, mode, or model type)
+    # Check if we need to recompute (new video or mode not cached)
     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 = (cache_key not in visualization_cache or mode not in visualization_cache[cache_key])
+    need_to_compute_mode = (video_path_str not in visualization_cache or mode not in visualization_cache[video_path_str])
     
     if need_to_load_video:
-        # 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
+        # Load video frames
         print(f"Loading video: {video_path_str}")
-        video_frames = load_video(video_path, num_frames=num_frames, sample_rate=sample_rate)
+        video_frames = load_video(video_path)
         stored_frames = video_frames
         current_video_path = video_path_str
     else:
+        # Reuse already loaded frames
         video_frames = stored_frames
     
-    # Initialize cache for this video+model combination
-    if cache_key not in visualization_cache:
-        visualization_cache[cache_key] = {}
+    # Initialize cache for this video
+    if video_path_str not in visualization_cache:
+        visualization_cache[video_path_str] = {}
     
     if need_to_compute_mode:
-        print(f"Computing {mode} visualization for all frames using {model_type}...")
+        # Compute all visualizations and cache them
+        print(f"Computing {mode} visualization for all frames...")
         num_frames = len(stored_frames)
         tubelet_size = model.config.tubelet_size
         
         if mode == "reconstruction":
-            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)
+            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)
         
         elif mode == "attention":
-            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] = {}
+            original_frames, attention_maps = compute_attention_all_frames(video_frames, model, processor)
+            visualization_cache[video_path_str][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):
@@ -659,16 +648,13 @@ def on_upload(video_path, mode, model_type):
                     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[cache_key][mode][i] = fig_to_image(fig)
+                    visualization_cache[video_path_str][mode][i] = fig_to_image(fig)
         
         elif mode == "latent":
-            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] = {}
+            original_frames, pca_images = compute_latent_all_frames(video_frames, model, processor)
+            visualization_cache[video_path_str][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))
@@ -676,30 +662,34 @@ def on_upload(video_path, mode, model_type):
                     ax.imshow(pca_images[temporal_patch_idx])
                     ax.set_title(f"PCA Components - Frame {i}")
                     ax.axis('off')
-                    visualization_cache[cache_key][mode][i] = fig_to_image(fig)
+                    visualization_cache[video_path_str][mode][i] = fig_to_image(fig)
         
-        print(f"Caching complete for {mode} mode using {model_type}")
+        print(f"Caching complete for {mode} mode")
     
     # 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]
     
-    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]
+    # 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]
         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, model_type):
+def update_frame(idx, mode):
     global stored_frames, visualization_cache, current_video_path
     if not stored_frames:
         return None, None
@@ -708,57 +698,49 @@ def update_frame(idx, mode, model_type):
     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
-    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]
+    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]
             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, model_type):
+    def _load_example_video(  mode):
         """Load the predefined example video"""
         example_path = f"examples/{video_file}"
-        return on_upload(example_path, mode, model_type)
+        return on_upload(example_path, mode)
     
     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="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.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.Row():
         with gr.Column():
@@ -770,49 +752,47 @@ with gr.Blocks(title="Video 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, model_radio], 
-                    outputs=[frame_slider, orig_output, viz_output]
-                )
+                load_example_btn.click(load_example_video(video_file), inputs=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]
-    )
+        # 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])
     
-    def on_mode_change(video_path, mode, model_type):
-        """Handle mode change"""
+    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
         if video_path is None:
             return gr.update(maximum=0), None, None
-        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]
-    )
+        
+        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)
     
-    model_radio.change(
-        on_model_change, 
-        inputs=[model_radio, video_input, mode_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])
 
 if __name__ == "__main__":
-    # Initialize default model at startup
-    initialize_model('videomae')
-    demo.launch()
+    # Initialize model at startup
+    initialize_model()
+    demo.launch()