Modules/diffusion/modules.py

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

from .utils import *

import torch
import torch.nn as nn
from einops import rearrange, reduce, repeat
from einops.layers.torch import Rearrange
from einops_exts import rearrange_many
from torch import Tensor, einsum


"""
Utils + Transformer with RoPE (Rotary Positional Embeddings)
Extended to support RMSNorm or LayerNorm via `norm_type` parameter.
"""


class RMSNorm(nn.Module):
    """RMSNorm: normalizacja po RMS (nie odejmuje średniej)."""

    def __init__(self, dim: int, eps: float = 1e-8, elementwise_affine: bool = True):
        super().__init__()
        self.dim = dim
        self.eps = eps
        self.elementwise_affine = elementwise_affine
        if elementwise_affine:
            self.weight = nn.Parameter(torch.ones(dim))
        else:
            self.register_buffer("weight", torch.ones(dim))

    def forward(self, x: Tensor) -> Tensor:
        # x: (..., dim)
        rms = x.pow(2).mean(dim=-1, keepdim=True).add(self.eps).sqrt()
        x_normed = x / rms
        return x_normed * self.weight


def _make_norm(norm_type: str, dim: int) -> nn.Module:
    """Helper: create norm module by name."""
    if norm_type is None or norm_type == "none":
        return nn.Identity()
    if norm_type == "layer":
        return nn.LayerNorm(dim)
    if norm_type == "rms":
        return RMSNorm(dim)
    raise ValueError(f"Unknown norm_type: {norm_type}")


class Transformer1d(nn.Module):
    def __init__(
        self,
        num_layers: int,
        channels: int,
        num_heads: int,
        head_features: int,
        multiplier: int,
        use_context_time: bool = True,
        use_rope: bool = False,
        rope_max_seq_len: int = 512,
        context_embedding_features: Optional[int] = None,
        embedding_max_length: int = 512,
        # Dropout params
        dropout: float = 0.0,  # general dropout (used as input/out dropout)
        attn_dropout: float = 0.0,  # dropout on attention weights
        ff_dropout: float = 0.0,  # dropout after feed-forward
        # norm type: "layer" (default) or "rms" or "none"
        norm_type: str = "layer",
    ):
        """
        Transformer1d simplified for the case where cross-attention/context
        features are never used.

        - context_features and context_features_multiplier removed.
        - All attention is self-attention.
        """
        super().__init__()

        # if context_embedding_features is None, treat as 0 (no extra embedding dim)
        context_embedding_features = context_embedding_features or 0

        self.context_embedding_features = context_embedding_features
        total_features = channels + context_embedding_features

        # Save dropout modules and params
        self.input_dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()
        self._dropout = dropout
        self._attn_dropout = attn_dropout
        self._ff_dropout = ff_dropout

        self.blocks = nn.ModuleList(
            [
                TransformerBlock(
                    features=total_features,
                    head_features=head_features,
                    num_heads=num_heads,
                    multiplier=multiplier,
                    use_rope=use_rope,
                    rope_max_seq_len=rope_max_seq_len,
                    # pass dropout settings into blocks
                    dropout=dropout,
                    attn_dropout=attn_dropout,
                    ff_dropout=ff_dropout,
                    norm_type=norm_type,
                )
                for _ in range(num_layers)
            ]
        )

        self.to_out = nn.Sequential(
            Rearrange("b t c -> b c t"),
            nn.Conv1d(in_channels=total_features, out_channels=channels, kernel_size=1),
        )

        # We assume `features: Optional[Tensor]` will never be passed.
        # So we only keep time-based context (if enabled).
        self.use_context_time = use_context_time

        if use_context_time:
            context_mapping_features = total_features
            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(),
            )

            self.to_time = nn.Sequential(
                TimePositionalEmbedding(dim=channels, out_features=context_mapping_features),
                nn.GELU(),
            )
            self.mapping_features = context_mapping_features
        else:
            self.to_mapping = None
            self.to_time = None
            self.mapping_features = None

        self.fixed_embedding = FixedEmbedding(
            max_length=embedding_max_length, features=context_embedding_features
        )

    def get_mapping(self, time: Optional[Tensor] = None) -> Optional[Tensor]:
        """Compute mapping solely from time. `features` is intentionally removed."""
        if not self.use_context_time:
            return None
        assert exists(time), "use_context_time=True but no time features provided"
        mapping = self.to_time(time)
        mapping = self.to_mapping(mapping)
        return mapping

    def run(self, x: Tensor, time: Tensor, embedding: Tensor) -> Tensor:
        # x: (b, seq_len_x, channels)
        # embedding: (b, seq_len_e, context_embedding_features)
        mapping = self.get_mapping(time)
        # Concatenate fixed embedding channels (if embedding features == 0, this is a no-op)
        x = torch.cat([x.expand(-1, embedding.size(1), -1), embedding], dim=-1)

        if mapping is not None:
            mapping = mapping.unsqueeze(1).expand(-1, embedding.size(1), -1)

        # Apply input dropout once to the inputs to the transformer blocks.
        x = self.input_dropout(x)

        for block in self.blocks:
            if mapping is not None:
                x = x + mapping
            x = block(x)

        x = x.mean(dim=1).unsqueeze(1)
        x = self.to_out(x)
        x = x.transpose(-1, -2)

        return x

    def forward(
        self,
        x: Tensor,
        time: Tensor,
        embedding_mask_proba: float = 0.1,
        embedding: Optional[Tensor] = None,
        embedding_scale: float = 1.0,
    ) -> Tensor:
        """
        Note: `features` tensor argument has been removed intentionally because it will
        never be provided.
        """
        assert exists(embedding), "embedding must be provided"
        b, device = embedding.shape[0], embedding.device
        fixed_embedding = self.fixed_embedding(embedding)
        if embedding_mask_proba > 0.0:
            # Randomly mask embedding per-batch
            batch_mask = rand_bool(
                shape=(b, 1, 1), proba=embedding_mask_proba, device=device
            )
            embedding = torch.where(batch_mask, fixed_embedding, embedding)

        if embedding_scale != 1.0:
            # Compute both normal and fixed embedding outputs (classifier-free guidance)
            out = self.run(x, time, embedding=embedding)
            out_masked = self.run(x, time, embedding=fixed_embedding)
            return out_masked + (out - out_masked) * embedding_scale
        else:
            return self.run(x, time, embedding=embedding)


