Modules/diffusion/module_dit.py

import math
from math import floor, log, pi
from typing import Any, List, Optional, Sequence, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, reduce, repeat
from einops.layers.torch import Rearrange
from torch import Tensor, einsum

# --- Utils ---

def exists(val):
    return val is not None

def default(val, d):
    return val if exists(val) else d

def rand_bool(shape, proba, device=None):
    if proba == 1:
        return torch.ones(shape, device=device, dtype=torch.bool)
    elif proba == 0:
        return torch.zeros(shape, device=device, dtype=torch.bool)
    return torch.bernoulli(torch.full(shape, proba, device=device)).to(torch.bool)

def rearrange_many(tensors, pattern, **kwargs):
    return tuple(rearrange(tensor, pattern, **kwargs) for tensor in tensors)

# --- RoPE (Rotary Positional Embeddings) ---

class RotaryEmbedding(nn.Module):
    def __init__(self, dim, theta=10000):
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)

    def forward(self, seq_len, device):
        t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
        freqs = torch.einsum("i , j -> i j", t, self.inv_freq)
        return torch.cat((freqs, freqs), dim=-1)

def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(pos, t):
    # t: [b, h, n, d]
    # pos: [n, d]
    pos = rearrange(pos, "n d -> 1 1 n d")
    return t * pos.cos() + rotate_half(t) * pos.sin()

# --- Core Layers ---

class AdaLayerNorm(nn.Module):
    def __init__(self, style_dim, channels, eps=1e-5):
        super().__init__()
        self.channels = channels
        self.eps = eps
        self.fc = nn.Linear(style_dim, channels * 2)

    def forward(self, x, s):
        # x: [b, c, t] -> needs transpose for LayerNorm
        x = x.transpose(-1, -2).transpose(1, -1) # [b, t, c]
        
        h = self.fc(s)
        h = h.view(h.size(0), 1, h.size(1))
        gamma, beta = torch.chunk(h, chunks=2, dim=2)
        
        x = F.layer_norm(x, (self.channels,), eps=self.eps)
        x = (1 + gamma) * x + beta
        
        return x.transpose(1, -1).transpose(-1, -2) # Back to [b, c, t]

def FeedForward(features: int, multiplier: int) -> nn.Module:
    mid_features = features * multiplier
    return nn.Sequential(
        nn.Conv1d(features, mid_features, 1),
        nn.GELU(),
        nn.Conv1d(mid_features, features, 1),
    )

# --- Attention with RoPE ---

class Attention(nn.Module):
    def __init__(
        self,
        features: int,
        *,
        head_features: int,
        num_heads: int,
        context_features: Optional[int] = None,
        use_rope: bool = True
    ):
        super().__init__()
        self.scale = head_features ** -0.5
        self.num_heads = num_heads
        mid_features = head_features * num_heads
        self.use_rope = use_rope

        # RoPE Generator
        if use_rope:
            self.rotary_emb = RotaryEmbedding(head_features)

        self.to_q = nn.Conv1d(features, mid_features, 1, bias=False)
        
        # Determine Key/Value input dimension (self vs cross attention)
        kv_dim = default(context_features, features)
        self.to_kv = nn.Conv1d(kv_dim, mid_features * 2, 1, bias=False)

        self.to_out = nn.Conv1d(mid_features, features, 1)

    def forward(self, x: Tensor, context: Optional[Tensor] = None) -> Tensor:
        b, c, n = x.shape
        # If context is None, it's self-attention
        kv_input = default(context, x)
        
        # Projections [b, (h d), n]
        q = self.to_q(x)
        k, v = self.to_kv(kv_input).chunk(2, dim=1)

        # Split heads: [b, h, n, d]
        q, k, v = map(
            lambda t: rearrange(t, "b (h d) n -> b h n d", h=self.num_heads), 
            (q, k, v)
        )

        # Apply RoPE (Rotary Positional Embeddings)
        if self.use_rope:
            # Generate frequencies for query length
            q_freqs = self.rotary_emb(n, q.device)
            q = apply_rotary_pos_emb(q_freqs, q)
            
            # If self-attention, k length is same as q. 
            # If cross-attention, we usually don't apply rotary to K unless time-aligned.
            # Assuming self-attention for RoPE here strictly or aligned cross.
            if context is None:
                k_freqs = q_freqs
                k = apply_rotary_pos_emb(k_freqs, k)

        # Attention
        sim = einsum("b h i d, b h j d -> b h i j", q, k) * self.scale
        attn = sim.softmax(dim=-1)
        
        out = einsum("b h i j, b h j d -> b h i d", attn, v)
        
        # Merge heads
        out = rearrange(out, "b h n d -> b (h d) n")
        return self.to_out(out)

