diff --git a/examples/inference/basic/basic_turbodiffusion.py b/examples/inference/basic/basic_turbodiffusion.py new file mode 100644 index 000000000..9b4a2a74f --- /dev/null +++ b/examples/inference/basic/basic_turbodiffusion.py @@ -0,0 +1,57 @@ +import os + +# Set SLA attention backend BEFORE fastvideo imports +os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLA_ATTN" + +from fastvideo import VideoGenerator + +OUTPUT_PATH = "video_samples_turbodiffusion" + + +def main() -> None: + # TurboDiffusion: 1-4 step video generation using RCM scheduler + SLA attention + # FastVideo will automatically use TurboDiffusionPipeline when specified + generator = VideoGenerator.from_pretrained( + "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers", + # FastVideo will automatically handle distributed setup + num_gpus=1, + # TurboDiffusion uses a custom pipeline with RCM scheduler + override_pipeline_cls_name="TurboDiffusionPipeline", + ) + + # Generate videos with the same simple API, regardless of GPU count + # TurboDiffusion uses guidance_scale=1.0 (no CFG) and only 4 steps + prompt = ( + "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones." + ) + video = generator.generate_video( + prompt, + output_path=OUTPUT_PATH, + save_video=True, + num_inference_steps=4, + seed=42, + guidance_scale=1.0, + ) + + # Generate another video with a different prompt, without reloading the model! + prompt2 = ( + "A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic." + ) + video2 = generator.generate_video( + prompt2, + output_path=OUTPUT_PATH, + save_video=True, + num_inference_steps=4, + seed=42, + guidance_scale=1.0, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/basic_turbodiffusion_14b.py b/examples/inference/basic/basic_turbodiffusion_14b.py new file mode 100644 index 000000000..b22bbbb5e --- /dev/null +++ b/examples/inference/basic/basic_turbodiffusion_14b.py @@ -0,0 +1,55 @@ +import os + +# Set SLA attention backend BEFORE fastvideo imports +os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "SLA_ATTN" + +from fastvideo import VideoGenerator + +OUTPUT_PATH = "video_samples_turbodiffusion_14B" + + +def main() -> None: + # TurboDiffusion 14B: 1-4 step video generation using RCM scheduler + SLA attention + # FastVideo will automatically use TurboDiffusionPipeline when specified + generator = VideoGenerator.from_pretrained( + "loayrashid/TurboWan2.1-T2V-14B-Diffusers", + # 14B model needs more GPUs + num_gpus=2, + # TurboDiffusion uses a custom pipeline with RCM scheduler + override_pipeline_cls_name="TurboDiffusionPipeline", + ) + + prompt = ( + "A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes " + "wide with interest. The playful yet serene atmosphere is complemented by soft " + "natural light filtering through the petals. Mid-shot, warm and cheerful tones." + ) + video = generator.generate_video( + prompt, + output_path=OUTPUT_PATH, + save_video=True, + num_inference_steps=4, + seed=42, + guidance_scale=1.0, + ) + + # Generate another video with a different prompt, without reloading the model! + prompt2 = ( + "A majestic lion strides across the golden savanna, its powerful frame " + "glistening under the warm afternoon sun. The tall grass ripples gently in " + "the breeze, enhancing the lion's commanding presence. The tone is vibrant, " + "embodying the raw energy of the wild. Low angle, steady tracking shot, " + "cinematic." + ) + video2 = generator.generate_video( + prompt2, + output_path=OUTPUT_PATH, + save_video=True, + num_inference_steps=4, + seed=42, + guidance_scale=1.0, + ) + + +if __name__ == "__main__": + main() diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/sla_triton.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/sla_triton.py new file mode 100644 index 000000000..64f2f6dd5 --- /dev/null +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/sla_triton.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from TurboDiffusion SLA implementation +# Copyright (c) 2025 by SLA team. +# +# Citation: +# @article{zhang2025sla, +# title={SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse-Linear Attention}, +# author={Jintao Zhang and Haoxu Wang and Kai Jiang and Shuo Yang and Kaiwen Zheng and +# Haocheng Xi and Ziteng Wang and Hongzhou Zhu and Min Zhao and Ion Stoica and +# Joseph E. Gonzalez and Jun Zhu and Jianfei Chen}, +# journal={arXiv preprint arXiv:2509.24006}, +# year={2025} +# } + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _attn_fwd( + Q, K, V, + qk_scale: tl.constexpr, + topk: tl.constexpr, + LUT, LSE, OS, + L: tl.constexpr, + M_BLOCKS: tl.constexpr, + D: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + idx_m = tl.program_id(0).to(tl.int64) + idx_bh = tl.program_id(1).to(tl.int64) + + qkv_offset = idx_bh * L * D + lut_offset = (idx_bh * M_BLOCKS + idx_m) * topk + lse_offset = idx_bh * L + offs_m = idx_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, D) + + Q_ptrs = Q + qkv_offset + offs_m[:, None] * D + offs_d[None, :] + K_ptrs = K + qkv_offset + offs_n[None, :] * D + offs_d[:, None] + V_ptrs = V + qkv_offset + offs_n[:, None] * D + offs_d[None, :] + OS_ptrs = OS + qkv_offset + offs_m[:, None] * D + offs_d[None, :] + LUT_ptr = LUT + lut_offset + LSE_ptrs = LSE + lse_offset + offs_m + + m_i = tl.full([BLOCK_M], -float('inf'), dtype=tl.float32) + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + o_s = tl.zeros([BLOCK_M, D], dtype=tl.float32) + + q = tl.load(Q_ptrs, mask=offs_m[:, None] < L) + for block_idx in tl.range(topk): + idx_n = tl.load(LUT_ptr + block_idx) + n_mask = offs_n < L - idx_n * BLOCK_N + + k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[None, :]) + qk = tl.dot(q, k) * (qk_scale * 1.4426950408889634) # = 1 / ln(2) + if L - idx_n * BLOCK_N < BLOCK_N: + qk = tl.where(n_mask[None, :], qk, float("-inf")) + + v = tl.load(V_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None]) + local_m = tl.max(qk, 1) + new_m = tl.maximum(m_i, local_m) + qk = qk - new_m[:, None] + + p = tl.math.exp2(qk) + l_ij = tl.sum(p, 1) + alpha = tl.math.exp2(m_i - new_m) + o_s = o_s * alpha[:, None] + o_s += tl.dot(p.to(v.dtype), v) + + l_i = l_i * alpha + l_ij + m_i = new_m + + o_s = o_s / l_i[:, None] + tl.store(OS_ptrs, o_s.to(OS.type.element_ty), mask=offs_m[:, None] < L) + + m_i += tl.math.log2(l_i) + tl.store(LSE_ptrs, m_i, mask=offs_m < L) + + +@triton.jit +def _attn_bwd_preprocess( + OS, DOS, DELTAS, + L, + D: tl.constexpr, + BLOCK_M: tl.constexpr, +): + idx_m = tl.program_id(0).to(tl.int64) + idx_bh = tl.program_id(1).to(tl.int64) + + OS += idx_bh * L * D + DOS += idx_bh * L * D + DELTAS += idx_bh * L + + offs_m = idx_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, D) + + o_s = tl.load(OS + offs_m[:, None] * D + offs_d[None, :], mask=offs_m[:, None] < L) + do_s = tl.load(DOS + offs_m[:, None] * D + offs_d[None, :], mask=offs_m[:, None] < L) + + delta_s = tl.sum(o_s * do_s, axis=1).to(DELTAS.type.element_ty) + tl.store(DELTAS + offs_m, delta_s, mask=offs_m < L) + + +@triton.jit +def _attn_bwd_dq( + Q, K, V, LSE, DELTAS, + DOS, DQ, LUT, + qk_scale: tl.constexpr, + topk: tl.constexpr, + L: tl.constexpr, + M_BLOCKS: tl.constexpr, + D: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + idx_m = tl.program_id(0).to(tl.int64) + idx_bh = tl.program_id(1).to(tl.int64) + + offs_m = idx_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, D) + + qkv_offset = idx_bh * L * D + lse_offset = idx_bh * L + lut_offset = (idx_bh * M_BLOCKS + idx_m) * topk + + Q_ptrs = Q + qkv_offset + offs_m[:, None] * D + offs_d[None, :] + K_ptrs = K + qkv_offset + offs_n[:, None] * D + offs_d[None, :] + V_ptrs = V + qkv_offset + offs_n[:, None] * D + offs_d[None, :] + DQ_ptrs = DQ + qkv_offset + offs_m[:, None] * D + offs_d[None, :] + DOS_ptrs = DOS + qkv_offset + offs_m[:, None] * D + offs_d[None, :] + LSE_ptrs = LSE + lse_offset + offs_m + DELTAS_ptrs = DELTAS + lse_offset + offs_m + LUT_ptr = LUT + lut_offset + + q = tl.load(Q_ptrs, mask=offs_m[:, None] < L) + do_s = tl.load(DOS_ptrs, mask=offs_m[:, None] < L) + delta_s = tl.load(DELTAS_ptrs, mask=offs_m < L) + lse = tl.load(LSE_ptrs, mask=offs_m < L, other=float("inf")) + + dq = tl.zeros([BLOCK_M, D], dtype=tl.float32) + for block_idx in tl.range(topk, num_stages=2): + idx_n = tl.load(LUT_ptr + block_idx) + n_mask = offs_n < L - idx_n * BLOCK_N + + k = tl.load(K_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None]) + v = tl.load(V_ptrs + idx_n * BLOCK_N * D, mask=n_mask[:, None]) + qk = tl.dot(q, k.T) * (qk_scale * 1.4426950408889634) + p = tl.math.exp2(qk - lse[:, None]) + p = tl.where(n_mask[None, :], p, 0.0) + + dp = tl.dot(do_s, v.T).to(tl.float32) + ds = p * (dp - delta_s[:, None]) + dq += tl.dot(ds.to(k.dtype), k) + tl.store(DQ_ptrs, dq * qk_scale, mask=offs_m[:, None] < L) + + +@triton.jit +def _attn_bwd_dkdv( + Q, K, V, DOS, DK, DV, + qk_scale, KBID, LSE, DELTAS, + L: tl.constexpr, + M_BLOCKS: tl.constexpr, + N_BLOCKS: tl.constexpr, + D: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_SLICE_FACTOR: tl.constexpr, +): + BLOCK_M2: tl.constexpr = BLOCK_M // BLOCK_SLICE_FACTOR + + idx_n = tl.program_id(0).to(tl.int64) + idx_bh = tl.program_id(1).to(tl.int64) + + offs_n = idx_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_m = tl.arange(0, BLOCK_M2) + offs_d = tl.arange(0, D) + + qkv_offset = idx_bh * L * D + kbid_offset = idx_bh * M_BLOCKS * N_BLOCKS + lse_offset = idx_bh * L + + Q_ptrs = Q + qkv_offset + offs_m[:, None] * D + offs_d[None, :] + K_ptrs = K + qkv_offset + offs_n[:, None] * D + offs_d[None, :] + V_ptrs = V + qkv_offset + offs_n[:, None] * D + offs_d[None, :] + DOS_ptrs = DOS + qkv_offset + offs_m[:, None] * D + offs_d[None, :] + DK_ptrs = DK + qkv_offset + offs_n[:, None] * D + offs_d[None, :] + DV_ptrs = DV + qkv_offset + offs_n[:, None] * D + offs_d[None, :] + LSE_ptrs = LSE + lse_offset + offs_m + DELTAS_ptrs = DELTAS + lse_offset + offs_m + KBID_ptr = KBID + kbid_offset + idx_n + + k = tl.load(K_ptrs, mask=offs_n[:, None] < L) + v = tl.load(V_ptrs, mask=offs_n[:, None] < L) + + dk = tl.zeros([BLOCK_N, D], dtype=tl.float32) + dv = tl.zeros([BLOCK_N, D], dtype=tl.float32) + for idx_m in tl.range(0, L, BLOCK_M2): + kbid = tl.load(KBID_ptr) + if kbid == 1: + m_mask = offs_m < L - idx_m + q = tl.load(Q_ptrs, mask=m_mask[:, None]) + lse = tl.load(LSE_ptrs, mask=m_mask, other=float("inf")) + qkT = tl.dot(k, q.T) * (qk_scale * 1.4426950408889634) + pT = tl.math.exp2(qkT - lse[None, :]) + pT = tl.where(offs_n[:, None] < L, pT, 0.0) + + do = tl.load(DOS_ptrs, mask=m_mask[:, None]) + dv += tl.dot(pT.to(do.dtype), do) + delta = tl.load(DELTAS_ptrs, mask=m_mask) + dpT = tl.dot(v, tl.trans(do)) + dsT = pT * (dpT - delta[None, :]) + dk += tl.dot(dsT.to(q.dtype), q) + + Q_ptrs += BLOCK_M2 * D + DOS_ptrs += BLOCK_M2 * D + LSE_ptrs += BLOCK_M2 + DELTAS_ptrs += BLOCK_M2 + if (idx_m + BLOCK_M2) % BLOCK_M == 0: + KBID_ptr += N_BLOCKS + + tl.store(DK_ptrs, dk * qk_scale, mask=offs_n[:, None] < L) + tl.store(DV_ptrs, dv, mask=offs_n[:, None] < L) + + +class _attention(torch.autograd.Function): + """Sparse attention forward/backward with autograd support.""" + + @staticmethod + def forward(ctx, q, k, v, k_block_id, lut, topk, BLOCK_M, BLOCK_N, qk_scale=None): + assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous() + assert k_block_id.is_contiguous() and lut.is_contiguous() + + assert BLOCK_M == 64 or BLOCK_M == 128 + assert BLOCK_N == 64 + + B, H, L, D = q.shape + if qk_scale is None: + qk_scale = D**-0.5 + + M_BLOCKS = triton.cdiv(L, BLOCK_M) + + o_s = torch.empty_like(v) + lse = torch.empty(q.shape[:-1], device=q.device, dtype=torch.float32) + + grid = (M_BLOCKS, B * H) + _attn_fwd[grid]( + q, k, v, qk_scale, topk, + lut, lse, o_s, + L, M_BLOCKS, + D, BLOCK_M, BLOCK_N, + num_warps=4 if q.shape[-1] == 64 else 8, + num_stages=3 + ) + + ctx.save_for_backward(q, k, v, k_block_id, lut, lse, o_s) + ctx.qk_scale = qk_scale + ctx.topk = topk + ctx.BLOCK_M = BLOCK_M + ctx.BLOCK_N = BLOCK_N + return o_s + + @staticmethod + def backward(ctx, do_s): + q, k, v, k_block_id, lut, lse, o_s = ctx.saved_tensors + do_s = do_s.contiguous() + + BLOCK_M, BLOCK_N = ctx.BLOCK_M, ctx.BLOCK_N + B, H, L, D = q.shape + + M_BLOCKS = triton.cdiv(L, BLOCK_M) + N_BLOCKS = triton.cdiv(L, BLOCK_N) + + dq = torch.empty_like(q) + dk = torch.empty_like(k) + dv = torch.empty_like(v) + delta_s = torch.empty_like(lse) + + grid = (M_BLOCKS, B * H) + _attn_bwd_preprocess[grid]( + o_s, do_s, delta_s, + L, D, BLOCK_M, + ) + + grid = (M_BLOCKS, B * H) + _attn_bwd_dq[grid]( + q, k, v, lse, delta_s, + do_s, dq, lut, + ctx.qk_scale, ctx.topk, + L, M_BLOCKS, + D, BLOCK_M, BLOCK_N, + num_warps=4 if q.shape[-1] == 64 else 8, + num_stages=4 if q.shape[-1] == 64 else 5 + ) + + grid = (N_BLOCKS, B * H) + _attn_bwd_dkdv[grid]( + q, k, v, do_s, dk, dv, + ctx.qk_scale, k_block_id, lse, delta_s, + L, M_BLOCKS, N_BLOCKS, + D, BLOCK_M, BLOCK_N, + BLOCK_SLICE_FACTOR=BLOCK_M // 64, + num_warps=4 if q.shape[-1] == 64 else 8, + num_stages=4 if q.shape[-1] == 64 else 5 + ) + + return dq, dk, dv, None, None, None, None, None, None diff --git a/fastvideo/attention/backends/sla.py b/fastvideo/attention/backends/sla.py new file mode 100644 index 000000000..726189536 --- /dev/null +++ b/fastvideo/attention/backends/sla.py @@ -0,0 +1,587 @@ +# SPDX-License-Identifier: Apache-2.0 +# SLA (Sparse-Linear Attention) backend for FastVideo +# Adapted from TurboDiffusion SLA implementation +# +# Copyright (c) 2025 by SLA team. +# Citation: +# @article{zhang2025sla, +# title={SLA: Beyond Sparsity in Diffusion Transformers via Fine-Tunable Sparse-Linear Attention}, +# author={Jintao Zhang and Haoxu Wang and Kai Jiang and Shuo Yang and Kaiwen Zheng and +# Haocheng Xi and Ziteng Wang and Hongzhou Zhu and Min Zhao and Ion Stoica and +# Joseph E. Gonzalez and Jun Zhu and Jianfei Chen}, +# journal={arXiv preprint arXiv:2509.24006}, +# year={2025} +# } + +from dataclasses import dataclass +from typing import Any +from collections.abc import Callable + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl + +from fastvideo.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, + AttentionMetadataBuilder, +) +from fastvideo_kernel.triton_kernels.sla_triton import _attention +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +# ============================================================================ +# SLA Utility functions (moved from sla_kernels/utils.py) +# ============================================================================ + + +@triton.jit +def compress_kernel( + X, + XM, + L: tl.constexpr, + D: tl.constexpr, + BLOCK_L: tl.constexpr, +): + idx_l = tl.program_id(0) + idx_bh = tl.program_id(1) + + offs_l = idx_l * BLOCK_L + tl.arange(0, BLOCK_L) + offs_d = tl.arange(0, D) + + x_offset = idx_bh * L * D + xm_offset = idx_bh * ((L + BLOCK_L - 1) // BLOCK_L) * D + x = tl.load(X + x_offset + offs_l[:, None] * D + offs_d[None, :], + mask=offs_l[:, None] < L) + + nx = min(BLOCK_L, L - idx_l * BLOCK_L) + x_mean = tl.sum(x, axis=0, dtype=tl.float32) / nx + tl.store(XM + xm_offset + idx_l * D + offs_d, + x_mean.to(XM.dtype.element_ty)) + + +def mean_pool(x: torch.Tensor, BLK: int) -> torch.Tensor: + """Mean pool tensor along sequence dimension with block size BLK.""" + assert x.is_contiguous() + + B, H, L, D = x.shape + L_BLOCKS = (L + BLK - 1) // BLK + x_mean = torch.empty((B, H, L_BLOCKS, D), device=x.device, dtype=x.dtype) + + grid = (L_BLOCKS, B * H) + compress_kernel[grid](x, x_mean, L, D, BLK) + return x_mean + + +def get_block_map( + q: torch.Tensor, + k: torch.Tensor, + topk_ratio: float, + BLKQ: int = 64, + BLKK: int = 64, +) -> tuple[torch.Tensor, torch.Tensor, int]: + """Compute sparse block map for attention based on QK similarity. + + Args: + q: Query tensor of shape (B, H, L, D) + k: Key tensor of shape (B, H, L, D) + topk_ratio: Ratio of key blocks to attend to (0-1) + BLKQ: Query block size + BLKK: Key block size + + Returns: + sparse_map: Binary mask of shape (B, H, num_q_blocks, num_k_blocks) + lut: Top-k indices of shape (B, H, num_q_blocks, topk) + topk: Number of key blocks selected + """ + arg_k = k - torch.mean( + k, dim=-2, keepdim=True) # smooth-k technique from SageAttention + pooled_qblocks = mean_pool(q, BLKQ) + pooled_kblocks = mean_pool(arg_k, BLKK) + pooled_score = pooled_qblocks @ pooled_kblocks.transpose(-1, -2) + + K = pooled_score.shape[-1] + topk = min(K, int(topk_ratio * K)) + lut = torch.topk(pooled_score, topk, dim=-1, sorted=False).indices + + sparse_map = torch.zeros_like(pooled_score, dtype=torch.int8) + sparse_map.scatter_(-1, lut, 1) + return sparse_map, lut, topk + + +# ============================================================================ +# SLA Backend classes +# ============================================================================ + + +class SLAAttentionBackend(AttentionBackend): + """Sparse-Linear Attention backend.""" + + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [64, 128] + + @staticmethod + def get_name() -> str: + return "SLA_ATTN" + + @staticmethod + def get_impl_cls() -> type["SLAAttentionImpl"]: + return SLAAttentionImpl + + @staticmethod + def get_metadata_cls() -> type["SLAAttentionMetadata"]: + return SLAAttentionMetadata + + @staticmethod + def get_builder_cls() -> type["SLAAttentionMetadataBuilder"]: + return SLAAttentionMetadataBuilder + + +@dataclass +class SLAAttentionMetadata(AttentionMetadata): + """Metadata for SLA attention.""" + current_timestep: int + topk_ratio: float = 0.5 # Ratio of key blocks to attend to + + +class SLAAttentionMetadataBuilder(AttentionMetadataBuilder): + """Builder for SLA attention metadata.""" + + def __init__(self) -> None: + pass + + def prepare(self) -> None: + pass + + def build( + self, + current_timestep: int, + topk_ratio: float = 0.5, + **kwargs: dict[str, Any], + ) -> SLAAttentionMetadata: + return SLAAttentionMetadata( + current_timestep=current_timestep, + topk_ratio=topk_ratio, + ) + + +class SLAAttentionImpl(AttentionImpl, nn.Module): + """SLA attention implementation with learnable linear projection. + + This implementation combines sparse attention with linear attention, + using a learnable projection to blend the outputs. The sparse attention + uses a block-sparse pattern determined by QK similarity. + + Args: + num_heads: Number of attention heads + head_size: Dimension of each head + topk_ratio: Ratio of key blocks to attend to (0-1), default 0.5 + feature_map: Feature map for linear attention ('softmax', 'elu', 'relu') + BLKQ: Query block size for sparse attention + BLKK: Key block size for sparse attention + use_bf16: Whether to use bfloat16 for computation + """ + + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool = False, + softmax_scale: float | None = None, + num_kv_heads: int | None = None, + prefix: str = "", + # SLA-specific parameters - matched to TurboDiffusion defaults + topk_ratio: float = 0.1, # TurboDiffusion uses topk=0.1 + feature_map: str = "softmax", + BLKQ: int = 128, # TurboDiffusion uses BLKQ=128 + BLKK: int = 64, # TurboDiffusion uses BLKK=64 + use_bf16: bool = True, + **extra_impl_args, + ) -> None: + nn.Module.__init__(self) + + self.num_heads = num_heads + self.head_size = head_size + self.softmax_scale = softmax_scale if softmax_scale else head_size**-0.5 + self.causal = causal + self.prefix = prefix + + # SLA-specific config + self.topk_ratio = topk_ratio + self.BLKQ = BLKQ + self.BLKK = BLKK + self.dtype = torch.bfloat16 if use_bf16 else torch.float16 + + # Learnable linear projection for combining sparse + linear attention + self.proj_l = nn.Linear(head_size, head_size, dtype=torch.float32) + + # Feature map for linear attention + # Type annotation for callables + self.feature_map_q: Callable[[torch.Tensor], torch.Tensor] + self.feature_map_k: Callable[[torch.Tensor], torch.Tensor] + if feature_map == "elu": + self.feature_map_q = lambda x: F.elu(x) + 1 + self.feature_map_k = lambda x: F.elu(x) + 1 + elif feature_map == "relu": + self.feature_map_q = F.relu + self.feature_map_k = F.relu + elif feature_map == "softmax": + self.feature_map_q = lambda x: F.softmax(x, dim=-1) + self.feature_map_k = lambda x: F.softmax(x, dim=-1) + else: + raise ValueError(f"Unknown feature map: {feature_map}") + + self._init_weights() + + def _init_weights(self) -> None: + """Initialize projection weights to zero for residual-like behavior.""" + with torch.no_grad(): + nn.init.zeros_(self.proj_l.weight) + nn.init.zeros_(self.proj_l.bias) # type: ignore[arg-type] + + def _calc_linear_attention( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: + """Compute linear attention: (Q @ K^T @ V) / normalizer. + + Args: + q: Query tensor (B, H, L, D) after feature map + k: Key tensor (B, H, L, D) after feature map + v: Value tensor (B, H, L, D) + + Returns: + Linear attention output (B, H, L, D) + """ + kvsum = k.transpose(-1, -2) @ v # (B, H, D, D) + ksum = torch.sum(k, dim=-2, keepdim=True) # (B, H, 1, D) + return (q @ kvsum) / (1e-5 + (q * ksum).sum(dim=-1, keepdim=True)) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + """Forward pass for SLA attention. + + Input tensors are in FastVideo format: (B, L, H, D) + Internally converted to SLA format: (B, H, L, D) + + Args: + query: Query tensor (B, L, H, D) + key: Key tensor (B, L, H, D) + value: Value tensor (B, L, H, D) + attn_metadata: Attention metadata + + Returns: + Output tensor (B, L, H, D) + """ + original_dtype = query.dtype + + # Convert from FastVideo format (B, L, H, D) to SLA format (B, H, L, D) + q = query.transpose(1, 2).contiguous() + k = key.transpose(1, 2).contiguous() + v = value.transpose(1, 2).contiguous() + + # Get topk ratio from metadata if available + topk_ratio = self.topk_ratio + if hasattr(attn_metadata, 'topk_ratio'): + topk_ratio = attn_metadata.topk_ratio # type: ignore[union-attr] + + # Compute block-sparse attention pattern + sparse_map, lut, real_topk = get_block_map(q, + k, + topk_ratio=topk_ratio, + BLKQ=self.BLKQ, + BLKK=self.BLKK) + + # Convert to compute dtype + q = q.to(self.dtype) + k = k.to(self.dtype) + v = v.to(self.dtype) + + # Sparse attention + o_s = _attention.apply(q, k, v, sparse_map, lut, real_topk, self.BLKQ, + self.BLKK) + + # Linear attention with feature maps + q_linear = self.feature_map_q(q).contiguous().to(self.dtype) + k_linear = self.feature_map_k(k).contiguous().to(self.dtype) + o_l = self._calc_linear_attention(q_linear, k_linear, v) + + # Project linear attention output and combine + with torch.amp.autocast('cuda', dtype=self.dtype): + o_l = self.proj_l(o_l) + + # Combine sparse and linear outputs + output = (o_s + o_l).to(original_dtype) + + # Convert back to FastVideo format (B, L, H, D) + output = output.transpose(1, 2) + + return output + + +# Check if spas_sage_attn is available for SageSLA +SAGESLA_ENABLED = True +try: + import spas_sage_attn._qattn as qattn + import spas_sage_attn._fused as fused + from spas_sage_attn.utils import get_vanilla_qk_quant, block_map_lut_triton +except ImportError: + SAGESLA_ENABLED = False + +SAGE2PP_ENABLED = True +try: + from spas_sage_attn._qattn import qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold +except ImportError: + SAGE2PP_ENABLED = False + + +class SageSLAAttentionBackend(AttentionBackend): + """Quantized Sparse-Linear Attention backend using SageAttention kernels.""" + + accept_output_buffer: bool = True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [64, 128] + + @staticmethod + def get_name() -> str: + return "SAGE_SLA_ATTN" + + @staticmethod + def get_impl_cls() -> type["SageSLAAttentionImpl"]: + return SageSLAAttentionImpl + + @staticmethod + def get_metadata_cls() -> type["SLAAttentionMetadata"]: + return SLAAttentionMetadata + + @staticmethod + def get_builder_cls() -> type["SLAAttentionMetadataBuilder"]: + return SLAAttentionMetadataBuilder + + +def _get_cuda_arch(device_index: int) -> str: + """Get CUDA architecture string for the given device.""" + major, minor = torch.cuda.get_device_capability(device_index) + return f"sm{major}{minor}" + + +class SageSLAAttentionImpl(AttentionImpl, nn.Module): + """SageSLA attention implementation using quantized SageAttention kernels. + + This uses INT8 quantization for Q/K and FP8 for V to achieve better performance + while maintaining accuracy. Requires spas_sage_attn package. + + Args: + num_heads: Number of attention heads + head_size: Dimension of each head (must be 64 or 128) + topk_ratio: Ratio of key blocks to attend to (0-1), default 0.5 + feature_map: Feature map for linear attention ('softmax', 'elu', 'relu') + use_bf16: Whether to use bfloat16 for computation + """ + + def __init__( + self, + num_heads: int, + head_size: int, + causal: bool = False, + softmax_scale: float | None = None, + num_kv_heads: int | None = None, + prefix: str = "", + # SageSLA-specific parameters + topk_ratio: float = 0.5, + feature_map: str = "softmax", + use_bf16: bool = True, + **extra_impl_args, + ) -> None: + nn.Module.__init__(self) + + if not SAGESLA_ENABLED: + raise ImportError( + "SageSLA requires spas_sage_attn. " + "Install with: pip install git+https://github.com/thu-ml/SpargeAttn.git" + ) + + assert head_size in [ + 64, 128 + ], f"SageSLA requires head_size in [64, 128], got {head_size}" + + self.num_heads = num_heads + self.head_size = head_size + self.softmax_scale = softmax_scale if softmax_scale else head_size**-0.5 + self.causal = causal + self.prefix = prefix + + # SageSLA-specific config + self.topk_ratio = topk_ratio + self.dtype = torch.bfloat16 if use_bf16 else torch.float16 + + # Learnable linear projection for combining sparse + linear attention + self.proj_l = nn.Linear(head_size, head_size, dtype=torch.float32) + + # Feature map for linear attention + # Type annotation for callables + self.feature_map_q: Callable[[torch.Tensor], torch.Tensor] + self.feature_map_k: Callable[[torch.Tensor], torch.Tensor] + if feature_map == "elu": + self.feature_map_q = lambda x: F.elu(x) + 1 + self.feature_map_k = lambda x: F.elu(x) + 1 + elif feature_map == "relu": + self.feature_map_q = F.relu + self.feature_map_k = F.relu + elif feature_map == "softmax": + self.feature_map_q = lambda x: F.softmax(x, dim=-1) + self.feature_map_k = lambda x: F.softmax(x, dim=-1) + else: + raise ValueError(f"Unknown feature map: {feature_map}") + + self._init_weights() + + def _init_weights(self) -> None: + """Initialize projection weights to zero for residual-like behavior.""" + with torch.no_grad(): + nn.init.zeros_(self.proj_l.weight) + nn.init.zeros_(self.proj_l.bias) # type: ignore[arg-type] + + def _calc_linear_attention( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: + """Compute linear attention: (Q @ K^T @ V) / normalizer.""" + kvsum = k.transpose(-1, -2) @ v + ksum = torch.sum(k, dim=-2, keepdim=True) + return (q @ kvsum) / (1e-5 + (q * ksum).sum(dim=-1, keepdim=True)) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + """Forward pass for SageSLA attention with quantized kernels. + + Input tensors are in FastVideo format: (B, L, H, D) + + Args: + query: Query tensor (B, L, H, D) + key: Key tensor (B, L, H, D) + value: Value tensor (B, L, H, D) + attn_metadata: Attention metadata + + Returns: + Output tensor (B, L, H, D) + """ + original_dtype = query.dtype + + # Convert from FastVideo format (B, L, H, D) to SLA format (B, H, L, D) + q = query.transpose(1, 2).contiguous() + k = key.transpose(1, 2).contiguous() + v = value.transpose(1, 2).contiguous() + + # Get topk ratio from metadata if available + topk_ratio = self.topk_ratio + if hasattr(attn_metadata, 'topk_ratio'): + topk_ratio = attn_metadata.topk_ratio # type: ignore[union-attr] + + # Determine block sizes based on GPU architecture + arch = _get_cuda_arch(q.device.index) + if arch == "sm90": + BLKQ, BLKK = 64, 128 + else: + BLKQ, BLKK = 128, 64 + + # Compute block-sparse attention pattern + sparse_map, lut, real_topk = get_block_map(q, + k, + topk_ratio=topk_ratio, + BLKQ=BLKQ, + BLKK=BLKK) + + # Convert to compute dtype + q = q.to(self.dtype) + k = k.to(self.dtype) + v = v.to(self.dtype) + + # ========== SPARGE QUANTIZED ATTENTION ========== + km = k.mean(dim=-2, keepdim=True) + headdim = q.size(-1) + scale = 1.0 / (headdim**0.5) + + # Quantize Q, K to INT8 + q_int8, q_scale, k_int8, k_scale = get_vanilla_qk_quant( + q, k, km, BLKQ, BLKK) + lut_triton, valid_block_num = block_map_lut_triton(sparse_map) + + # Quantize V to FP8 + b, h_kv, kv_len, head_dim = v.shape + padded_len = (kv_len + 127) // 128 * 128 + v_transposed_permutted = torch.empty((b, h_kv, head_dim, padded_len), + dtype=v.dtype, + device=v.device) + fused.transpose_pad_permute_cuda(v, v_transposed_permutted, 1) + v_fp8 = torch.empty(v_transposed_permutted.shape, + dtype=torch.float8_e4m3fn, + device=v.device) + v_scale = torch.empty((b, h_kv, head_dim), + dtype=torch.float32, + device=v.device) + fused.scale_fuse_quant_cuda(v_transposed_permutted, v_fp8, v_scale, + kv_len, 2.25, 1) + + # Sparse attention with quantized kernels + o_s = torch.empty_like(q) + if arch == "sm90": + qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_sm90( + q_int8, k_int8, v_fp8, o_s, lut_triton, valid_block_num, + q_scale, k_scale, v_scale, 1, False, 1, scale) + else: + pvthreshold = torch.full((q.shape[-3], ), + 1e6, + dtype=torch.float32, + device=q.device) + if SAGE2PP_ENABLED: + qk_int8_sv_f8_accum_f16_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold( + q_int8, k_int8, v_fp8, o_s, lut_triton, valid_block_num, + pvthreshold, q_scale, k_scale, v_scale, 1, False, 1, scale, + 0) + else: + qattn.qk_int8_sv_f8_accum_f32_block_sparse_attn_inst_buf_fuse_v_scale_with_pv_threshold( + q_int8, k_int8, v_fp8, o_s, lut_triton, valid_block_num, + pvthreshold, q_scale, k_scale, v_scale, 1, False, 1, scale, + 0) + # ========== END SPARGE ========== + + # Linear attention with feature maps + q_linear = self.feature_map_q(q).contiguous().to(self.dtype) + k_linear = self.feature_map_k(k).contiguous().to(self.dtype) + o_l = self._calc_linear_attention(q_linear, k_linear, v) + + # Project linear attention output and combine + with torch.amp.autocast('cuda', dtype=self.dtype): + o_l = self.proj_l(o_l) + + # Combine sparse and linear outputs + output = (o_s + o_l).to(original_dtype) + + # Convert back to FastVideo format (B, L, H, D) + output = output.transpose(1, 2) + + return output diff --git a/fastvideo/attention/layer.py b/fastvideo/attention/layer.py index e7985e58e..f1177dadd 100644 --- a/fastvideo/attention/layer.py +++ b/fastvideo/attention/layer.py @@ -50,6 +50,10 @@ class DistributedAttention(nn.Module): num_kv_heads=num_kv_heads, prefix=f"{prefix}.impl", **extra_impl_args) + # Register attn_impl as submodule if it has learnable parameters (e.g., SLA's proj_l) + # This ensures its parameters are included in state_dict() for saving/loading + if isinstance(self.attn_impl, nn.Module): + self.add_module('attn_impl', self.attn_impl) self.num_heads = num_heads self.head_size = head_size self.num_kv_heads = num_kv_heads diff --git a/fastvideo/configs/models/dits/base.py b/fastvideo/configs/models/dits/base.py index 818e0bee3..3738b1188 100644 --- a/fastvideo/configs/models/dits/base.py +++ b/fastvideo/configs/models/dits/base.py @@ -18,7 +18,8 @@ class DiTArchConfig(ArchConfig): AttentionBackendEnum.SLIDING_TILE_ATTN, AttentionBackendEnum.SAGE_ATTN, AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN, - AttentionBackendEnum.SAGE_ATTN_THREE) + AttentionBackendEnum.SAGE_ATTN_THREE, AttentionBackendEnum.SLA_ATTN, + AttentionBackendEnum.SAGE_SLA_ATTN) hidden_size: int = 0 num_attention_heads: int = 0 diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index ec1987efb..3ea13a0e3 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -326,7 +326,7 @@ def load_model_from_full_model_state_dict( unused_keys) # List of allowed parameter name patterns - ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress"] # Can be extended as needed + ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress", "proj_l"] # Can be extended as needed for new_param_name in unused_keys: if not any(pattern in new_param_name for pattern in ALLOWED_NEW_PARAM_PATTERNS): diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 4f01d261c..0cc161006 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -75,6 +75,8 @@ _SCHEDULERS = { "SelfForcingFlowMatchScheduler": ("schedulers", "scheduling_self_forcing_flow_match", "SelfForcingFlowMatchScheduler"), + "RCMScheduler": + ("schedulers", "scheduling_rcm", "RCMScheduler"), } _FAST_VIDEO_MODELS = { diff --git a/fastvideo/models/schedulers/scheduling_rcm.py b/fastvideo/models/schedulers/scheduling_rcm.py new file mode 100644 index 000000000..cf96d1c05 --- /dev/null +++ b/fastvideo/models/schedulers/scheduling_rcm.py @@ -0,0 +1,323 @@ +# SPDX-License-Identifier: Apache-2.0 +# rCM (recurrent Consistency Model) Scheduler for TurboDiffusion support +# +# This scheduler implements the rCM sampling method from TurboDiffusion, +# enabling 1-4 step video generation with distilled checkpoints. +# +# Reference: +# TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times +# https://arxiv.org/pdf/2512.16093 + +import math +from dataclasses import dataclass +from typing import Any + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.base import BaseScheduler + +logger = init_logger(__name__) + + +@dataclass +class RCMSchedulerOutput(BaseOutput): + """ + Output class for the RCM scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, ...)`): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used + as next model input in the denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class RCMScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): + """ + rCM (recurrent Consistency Model) scheduler for TurboDiffusion. + + This scheduler implements the rCM sampling method which enables 1-4 step + video generation using distilled checkpoints. It uses: + 1. TrigFlow → RectifiedFlow timestep conversion + 2. SDE sampling formula: x = (1 - t_next) * (x - t_cur * v_pred) + t_next * noise + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps used to train the model. + sigma_max (`float`, defaults to 80.0): + The initial sigma value for rCM sampling. Controls the noise level + at the start of sampling. + mid_timesteps (`list[float]`, *optional*): + Custom intermediate timesteps. If None, uses optimized defaults + [1.5, 1.4, 1.0] for better visual quality. + """ + + _compatibles: list[Any] = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + sigma_max: float = 80.0, + mid_timesteps: list[float] | None = None, + ): + # Default mid timesteps optimized for visual quality + if mid_timesteps is None: + mid_timesteps = [1.5, 1.4, 1.0] + self._mid_timesteps = mid_timesteps + + self.num_train_timesteps = num_train_timesteps + self.sigma_max = sigma_max + + # Initialize with default timesteps (will be set properly via set_timesteps) + self.timesteps = torch.tensor([1.0, 0.0], dtype=torch.float64) + self.sigmas = self.timesteps.clone() + + self._step_index: int | None = None + self._begin_index: int | None = None + + BaseScheduler.__init__(self) + + @property + def step_index(self) -> int | None: + """ + The index counter for current timestep. Increases by 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self) -> int | None: + """ + The index for the first timestep. + """ + return self._begin_index + + @property + def init_noise_sigma(self) -> float: + """ + Initial noise sigma for scaling latents. + + In rCM, initial noise is scaled by the first sigma value: + x_0 = noise * sigmas[0] + + This property is used by LatentPreparationStage to scale initial latents. + """ + return float(self.sigmas[0]) + + def set_begin_index(self, begin_index: int = 0) -> None: + """ + Sets the begin index for the scheduler. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def set_shift(self, shift: float) -> None: + """ + rCM doesn't use shift parameter, but required by BaseScheduler. + """ + pass + + def set_timesteps( + self, + num_inference_steps: int, + device: str | torch.device | None = None, + sigma_max: float | None = None, + ) -> None: + """ + Sets the discrete timesteps used for the rCM sampling process. + + The timesteps are computed using TrigFlow → RectifiedFlow conversion: + 1. Start with atan(sigma_max) and intermediate values + 2. Convert via: t = sin(t) / (cos(t) + sin(t)) + + Args: + num_inference_steps (`int`): + The number of diffusion steps (1-4 for rCM). + device (`str` or `torch.device`, *optional*): + The device to move timesteps to. + sigma_max (`float`, *optional*): + Override the initial sigma value. + """ + if num_inference_steps < 1 or num_inference_steps > 4: + logger.warning( + "rCM is optimized for 1-4 steps, got %d steps. " + "Performance may be suboptimal.", num_inference_steps + ) + + self.num_inference_steps = num_inference_steps + + if sigma_max is not None: + self.sigma_max = sigma_max + + # Build timestep schedule + mid_t = self._mid_timesteps[:num_inference_steps - 1] + + # TrigFlow timesteps: [atan(sigma_max), mid_t..., 0] + t_steps = torch.tensor( + [math.atan(self.sigma_max), *mid_t, 0], + dtype=torch.float64, + device=device, + ) + + # Convert TrigFlow → RectifiedFlow: t = sin(t) / (cos(t) + sin(t)) + t_steps = torch.sin(t_steps) / (torch.cos(t_steps) + torch.sin(t_steps)) + + # Store raw sigmas for use in step() formula + self.sigmas = t_steps.clone() + + # Scale timesteps by 1000 for model input (as per TurboDiffusion) + self.timesteps = t_steps * 1000 + + self._step_index = None + self._begin_index = None + + logger.debug("rCM timesteps (scaled): %s", self.timesteps.tolist()) + logger.debug("rCM sigmas (raw): %s", self.sigmas.tolist()) + + def _init_step_index(self, timestep: torch.FloatTensor | None = None) -> None: + """Initialize step index at the beginning of sampling.""" + if self._begin_index is None: + self._step_index = 0 + else: + self._step_index = self._begin_index + + def scale_model_input( + self, + sample: torch.Tensor, + timestep: int | None = None, + ) -> torch.Tensor: + """ + rCM doesn't scale model input, returns sample as-is. + """ + return sample + + def scale_noise( + self, + sample: torch.FloatTensor, + timestep: torch.FloatTensor | None = None, + noise: torch.FloatTensor | None = None, + ) -> torch.FloatTensor: + """ + Scale initial noise for rCM sampling. + + In rCM, initial noise is scaled by the first timestep (raw sigma): + x_0 = noise * t_steps[0] + + Args: + sample: Not used (for API compatibility) + timestep: Not used (for API compatibility) + noise: The noise tensor to scale + + Returns: + Scaled noise tensor ready for sampling + """ + if noise is None: + raise ValueError("noise must be provided for rCM scale_noise") + + # Use raw sigma (not scaled timestep) for initial noise scaling + t_initial = self.sigmas[0] + return noise.to(torch.float64) * t_initial + + def step( + self, + model_output: torch.FloatTensor, + timestep: int | torch.Tensor, + sample: torch.FloatTensor, + generator: torch.Generator | None = None, + return_dict: bool = True, + ) -> RCMSchedulerOutput | tuple[torch.FloatTensor, ...]: + """ + Predict the sample from the previous timestep using rCM update rule. + + The rCM update formula is: + x_{t+1} = (1 - t_next) * (x_t - t_cur * v_pred) + t_next * noise + + Args: + model_output (`torch.FloatTensor`): + The velocity prediction from the model (v_pred). + timestep (`int` or `torch.Tensor`): + The current timestep index (not the actual timestep value). + For rCM, this should be the index into self.timesteps. + sample (`torch.FloatTensor`): + Current sample x_t. + generator (`torch.Generator`, *optional*): + Random number generator for noise. + return_dict (`bool`): + Whether to return RCMSchedulerOutput or tuple. + + Returns: + `RCMSchedulerOutput` or `tuple`: + The denoised sample for the next step. + """ + if self._step_index is None: + self._init_step_index() + + assert self._step_index is not None + + # Get current and next sigma values (raw, unscaled) for rCM formula + # Note: self.timesteps is scaled by 1000 for model input, + # but we need raw values for the step formula + t_cur = self.sigmas[self._step_index] + # On the final step, t_next should be 0 (fully denoised) + if self._step_index + 1 < len(self.sigmas): + t_next = self.sigmas[self._step_index + 1] + else: + t_next = torch.tensor(0.0, device=sample.device, dtype=torch.float64) + + # Ensure we're working in float64 for precision + sample = sample.to(torch.float64) + model_output = model_output.to(torch.float64) + + # rCM update: x = (1 - t_next) * (x - t_cur * v_pred) + t_next * noise + x_denoised = sample - t_cur * model_output + + # Generate noise for SDE sampling + if isinstance(generator, list): + generator = generator[0] + noise = torch.randn( + sample.shape, + dtype=torch.float32, + device="cpu", + generator=generator, + ).to(sample.device).to(torch.float64) + + prev_sample = (1 - t_next) * x_denoised + t_next * noise + + # Increment step counter + self._step_index += 1 + + # Cast back to model output dtype + prev_sample = prev_sample.to(model_output.dtype) + + if not return_dict: + return (prev_sample,) + + return RCMSchedulerOutput(prev_sample=prev_sample) + + def add_noise( + self, + original_samples: torch.Tensor, + noise: torch.Tensor, + timesteps: torch.IntTensor, + ) -> torch.Tensor: + """ + Add noise to samples (forward diffusion process). + + Not typically used for rCM inference, but provided for API compatibility. + """ + raise NotImplementedError( + "add_noise is not implemented for RCMScheduler. " + "Use scale_noise for initializing the sampling process." + ) + + def __len__(self) -> int: + return self.config.num_train_timesteps diff --git a/fastvideo/pipelines/basic/turbodiffusion/__init__.py b/fastvideo/pipelines/basic/turbodiffusion/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py b/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py new file mode 100644 index 000000000..c0918080c --- /dev/null +++ b/fastvideo/pipelines/basic/turbodiffusion/turbodiffusion_pipeline.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +TurboDiffusion Video Pipeline Implementation. + +This module contains an implementation of the TurboDiffusion video diffusion pipeline +for 1-4 step video generation using rCM (recurrent Consistency Model) sampling +with SLA (Sparse-Linear Attention). +""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_rcm import RCMScheduler +from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, + DenoisingStage, InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage) + +logger = init_logger(__name__) + + +class TurboDiffusionPipeline(LoRAPipeline, ComposedPipelineBase): + """ + TurboDiffusion video pipeline for 1-4 step generation. + + Uses RCM scheduler and SLA attention for fast, high-quality video generation. + """ + + _required_config_modules = [ + "text_encoder", "tokenizer", "vae", "transformer", "scheduler" + ] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + # Use RCM scheduler for TurboDiffusion + logger.info("Initializing RCM scheduler for TurboDiffusion") + self.modules["scheduler"] = RCMScheduler(sigma_max=80.0) + + # Store checkpoint path for later loading + self._turbodiffusion_checkpoint = getattr(fastvideo_args, + 'turbodiffusion_checkpoint', + None) + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + """Set up pipeline stages with proper dependency injection.""" + + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + self.add_stage(stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + )) + + self.add_stage(stage_name="conditioning_stage", + stage=ConditioningStage()) + + self.add_stage(stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="latent_preparation_stage", + stage=LatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer", None))) + + self.add_stage(stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + transformer_2=self.get_module("transformer_2", None), + scheduler=self.get_module("scheduler"), + vae=self.get_module("vae"), + pipeline=self)) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"), + pipeline=self)) + + +EntryClass = TurboDiffusionPipeline diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index f5ac67b9f..4152c5819 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/pipelines/pipeline_registry.py @@ -23,6 +23,7 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = { "WanImageToVideoPipeline": "wan", "WanVideoToVideoPipeline": "wan", "WanCausalDMDPipeline": "wan", + "TurboDiffusionPipeline": "turbodiffusion", "StepVideoPipeline": "stepvideo", "HunyuanVideoPipeline": "hunyuan", "HunyuanVideo15Pipeline": "hunyuan15", diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 4ab4e1b76..dd748c95d 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -442,7 +442,6 @@ class DenoisingStage(PipelineStage): and progress_bar is not None): progress_bar.update() - # Gather results if using sequence parallelism trajectory_tensor: torch.Tensor | None = None if trajectory_latents: trajectory_tensor = torch.stack(trajectory_latents, dim=1) diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index c93be7f03..adbdec831 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/platforms/cuda.py @@ -197,6 +197,34 @@ class CudaPlatformBase(Platform): "Failed to import Video MoBA Attention backend: %s", str(e)) raise ImportError( "Video MoBA Attention backend is not installed. ") from e + elif selected_backend == AttentionBackendEnum.SLA_ATTN: + try: + from fastvideo.attention.backends.sla import ( # noqa: F401 + SLAAttentionBackend) + logger.info("Using SLA (Sparse-Linear Attention) backend.") + + return "fastvideo.attention.backends.sla.SLAAttentionBackend" + except ImportError as e: + logger.error("Failed to import SLA Attention backend: %s", + str(e)) + raise ImportError( + "SLA Attention backend is not available. ") from e + elif selected_backend == AttentionBackendEnum.SAGE_SLA_ATTN: + try: + from fastvideo.attention.backends.sla import ( # noqa: F401 + SageSLAAttentionBackend) + logger.info( + "Using SageSLA (Quantized Sparse-Linear Attention) backend." + ) + + return "fastvideo.attention.backends.sla.SageSLAAttentionBackend" + except ImportError as e: + logger.error("Failed to import SageSLA Attention backend: %s", + str(e)) + raise ImportError( + "SageSLA Attention backend requires spas_sage_attn. " + "Install with: pip install git+https://github.com/thu-ml/SpargeAttn.git" + ) from e elif selected_backend == AttentionBackendEnum.TORCH_SDPA: logger.info("Using Torch SDPA backend.") return "fastvideo.attention.backends.sdpa.SDPABackend" diff --git a/fastvideo/platforms/interface.py b/fastvideo/platforms/interface.py index f6d47f8ff..1f9e8a808 100644 --- a/fastvideo/platforms/interface.py +++ b/fastvideo/platforms/interface.py @@ -18,6 +18,8 @@ class AttentionBackendEnum(enum.Enum): SAGE_ATTN_THREE = enum.auto() VIDEO_SPARSE_ATTN = enum.auto() VMOBA_ATTN = enum.auto() + SLA_ATTN = enum.auto() + SAGE_SLA_ATTN = enum.auto() NO_ATTENTION = enum.auto() diff --git a/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 new file mode 100644 index 000000000..ff53f78e8 Binary files /dev/null and b/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4 differ diff --git a/fastvideo/tests/ssim/generated_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/steps4_Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of _ssim.json b/fastvideo/tests/ssim/generated_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/steps4_Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of _ssim.json new file mode 100644 index 000000000..c2d82a787 --- /dev/null +++ b/fastvideo/tests/ssim/generated_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/steps4_Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of _ssim.json @@ -0,0 +1,11 @@ +{ + "mean_ssim": 0.9917645101194028, + "min_ssim": 0.9908460974693298, + "max_ssim": 0.9925968050956726, + "reference_video": "/FastVideo/fastvideo/tests/ssim/L40S_reference_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4", + "generated_video": "/FastVideo/fastvideo/tests/ssim/generated_videos/TurboWan2.1-T2V-1.3B-Diffusers/SLA_ATTN/Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of.mp4", + "parameters": { + "num_inference_steps": 4, + "prompt": "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting." + } +} \ No newline at end of file diff --git a/fastvideo/tests/ssim/test_turbodiffusion_similarity.py b/fastvideo/tests/ssim/test_turbodiffusion_similarity.py new file mode 100644 index 000000000..b1ae13ef8 --- /dev/null +++ b/fastvideo/tests/ssim/test_turbodiffusion_similarity.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +SSIM-based similarity test for TurboDiffusion inference. + +TurboDiffusion uses the SLA (Sparse-Linear Attention) backend and RCM scheduler +for 1-4 step video generation. +""" +import os + +import torch +import pytest + +from fastvideo import VideoGenerator +from fastvideo.logger import init_logger +from fastvideo.tests.utils import compute_video_ssim_torchvision, write_ssim_results +from fastvideo.worker.multiproc_executor import MultiprocExecutor + +logger = init_logger(__name__) + +device_name = torch.cuda.get_device_name() +device_reference_folder_suffix = '_reference_videos' + +if "A40" in device_name: + device_reference_folder = "A40" + device_reference_folder_suffix +elif "L40S" in device_name: + device_reference_folder = "L40S" + device_reference_folder_suffix +else: + # device_reference_folder = "L40S" + device_reference_folder_suffix + logger.warning(f"Unsupported device for ssim tests: {device_name}, using L40S references") + raise ValueError(f"Unsupported device for ssim tests: {device_name}") + + +# TurboDiffusion parameters (1-4 step generation with RCM scheduler + SLA attention) +TURBODIFFUSION_PARAMS = { + "num_gpus": 2, + "model_path": "loayrashid/TurboWan2.1-T2V-1.3B-Diffusers", + "height": 480, + "width": 832, + "num_frames": 81, + "num_inference_steps": 4, # TurboDiffusion uses 1-4 steps + "guidance_scale": 1.0, # No CFG for TurboDiffusion + "seed": 42, + "sp_size": 2, + "tp_size": 1, + "fps": 24, +} + +TURBODIFFUSION_MODEL_TO_PARAMS = { + "TurboWan2.1-T2V-1.3B-Diffusers": TURBODIFFUSION_PARAMS, +} + +TURBODIFFUSION_TEST_PROMPTS = [ + "Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting.", +] + + +@pytest.mark.parametrize("prompt", TURBODIFFUSION_TEST_PROMPTS) +@pytest.mark.parametrize("model_id", list(TURBODIFFUSION_MODEL_TO_PARAMS.keys())) +def test_turbodiffusion_inference_similarity(prompt, model_id): + """ + Test that runs TurboDiffusion inference with SLA attention and RCM scheduler, + then compares the output to reference videos using SSIM. + """ + # TurboDiffusion requires SLA attention backend + ATTENTION_BACKEND = "SLA_ATTN" + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + + script_dir = os.path.dirname(os.path.abspath(__file__)) + + base_output_dir = os.path.join(script_dir, 'generated_videos', model_id) + output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND) + output_video_name = f"{prompt[:100].strip()}.mp4" + + os.makedirs(output_dir, exist_ok=True) + + BASE_PARAMS = TURBODIFFUSION_MODEL_TO_PARAMS[model_id] + num_inference_steps = BASE_PARAMS["num_inference_steps"] + + init_kwargs = { + "num_gpus": BASE_PARAMS["num_gpus"], + "sp_size": BASE_PARAMS["sp_size"], + "tp_size": BASE_PARAMS["tp_size"], + "override_pipeline_cls_name": "TurboDiffusionPipeline", + } + + generation_kwargs = { + "num_inference_steps": num_inference_steps, + "output_path": output_dir, + "height": BASE_PARAMS["height"], + "width": BASE_PARAMS["width"], + "num_frames": BASE_PARAMS["num_frames"], + "guidance_scale": BASE_PARAMS["guidance_scale"], + "seed": BASE_PARAMS["seed"], + "fps": BASE_PARAMS["fps"], + } + + generator = VideoGenerator.from_pretrained( + model_path=BASE_PARAMS["model_path"], + **init_kwargs + ) + generator.generate_video(prompt, **generation_kwargs) + + if isinstance(generator.executor, MultiprocExecutor): + generator.executor.shutdown() + + assert os.path.exists(output_dir), f"Output video was not generated at {output_dir}" + + reference_folder = os.path.join( + script_dir, device_reference_folder, model_id, ATTENTION_BACKEND + ) + + if not os.path.exists(reference_folder): + logger.error("Reference folder missing") + raise FileNotFoundError( + f"Reference video folder does not exist: {reference_folder}" + ) + + # Find the matching reference video based on the prompt + reference_video_name = None + + for filename in os.listdir(reference_folder): + if filename.endswith('.mp4') and prompt[:100].strip() in filename: + reference_video_name = filename + break + + if not reference_video_name: + logger.error( + f"Reference video not found for prompt: {prompt} with backend: {ATTENTION_BACKEND}" + ) + raise FileNotFoundError(f"Reference video missing") + + reference_video_path = os.path.join(reference_folder, reference_video_name) + generated_video_path = os.path.join(output_dir, output_video_name) + + logger.info( + f"Computing SSIM between {reference_video_path} and {generated_video_path}" + ) + ssim_values = compute_video_ssim_torchvision( + reference_video_path, generated_video_path, use_ms_ssim=True + ) + + mean_ssim = ssim_values[0] + logger.info(f"SSIM mean value: {mean_ssim}") + logger.info(f"Writing SSIM results to directory: {output_dir}") + + success = write_ssim_results( + output_dir, ssim_values, reference_video_path, + generated_video_path, num_inference_steps, prompt + ) + + if not success: + logger.error("Failed to write SSIM results to file") + + # TurboDiffusion uses fewer steps, may have slightly lower SSIM + min_acceptable_ssim = 0.90 + assert mean_ssim >= min_acceptable_ssim, ( + f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} " + f"for {model_id} with backend {ATTENTION_BACKEND}" + ) diff --git a/fastvideo/tests/transformers/test_wanvideo.py b/fastvideo/tests/transformers/test_wanvideo.py index 1901a1811..1028138b5 100644 --- a/fastvideo/tests/transformers/test_wanvideo.py +++ b/fastvideo/tests/transformers/test_wanvideo.py @@ -121,4 +121,4 @@ def test_wan_transformer(): assert output1.dtype == output2.dtype, f"Output dtype don't match: {output1.dtype} vs {output2.dtype}" # Check if outputs are similar (allowing for small numerical differences) - assert_close(output1, output2, atol=1e-1, rtol=1e-2) \ No newline at end of file + assert_close(output1, output2, atol=1e-1, rtol=1e-2) diff --git a/scripts/checkpoint_conversion/convert_turbodiffusion_to_diffusers.py b/scripts/checkpoint_conversion/convert_turbodiffusion_to_diffusers.py new file mode 100644 index 000000000..077add50e --- /dev/null +++ b/scripts/checkpoint_conversion/convert_turbodiffusion_to_diffusers.py @@ -0,0 +1,205 @@ +#!/usr/bin/env python3 +""" +Convert TurboDiffusion .pth checkpoint to Diffusers safetensors format. + +This script: +1. Loads the TurboDiffusion .pth checkpoint +2. Applies weight key renaming to match FastVideo/Diffusers format +3. Reshapes patch_embedding from [D, C*P] to [D, C, P_t, P_h, P_w] +4. Saves as sharded safetensors files compatible with diffusers + +Usage: + python convert_turbodiffusion_to_diffusers.py \ + --input_path /path/to/TurboWan2.1-T2V-1.3B-480P.pth \ + --output_dir /path/to/output/transformer \ + --reference_repo Wan-AI/Wan2.1-T2V-1.3B-Diffusers +""" + +import argparse +import os +import re +import json +import torch +import shutil +import glob +from safetensors import safe_open +from safetensors.torch import save_file +from huggingface_hub import snapshot_download + +# Weight mapping from TurboDiffusion -> Diffusers/FastVideo format +TURBODIFFUSION_WEIGHT_MAPPING = { + # Self attention + r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.to_q.\2", + r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.to_k.\2", + r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.to_v.\2", + r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.to_out.\2", + r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2", + r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2", + # Cross attention + r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2", + r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2", + r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2", + r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2", + # Norms and FFN + r"^blocks\.(\d+)\.norm1\.(.*)$": r"blocks.\1.norm1.\2", + r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2", + r"^blocks\.(\d+)\.norm2\.(.*)$": r"blocks.\1.norm3.\2", + r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", + r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", + r"^blocks\.(\d+)\.modulation$": r"blocks.\1.scale_shift_table", + # Embeddings - DON'T add .proj here! WanVideoArchConfig's param_names_mapping will add it + # patch_embedding.weight stays as patch_embedding.weight (HF format needs this) + r"^text_embedding\.0\.(.*)$": r"condition_embedder.text_embedder.fc_in.\1", + r"^text_embedding\.2\.(.*)$": r"condition_embedder.text_embedder.fc_out.\1", + r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1", + r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1", + r"^time_projection\.1\.(.*)$": r"condition_embedder.time_modulation.linear.\1", + # Head + r"^head\.head\.(.*)$": r"proj_out.\1", + r"^head\.norm\.(.*)$": r"norm_out.\1", + r"^head\.modulation$": r"scale_shift_table", + # SLA proj_l weights - include them! They're the distilled attention weights + r"^blocks\.(\d+)\.self_attn\.attn_op\.local_attn\.proj_l\.(.*)$": r"blocks.\1.attn1.attn_impl.proj_l.\2", +} + +# No keys to skip - we want all weights including proj_l +SKIP_PATTERNS = [] + + +def should_skip_key(key: str) -> bool: + """Check if a key should be skipped (SLA-specific weights).""" + for pattern in SKIP_PATTERNS: + if re.match(pattern, key): + return True + return False + + +def convert_key(turbo_key: str) -> str: + """Convert TurboDiffusion key to Diffusers format.""" + for pattern, replacement in TURBODIFFUSION_WEIGHT_MAPPING.items(): + if re.match(pattern, turbo_key): + return re.sub(pattern, replacement, turbo_key) + return turbo_key # Return unchanged if no pattern matches + + +def reshape_patch_embedding(tensor: torch.Tensor, target_shape: tuple) -> torch.Tensor: + """Reshape patch_embedding from [D, C*P_t*P_h*P_w] to [D, C, P_t, P_h, P_w].""" + if len(tensor.shape) == 2 and len(target_shape) == 5: + return tensor.view(target_shape) + return tensor + + +def get_reference_shapes(reference_repo: str) -> dict: + """Download reference model and get expected shapes for patch_embedding.""" + print(f"Downloading reference model shapes from {reference_repo}...") + + # Download just the transformer config and a weight file to get shapes + local_dir = snapshot_download( + repo_id=reference_repo, + allow_patterns=["transformer/config.json", "transformer/diffusion_pytorch_model*.safetensors"], + local_dir_use_symlinks=False + ) + + # Load the first safetensors file to get shapes + weight_files = glob.glob(os.path.join(local_dir, "transformer", "*.safetensors")) + shapes = {} + + for wf in weight_files: + with safe_open(wf, framework="pt") as f: + for key in f.keys(): + shapes[key] = f.get_tensor(key).shape + + return shapes + + +def main(): + parser = argparse.ArgumentParser(description="Convert TurboDiffusion checkpoint to Diffusers format") + parser.add_argument("--input_path", type=str, required=True, + help="Path to TurboDiffusion .pth checkpoint") + parser.add_argument("--output_dir", type=str, required=True, + help="Output directory for converted safetensors") + parser.add_argument("--reference_repo", type=str, default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + help="Reference HF repo to get expected tensor shapes") + parser.add_argument("--skip_sla_weights", action="store_true", default=False, + help="Skip SLA-specific weights (proj_l) that aren't in base model") + + args = parser.parse_args() + + # Load TurboDiffusion checkpoint + print(f"Loading TurboDiffusion checkpoint from {args.input_path}...") + turbo_state_dict = torch.load(args.input_path, map_location="cpu", weights_only=True) + print(f"Loaded {len(turbo_state_dict)} keys") + + # Get reference shapes for reshaping + ref_shapes = get_reference_shapes(args.reference_repo) + print(f"Got {len(ref_shapes)} reference shapes") + + # Convert keys and reshape tensors + converted_state_dict = {} + skipped_keys = [] + + for turbo_key, tensor in turbo_state_dict.items(): + # Skip SLA-specific weights if requested + if args.skip_sla_weights and should_skip_key(turbo_key): + skipped_keys.append(turbo_key) + continue + + # Convert key name + new_key = convert_key(turbo_key) + + # Reshape patch_embedding if needed + if "patch_embedding" in new_key and new_key in ref_shapes: + target_shape = ref_shapes[new_key] + if tensor.shape != target_shape: + print(f"Reshaping {new_key}: {tensor.shape} -> {target_shape}") + tensor = reshape_patch_embedding(tensor, target_shape) + + # Verify shape matches reference if available + if new_key in ref_shapes: + if tensor.shape != ref_shapes[new_key]: + print(f"WARNING: Shape mismatch for {new_key}: got {tensor.shape}, expected {ref_shapes[new_key]}") + + converted_state_dict[new_key] = tensor + + print(f"\nConversion summary:") + print(f" Converted: {len(converted_state_dict)} keys") + print(f" Skipped (SLA): {len(skipped_keys)} keys") + + if skipped_keys: + print(f"\nSkipped SLA keys (first 5):") + for k in skipped_keys[:5]: + print(f" - {k}") + + # Create output directory + os.makedirs(args.output_dir, exist_ok=True) + + # Save as safetensors + output_path = os.path.join(args.output_dir, "diffusion_pytorch_model.safetensors") + print(f"\nSaving to {output_path}...") + save_file(converted_state_dict, output_path) + + # Copy config.json from reference + ref_local = snapshot_download( + repo_id=args.reference_repo, + allow_patterns=["transformer/config.json"], + local_dir_use_symlinks=False + ) + src_config = os.path.join(ref_local, "transformer", "config.json") + dst_config = os.path.join(args.output_dir, "config.json") + shutil.copy(src_config, dst_config) + print(f"Copied config.json") + + print(f"\nDone! Converted weights saved to: {args.output_dir}") + print(f"\nNext steps:") + print(f" 1. Use create_hf_repo.py to create a complete diffusers repo:") + print(f" python scripts/checkpoint_conversion/create_hf_repo.py \\") + print(f" --repo_id Wan-AI/Wan2.1-T2V-1.3B-Diffusers \\") + print(f" --local_dir /tmp/turbodiffusion-wan \\") + print(f" --checkpoint_dir {args.output_dir} \\") + print(f" --push_to_hub --upload_repo_id YOUR_USERNAME/TurboWan-Diffusers") + + +if __name__ == "__main__": + main()