src/r4t/diffusion.py

from __future__ import annotations

import math

import torch
from torch import Tensor, nn

from .config import DiffusionConfig


def _broadcast(values: Tensor, target: Tensor) -> Tensor:
    return values.reshape(values.shape[0], *([1] * (target.ndim - 1)))


def sinusoidal_embedding(values: Tensor, dimension: int) -> Tensor:
    half = dimension // 2
    if half == 0:
        return values[:, None]
    frequencies = torch.exp(
        -math.log(10_000) * torch.arange(half, device=values.device) / max(half - 1, 1)
    )
    angles = values[:, None] * frequencies[None]
    embedding = torch.cat((angles.sin(), angles.cos()), dim=-1)
    if dimension % 2:
        embedding = torch.nn.functional.pad(embedding, (0, 1))
    return embedding


class CoherentTransformer(nn.Module):
    """Continuous-input transformer with query cross-attention conditioning."""

    def __init__(self, config: DiffusionConfig) -> None:
        super().__init__()
        if config.hidden_dim % config.heads:
            raise ValueError("hidden_dim must be divisible by heads")
        self.config = config
        self.input_projection = nn.Linear(config.embedding_dim, config.hidden_dim)
        self.output_projection = nn.Linear(config.hidden_dim, config.embedding_dim)
        self.query_projection = nn.Linear(config.embedding_dim, config.hidden_dim)
        self.position = nn.Parameter(torch.zeros(1, config.sequence_length, config.hidden_dim))
        self.time_mlp = nn.Sequential(
            nn.Linear(config.hidden_dim, config.hidden_dim),
            nn.SiLU(),
            nn.Linear(config.hidden_dim, config.hidden_dim),
        )
        layer = nn.TransformerDecoderLayer(
            d_model=config.hidden_dim,
            nhead=config.heads,
            dim_feedforward=config.mlp_dim,
            dropout=config.dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        # PyTorch exposes one dropout value for both attention and residual paths.
        # Keep the separately reported paper value in config for provenance.
        self.decoder = nn.TransformerDecoder(layer, num_layers=config.layers)
        self.final_norm = nn.LayerNorm(config.hidden_dim)
        nn.init.normal_(self.position, std=0.02)

    def forward(
        self,
        noisy_targets: Tensor,
        noise_embedding: Tensor,
        query_embedding: Tensor,
        condition_drop_mask: Tensor | None = None,
    ) -> Tensor:
        batch, length, _ = noisy_targets.shape
        if length > self.config.sequence_length:
            raise ValueError("input sequence exceeds configured sequence_length")
        if condition_drop_mask is not None:
            query_embedding = query_embedding.masked_fill(condition_drop_mask[:, None], 0.0)
        tokens = self.input_projection(noisy_targets) + self.position[:, :length]
        tokens = tokens + self.time_mlp(noise_embedding)[:, None]
        memory = self.query_projection(query_embedding)[:, None]
        decoded = self.decoder(tokens, memory)
        return self.output_projection(self.final_norm(decoded))


class EDMDenoiser(nn.Module):
    """EDM-preconditioned denoiser predicting clean set-valued targets."""

    def __init__(self, config: DiffusionConfig) -> None:
        super().__init__()
        self.config = config
        self.backbone = CoherentTransformer(config)

    def forward(
        self,
        noisy_targets: Tensor,
        sigma: Tensor,
        query_embedding: Tensor,
        condition_drop_mask: Tensor | None = None,
    ) -> Tensor:
        sigma_data = self.config.sigma_data
        sigma_sq = sigma.square()
        denominator = sigma_sq + sigma_data**2
        c_skip = sigma_data**2 / denominator
        c_out = sigma * sigma_data / denominator.sqrt()
        c_in = denominator.rsqrt()
        c_noise = sigma.clamp_min(1e-12).log() / 4.0
        network_output = self.backbone(
            _broadcast(c_in, noisy_targets) * noisy_targets,
            sinusoidal_embedding(c_noise, self.config.hidden_dim),
            query_embedding,
            condition_drop_mask=condition_drop_mask,
        )
        return (
            _broadcast(c_skip, noisy_targets) * noisy_targets
            + _broadcast(c_out, noisy_targets) * network_output
        )


def sample_tangent_sigmas(
    batch_size: int,
    sigma_min: float,
    sigma_max: float,
    device: torch.device,
    generator: torch.Generator | None = None,
) -> Tensor:
    theta_min = math.atan(sigma_min)
    theta_max = math.atan(sigma_max)
    uniform = torch.rand(batch_size, device=device, generator=generator)
    return torch.tan(theta_min + uniform * (theta_max - theta_min))


def sample_lognormal_sigmas(
    batch_size: int,
    sigma_data: float = 0.0361,
    p_mean: float | None = None,
    p_std: float = 1.2,
    device: torch.device = torch.device("cpu"),
    generator: torch.Generator | None = None,
) -> Tensor:
    """Sample noise levels sigma from log-normal distribution (Karras et al. 2022 EDM)."""
    mean = math.log(sigma_data) if p_mean is None else p_mean
    normal = torch.randn(batch_size, device=device, generator=generator)
    return (mean + normal * p_std).exp()


def diffusion_loss(
    model: EDMDenoiser,
    targets: Tensor,
    query_embeddings: Tensor,
    generator: torch.Generator | None = None,
) -> Tensor:
    config = model.config
    sigma = sample_lognormal_sigmas(
        len(targets),
        sigma_data=config.sigma_data,
        p_std=1.2,
        device=targets.device,
        generator=generator,
    )
    noise = torch.randn(targets.shape, device=targets.device, dtype=targets.dtype, generator=generator)
    noisy = targets + _broadcast(sigma, targets) * noise
    drop_mask = torch.rand(len(targets), device=targets.device, generator=generator) < config.condition_drop_probability
    prediction = model(noisy, sigma, query_embeddings, condition_drop_mask=drop_mask).float()
    targets_f = targets.float()
    sigma_f = sigma.float()
    weight = ((sigma_f.square() + config.sigma_data**2) / (
        sigma_f * config.sigma_data
    ).square().clamp_min(1e-12)).clamp_max(25_000.0)
    per_example = (prediction - targets_f).square().flatten(1).mean(dim=-1)
    return (weight * per_example).mean()


def _guided_denoise(
    model: EDMDenoiser,
    values: Tensor,
    sigma: Tensor,
    query_embeddings: Tensor,
    cfg_strength: float,
) -> Tensor:
    conditional = model(values, sigma, query_embeddings)
    if cfg_strength == 0:
        return conditional
    unconditional = model(
        values,
        sigma,
        query_embeddings,
        condition_drop_mask=torch.ones(len(values), dtype=torch.bool, device=values.device),
    )
    return conditional + cfg_strength * (conditional - unconditional)


@torch.no_grad()
def sample_edm(
    model: EDMDenoiser,
    query_embeddings: Tensor,
    sampling_steps: int | None = None,
    cfg_strength: float | None = None,
    generator: torch.Generator | None = None,
    rho: float = 7.0,
) -> Tensor:
    """Deterministic Heun solve of the EDM probability-flow ODE using Karras polynomial schedule."""

    config = model.config
    steps = sampling_steps or config.sampling_steps
    guidance = config.cfg_strength if cfg_strength is None else cfg_strength

    # Karras polynomial schedule (Karras et al. 2022)
    ramp = torch.linspace(0, 1, steps, device=query_embeddings.device)
    min_inv_rho = config.sigma_min ** (1 / rho)
    max_inv_rho = config.sigma_max ** (1 / rho)
    sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
    schedule = torch.cat([sigmas, sigmas.new_zeros(1)])
    shape = (
        len(query_embeddings),
        config.sequence_length,
        config.embedding_dim,
    )
    values = torch.randn(
        shape,
        device=query_embeddings.device,
        dtype=query_embeddings.dtype,
        generator=generator,
    ) * config.sigma_max
    for index in range(len(schedule) - 1):
        current = schedule[index]
        following = schedule[index + 1]
        sigma = current.expand(len(values))
        denoised = _guided_denoise(model, values, sigma, query_embeddings, guidance)
        derivative = (values - denoised) / current.clamp_min(1e-12)
        proposal = values + (following - current) * derivative
        if following > 0:
            next_sigma = following.expand(len(values))
            next_denoised = _guided_denoise(
                model, proposal, next_sigma, query_embeddings, guidance
            )
            next_derivative = (proposal - next_denoised) / following
            values = values + (following - current) * 0.5 * (
                derivative + next_derivative
            )
        else:
            values = proposal
    return values


class ExponentialMovingAverage:
    def __init__(self, model: nn.Module, decay: float) -> None:
        self.decay = decay
        self.shadow = {
            name: parameter.detach().clone()
            for name, parameter in model.named_parameters()
            if parameter.requires_grad
        }

    @torch.no_grad()
    def update(self, model: nn.Module) -> None:
        for name, parameter in model.named_parameters():
            if name in self.shadow:
                self.shadow[name].lerp_(parameter.detach(), 1.0 - self.decay)

    @torch.no_grad()
    def copy_to(self, model: nn.Module) -> None:
        for name, parameter in model.named_parameters():
            if name in self.shadow:
                parameter.copy_(self.shadow[name])

    def state_dict(self) -> dict[str, object]:
        return {"decay": self.decay, "shadow": self.shadow}

    def load_state_dict(self, state_dict: dict[str, object]) -> None:
        self.decay = float(state_dict.get("decay", self.decay))
        shadow = state_dict.get("shadow", {})
        self.shadow = {name: tensor.clone() for name, tensor in shadow.items()}