class StyleAttention(nn.Module):
    """
    Wrapper that handles AdaLayerNorm before passing to the main Attention mechanism.
    """
    def __init__(
        self,
        features: int,
        *,
        style_dim: int,
        head_features: int,
        num_heads: int,
        context_features: Optional[int] = None,
        use_rope: bool = True,
    ):
        super().__init__()
        self.context_features = context_features
        context_dim = default(context_features, features)

        self.norm = AdaLayerNorm(style_dim, features)
        self.norm_context = AdaLayerNorm(style_dim, context_dim)
        
        self.attention = Attention(
            features=features,
            head_features=head_features,
            num_heads=num_heads,
            context_features=context_features,
            use_rope=use_rope
        )

    def forward(self, x: Tensor, s: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
        if self.context_features is not None:
             assert exists(context), "Context required for cross attention"

        context_input = default(context, x)
        
        # Adaptive Norm
        x = self.norm(x, s)
        context_input = self.norm_context(context_input, s)
        
        return self.attention(x, context=context_input if self.context_features else None)

# --- Transformer Blocks ---

class StyleTransformerBlock(nn.Module):
    def __init__(
        self,
        features: int,
        num_heads: int,
        head_features: int,
        style_dim: int,
        multiplier: int,
        context_features: Optional[int] = None,
    ):
        super().__init__()
        self.use_cross_attention = exists(context_features) and context_features > 0

        self.attention = StyleAttention(
            features=features,
            style_dim=style_dim,
            num_heads=num_heads,
            head_features=head_features,
            use_rope=True  # Force RoPE
        )

        if self.use_cross_attention:
            self.cross_attention = StyleAttention(
                features=features,
                style_dim=style_dim,
                num_heads=num_heads,
                head_features=head_features,
                context_features=context_features,
                use_rope=False # Usually no positional embedding on cross-attn to generic context
            )

        self.feed_forward = FeedForward(features=features, multiplier=multiplier)

    def forward(self, x: Tensor, s: Tensor, *, context: Optional[Tensor] = None) -> Tensor:
        x = self.attention(x, s) + x
        if self.use_cross_attention:
            x = self.cross_attention(x, s, context=context) + x
        x = self.feed_forward(x) + x
        return x

# --- Embeddings & Time ---

class LearnedPositionalEmbedding(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        assert (dim % 2) == 0
        half_dim = dim // 2
        self.weights = nn.Parameter(torch.randn(half_dim))

    def forward(self, x: Tensor) -> Tensor:
        x = rearrange(x, "b -> b 1")
        freqs = x * rearrange(self.weights, "d -> 1 d") * 2 * pi
        fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1)
        fouriered = torch.cat((x, fouriered), dim=-1)
        return fouriered

def TimePositionalEmbedding(dim: int, out_features: int) -> nn.Module:
    return nn.Sequential(
        LearnedPositionalEmbedding(dim),
        nn.Linear(in_features=dim + 1, out_features=out_features),
    )

class FixedEmbedding(nn.Module):
    def __init__(self, max_length: int, features: int):
        super().__init__()
        self.max_length = max_length
        self.embedding = nn.Embedding(max_length, features)

    def forward(self, x: Tensor) -> Tensor:
        batch_size, length, device = *x.shape[0:2], x.device
        assert length <= self.max_length
        position = torch.arange(length, device=device)
        fixed_embedding = self.embedding(position)
        fixed_embedding = repeat(fixed_embedding, "n d -> b n d", b=batch_size)
        return fixed_embedding

# --- Main Backbone: StyleTransformer1d ---

class StyleTransformer1d(nn.Module):
    def __init__(
        self,
        num_layers: int,
        channels: int,
        num_heads: int,
        head_features: int,
        multiplier: int,
        use_context_time: bool = True,
        context_features: Optional[int] = None,
        context_embedding_features: Optional[int] = None,
        embedding_max_length: int = 512,
    ):
        super().__init__()

        self.channels = channels
        self.context_embedding_features = default(context_embedding_features, 0)
        
        # Input projection
        input_dim = channels + self.context_embedding_features
        
        self.blocks = nn.ModuleList([
            StyleTransformerBlock(
                features=input_dim,
                head_features=head_features,
                num_heads=num_heads,
                multiplier=multiplier,
                style_dim=channels + self.context_embedding_features, # Mapping dim
                context_features=context_features,
            )
            for _ in range(num_layers)
        ])

        self.to_out = nn.Sequential(
            nn.Conv1d(input_dim, channels, 1)
        )
        
        use_context_features = exists(context_features)
        self.use_context_features = use_context_features
        self.use_context_time = use_context_time

        # Mapping Network (Conditioning)
        # We project time/features into a shared 'style' vector 's'
        if use_context_time or use_context_features:
            context_mapping_features = input_dim # Must match block width for AdaLN

            self.to_mapping = nn.Sequential(
                nn.Linear(context_mapping_features, context_mapping_features),
                nn.GELU(),
                nn.Linear(context_mapping_features, context_mapping_features),
                nn.GELU(),
            )
        
        if use_context_time:
            self.to_time = nn.Sequential(
                TimePositionalEmbedding(dim=channels, out_features=context_mapping_features),
                nn.GELU(),
            )

        if use_context_features:
            self.to_features = nn.Sequential(
                nn.Linear(in_features=context_features, out_features=context_mapping_features),
                nn.GELU(),
            )
            
        if self.context_embedding_features > 0:
            self.fixed_embedding = FixedEmbedding(
                max_length=embedding_max_length, features=self.context_embedding_features
            )
        

    def get_mapping(
        self, time: Optional[Tensor] = None, features: Optional[Tensor] = None
    ) -> Optional[Tensor]:
        """Combines time and global features into a style vector s"""
        items = []
        if self.use_context_time:
            assert exists(time)
            items += [self.to_time(time)]
        
        if self.use_context_features:
            assert exists(features)
            items += [self.to_features(features)]

        mapping = None
        if items:
            # sum aggregation of conditions
            mapping = reduce(torch.stack(items), "n b m -> b m", "sum")
            mapping = self.to_mapping(mapping)

        return mapping
            
    def run_network(self, x, time, embedding=None, features=None):
        # x: [b, c, t] (expecting channels first for convs, but transposing internally if needed)
        # Note: Previous blocks expect [b, c, t]
        
        mapping = self.get_mapping(time, features) # [b, style_dim]
        
        # Concatenate x with learned embeddings if they exist
        if exists(embedding):
            # embedding: [b, t, c_emb] -> [b, c_emb, t]
            emb_t = embedding.transpose(1, 2)
            x = torch.cat([x, emb_t], dim=1)
        
        # Forward through blocks
        for block in self.blocks:
            x = block(x, mapping, context=features) # context used for cross-attn if enabled
        
        x = self.to_out(x)
        return x
        
    def forward(self, 
                x: Tensor, 
                time: Tensor, 
                embedding_mask_proba: float = 0.0,
                embedding: Optional[Tensor] = None, 
                features: Optional[Tensor] = None,
                embedding_scale: float = 1.0) -> Tensor:
        
        b, device = x.shape[0], x.device
        
        # Handle Discrete Embedding (e.g. text/class tokens)
        final_embedding = None
        if self.context_embedding_features > 0 and exists(embedding):
            fixed_emb = self.fixed_embedding(embedding) # [b, t, d]
            
            if embedding_mask_proba > 0.0:
                batch_mask = rand_bool(shape=(b, 1, 1), proba=embedding_mask_proba, device=device)
                # Null embedding is assumed to be the fixed positional one here, 
                # or typically a separate learnable null token.
                # Here we just mask with the fixed pos embedding as 'unconditional' proxy 
                # or keep it simple.
                final_embedding = torch.where(batch_mask, fixed_emb, fixed_emb) # Placeholder logic
                # Ideally: if masked, replace with a learnable null parameter. 
                # For simplicity in this structure: we assume unmasked is passed explicitly or we use 0.
                if embedding_mask_proba > 0:
                     mask = rand_bool((b, 1, 1), 1 - embedding_mask_proba, device)
                     final_embedding = fixed_emb * mask 
            else:
                final_embedding = fixed_emb

        # Classifier-Free Guidance (CFG) Logic
        if embedding_scale != 1.0:
            # 1. Conditional Forward Pass
            out = self.run_network(x, time, embedding=final_embedding, features=features)
            
            # 2. Unconditional Forward Pass
            # To make it unconditional, we zero out the condition or use a null token.
            # Here we simulate unconditional by zeroing the explicit embedding/features if they exist.
            null_embedding = torch.zeros_like(final_embedding) if exists(final_embedding) else None
            null_features = torch.zeros_like(features) if exists(features) else None
            
            out_uncond = self.run_network(x, time, embedding=null_embedding, features=null_features)
            
            # CFG Formula: uncond + scale * (cond - uncond)
            return out_uncond + (out - out_uncond) * embedding_scale
            
        else:
            return self.run_network(x, time, embedding=final_embedding, features=features)


# --- Flow Matching Wrapper ---

class FlowMatching(nn.Module):
    def __init__(self, model: nn.Module, sigma_min=0.0):
        super().__init__()
        self.model = model
        self.sigma_min = sigma_min

    def forward(self, x1, condition=None):
        """
        Training Step:
        x1: Real data sample [b, c, n]
        condition: Optional condition tensor
        """
        b, c, n = x1.shape
        device = x1.device

        # 1. Sample Noise x0
        x0 = torch.randn_like(x1)

        # 2. Sample Time t
        t = torch.rand((b,), device=device)
        
        # 3. Compute Conditional Flow (Optimal Transport path)
        # x_t = (1 - (1 - sigma_min) * t) * x0 + t * x1
        # Simplified (assuming sigma_min approx 0 for standard OT): x_t = (1 - t) * x0 + t * x1
        t_expand = t.view(b, 1, 1)
        x_t = (1 - (1 - self.sigma_min) * t_expand) * x0 + t_expand * x1

        # 4. Compute Target Vector Field
        # u_t(x|x1) = x1 - (1 - sigma_min) * x0
        target_v = x1 - (1 - self.sigma_min) * x0

        # 5. Predict Vector Field
        # We pass t, x_t, and conditions to the backbone
        # We reshape features if necessary depending on how they are passed
        v_pred = self.model(x_t, t, features=condition, embedding_mask_proba=0.1) # 10% dropout for CFG training

        # 6. MSE Loss
        loss = F.mse_loss(v_pred, target_v)
        return loss

    @torch.no_grad()
    def sample(self, shape, steps=50, condition=None, cfg_scale=1.0):
        """
        Euler ODE Solver for generation
        """
        b, c, n = shape
        device = next(self.model.parameters()).device
        
        # Start from Normal distribution (x0)
        x = torch.randn(shape, device=device)
        
        dt = 1.0 / steps
        traj = [] # keep track of trajectory if needed

        for i in range(steps):
            t_val = i / steps
            t = torch.full((b,), t_val, device=device)
            
            # Predict velocity
            v = self.model(x, t, features=condition, embedding_scale=cfg_scale)
            
            # Euler step
            x = x + v * dt
            
        return x

# --- Example Usage ---

if __name__ == "__main__":
    # Dimensions
    B, Channels, Length = 2, 32, 128
    Cond_Dim = 64
    
    # Instantiate Backbone
    backbone = StyleTransformer1d(
        num_layers=4,
        channels=Channels,
        num_heads=4,
        head_features=32,
        multiplier=2,
        use_context_time=True,
        context_features=Cond_Dim, # Global condition (e.g. class embedding mapped)
        context_embedding_features=0, # Sequence condition
        embedding_max_length=Length
    )

    # Wrap in Flow Matching
    flow_model = FlowMatching(backbone)
    
    # --- Training Step ---
    data = torch.randn(B, Channels, Length)
    cond = torch.randn(B, Cond_Dim)
    
    loss = flow_model(data, condition=cond)
    print(f"Training Loss: {loss.item()}")

    # --- Sampling Step ---
    # Generate new samples with CFG
    generated = flow_model.sample(
        shape=(B, Channels, Length), 
        steps=20, 
        condition=cond, 
        cfg_scale=3.0
    )
    print(f"Generated Shape: {generated.shape}")