"""
Rotary Positional Embedding implementation
"""


class RotaryEmbedding(nn.Module):
    """
    RoPE implementation that caches sin/cos for up to max_seq_len.
    Works on head dimension d (must be even).
    """

    def __init__(self, dim: int, max_seq_len: int = 512, base: int = 10000):
        super().__init__()
        assert dim % 2 == 0, "Rotary embedding dim must be even"
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.base = base
        # inv_freq = 1.0 / (base ** (i/dim)) for even positions
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        # No precomputed cos/sin buffer here — generate on demand (optionally could cache)

    def _build_sin_cos(self, seq_len: int, device: torch.device, dtype: torch.dtype):
        # seq_len x (dim/2)
        positions = torch.arange(seq_len, device=device, dtype=dtype).unsqueeze(1)
        angles = positions * rearrange(self.inv_freq.to(device=device, dtype=dtype), "d -> 1 d")
        sin = torch.sin(angles)  # seq_len x (dim/2)
        cos = torch.cos(angles)  # seq_len x (dim/2)
        # interleave to match original (d) layout: [cos0, cos0, cos1, cos1, ...] when applied
        sin = torch.stack([sin, sin], dim=-1).reshape(seq_len, self.dim)
        cos = torch.stack([cos, cos], dim=-1).reshape(seq_len, self.dim)
        return sin, cos

    @staticmethod
    def rotate_half(x: Tensor) -> Tensor:
        # x: (..., d) where d is even
        x1 = x[..., ::2]
        x2 = x[..., 1::2]
        # rotate: (-x2, x1) interleaved
        x_rotated = torch.stack((-x2, x1), dim=-1).reshape_as(x)
        return x_rotated

    def apply_rotary(self, q: Tensor, k: Tensor) -> Tuple[Tensor, Tensor]:
        """
        Apply RoPE to q and k.

        q,k shapes: (b, h, n, d) with d == self.dim
        """
        assert q.shape[-1] == self.dim and k.shape[-1] == self.dim

        seq_len = q.shape[-2]
        device = q.device
        dtype = q.dtype

        sin, cos = self._build_sin_cos(seq_len, device=device, dtype=dtype)
        # make broadcastable: (1, 1, n, d)
        sin = sin.unsqueeze(0).unsqueeze(0)
        cos = cos.unsqueeze(0).unsqueeze(0)

        q_out = (q * cos) + (self.rotate_half(q) * sin)
        k_out = (k * cos) + (self.rotate_half(k) * sin)
        return q_out, k_out


"""
Attention Components (self-attention only) with RoPE
"""


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


