models/hyperfeature_fusion.py

"""
Aggregates features across multiple timesteps and layers using learned attention,
allowing the model to adaptively combine information from different denoising stages.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Dict, Tuple


class FeatureNormalizer(nn.Module):
    """
    Normalizes features from different timesteps/layers to comparable scales.
    Uses learnable normalization parameters per timestep.
    """
    def __init__(self, feature_dim: int = 3072, num_timesteps: int = 4):
        super().__init__()
        self.feature_dim = feature_dim
        self.num_timesteps = num_timesteps
        
        # Learnable normalization parameters (per timestep)
        self.layer_norms = nn.ModuleList([
            nn.LayerNorm(feature_dim) for _ in range(num_timesteps)
        ])
    
    def forward(self, features_per_timestep: List[torch.Tensor]) -> List[torch.Tensor]:
        normalized = []
        for i, feat in enumerate(features_per_timestep):
            # LayerNorm expects channels last: [B, C, H, W] -> [B, H, W, C]
            B, C, H, W = feat.shape
            feat_reshaped = feat.permute(0, 2, 3, 1).contiguous()  # [B, H, W, C]
            
            normalized_feat = self.layer_norms[i](feat_reshaped)
            
            normalized_feat = normalized_feat.permute(0, 3, 1, 2).contiguous()  # back to [B, C, H, W]
            normalized.append(normalized_feat)
        
        return normalized


class FeatureProjector(nn.Module):
    """
    Projects features from different layers/timesteps to a common dimension.
    Uses 1x1 convolutions for efficient channel reduction.
    """
    def __init__(self, in_channels: int = 3072, out_channels: int = 512, num_timesteps: int = 4):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        
        # Separate linear projection for each timestep (allows timestep-specific transformations)
        # No activation here: the transformer and output_proj provide nonlinearities,
        # and keeping this linear ensures the skip connection can correct in both directions.
        self.projectors = nn.ModuleList([
            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=True)
            for _ in range(num_timesteps)
        ])
    
    def forward(self, features_per_timestep: List[torch.Tensor]) -> List[torch.Tensor]:
        return [self.projectors[i](feat) for i, feat in enumerate(features_per_timestep)]


class CrossTimestepAttention(nn.Module):
    """
    Pixel-wise cross-timestep attention: for each spatial location, attends across
    features from different timesteps to learn per-pixel timestep importance.
    """
    def __init__(self, feature_dim: int = 512, num_timesteps: int = 4, num_heads: int = 4):
        super().__init__()
        self.feature_dim = feature_dim
        self.num_timesteps = num_timesteps
        self.num_heads = num_heads
        self.head_dim = feature_dim // num_heads
        
        assert feature_dim % num_heads == 0, "feature_dim must be divisible by num_heads"
        
        self.q_proj = nn.Linear(feature_dim, feature_dim)
        self.k_proj = nn.Linear(feature_dim, feature_dim)
        self.v_proj = nn.Linear(feature_dim, feature_dim)
        
        self.out_proj = nn.Linear(feature_dim, feature_dim)
        
        self.scale = self.head_dim ** -0.5
    
    def forward(self, features_per_timestep: List[torch.Tensor]) -> torch.Tensor:
        B, C, H, W = features_per_timestep[0].shape
        T = len(features_per_timestep)
        
        stacked = torch.stack(features_per_timestep, dim=1)  # [B, T, C, H, W]
        
        stacked = stacked.permute(0, 3, 4, 1, 2).contiguous()  # [B, H, W, T, C]
        stacked = stacked.view(B * H * W, T, C)
        
        # Query from last (most refined) timestep; keys/values from all timesteps
        q = self.q_proj(stacked[:, -1:, :])  # [B*H*W, 1, C]
        k = self.k_proj(stacked)  # [B*H*W, T, C]
        v = self.v_proj(stacked)  # [B*H*W, T, C]
        
        q = q.view(B * H * W, 1, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B * H * W, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(B * H * W, T, self.num_heads, self.head_dim).transpose(1, 2)
        
        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale  # [B*H*W, num_heads, 1, T]
        attn = F.softmax(attn, dim=-1)
        
        out = torch.matmul(attn, v)  # [B*H*W, num_heads, 1, head_dim]
        
        out = out.transpose(1, 2).contiguous().view(B * H * W, 1, C)
        
        out = self.out_proj(out)
        
        out = out.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()  # [B, C, H, W]
        
        return out


class LayerScaleTransformerLayer(nn.Module):
    """
    Transformer layer with LayerScale (from CaiT: "Going deeper with Image Transformers").
    LayerScale helps stabilize training of deeper transformers by adding learnable
    diagonal scaling matrices initialized to small values (e.g., 1e-4).
    """
    def __init__(self, feature_dim: int, nhead: int = 8, dropout: float = 0.1, 
                 layer_scale_init: float = 1e-4):
        super().__init__()
        self.feature_dim = feature_dim
        
        self.self_attn = nn.MultiheadAttention(
            feature_dim, nhead, dropout=dropout, batch_first=True
        )
        
        self.ffn = nn.Sequential(
            nn.Linear(feature_dim, feature_dim * 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(feature_dim * 2, feature_dim),
            nn.Dropout(dropout)
        )
        
        self.norm1 = nn.LayerNorm(feature_dim)
        self.norm2 = nn.LayerNorm(feature_dim)
        
        # LayerScale diagonal scaling, initialized small to stabilize deep transformers
        self.layer_scale_1 = nn.Parameter(
            torch.ones(feature_dim) * layer_scale_init
        )
        self.layer_scale_2 = nn.Parameter(
            torch.ones(feature_dim) * layer_scale_init
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Pre-norm self-attention with LayerScale (CaiT style)
        x_normed = self.norm1(x)
        attn_out, _ = self.self_attn(x_normed, x_normed, x_normed)
        x = x + self.layer_scale_1 * attn_out
        
        ffn_out = self.ffn(self.norm2(x))
        x = x + self.layer_scale_2 * ffn_out
        
        return x


class TemporalTransformerAggregator(nn.Module):
    """
    Transformer that aggregates features across timesteps with LayerScale.
    Predicts content-adaptive per-pixel fusion weights rather than global weights.
    """
    def __init__(self, feature_dim: int = 512, num_timesteps: int = 4, num_layers: int = 2,
                 layer_scale_init: float = 1e-4):
        super().__init__()
        self.feature_dim = feature_dim
        self.num_timesteps = num_timesteps
        self.num_layers = num_layers
        
        # Learnable temporal positional encoding, small init to avoid early dominance
        self.temporal_pos_embed = nn.Parameter(torch.randn(1, num_timesteps, feature_dim))
        nn.init.trunc_normal_(self.temporal_pos_embed, std=0.02)
        
        self.layers = nn.ModuleList([
            LayerScaleTransformerLayer(
                feature_dim=feature_dim,
                nhead=8,
                dropout=0.1,
                layer_scale_init=layer_scale_init
            ) for _ in range(num_layers)
        ])
        
        # Per-pixel timestep weights: [B*H*W, C, T] -> [B*H*W, 1, T]
        self.alpha_head = nn.Conv1d(self.feature_dim, 1, kernel_size=1)
        # Small init so weighting starts nearly uniform across timesteps
        nn.init.normal_(self.alpha_head.weight, mean=0.0, std=0.01)
        nn.init.zeros_(self.alpha_head.bias)
    
    def forward(self, features_per_timestep: List[torch.Tensor], return_alpha: bool = False):
        """If return_alpha, also returns the per-pixel timestep weight map [B, T, H, W]."""
        B, C, H, W = features_per_timestep[0].shape
        T = len(features_per_timestep)
        
        stacked = torch.stack(features_per_timestep, dim=1)  # [B, T, C, H, W]
        stacked = stacked.permute(0, 3, 4, 1, 2).contiguous()  # [B, H, W, T, C]
        stacked = stacked.view(B * H * W, T, C)
        
        stacked = stacked + self.temporal_pos_embed
        
        transformed = stacked
        for layer in self.layers:
            transformed = layer(transformed)  # [B*H*W, T, C]
        
        transformed_t = transformed.transpose(1, 2)  # [B*H*W, C, T]
        
        logits_alpha = self.alpha_head(transformed_t)  # [B*H*W, 1, T]
        logits_alpha = logits_alpha.squeeze(1)  # [B*H*W, T]
        
        # Softmax over timesteps per pixel
        alpha = torch.softmax(logits_alpha, dim=-1)  # [B*H*W, T]
        
        alpha_expanded = alpha.unsqueeze(1)  # [B*H*W, 1, T]
        pooled = (alpha_expanded * transformed_t).sum(dim=-1)  # [B*H*W, C]
        
        out = pooled.view(B, H, W, C).permute(0, 3, 1, 2).contiguous()  # [B, C, H, W]
        
        if return_alpha:
            alpha_map = alpha.view(B, H, W, T).permute(0, 3, 1, 2).contiguous()  # [B, T, H, W]
            return out, alpha_map
        
        return out


class CBAM(nn.Module):
    """
    Convolutional Block Attention Module (CBAM, ECCV 2018).
    Sequentially applies channel then spatial attention to refine features.
    """
    def __init__(self, channels: int, reduction: int = 16, kernel_size: int = 7):
        super().__init__()
        self.channels = channels
        
        self.channel_mlp = nn.Sequential(
            nn.Linear(channels, channels // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(channels // reduction, channels)
        )
        
        self.spatial_conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size // 2)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, C, H, W = x.shape
        
        avg_pool = x.mean(dim=(2, 3), keepdim=False)  # [B, C]
        max_pool = x.amax(dim=(2, 3), keepdim=False)  # [B, C]
        
        # Shared MLP applied to both pooled descriptors
        channel_weight = torch.sigmoid(
            self.channel_mlp(avg_pool) + self.channel_mlp(max_pool)
        ).view(B, C, 1, 1)
        
        x = x * channel_weight
        
        avg_spatial = x.mean(dim=1, keepdim=True)  # [B, 1, H, W]
        max_spatial = x.amax(dim=1, keepdim=True)  # [B, 1, H, W]
        
        spatial_weight = torch.sigmoid(
            self.spatial_conv(torch.cat([avg_spatial, max_spatial], dim=1))
        )
        
        return x * spatial_weight


class HyperfeatureFusion(nn.Module):
    """
    Complete Hyperfeature Fusion module.
    Implements the full pipeline: normalize -> project -> attend -> fuse -> refine (CBAM).
    """
    def __init__(
        self,
        in_channels: int = 3072,
        out_channels: int = 3072,
        hidden_dim: int = 512,
        num_timesteps: int = 4,
        fusion_type: str = 'attention',  # 'attention' or 'transformer'
        num_attention_heads: int = 4,
        num_transformer_layers: int = 2,
        layer_scale_init: float = 1e-4,
        return_alpha: bool = False
    ):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.hidden_dim = hidden_dim
        self.num_timesteps = num_timesteps
        self.fusion_type = fusion_type
        self.return_alpha = return_alpha
        
        self.normalizer = FeatureNormalizer(in_channels, num_timesteps)
        
        self.projector = FeatureProjector(in_channels, hidden_dim, num_timesteps)
        
        if fusion_type == 'attention':
            self.aggregator = CrossTimestepAttention(hidden_dim, num_timesteps, num_attention_heads)
        elif fusion_type == 'transformer':
            self.aggregator = TemporalTransformerAggregator(
                hidden_dim, num_timesteps, num_transformer_layers, layer_scale_init
            )
        else:
            raise ValueError(f"Unknown fusion_type: {fusion_type}")
        
        self.output_proj = nn.Sequential(
            nn.Conv2d(hidden_dim, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=1)
        )
        
        # Align hidden_dim to out_channels for the skip from the latest timestep
        if hidden_dim != out_channels:
            self.skip_align = nn.Conv2d(hidden_dim, out_channels, kernel_size=1, bias=False)
        else:
            self.skip_align = nn.Identity()
        
        # Learnable skip gate initialized near zero so the aggregated path dominates
        # early and the skip is gradually introduced during training.
        self.skip_gate = nn.Parameter(torch.tensor(0.1))
        
        self.cbam = CBAM(out_channels, reduction=16, kernel_size=7)
    
    def forward(self, features_per_timestep: List[torch.Tensor]):
        """features_per_timestep: list of [B, C_in, H, W] tensors ordered early to late."""
        normalized = self.normalizer(features_per_timestep)
        
        projected = self.projector(normalized)
        
        if self.return_alpha and self.fusion_type == 'transformer':
            aggregated, alpha_map = self.aggregator(projected, return_alpha=True)
        else:
            aggregated = self.aggregator(projected)
            alpha_map = None
        
        output = self.output_proj(aggregated)
        
        # Gated skip connection from the latest timestep preserves refined features
        latest_skip = self.skip_align(projected[-1])
        output = output + self.skip_gate * latest_skip
        
        output = self.cbam(output)
        
        if self.return_alpha and alpha_map is not None:
            return output, alpha_map
        
        return output


class MultiLayerHyperfeatureFusion(nn.Module):
    """
    Applies Hyperfeature fusion independently for each FLUX feature layer.
    This allows different layers to learn different temporal aggregation strategies.
    """
    def __init__(
        self,
        in_channels: int = 3072,
        out_channels: int = 3072,
        hidden_dim: int = 512,
        num_layers: int = 4,
        num_timesteps: int = 4,
        fusion_type: str = 'attention',
        num_attention_heads: int = 4,
        num_transformer_layers: int = 2,
        layer_scale_init: float = 1e-4,
        return_alpha: bool = False
    ):
        super().__init__()
        self.num_layers = num_layers
        self.num_timesteps = num_timesteps
        self.return_alpha = return_alpha
        
        self.layer_fusions = nn.ModuleList([
            HyperfeatureFusion(
                in_channels=in_channels,
                out_channels=out_channels,
                hidden_dim=hidden_dim,
                num_timesteps=num_timesteps,
                fusion_type=fusion_type,
                num_attention_heads=num_attention_heads,
                num_transformer_layers=num_transformer_layers,
                layer_scale_init=layer_scale_init,
                return_alpha=return_alpha
            ) for _ in range(num_layers)
        ])
    
    def forward(self, multi_timestep_features: Dict[int, List[torch.Tensor]]):
        """multi_timestep_features: dict mapping timestep -> list of per-layer maps [B, C, H, W]."""
        # Regroup features by layer instead of by timestep
        features_per_layer = []
        timesteps = sorted(multi_timestep_features.keys())
        
        for layer_idx in range(self.num_layers):
            layer_features_across_timesteps = [
                multi_timestep_features[t][layer_idx] for t in timesteps
            ]
            features_per_layer.append(layer_features_across_timesteps)
        
        fused_layers = []
        alpha_layers = []
        for layer_idx, layer_features in enumerate(features_per_layer):
            result = self.layer_fusions[layer_idx](layer_features)
            if self.return_alpha:
                fused, alpha = result
                fused_layers.append(fused)
                alpha_layers.append(alpha)
            else:
                fused_layers.append(result)
        
        if self.return_alpha:
            return fused_layers, alpha_layers
        
        return fused_layers


def create_hyperfeature_fusion(
    num_timesteps: int = 4,
    num_layers: int = 4,
    fusion_type: str = 'attention',
    hidden_dim: int = 512,
    num_transformer_layers: int = 2,
    layer_scale_init: float = 1e-4,
    return_alpha: bool = False,
    feature_dim: int = 3072
) -> MultiLayerHyperfeatureFusion:
    """Factory for MultiLayerHyperfeatureFusion (feature_dim 3072 for FLUX, 1536 for SD3.5)."""
    return MultiLayerHyperfeatureFusion(
        in_channels=feature_dim,
        out_channels=feature_dim,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        num_timesteps=num_timesteps,
        fusion_type=fusion_type,
        num_attention_heads=8,
        num_transformer_layers=num_transformer_layers,
        layer_scale_init=layer_scale_init,
        return_alpha=return_alpha
    )