SANA-Sprint Runs Text-to-Image in One Step. The Reason It Works Is Not the Distillation. It Is the Training Stability Fix Nobody Talks About.

# GitHub: NVlabs/Sana
git clone https://github.com/NVlabs/Sana

# HuggingFace demo
# https://huggingface.co/spaces/Efficient-Large-Model/SanaSprint

# Self-hosted demo
# https://sana.hanlab.ai/sprint

# Install
pip install sana-sprint

# SANA-Sprint: TrigFlow transformation and continuous-time consistency distillation
# Reconstructed from arXiv:2503.09641 methodology
# The training-free reparameterization that converts flow-matching to sCM

import torch
import torch.nn as nn
import math

class TrigFlowTransform:
    """
    Training-free transformation of a flow-matching model into a sCM-compatible student.

    Flow-matching uses linear interpolation between noise and data:
      x_t = t * x_0 + (1-t) * x_noise   for t ∈ [0, 1]

    Continuous-time consistency models use cosine schedule:
      sigma(s) = tan(s)  for s ∈ [0, π/2)

    TrigFlow maps between these by:
      s = arctan(t)   ← the "TrigFlow" transformation

    ← WHY THIS IS TRAINING-FREE:
      The teacher model was trained on t ∈ [0, 1].
      After TrigFlow, the SAME teacher can be queried at corresponding s values.
      No retraining. No new weights. Just reparameterize time.
      This eliminates the most expensive part of prior sCM work:
      training a teacher from scratch in the sCM time parameterization.
    """
    @staticmethod
    def t_to_s(t: torch.Tensor) -> torch.Tensor:
        """Map flow-matching time t ∈ [0,1] to sCM time s ∈ [0, π/2)"""
        return torch.arctan(t)  # ← THIS is the trick: the TrigFlow mapping

    @staticmethod
    def s_to_t(s: torch.Tensor) -> torch.Tensor:
        """Inverse: sCM time back to flow-matching time"""
        return torch.tan(s)

    @staticmethod
    def sigma(s: torch.Tensor) -> torch.Tensor:
        """Noise schedule in sCM parameterization"""
        return torch.tan(s)   # sigma(s) = tan(s) in the TrigFlow space


def scm_consistency_loss(
    student: nn.Module,
    teacher: nn.Module,
    x_0: torch.Tensor,        # real images [B, C, H, W]
    noise: torch.Tensor,      # Gaussian noise [B, C, H, W]
    transform: TrigFlowTransform,
) -> torch.Tensor:
    """
    Continuous-time consistency distillation loss.
    At each training step, sample a random time pair (s_i, s_{i+1})
    and require student(x_{s_i}, s_i) ≈ teacher_rollout(x_{s_i}, s_i → s_{i+1}).

    This is "local" supervision: the student learns from adjacent timestep pairs.
    CTM (prior work) analyzed that this local learning introduces implicit
    extrapolation, which slows single-step quality convergence.
    LADD (below) addresses this by adding global supervision.
    """
    B = x_0.shape[0]
    device = x_0.device

    # Sample adjacent time pairs for consistency constraint
    # s_i is current, s_{i+eps} is teacher target time
    s_i = torch.rand(B, device=device) * (math.pi / 2 - 0.01)  # ← sample in [0, π/2)
    eps = 0.01  # small consistency window
    s_next = torch.clamp(s_i + eps, max=math.pi / 2 - 0.001)

    # Noisy input at time s_i
    sigma_i = transform.sigma(s_i).view(B, 1, 1, 1)
    x_si = x_0 + sigma_i * noise     # interpolation in sCM space

    # Student prediction: jump directly from x_{s_i} to denoised output
    # ← Student predicts x_0 directly from noisy x_si
    x0_student = student(x_si, s_i)   # one-step denoised prediction

    # Teacher prediction at s_next (small step from current position)
    # ← Convert s → t for teacher (teacher lives in flow-matching time)
    t_next = transform.s_to_t(s_next)
    sigma_next = transform.sigma(s_next).view(B, 1, 1, 1)
    x_snext = x_0 + sigma_next * noise  # noisy input at s_next

    with torch.no_grad():
        x0_teacher = teacher(x_snext, t_next)  # teacher's one-step prediction

    # Consistency loss: student at s_i should match teacher at s_{i+1}
    # ← Both should predict the same x_0 if model is internally consistent
    loss = nn.functional.mse_loss(x0_student, x0_teacher.detach())
    return loss


