Modules/diffusion/sampler.py

from math import atan, cos, pi, sin, sqrt, log
from typing import Any, Callable, List, Optional, Tuple, Type, Literal
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, reduce as einops_reduce
from torch import Tensor
import math

# -----------------------------------------------------------------------------
# Helpers
# -----------------------------------------------------------------------------

def exists(x):
    return x is not None

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

def pad_dims(x: Tensor, ndim: int) -> Tensor:
    # Pads additional ndims to the right of the tensor
    return x.view(*x.shape, *((1,) * ndim))

def clip(x: Tensor, dynamic_threshold: float = 0.0):
    if dynamic_threshold == 0.0:
        return x.clamp(-1.0, 1.0)
    else:
        # Dynamic thresholding
        x_flat = rearrange(x, "b ... -> b (...)")
        scale = torch.quantile(x_flat.abs(), dynamic_threshold, dim=-1)
        scale.clamp_(min=1.0)
        scale = pad_dims(scale, ndim=x.ndim - scale.ndim)
        x = x.clamp(-scale, scale) / scale
        return x

def _to_batch_sigma(
    batch_size: int,
    device: torch.device,
    x: Optional[float] = None,
    xs: Optional[Tensor] = None,
) -> Tensor:
    if (x is None) and (xs is None):
        raise ValueError("Either x (scalar) or xs (tensor) must be provided")
    if (x is not None) and (xs is not None):
        raise ValueError("Only one of x or xs may be provided")

    if xs is None:
        return torch.full((batch_size,), float(x), device=device)

    if not isinstance(xs, torch.Tensor):
        xs = torch.tensor(xs, device=device, dtype=torch.get_default_dtype())
    else:
        xs = xs.to(device)

    xs = xs.reshape(-1)
    if xs.numel() == 1:
        xs = xs.expand(batch_size).contiguous()
    elif xs.numel() != batch_size:
        raise ValueError(f"xs has {xs.numel()} values but batch_size is {batch_size}")
    return xs

# -----------------------------------------------------------------------------
# Distributions
# -----------------------------------------------------------------------------

class Distribution:
    def __call__(self, num_samples: int, device: torch.device):
        raise NotImplementedError()

class LogNormalDistribution(Distribution):
    def __init__(self, mean: float, std: float):
        self.mean = mean
        self.std = std

    def __call__(
        self, num_samples: int, device: torch.device = torch.device("cpu")
    ) -> Tensor:
        normal = self.mean + self.std * torch.randn((num_samples,), device=device)
        return normal.exp()

# -----------------------------------------------------------------------------
# Base Diffusion
# -----------------------------------------------------------------------------

class Diffusion(nn.Module):
    alias: str = ""
    def denoise_fn(self, x_noisy: Tensor, sigmas: Optional[Tensor] = None, sigma: Optional[float] = None, **kwargs) -> Tensor:
        raise NotImplementedError("Diffusion class missing denoise_fn")
    def forward(self, x: Tensor, noise: Tensor = None, **kwargs) -> Tensor:
        raise NotImplementedError("Diffusion class missing forward function")

# -----------------------------------------------------------------------------
# VKDiffusion (Corrected)
# -----------------------------------------------------------------------------
class UniformDistribution(Distribution):
    def __call__(self, num_samples: int, device: torch.device = torch.device("cpu")):
        return torch.rand(num_samples, device=device)
