src/models/dit.py

"""Conditional Diffusion Transformer used as the RiboSphere denoiser."""

from __future__ import annotations

import math
from typing import Any

import torch
from torch import Tensor, nn

from .attention import SelfAttention
from .layers import FeedForward


def apply_adaptive_modulation(
    inputs: Tensor,
    shift: Tensor,
    scale: Tensor,
) -> Tensor:
    """Apply adaptive affine modulation over the feature dimension."""
    return inputs * (1 + scale.unsqueeze(-2)) + shift.unsqueeze(-2)


class DiffusionTransformer(nn.Module):
    """Conditional diffusion transformer that predicts coordinate flow."""

    def __init__(
        self,
        *,
        num_channels: int,
        input_channels: int,
        num_layers: int,
        num_heads: int,
        conditioning_type: str = "cat",
        mlp_factor: int = 4,
        normalize_queries_and_keys: bool = False,
        share_adaln: bool = True,
        attention_backend: str = "sdpa",
    ) -> None:
        super().__init__()
        if min(num_channels, input_channels, num_layers, num_heads) <= 0:
            raise ValueError("All dimensions and layer counts must be positive.")
        if conditioning_type != "cat":
            raise ValueError("Only 'cat' conditioning is currently supported.")

        self.input_projection = AdaptiveInputProjection(
            input_channels,
            num_channels,
        )
        self.share_adaln = share_adaln

        if share_adaln:
            self.shared_adaln_modulation = nn.Sequential(
                nn.SiLU(),
                nn.Linear(num_channels, num_channels * 6, bias=True),
            )

        self.blocks = nn.ModuleList(
            [
                DiffusionTransformerBlock(
                    num_channels=num_channels,
                    num_heads=num_heads,
                    mlp_factor=mlp_factor,
                    normalize_queries_and_keys=normalize_queries_and_keys,
                    attention_backend=attention_backend,
                    shared_adaln=(
                        self.shared_adaln_modulation if share_adaln else None
                    ),
                )
                for _ in range(num_layers)
            ]
        )
        self.timestep_embedding = SinusoidalTimestepEmbedding(num_channels)
        self.conditioning_type = conditioning_type
        self.condition_embedding = nn.Embedding(2, num_channels)
        self.output_projection = AdaptiveOutputProjection(
            num_channels,
            input_channels,
        )

    def forward(
        self,
        input_states: Tensor,
        times: Tensor,
        conditioning_states: Tensor | None = None,
    ) -> Tensor:
        """Predict a vector field for ``[B, L, input_channels]`` states."""
        if input_states.ndim != 3:
            raise ValueError("input_states must have shape [B, L, D].")
        if times.ndim == 0:
            times = times.expand(input_states.shape[0])
        if times.shape != (input_states.shape[0],):
            raise ValueError("times must have shape [B].")
        if conditioning_states is None:
            raise ValueError("conditioning_states must be provided.")
        if conditioning_states.shape[:2] != input_states.shape[:2]:
            raise ValueError(
                "conditioning_states must match input batch and sequence dimensions."
            )

        time_conditioning = self.timestep_embedding(times)
        hidden_states = self.input_projection(
            input_states,
            time_conditioning,
        )

        condition_shape = conditioning_states.shape[:-1]
        device = conditioning_states.device
        condition_type_ids = torch.cat(
            (
                torch.zeros(condition_shape, dtype=torch.long, device=device),
                torch.ones(condition_shape, dtype=torch.long, device=device),
            ),
            dim=-1,
        )
        hidden_states = torch.cat(
            [hidden_states, conditioning_states],
            dim=-2,
        )
        hidden_states = (
            hidden_states + self.condition_embedding(condition_type_ids)
        )

        for block in self.blocks:
            hidden_states = block(hidden_states, time_conditioning)

        sequence_length = input_states.size(1)
        hidden_states = hidden_states[:, :sequence_length, :]
        return self.output_projection(hidden_states, time_conditioning)