def ladd_loss(
    discriminator: nn.Module,  # discriminator in LATENT space
    vae_encoder,               # VAE encoder to get latents
    student: nn.Module,
    x_0: torch.Tensor,         # real images [B, C, H, W]
    text_embedding: torch.Tensor,
) -> torch.Tensor:
    """
    Latent Adversarial Distillation loss.
    Discriminator operates in VAE latent space, not pixel space.

    ← WHY LATENT SPACE:
      Pixel-space GAN is computationally expensive at 1024×1024.
      Latent space is 8x-16x smaller, so discriminator is fast.
      Latent-space features capture semantic quality better than pixels.

    ← WHY GAN HELPS:
      sCM provides LOCAL consistency (step-to-step alignment with teacher).
      But it doesn't directly optimize "does the final output look realistic?"
      The discriminator provides GLOBAL quality feedback:
      "Is this generated latent distinguishable from real image latents?"
      This is what prevents the step-to-step local learning from
      missing global distribution quality.
    """
    # One-step student generation from pure noise
    noise = torch.randn_like(x_0)
    s_start = torch.ones(x_0.shape[0]) * (math.pi / 2 - 0.01)  # near max noise
    x0_generated = student(noise, s_start, text_embedding)   # single-step generation

    # Encode both real and generated to latent space
    z_real = vae_encoder(x_0).detach()           # real image latents
    z_fake = vae_encoder(x0_generated)            # generated image latents

    # ← THIS is the trick: adversarial loss on LATENTS, not pixels
    # Real latents should be classified as real, fake as fake
    d_real = discriminator(z_real)
    d_fake = discriminator(z_fake)

    # Standard non-saturating GAN loss
    loss_d = -torch.log(torch.sigmoid(d_real)).mean() \
             -torch.log(1 - torch.sigmoid(d_fake)).mean()
    loss_g = -torch.log(torch.sigmoid(d_fake)).mean()   # generator wants D(fake) = 1

    return loss_g   # return generator loss for student training

# SANA-Sprint: Stabilization techniques and step-adaptive inference
# arXiv:2503.09641, stabilization analysis Section 3.2
# The fixes that make continuous-time distillation at scale work

import torch
import torch.nn as nn
import math

class StabilizedAttention(nn.Module):
    """
    Attention with QK-Normalization: the fix for gradient explosion at scale.

    Standard attention: softmax(QK^T / sqrt(d)) · V
    When model is large (1.6B) or resolution is high (1024×1024):
    - Q and K values can grow large during training
    - QK^T entries explode → softmax saturates → gradients vanish or explode
    - Model collapse: training diverges before convergence

    QK-Normalization: normalize Q and K before dot product
    - ||Q||₂ = ||K||₂ = 1 after normalization
    - Attention logits bounded by geometry: max(QK^T) ≤ 1
    - Gradient norm stays controlled throughout training

    ← Applied in BOTH self-attention (across image tokens)
       AND cross-attention (between image tokens and text embeddings)
       Both types fail independently at 1.6B without normalization.
    """
    def __init__(self, dim: int, num_heads: int):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.out = nn.Linear(dim, dim, bias=False)
        self.q_norm = nn.LayerNorm(self.head_dim)   # ← QK-Norm
        self.k_norm = nn.LayerNorm(self.head_dim)   # ← QK-Norm

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, N, C = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)

        # Reshape to multi-head format
        q = q.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, N, self.num_heads, self.head_dim).transpose(1, 2)

        # ← THIS is the trick: normalize Q and K before dot product
        # Without this: at 1.6B parameters, gradient norms explode
        # With this: training remains stable through the full distillation
        q = self.q_norm(q)   # normalize per head dimension
        k = self.k_norm(k)   # normalize per head dimension

        attn = torch.softmax(q @ k.transpose(-2, -1) / math.sqrt(self.head_dim), dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(B, N, C)
        return self.out(out)


class DenseTimeEmbedding(nn.Module):
    """
    Dense time-embedding: inject time at EVERY transformer block.

    Standard time conditioning: embed time t once, add to layer 0
    Problem: by layer 32 in a deep transformer, the time signal
    has been diluted through residual additions and normalizations.
    The model "forgets" which timestep it is supposed to be denoising at.

    Dense injection: inject time embedding RESIDUALLY at every block.
    ← This is not a conditioning; it is a residual correction.
    ← Every layer independently knows what time it is.
    ← Critical for sCM because the model must have very precise
      time awareness to implement the consistency constraint correctly.
    """
    def __init__(self, embed_dim: int):
        super().__init__()
        self.time_proj = nn.Sequential(
            nn.Linear(256, embed_dim),
            nn.SiLU(),
            nn.Linear(embed_dim, embed_dim)
        )

    def get_timestep_embedding(self, t: torch.Tensor, dim: int = 256) -> torch.Tensor:
        """Sinusoidal position-style time embedding (standard)."""
        half = dim // 2
        freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / half)
        args = t[:, None] * freqs[None]
        return torch.cat([torch.cos(args), torch.sin(args)], dim=-1)

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        t_embed = self.get_timestep_embedding(t)
        return self.time_proj(t_embed)   # [B, embed_dim]