class VKDiffusion(Diffusion):
    """
    Velocity-prediction diffusion (v-prediction) specifically adapted for 
    Karras/Elucidated framework.
    
    Target: v = alpha * eps - beta * x0
    Reconstruction: x0 = alpha * x_noised - beta * v_pred
    """

    alias = "v"

    def __init__(
        self,
        net: nn.Module,
        *,
        sigma_distribution: Distribution,
        sigma_data: float = 1.0,
        dynamic_threshold: float = 0.0,
        loss_type: str = "mse",
        huber_delta: float = 0.1,
    ):
        super().__init__()
        self.net = net
        self.register_buffer("sigma_data", torch.tensor(float(sigma_data)))
        self.sigma_distribution = sigma_distribution
        self.dynamic_threshold = float(dynamic_threshold)
        self.loss_type = loss_type
        self.huber_delta = float(huber_delta)

    def get_scalings(self, sigmas: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
        """
        Returns scaling factors for Karras inputs and V-target calculation.
        """
        sigma_data = self.sigma_data
        
        # 1. Network Input Scalings (Karras et al.)
        # c_in: scales input to unit variance
        c_in = (sigmas ** 2 + sigma_data ** 2).rsqrt() 
        # c_noise: conditioning on noise level
        c_noise = sigmas.log() * 0.25

        # 2. Geometric factors for V-prediction
        # These represent the 'angle' of the noise vs data
        # z = sqrt(sigma^2 + sigma_data^2)
        # alpha = sigma_data / z
        # beta = sigma / z
        
        # We compute them via c_in to save ops: c_in = 1/z
        alpha = c_in * sigma_data
        beta = c_in * sigmas

        return c_in, c_noise, alpha, beta

    def denoise_fn(
        self,
        x_noisy: Tensor,
        sigmas: Optional[Tensor] = None,
        sigma: Optional[float] = None,
        **kwargs,
    ) -> Tensor:
        """
        Inference time: Reconstruct x0 from x_noisy and predicted v.
        """
        batch_size, device = x_noisy.shape[0], x_noisy.device
        sigmas = _to_batch_sigma(batch_size, device, sigma, sigmas)
        
        # Pad sigmas for broadcasting [B, 1, 1]
        sigmas_view = pad_dims(sigmas, x_noisy.ndim - 1)

        c_in, c_noise, alpha, beta = self.get_scalings(sigmas_view)

        # 1. Run Network
        # IMPORTANT: kwargs must contain 'embedding' and optionally 'mask' for the transformer
        v_pred = self.net(x_noisy * c_in, c_noise.flatten(), **kwargs)

        # 2. Reconstruct x0
        # x0 = alpha * x_noisy - beta * v_pred (scaled by sigma_data implicitly in definition)
        x_denoised = (alpha * x_noisy) - (beta * v_pred) * self.sigma_data

        if self.dynamic_threshold > 0.0:
            x_denoised = clip(x_denoised, dynamic_threshold=self.dynamic_threshold)

        return x_denoised

    def forward(
        self,
        x: Tensor,
        noise: Optional[Tensor] = None,
        mask: Optional[Tensor] = None,
        **kwargs,
    ) -> Tensor:
        """
        Training forward pass.
        x: [B, 2, T] (Pitch, Energy)
        mask: [B, T] (Valid frames) - Passed to both loss and attention
        """
        batch_size, device = x.shape[0], x.device
        
        # 1. Sample Sigmas
        sigmas = self.sigma_distribution(num_samples=batch_size, device=device)
        sigmas_view = pad_dims(sigmas, x.ndim - 1) # [B, 1, 1]

        # 2. Perturb Data
        noise = default(noise, lambda: torch.randn_like(x))
        x_noisy = x + sigmas_view * noise

        # 3. Get Scalings
        c_in, c_noise, alpha, beta = self.get_scalings(sigmas_view)

        # 4. Calculate Target V
        # v_target = alpha * noise - beta * (x / sigma_data)
        # Note: Inputs are normalized by dataloader, so x is effectively x/1.0
        v_target = (alpha * noise) - (beta * x)

        # 5. Network Prediction
        # Flatten c_noise for the linear layer in model [B]
        # CRITICAL: We pass 'mask' to the net here for ATTENTION MASKING
        v_pred = self.net(x_noisy * c_in, c_noise.flatten(), mask=mask, **kwargs)

        # 6. Loss Calculation
        if self.loss_type == "mse":
            loss = F.mse_loss(v_pred, v_target, reduction="none")
        elif self.loss_type == "huber":
            loss = F.smooth_l1_loss(v_pred, v_target, reduction="none", beta=self.huber_delta)

        # 7. Masking for Loss (CRITICAL for Seq2Seq)
        if mask is not None:
            # mask comes in as [B, T], we need [B, 1, T] for broadcasting to [B, 2, T]
            if mask.ndim == 2:
                mask_bc = mask.unsqueeze(1) 
            else: 
                mask_bc = mask
            
            loss = loss * mask_bc
            # Sum over time and channels, divide by sum of mask
            # Add epsilon to avoid divide by zero
            return loss.sum() / (mask_bc.sum() * x.shape[1] + 1e-8)
        
        return loss.mean()

# -----------------------------------------------------------------------------
# Samplers / Schedules
# -----------------------------------------------------------------------------

class Schedule(nn.Module):
    def forward(self, num_steps: int, device: torch.device) -> Tensor:
        raise NotImplementedError()

class KarrasSchedule(Schedule):
    def __init__(self, sigma_min: float, sigma_max: float, rho: float = 7.0):
        super().__init__()
        self.sigma_min = sigma_min
        self.sigma_max = sigma_max
        self.rho = rho

    def forward(self, num_steps: int, device: Any) -> Tensor:
        rho_inv = 1.0 / self.rho
        steps = torch.arange(num_steps, device=device, dtype=torch.float32)
        sigmas = (
            self.sigma_max ** rho_inv
            + (steps / (num_steps - 1))
            * (self.sigma_min ** rho_inv - self.sigma_max ** rho_inv)
        ) ** self.rho
        sigmas = F.pad(sigmas, pad=(0, 1), value=0.0)
        return sigmas

class Sampler(nn.Module):
    def forward(self, noise: Tensor, fn: Callable, sigmas: Tensor, num_steps: int) -> Tensor:
        raise NotImplementedError()

class DPMpp2MSampler(Sampler):
    """ DPM-Solver++(2M) """
    def __init__(self, s_churn: float = 0.0, s_tmin: float = 0.0,
                 s_tmax: float = float("inf"), s_noise: float = 1.0):
        super().__init__()
        self.s_churn = s_churn
        self.s_tmin = s_tmin
        self.s_tmax = s_tmax
        self.s_noise = s_noise

    def step(
        self,
        x: Tensor,
        fn: Callable,
        sigma: float,
        sigma_next: float,
        i: int,
        n: int,
        d_prev: Optional[Tensor],
        sigma_prev_hat: Optional[float],
    ) -> Tuple[Tensor, Tensor, float]:
        
        gamma = 0.0
        if self.s_tmin <= sigma <= self.s_tmax:
            gamma = min(self.s_churn / max(n - 1, 1), sqrt(2.0) - 1.0)
        sigma_hat = sigma * (1.0 + gamma)
        if gamma > 0.0:
            eps = torch.randn_like(x) * self.s_noise
            x = x + (sigma_hat**2 - sigma**2).sqrt() * eps

        denoised = fn(x, sigma=sigma_hat)
        d = (x - denoised) / sigma_hat
        h = sigma_next - sigma_hat

        if d_prev is None:
            # First step -> Euler
            x_next = x + h * d
        else:
            # 2M multistep
            h_prev = sigma_hat - sigma_prev_hat
            r = h / (h_prev + 1e-12)
            x_next = x + h * ((1 + r) * d - r * d_prev)

        return x_next, d, sigma_hat

    def forward(
        self, noise: Tensor, fn: Callable, sigmas: Tensor, num_steps: int
    ) -> Tensor:
        x = sigmas[0] * noise
        d_prev: Optional[Tensor] = None
        sigma_prev_hat: Optional[float] = None

        for i in range(num_steps - 1):
            x, d_prev, sigma_prev_hat = self.step(
                x,
                fn=fn,
                sigma=sigmas[i],
                sigma_next=sigmas[i + 1],
                i=i,
                n=num_steps,
                d_prev=d_prev,
                sigma_prev_hat=sigma_prev_hat,
            )
        return x

class DiffusionSampler(nn.Module):
    def __init__(
        self,
        diffusion: Diffusion,
        *,
        sampler: Sampler,
        sigma_schedule: Schedule,
        num_steps: Optional[int] = None,
        clamp: bool = True,
    ):
        super().__init__()
        self.denoise_fn = diffusion.denoise_fn
        self.sampler = sampler
        self.sigma_schedule = sigma_schedule
        self.num_steps = num_steps
        self.clamp = clamp

    def forward(
        self, noise: Tensor, num_steps: Optional[int] = None, **kwargs
    ) -> Tensor:
        device = noise.device
        num_steps = default(num_steps, self.num_steps)
        assert exists(num_steps), "Parameter `num_steps` must be provided"
        
        sigmas = self.sigma_schedule(num_steps, device)
        
        # We wrap denoise_fn to pass **kwargs (like embedding, mask)
        fn = lambda *a, **ka: self.denoise_fn(*a, **{**ka, **kwargs})
        
        x = self.sampler(noise, fn=fn, sigmas=sigmas, num_steps=sigmas.numel())
        
        if self.clamp:
             x = x.clamp(-1.0, 1.0)
        return x