class DiffusionTransformerBlock(nn.Module):
    """AdaLN-modulated transformer block."""

    def __init__(
        self,
        *,
        num_channels: int,
        num_heads: int,
        mlp_factor: int,
        normalize_queries_and_keys: bool = False,
        dropout: float = 0.1,
        shared_adaln: nn.Module | None = None,
        attention_backend: str = "sdpa",
    ) -> None:
        super().__init__()
        self.attention_backend = attention_backend
        self.attention = SelfAttention(
            num_channels,
            num_heads,
            attention_backend=attention_backend,
            dropout=dropout,
            normalize_queries_and_keys=normalize_queries_and_keys,
        )
        self.feed_forward = FeedForward(
            num_channels,
            num_channels * mlp_factor,
            num_channels,
            activation=nn.GELU,
            dropout=dropout,
        )
        self.norm1 = nn.LayerNorm(num_channels, elementwise_affine=False)
        self.norm2 = nn.LayerNorm(num_channels, elementwise_affine=False)
        # Retained for checkpoint compatibility with the training architecture.
        self.norm3 = nn.LayerNorm(num_channels, elementwise_affine=False)

        if shared_adaln is not None:
            self.adaptive_norm_modulation = shared_adaln
        else:
            self.adaptive_norm_modulation = nn.Sequential(
                nn.SiLU(),
                nn.Linear(num_channels, num_channels * 6, bias=True),
            )

    def _get_attention_options(self) -> dict[str, Any]:
        if self.attention_backend == "sdpa":
            return {"attn_mask": None}
        if self.attention_backend == "flex":
            return {"block_mask": None, "score_mod": None}
        raise RuntimeError(
            f"Unsupported attention backend: {self.attention_backend}"
        )

    def forward(
        self,
        hidden_states: Tensor,
        time_conditioning: Tensor,
    ) -> Tensor:
        """Transform hidden states conditioned on diffusion time."""
        adaptive_norm_parameters = self.adaptive_norm_modulation(
            time_conditioning
        )
        (
            attention_shift,
            attention_scale,
            attention_gate,
            feed_forward_shift,
            feed_forward_scale,
            feed_forward_gate,
        ) = adaptive_norm_parameters.chunk(6, dim=-1)

        hidden_states = hidden_states + attention_gate.unsqueeze(
            1
        ) * self.attention(
            apply_adaptive_modulation(
                self.norm1(hidden_states),
                attention_shift,
                attention_scale,
            ),
            **self._get_attention_options(),
        )
        hidden_states = hidden_states + feed_forward_gate.unsqueeze(
            1
        ) * self.feed_forward(
            apply_adaptive_modulation(
                self.norm2(hidden_states),
                feed_forward_shift,
                feed_forward_scale,
            )
        )
        return hidden_states


class AdaptiveInputProjection(nn.Module):
    """Project inputs and modulate them with a conditioning vector."""

    def __init__(self, input_channels: int, output_channels: int) -> None:
        super().__init__()
        self.projection = nn.Linear(
            input_channels,
            output_channels,
            bias=True,
        )
        self.norm = nn.LayerNorm(
            output_channels,
            elementwise_affine=False,
            eps=1e-6,
        )
        self.adaptive_norm_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(output_channels, 2 * output_channels, bias=True),
        )

    def forward(self, inputs: Tensor, conditioning: Tensor) -> Tensor:
        """Project ``inputs`` and apply conditioning-derived shift and scale."""
        shift, scale = self.adaptive_norm_modulation(conditioning).chunk(
            2,
            dim=-1,
        )
        outputs = self.projection(inputs)
        return apply_adaptive_modulation(self.norm(outputs), shift, scale)


class AdaptiveOutputProjection(nn.Module):
    """Final adaptive projection adopted from DiT."""

    def __init__(self, model_channels: int, output_channels: int) -> None:
        super().__init__()
        self.norm = nn.LayerNorm(
            model_channels,
            elementwise_affine=False,
            eps=1e-6,
        )
        self.projection = nn.Linear(
            model_channels,
            output_channels,
            bias=True,
        )
        self.adaptive_norm_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(model_channels, 2 * model_channels, bias=True),
        )

    def forward(self, inputs: Tensor, conditioning: Tensor) -> Tensor:
        """Modulate and project hidden states to the output dimension."""
        shift, scale = self.adaptive_norm_modulation(conditioning).chunk(
            2,
            dim=-1,
        )
        outputs = apply_adaptive_modulation(
            self.norm(inputs),
            shift,
            scale,
        )
        return self.projection(outputs)


class SinusoidalTimestepEmbedding(nn.Module):
    """Embed scalar timesteps into vector representations."""

    def __init__(
        self,
        hidden_size: int,
        frequency_embedding_size: int = 256,
    ) -> None:
        super().__init__()
        if hidden_size <= 0 or frequency_embedding_size <= 1:
            raise ValueError("Embedding dimensions must be positive.")
        self.projection = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size, bias=True),
        )
        self.frequency_embedding_size = frequency_embedding_size

    @staticmethod
    def create_sinusoidal_embedding(
        times: Tensor,
        embedding_dimension: int,
        max_period: int = 10_000,
    ) -> Tensor:
        """Create sinusoidal embeddings for one-dimensional timesteps."""
        if times.ndim != 1:
            raise ValueError("times must be one-dimensional.")
        half_dimension = embedding_dimension // 2
        frequencies = torch.exp(
            -math.log(max_period)
            * torch.arange(
                half_dimension,
                dtype=torch.float32,
                device=times.device,
            )
            / half_dimension
        )
        phase = times[:, None].float() * frequencies[None]
        embedding = torch.cat(
            [torch.cos(phase), torch.sin(phase)],
            dim=-1,
        )
        if embedding_dimension % 2:
            embedding = torch.cat(
                [embedding, torch.zeros_like(embedding[:, :1])],
                dim=-1,
            )
        return embedding

    def forward(self, times: Tensor) -> Tensor:
        """Embed a ``[B]`` tensor of timesteps."""
        frequency_embedding = self.create_sinusoidal_embedding(
            times,
            self.frequency_embedding_size,
        )
        return self.projection(frequency_embedding)