class SanaDiTBlock(nn.Module):
    """
    One transformer block with dense time-embedding injection.
    """
    def __init__(self, dim: int, num_heads: int):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = StabilizedAttention(dim, num_heads)
        self.norm2 = nn.LayerNorm(dim)
        self.ff = nn.Sequential(nn.Linear(dim, 4*dim), nn.GELU(), nn.Linear(4*dim, dim))
        self.time_proj = nn.Linear(dim, dim)  # per-block time adaptation

    def forward(self, x: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor:
        # ← Dense time-embedding: inject time residually at THIS block
        # Not just at the first layer, but here, for every block
        time_scale = self.time_proj(t_emb).unsqueeze(1)   # [B, 1, dim]
        x = x + self.attn(self.norm1(x)) + time_scale     # residual + time correction
        x = x + self.ff(self.norm2(x))
        return x


def step_adaptive_inference(
    model: nn.Module,
    noise: torch.Tensor,         # [B, C, H/8, W/8] latent noise
    prompt_embedding: torch.Tensor,
    num_steps: int = 1,           # 1, 2, or 4: same model, different quality/speed
    transform: TrigFlowTransform = None,
) -> torch.Tensor:
    """
    Unified step-adaptive inference with a single set of model weights.

    ← Why one model for all step counts:
      Standard multi-step distillation trains separate models for each step count.
      (1-step model, 4-step model, etc.)
      SANA-Sprint trains ONE model with sCM objective that enforces
      self-consistency at ALL timesteps simultaneously.
      At inference: choose trajectory length, same weights.

    For 1 step: single jump from t=T (pure noise) to t=0 (clean image)
    For 4 steps: 4 Euler steps with smaller intervals
    """
    if transform is None:
        transform = TrigFlowTransform()

    x = noise.clone()

    if num_steps == 1:
        # ← Single NFE: jump from max noise to clean image directly
        s = torch.ones(noise.shape[0], device=noise.device) * (math.pi / 2 - 0.01)
        x0_pred = model(x, s, prompt_embedding)   # 0.1s on H100

    else:
        # Multi-step Euler: divide the time interval evenly
        s_vals = torch.linspace(math.pi / 2 - 0.01, 0.01, num_steps + 1)

        for i in range(num_steps):
            s_t = s_vals[i].expand(noise.shape[0])
            s_next = s_vals[i + 1].expand(noise.shape[0])

            x0_pred = model(x, s_t, prompt_embedding)  # predict clean image

            # Move toward predicted x0 (Euler step in sCM space)
            sigma_t = transform.sigma(s_t).view(-1, 1, 1, 1)
            sigma_next = transform.sigma(s_next).view(-1, 1, 1, 1)
            # ← Linear interpolation between current noisy and predicted clean
            x = x0_pred + sigma_next * (x - x0_pred) / sigma_t

    return x0_pred   # [B, C, H/8, W/8] denoised latent → decode with VAE

prompt = "A majestic castle perched on a floating island above clouds, "
         "golden hour lighting, highly detailed, 8k resolution"
num_steps = 1         # single step inference
model_size = "0.6B"   # or "1.6B" for better GenEval

Inference starting time: s_start = π/2 - 0.01 ≈ 1.5608 (maximum noise level)
Corresponding sigma: σ = tan(s_start) ≈ 158 (very high noise, near-pure Gaussian)

Input to model: x ~ N(0, σ²·I) in latent space  [B=1, C=4, H=128, W=128]
                                                    (1024/8 = 128 for VAE with 8x compression)

Prompt → T5 encoder → text_embedding [B=1, L=128, D=768]
Time embedding: dense_time_emb = time_proj(sinusoidal(s_start))   [B=1, D=hidden]
Applied at EVERY DiT block (32 blocks for 0.6B model)

Input latent: x ∈ ℝ^{1×4×128×128} (4-channel VAE latent, 128×128 spatial)
After patch embedding: [1, 1024, hidden_dim]  (16×16 patches of 8×8 spatial)

32 transformer blocks, each:
  - QK-normalized self-attention (spatial, image tokens × image tokens)
  - QK-normalized cross-attention (image tokens × text embedding tokens)
  - Dense time-embedding residual injection
  - Feed-forward

Output: x0_pred ∈ ℝ^{1×4×128×128}   (predicted clean latent, one forward pass)

Timing:
  Attention per block: ~0.8ms
  32 blocks total: ~26ms
  VAE decode: ~35ms
  Overhead (tokenize, encode, scaffold): ~39ms
  Total: ~0.1s on H100 (matches paper)

Single NFE (Number of Function Evaluations): 1
(vs teacher SANA: 20 NFE at 3s total)

x0_pred [1, 4, 128, 128] → VAE decoder → image [1, 3, 1024, 1024]
Quality at 1-step:
  FID: 7.04  (vs FLUX-schnell: 7.94; lower is better)
  GenEval: 0.72  (vs FLUX-schnell: 0.71; higher is better)

If using ControlNet (depth map, edge map, sketch):
  ControlNet encoder processes condition → injects at each DiT block
  Additional compute: ~35ms → 0.25s total on H100
  Enables real-time interactive generation: user changes condition, image updates immediately

SANA-Sprint 0.6B, 1-step:  0.1s  H100  / 0.31s  RTX 4090  → 7.22 samples/s / 3.2 samples/s
SANA-Sprint 0.6B, 4-step:  0.32s H100  → 5.34 samples/s
SANA-Sprint 1.6B, 4-step:  ~0.5s H100  → 5.20 samples/s, GenEval 0.77
FLUX-schnell (12B), 4-step: 2.10s H100  → 0.5 samples/s (10x slower than SANA-Sprint 0.6B)
Teacher SANA (20-step):     ~3s   H100  → reference quality