class AttentionBase(nn.Module):
    def __init__(
        self,
        features: int,
        *,
        head_features: int,
        num_heads: int,
        use_rope: bool,
        rope_max_seq_len: int = 512,
        out_features: Optional[int] = None,
        # new dropout params
        attn_dropout: float = 0.0,
        out_dropout: float = 0.0,
    ):
        super().__init__()
        self.scale = head_features ** -0.5
        self.num_heads = num_heads
        self.use_rope = use_rope
        mid_features = head_features * num_heads

        if out_features is None:
            out_features = features

        self.to_out = nn.Linear(in_features=mid_features, out_features=out_features)

        # dropout modules
        self.attn_dropout = nn.Dropout(attn_dropout) if attn_dropout > 0.0 else nn.Identity()
        self.out_dropout = nn.Dropout(out_dropout) if out_dropout > 0.0 else nn.Identity()

        # Rotary embedding per-head-dim
        if use_rope:
            # head_features is the d per head that RoPE should be applied to
            self.rotary = RotaryEmbedding(dim=head_features, max_seq_len=rope_max_seq_len)
        else:
            self.rotary = None

    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
        # Split heads
        # q,k,v shape before: (b, n, h * d) -> after: (b, h, n, d)
        q, k, v = rearrange_many((q, k, v), "b n (h d) -> b h n d", h=self.num_heads)

        # Apply RoPE to q and k if enabled
        if self.rotary is not None:
            q, k = self.rotary.apply_rotary(q, k)

        # Compute similarity matrix
        sim = einsum("... n d, ... m d -> ... n m", q, k)
        sim = sim * self.scale
        # Get attention matrix with softmax
        attn = sim.softmax(dim=-1)
        # apply dropout to attention weights if configured
        attn = self.attn_dropout(attn)
        # Compute values
        out = einsum("... n m, ... m d -> ... n d", attn, v)
        out = rearrange(out, "b h n d -> b n (h d)")
        out = self.to_out(out)
        out = self.out_dropout(out)
        return out


class Attention(nn.Module):
    def __init__(
        self,
        features: int,
        *,
        head_features: int,
        num_heads: int,
        out_features: Optional[int] = None,
        use_rope: bool,
        rope_max_seq_len: int = 512,
        # propagate dropout params
        attn_dropout: float = 0.0,
        out_dropout: float = 0.0,
        # norm type: "layer", "rms", or "none"
        norm_type: str = "layer",
    ):
        """
        Self-attention only (context / cross-attention removed).
        """
        super().__init__()

        mid_features = head_features * num_heads

        self.norm = _make_norm(norm_type, features)
        # For self-attention we compute q from x and k,v from x as well.
        self.to_q = nn.Linear(in_features=features, out_features=mid_features, bias=False)
        self.to_kv = nn.Linear(in_features=features, out_features=mid_features * 2, bias=False)

        self.attention = AttentionBase(
            features,
            out_features=out_features,
            num_heads=num_heads,
            head_features=head_features,
            use_rope=use_rope,
            rope_max_seq_len=rope_max_seq_len,
            attn_dropout=attn_dropout,
            out_dropout=out_dropout,
        )

    def forward(self, x: Tensor) -> Tensor:
        # Pre-norm before computing q/k/v (recommended for stability)
        if not isinstance(self.norm, nn.Identity):
            x_norm = self.norm(x)
        else:
            x_norm = x
        q = self.to_q(x_norm)
        k, v = torch.chunk(self.to_kv(x_norm), 2, dim=-1)
        return self.attention(q, k, v)


"""
Transformer Blocks
"""


class TransformerBlock(nn.Module):
    def __init__(
        self,
        features: int,
        num_heads: int,
        head_features: int,
        multiplier: int,
        use_rope: bool,
        rope_max_seq_len: int = 512,
        # new: dropout params per-block
        dropout: float = 0.0,
        attn_dropout: float = 0.0,
        ff_dropout: float = 0.0,
        # norm type
        norm_type: str = "layer",
    ):
        super().__init__()

        # Only self-attention (no cross-attention)
        self.attention = Attention(
            features=features,
            num_heads=num_heads,
            head_features=head_features,
            use_rope=use_rope,
            rope_max_seq_len=rope_max_seq_len,
            attn_dropout=attn_dropout,
            out_dropout=dropout,
            norm_type=norm_type,
        )

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

        # LayerNorm or RMSNorm before feed-forward and dropout after FF
        self.norm_ff = _make_norm(norm_type, features)
        self.ff_dropout = nn.Dropout(ff_dropout) if ff_dropout > 0.0 else nn.Identity()

    def forward(self, x: Tensor) -> Tensor:
        x = self.attention(x) + x
        # Apply norm before feed-forward (pre-norm style) and apply FF + dropout then residual
        if not isinstance(self.norm_ff, nn.Identity):
            ff_in = self.norm_ff(x)
        else:
            ff_in = x
        ff_out = self.feed_forward(ff_in)
        ff_out = self.ff_dropout(ff_out)
        x = ff_out + x
        return x


"""
Time Embeddings (unchanged)
"""


class SinusoidalEmbedding(nn.Module):
    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim

    def forward(self, x: Tensor) -> Tensor:
        device, half_dim = x.device, self.dim // 2
        emb = torch.tensor(log(10000) / (half_dim - 1), device=device)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = rearrange(x, "i -> i 1") * rearrange(emb, "j -> 1 j")
        return torch.cat((emb.sin(), emb.cos()), dim=-1)


class LearnedPositionalEmbedding(nn.Module):
    """Used for continuous time"""

    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_message = "Input sequence length must be <= max_length"
        assert length <= self.max_length, assert_message
        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