Upload vil_tracker/models/film_temporal.py with huggingface_hub

vil_tracker/models/film_temporal.py

@@ -0,0 +1,188 @@
+"""
+FiLM (Feature-wise Linear Modulation) Temporal Module.
+
+Replaces DTPTrack's temporal prompt tokens (which are broken for bidirectional mLSTM
+scanning) with channel-wise affine modulation conditioned on temporal context.
+
+Architecture:
+1. TemporalReliabilityCalibrator: learns reliability weights for temporal features
+2. FiLMTemporalModulation: γ(t)·x + β(t) modulation per block
+3. TemporalModulationManager: manages FiLM layers across all blocks
+
+Reference: Perez et al., "FiLM: Visual Reasoning with a General Conditioning Layer"
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class TemporalReliabilityCalibrator(nn.Module):
+    """Learns a reliability score for temporal context.
+    
+    Takes temporal features (e.g., from previous frame's mLSTM states)
+    and produces a scalar reliability weight in [0, 1] for each token.
+    """
+    def __init__(self, dim: int = 384):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, dim // 4),
+            nn.GELU(),
+            nn.Linear(dim // 4, 1),
+            nn.Sigmoid(),
+        )
+    
+    def forward(self, temporal_feat: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            temporal_feat: (B, S, D) temporal context features
+        Returns:
+            reliability: (B, S, 1) reliability weights in [0, 1]
+        """
+        return self.net(temporal_feat)
+
+
+class FiLMTemporalModulation(nn.Module):
+    """Feature-wise Linear Modulation conditioned on temporal context.
+    
+    Computes: output = γ(temporal) · x + β(temporal)
+    where γ, β are learned from temporal features via small networks.
+    """
+    def __init__(self, dim: int = 384):
+        super().__init__()
+        # Generate scale (γ) and shift (β) from temporal context
+        self.gamma_net = nn.Sequential(
+            nn.Linear(dim, dim // 4),
+            nn.GELU(),
+            nn.Linear(dim // 4, dim),
+        )
+        self.beta_net = nn.Sequential(
+            nn.Linear(dim, dim // 4),
+            nn.GELU(),
+            nn.Linear(dim // 4, dim),
+        )
+        
+        # Initialize γ near 1 and β near 0 (identity modulation at init)
+        nn.init.zeros_(self.gamma_net[-1].weight)
+        nn.init.ones_(self.gamma_net[-1].bias)
+        nn.init.zeros_(self.beta_net[-1].weight)
+        nn.init.zeros_(self.beta_net[-1].bias)
+    
+    def forward(
+        self,
+        x: torch.Tensor,
+        temporal_context: torch.Tensor,
+        reliability: torch.Tensor = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: (B, S, D) input features from current frame
+            temporal_context: (B, S, D) temporal features (prev frame, pooled states, etc.)
+            reliability: (B, S, 1) optional reliability weights
+        Returns:
+            (B, S, D) modulated features
+        """
+        gamma = self.gamma_net(temporal_context)  # (B, S, D)
+        beta = self.beta_net(temporal_context)     # (B, S, D)
+        
+        if reliability is not None:
+            # Blend between identity (no modulation) and full modulation based on reliability
+            gamma = reliability * gamma + (1 - reliability) * torch.ones_like(gamma)
+            beta = reliability * beta
+        
+        return gamma * x + beta
+
+
+class TemporalModulationManager(nn.Module):
+    """Manages FiLM modulation across multiple backbone blocks.
+    
+    Applies FiLM modulation after every N-th block, using temporal context
+    from the previous frame's features (or running average).
+    """
+    def __init__(
+        self,
+        dim: int = 384,
+        num_blocks: int = 24,
+        modulation_interval: int = 6,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_blocks = num_blocks
+        self.modulation_interval = modulation_interval
+        
+        # FiLM layers at intervals
+        num_film = num_blocks // modulation_interval
+        self.film_layers = nn.ModuleList([
+            FiLMTemporalModulation(dim=dim)
+            for _ in range(num_film)
+        ])
+        
+        # Reliability calibrator
+        self.reliability = TemporalReliabilityCalibrator(dim=dim)
+        
+        # Temporal context projection (map prev features to context)
+        self.context_proj = nn.Linear(dim, dim)
+        
+        # Running temporal context (registered as buffer, not parameter)
+        self.register_buffer('_temporal_context', None)
+    
+    def should_modulate(self, block_idx: int) -> bool:
+        """Check if this block index should apply FiLM modulation."""
+        return (block_idx + 1) % self.modulation_interval == 0
+    
+    def get_film_layer(self, block_idx: int) -> FiLMTemporalModulation:
+        """Get the FiLM layer for a given block index."""
+        film_idx = (block_idx + 1) // self.modulation_interval - 1
+        return self.film_layers[film_idx]
+    
+    def update_temporal_context(self, features: torch.Tensor):
+        """Update temporal context from current frame features.
+        
+        Args:
+            features: (B, S, D) features from current frame processing
+        """
+        context = self.context_proj(features.detach())
+        if self._temporal_context is None:
+            self._temporal_context = context
+        else:
+            # EMA update
+            self._temporal_context = 0.7 * self._temporal_context + 0.3 * context
+    
+    def modulate(
+        self,
+        x: torch.Tensor,
+        block_idx: int,
+    ) -> torch.Tensor:
+        """Apply FiLM modulation at the appropriate block.
+        
+        Args:
+            x: (B, S, D) features at block_idx
+            block_idx: current block index
+        Returns:
+            (B, S, D) modulated features (or unchanged if not a modulation block)
+        """
+        if not self.should_modulate(block_idx):
+            return x
+        
+        if self._temporal_context is None:
+            return x  # No temporal context yet (first frame)
+        
+        film = self.get_film_layer(block_idx)
+        
+        # Ensure temporal context matches spatial dimension
+        tc = self._temporal_context
+        if tc.shape[1] != x.shape[1]:
+            # Interpolate or pad temporal context
+            tc = F.interpolate(
+                tc.transpose(1, 2),
+                size=x.shape[1],
+                mode='linear',
+                align_corners=False,
+            ).transpose(1, 2)
+        
+        reliability = self.reliability(tc)
+        return film(x, tc, reliability)
+    
+    def reset(self):
+        """Reset temporal context (e.g., for new tracking sequence)."""
+        self._temporal_context = None