From fe8e83a6dd1c0b62eb21ce6e5df438a488b966fe Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Sun, 9 Aug 2026 16:20:11 +0000 Subject: [PATCH 1/3] feat(swiftvr): add real-time restoration pipeline Port the released SwiftVR ReAE and transformer into a causal single-device pipeline with stateful chunk processing and pinned output conversion. Add a FlashVSR-style H100 example, documentation, and focused model/session tests. The example warms full and tail shapes and sustains 1080p generation on one H100 without keeping the development validation scripts. Verified with 16 SwiftVR unit tests, Ruff check/format, diff checks, and an 81-frame 1080p H100 run at 34.78 generation FPS and 22.18 FPS including H.264 encoding. --- README.md | 2 + docs/en/swiftvr.md | 103 ++ examples/swiftvr/README.md | 86 ++ examples/swiftvr/swiftvr_restore_h100.py | 154 +++ telefuser/models/swiftvr_reae.py | 171 +++ telefuser/models/swiftvr_transformer.py | 984 ++++++++++++++++++ telefuser/pipelines/swiftvr/__init__.py | 9 + telefuser/pipelines/swiftvr/chunk.py | 49 + telefuser/pipelines/swiftvr/io.py | 56 + telefuser/pipelines/swiftvr/pipeline.py | 369 +++++++ telefuser/pipelines/swiftvr/streaming_dit.py | 207 ++++ telefuser/pipelines/swiftvr/streaming_tae.py | 165 +++ tests/unit/models/test_swiftvr_transformer.py | 72 ++ tests/unit/pipelines/swiftvr/test_example.py | 95 ++ .../unit/pipelines/swiftvr/test_streaming.py | 157 +++ 15 files changed, 2679 insertions(+) create mode 100644 docs/en/swiftvr.md create mode 100644 examples/swiftvr/README.md create mode 100644 examples/swiftvr/swiftvr_restore_h100.py create mode 100644 telefuser/models/swiftvr_reae.py create mode 100644 telefuser/models/swiftvr_transformer.py create mode 100644 telefuser/pipelines/swiftvr/__init__.py create mode 100644 telefuser/pipelines/swiftvr/chunk.py create mode 100644 telefuser/pipelines/swiftvr/io.py create mode 100644 telefuser/pipelines/swiftvr/pipeline.py create mode 100644 telefuser/pipelines/swiftvr/streaming_dit.py create mode 100644 telefuser/pipelines/swiftvr/streaming_tae.py create mode 100644 tests/unit/models/test_swiftvr_transformer.py create mode 100644 tests/unit/pipelines/swiftvr/test_example.py create mode 100644 tests/unit/pipelines/swiftvr/test_streaming.py diff --git a/README.md b/README.md index 39fe20d0..2f9d346b 100644 --- a/README.md +++ b/README.md @@ -237,6 +237,7 @@ telefuser/ | `ABot-World-0-5B-LF` | Single-GPU interactive world model | Direct HTTP or shared LiveKit controller via [examples/abot_world/README.md](examples/abot_world/README.md) | | `LiveAct` | S2V | Speech-driven talking head generation via [examples/liveact/liveact_s2v_h100.py](examples/liveact/liveact_s2v_h100.py) | | `FlashVSR` | VSR | Streaming video super-resolution via [examples/flashvsr/README.md](examples/flashvsr/README.md) | +| `SwiftVR` | Causal video restoration | Single-GPU BF16 offline and direct streaming restoration via [examples/swiftvr/README.md](examples/swiftvr/README.md) | ### Video Generation @@ -274,6 +275,7 @@ See [examples/README.md](examples/README.md) for the example runner and baseline - [docs/en/adding_new_model.md](docs/en/adding_new_model.md): integrating new models - [docs/en/adding_new_example.md](docs/en/adding_new_example.md): authoring examples and pipeline contracts - [docs/en/abot_world.md](docs/en/abot_world.md): ABot-World single-GPU interactive pipeline, controls, and tests +- [docs/en/swiftvr.md](docs/en/swiftvr.md): SwiftVR checkpoint loading, streaming usage, performance, and service limits ## Known Limitations diff --git a/docs/en/swiftvr.md b/docs/en/swiftvr.md new file mode 100644 index 00000000..1b00dab2 --- /dev/null +++ b/docs/en/swiftvr.md @@ -0,0 +1,103 @@ +# SwiftVR + +TeleFuser integrates the released `H-oliday/SwiftVR` video-restoration model as +a faithful, sequential, single-GPU pipeline. The default path preserves the +upstream BF16 model, one-step DiT, mask-free shifted-window attention, causal +ReAE state, and fixed chunk/flush behavior. Dense attention is dispatched +through `telefuser.ops.attention` with Torch SDPA. + +## Provenance and checkpoint + +The reference source is SwiftVR commit +`5ca168cef6ca7200f135fdfea85e5e13d12c5b53`. The local checkpoint is the model +revision `743ed2530c550764905400f38eb6cc41af5abc80` under `/data/SwiftVR`. + +The implementation loads ReAE and the Diffusers transformer through the +existing `ModuleManager`. It keeps checkpoint keys unchanged. Both models are +loaded on CPU and moved ReAE-first to the target GPU, matching upstream +allocation and cuDNN plan selection. The prompt embedding remains on CPU until +the DiT condition cache is created, also matching upstream behavior. + +## Offline usage + +Install the project dependencies and run: + +```bash +python examples/swiftvr/swiftvr_restore_h100.py \ + --model_root /data/SwiftVR \ + --height 360 --width 640 \ + --scale 3 \ + --output restored_1080p.mp4 +``` + +The default input is FlashVSR's versioned `examples/data/dag.mp4` test video. +The CLI follows FlashVSR's common options: `input_video`, `scale`, `height`, +`width`, `gpu_num`, `model_root`, and `output`. Dimensions are internally padded +to multiples of 32 and cropped back exactly. The output preserves `4k+1` source +frames because that is the released temporal contract. BF16 and one GPU are the +supported parity path. `get_pipeline()` owns model loading, `run()` accepts +loaded PIL frames, and TeleFuser's shared video utilities own file I/O. `run()` +uses the stateful 24-frame session path. A throwaway session warms two full +chunks plus the actual tail shape before measuring the real request, then the +example reports both generation FPS and end-to-end FPS including H.264 +encoding. + +## Streaming behavior + +`SwiftVRPipeline.stream()` creates independent ReAE encoder/decoder boundary +state, DiT overlap state, RoPE position state, and condition cache. `step()` +accepts arbitrary uint8 frame counts, buffers non-four-aligned tails, and +`flush()` pads only the final encoder group while preserving output frame +count. GPU execution is serialized by the shared pipeline lock. + +Direct sessions are isolated and tested with interleaved inputs so ReAE +boundary state, DiT overlap, RoPE offsets, and frames cannot cross sessions. +`close()` releases all retained causal state. The example intentionally does +not expose `get_service()`: the stock LiveKit protocol has no inbound video +track or frame payload, so a local queue adapter would not provide a usable +`stream-serve` transport. + +## H100 baseline + +The target-side compute measurements exclude model loading, video encoding, +LiveKit delivery, and client playback. Each measured chunk contains 24 output +frames after two warmup chunks. + +| Resolution | Compute FPS | TTFC | P50 / P95 chunk | Peak allocated | Retained session state | +| --- | ---: | ---: | ---: | ---: | ---: | +| 1920x1080 | 47.28 | 15.01 s | 0.506 / 0.511 s | 27.08 GiB | 235.5 MiB | +| 2560x1440 | 26.50 | 16.31 s | 0.895 / 0.960 s | 35.93 GiB | 411.8 MiB | + +The complete example path was separately accepted on one H100 using all 81 +frames of `examples/data/dag.mp4`, resized to `640x360` and restored at 3x. The +measured session produced the `1920x1080` frames in 2.33 seconds (34.78 FPS), +and H.264 encoding took 1.32 seconds, for 22.18 end-to-end FPS. The resulting +video contains 81 frames at 16 FPS. These example figures exclude one-time +model loading, input decoding, and shape warmup; they include output conversion +and device-to-host transfer, and the end-to-end figure additionally includes +video encoding. + +TTFC includes cold cuDNN benchmarking and kernel-plan setup after model loading. +Retained memory in the table is the state of one direct causal session. No +multi-session service capacity is claimed because there is no supported +SwiftVR video transport in the shared streaming server. + +The 1080p long-stream run completed 1176 measured chunks and 28,224 frames over +600.46 seconds at 47.01 compute FPS. P50/P95/max chunk latency was +0.508/0.523/0.654 seconds. The first and last 100 chunks had P50 values of +0.50810 and 0.50796 seconds respectively, so no latency growth was observed. +Session cleanup completed and the benchmark process returned all device memory +to the driver. + +## Current constraints + +- No quantization, feature cache, sparse-attention substitution, or + `torch.compile` path is enabled. +- No stage splitting or pipeline overlap is enabled; execution remains in the + upstream order. +- The public example uses eager Torch SDPA and does not expose alternative + attention backends. +- The supported online surface is the direct causal session API; the example + does not expose a partial `stream-serve` adapter. +- RTX 5090 measurements require that hardware and are not represented by H100 + results. diff --git a/examples/swiftvr/README.md b/examples/swiftvr/README.md new file mode 100644 index 00000000..01725cdb --- /dev/null +++ b/examples/swiftvr/README.md @@ -0,0 +1,86 @@ +# SwiftVR Video Restoration + +This example restores videos with the released `H-oliday/SwiftVR` checkpoint on +one CUDA GPU. The parity path uses BF16, dense Torch SDPA, the upstream fixed +timestep, and the original 24-frame causal chunk protocol. + +The checkpoint root must contain: + +```text +reae.safetensors +prompt_embedding.safetensors +transformer/config.json +transformer/diffusion_pytorch_model.safetensors +``` + +Run offline restoration with: + +```bash +python examples/swiftvr/swiftvr_restore_h100.py \ + --model_root /data/SwiftVR \ + --height 360 \ + --width 640 \ + --scale 3 \ + --output restored_1080p.mp4 +``` + +Like the FlashVSR example, the file exposes `get_pipeline()` for model loading +and `run()` for loaded PIL frames. Its CLI accepts the same common video options: +`input_video`, `scale`, `height`, `width`, `gpu_num`, `model_root`, and `output`. +It defaults to FlashVSR's versioned `examples/data/dag.mp4` test video. The +SwiftVR-specific 24-frame causal protocol remains an internal pipeline default. +The command above resizes the low-resolution input to `640x360` and restores a +`1920x1080` video. A separate warmup session covers two full chunks and the +actual tail shape to absorb cold cuDNN plan selection. The reported processing +FPS then covers stateful inference and output transfer, while end-to-end FPS +also includes H.264 encoding. + +The direct causal API accepts uint8 `[T,H,W,3]` tensors and returns BF16 +`[1,T,3,H,W]` tensors in `[0,1]`: + +```python +import torch + +from examples.swiftvr.swiftvr_restore_h100 import get_pipeline + +pipeline = get_pipeline(model_root="/data/SwiftVR") +session = pipeline.stream(resolution=(1920, 1080), clip_len=24, dit_overlap=1) +try: + first = session.step(torch.zeros((24, 540, 960, 3), dtype=torch.uint8)) + tail = session.step(torch.zeros((5, 540, 960, 3), dtype=torch.uint8)) + flushed = session.flush() +finally: + session.close() +``` + +## H100 performance + +The following baseline was measured on one NVIDIA H100 80GB HBM3 with BF16, +eager Torch SDPA, and 24-frame chunks. The values are synchronized target-side +compute metrics; model loading, video encoding, network delivery, and client +playback are excluded. + +| Output resolution | Compute FPS | TTFC | P50 / P95 chunk latency | Peak allocated memory | +| --- | ---: | ---: | ---: | ---: | +| 1920x1080 | 47.28 | 15.01 s | 0.506 / 0.511 s | 27.08 GiB | +| 2560x1440 | 26.50 | 16.31 s | 0.895 / 0.960 s | 35.93 GiB | + +The full example command above was also verified on one H100 with all 81 frames +of `dag.mp4`. The measured session produced `1920x1080` frames at 34.78 FPS; +including H.264 encoding, end-to-end throughput was 22.18 FPS versus the source +rate of 16 FPS. The output contains all 81 frames at 16 FPS. These figures +exclude one-time model loading, input decoding, and shape warmup. + +TTFC includes cold cuDNN benchmarking and kernel-plan setup after model +loading. In the 1080p ten-minute stability run, SwiftVR processed 28,224 frames +in 1,176 chunks at 47.01 compute FPS. P50/P95/max chunk latency was +0.508/0.523/0.654 seconds, with no latency growth between the beginning and end +of the run. Session cleanup returned all device memory to the driver. + +`SwiftVRPipeline.stream()` is the supported stateful interface. The example +does not expose `get_service()` because the current LiveKit protocol has no +video-input transport; presenting a partial service adapter would not make the +model usable through `telefuser stream-serve`. + +See [the SwiftVR integration guide](../../docs/en/swiftvr.md) for checkpoint +loading, constraints, and measured results. diff --git a/examples/swiftvr/swiftvr_restore_h100.py b/examples/swiftvr/swiftvr_restore_h100.py new file mode 100644 index 00000000..09a454d8 --- /dev/null +++ b/examples/swiftvr/swiftvr_restore_h100.py @@ -0,0 +1,154 @@ +"""SwiftVR streaming video-restoration example.""" + +from __future__ import annotations + +import os +import time + +import click +import numpy as np +import torch +from PIL import Image + +from telefuser.core.config import AttentionConfig, AttnImplType +from telefuser.pipelines.swiftvr import SwiftVRPipeline +from telefuser.pipelines.swiftvr.io import ntchw_to_uint8_frames +from telefuser.utils.utils import get_example_name +from telefuser.utils.video import VideoData, save_video + +TF_MODEL_ZOO_PATH = os.environ.get("TF_MODEL_ZOO_PATH", "model_zoo") +PPL_CONFIG = dict( + name="swiftvr_restore_h100", + model_root=TF_MODEL_ZOO_PATH + "/SwiftVR", + scale=4, + fps=16, + video_quality=6, + chunk_size=24, + warmup_chunks=2, +) + + +def get_pipeline( + parallelism: int = 1, + model_root: str = PPL_CONFIG["model_root"], +) -> SwiftVRPipeline: + """Initialize the SwiftVR pipeline.""" + if parallelism != 1: + raise ValueError("SwiftVR supports exactly one GPU") + return SwiftVRPipeline.from_pretrained( + model_root, + device="cuda", + torch_dtype=torch.bfloat16, + attention_config=AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA), + ) + + +def run( + pipeline: SwiftVRPipeline, + input_video: list[Image.Image], + scale: int = PPL_CONFIG["scale"], +) -> list[Image.Image]: + """Restore loaded PIL frames through a stateful streaming session.""" + if not input_video: + raise ValueError("input_video must contain at least one frame") + width, height = input_video[0].size + crop_width, crop_height = width // 8 * 8, height // 8 * 8 + if crop_width <= 0 or crop_height <= 0: + raise ValueError(f"input frames are too small: {width}x{height}") + arrays = [ + np.asarray(frame.convert("RGB").crop((0, 0, crop_width, crop_height)), dtype=np.uint8) for frame in input_video + ] + frames = torch.from_numpy(np.stack(arrays)).contiguous() + restored_frames: list[Image.Image] = [] + session = pipeline.stream(upscale=scale) + try: + for start in range(0, len(frames), PPL_CONFIG["chunk_size"]): + output = ntchw_to_uint8_frames(session.step(frames[start : start + PPL_CONFIG["chunk_size"]])) + if output is not None: + restored_frames.extend(Image.fromarray(frame) for frame in output) + output = ntchw_to_uint8_frames(session.flush()) + if output is not None: + restored_frames.extend(Image.fromarray(frame) for frame in output) + finally: + session.close() + return restored_frames + + +def _warmup_frame_count(frame_count: int) -> int: + full_chunks = PPL_CONFIG["chunk_size"] * PPL_CONFIG["warmup_chunks"] + tail_frames = frame_count % PPL_CONFIG["chunk_size"] + return min(frame_count, full_chunks + tail_frames) + + +@click.command() +@click.option( + "--input_video", + "-i", + default=f"{os.path.dirname(__file__)}/../data/dag.mp4", + help="Path to input low-quality video", +) +@click.option("--scale", "-s", default=PPL_CONFIG["scale"], type=int, help="Upscaling factor (default: 4)") +@click.option("--height", "-h", default=None, type=int, help="Input video height (default: auto-detect)") +@click.option("--width", "-w", default=None, type=int, help="Input video width (default: auto-detect)") +@click.option("--gpu_num", default=1, type=int, help="Number of GPUs to use (default: 1)") +@click.option( + "--model_root", + default=PPL_CONFIG["model_root"], + help=f"Root directory containing model files (default: {PPL_CONFIG['model_root']})", +) +@click.option("--output", "-o", default=None, help="Output video path (default: auto-generated)") +def main( + input_video: str, + scale: int, + height: int | None, + width: int | None, + gpu_num: int, + model_root: str, + output: str | None, +) -> None: + """SwiftVR streaming video restoration.""" + if not os.path.exists(input_video): + raise FileNotFoundError(f"Input video not found: {input_video}") + + if output is None: + output_dir = os.getenv("TELEAI_EXAMPLE_OUTPUT_DIR", "./") + filename = get_example_name(__file__).replace(".mp4", f"_scale{scale}_{gpu_num}gpu.mp4") + output = os.path.join(output_dir, filename) + + click.echo(f"Input video: {input_video}") + click.echo(f"Input resolution: {width or 'auto'}x{height or 'auto'}") + click.echo(f"Scale: {scale}x") + click.echo(f"GPUs: {gpu_num}") + click.echo(f"Model root: {model_root}") + click.echo(f"Output: {output}") + + click.echo("Loading pipeline...") + pipeline = get_pipeline(gpu_num, model_root) + + click.echo("Loading video...") + input_frames = VideoData(video_file=input_video, height=height, width=width).raw_data() + click.echo(f"Total frames: {len(input_frames)}") + + warmup_frame_count = _warmup_frame_count(len(input_frames)) + click.echo(f"Warmup pass ({warmup_frame_count} frames)...") + run(pipeline, input_frames[:warmup_frame_count], scale=scale) + + click.echo("Processing video...") + start_time = time.perf_counter() + restored = run(pipeline, input_frames, scale=scale) + processing_seconds = time.perf_counter() - start_time + processing_fps = len(restored) / processing_seconds + click.echo(f"Processing time: {processing_seconds:.2f} seconds ({processing_fps:.2f} FPS)") + + click.echo(f"Saving to {output}...") + encoding_started = time.perf_counter() + save_video(restored, output, fps=PPL_CONFIG["fps"], quality=PPL_CONFIG["video_quality"]) + encoding_seconds = time.perf_counter() - encoding_started + end_to_end_fps = len(restored) / (processing_seconds + encoding_seconds) + click.echo(f"Encoding time: {encoding_seconds:.2f} seconds") + click.echo(f"End-to-end throughput: {end_to_end_fps:.2f} FPS") + click.echo("Done!") + + +if __name__ == "__main__": + main() diff --git a/telefuser/models/swiftvr_reae.py b/telefuser/models/swiftvr_reae.py new file mode 100644 index 00000000..99164f46 --- /dev/null +++ b/telefuser/models/swiftvr_reae.py @@ -0,0 +1,171 @@ +"""Restoration-aware Autoencoder (ReAE). + +A lightweight causal-streaming encoder/decoder used by SwiftVR to move between pixel +space and the DiT latent space. The memory-block topology is adapted from TAEHV +(https://github.com/madebyollin/taehv, MIT License). +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def conv(n_in: int, n_out: int, **kwargs: object) -> nn.Conv2d: + return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs) + + +class Clamp(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.tanh(x / 3) * 3 + + +class MemBlock(nn.Module): + """Residual block that fuses the current frame with the previous one.""" + + def __init__(self, n_in: int, n_out: int) -> None: + super().__init__() + self.conv = nn.Sequential( + conv(n_in * 2, n_out), + nn.ReLU(inplace=True), + conv(n_out, n_out), + nn.ReLU(inplace=True), + conv(n_out, n_out), + ) + self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() + self.act = nn.ReLU(inplace=True) + + def forward(self, x: torch.Tensor, past: torch.Tensor) -> torch.Tensor: + return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x)) + + +class TPool(nn.Module): + """Temporal pooling by a factor of ``stride`` via a 1x1 channel mix.""" + + def __init__(self, n_f: int, stride: int) -> None: + super().__init__() + self.stride = stride + self.conv = nn.Conv2d(n_f * stride, n_f, 1, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + _NT, C, H, W = x.shape + return self.conv(x.reshape(-1, self.stride * C, H, W)) + + +class TGrow(nn.Module): + """Temporal upsampling by a factor of ``stride``. + + ``stride == 1`` is a plain 1x1 projection; ``stride == 2`` upsamples in time + with nearest interpolation followed by a depth-only 3D convolution. ``conv`` + is kept solely for checkpoint compatibility and is not used at inference. + """ + + def __init__(self, n_f: int, stride: int) -> None: + super().__init__() + self.stride = stride + self.n_f = n_f + + if stride == 1: + self.proj = nn.Conv2d(n_f, n_f, 1, bias=False) + self.conv3d = None + else: + self.conv3d = nn.Conv3d(n_f, n_f, kernel_size=(3, 1, 1), padding=(1, 0, 0), bias=False) + self.proj = None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + NT, C, H, W = x.shape + if self.stride == 1: + return self.proj(x) + x = x.unsqueeze(2) + x = F.interpolate(x, size=(self.stride, H, W), mode="nearest") + x = self.conv3d(x) + return x.permute(0, 2, 1, 3, 4).reshape(NT * self.stride, C, H, W) + + +class ReAE(nn.Module): + def __init__( + self, + checkpoint_path: str | None = None, + width_mult: int = 2, + decoder_time_upscale: tuple[bool, bool] = (True, True), + decoder_space_upscale: tuple[bool, bool, bool] = (True, True, True), + patch_size: int = 2, + latent_channels: int = 48, + ) -> None: + super().__init__() + self.width_mult = width_mult + self.image_channels = 3 + self.patch_size = patch_size + self.latent_channels = latent_channels + + e_enc = 64 + + self.encoder = nn.Sequential( + conv(self.image_channels * self.patch_size**2, e_enc), + nn.ReLU(inplace=True), + TPool(e_enc, 2), + conv(e_enc, e_enc, stride=2, bias=False), + MemBlock(e_enc, e_enc), + MemBlock(e_enc, e_enc), + MemBlock(e_enc, e_enc), + TPool(e_enc, 2), + conv(e_enc, e_enc, stride=2, bias=False), + MemBlock(e_enc, e_enc), + MemBlock(e_enc, e_enc), + MemBlock(e_enc, e_enc), + TPool(e_enc, 1), + conv(e_enc, e_enc, stride=2, bias=False), + MemBlock(e_enc, e_enc), + MemBlock(e_enc, e_enc), + MemBlock(e_enc, e_enc), + conv(e_enc, self.latent_channels), + ) + + n_f = [256 * width_mult, 128 * width_mult, 64 * width_mult, 64] + self.frames_to_trim = 2 ** sum(decoder_time_upscale) - 1 + + self.decoder = nn.Sequential( + Clamp(), + conv(self.latent_channels, n_f[0]), + nn.ReLU(inplace=True), + MemBlock(n_f[0], n_f[0]), + MemBlock(n_f[0], n_f[0]), + MemBlock(n_f[0], n_f[0]), + nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), + TGrow(n_f[0], 1), + conv(n_f[0], n_f[1], bias=False), + MemBlock(n_f[1], n_f[1]), + MemBlock(n_f[1], n_f[1]), + MemBlock(n_f[1], n_f[1]), + nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), + TGrow(n_f[1], 2 if decoder_time_upscale[0] else 1), + conv(n_f[1], n_f[2], bias=False), + MemBlock(n_f[2], n_f[2]), + MemBlock(n_f[2], n_f[2]), + MemBlock(n_f[2], n_f[2]), + nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), + TGrow(n_f[2], 2 if decoder_time_upscale[1] else 1), + conv(n_f[2], n_f[3], bias=False), + nn.ReLU(inplace=True), + conv(n_f[3], self.image_channels * self.patch_size**2), + ) + + if checkpoint_path is not None: + raise ValueError("Load ReAE checkpoints through ModuleManager") + + @classmethod + def state_dict_converter(cls) -> SwiftVRReAEStateDictConverter: + return SwiftVRReAEStateDictConverter() + + +class SwiftVRReAEStateDictConverter: + """Keep the released ReAE checkpoint keys unchanged.""" + + @staticmethod + def from_official(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + return state_dict + + @staticmethod + def from_diffusers(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + return state_dict diff --git a/telefuser/models/swiftvr_transformer.py b/telefuser/models/swiftvr_transformer.py new file mode 100644 index 00000000..6c7d7daa --- /dev/null +++ b/telefuser/models/swiftvr_transformer.py @@ -0,0 +1,984 @@ +"""Shifted-window self-attention diffusion transformer for SwiftVR. + +Adapted from the ``WanTransformer3DModel`` (Wan2.2-TI2V) implementation in +Hugging Face ``diffusers`` (Apache-2.0). The mask-free shifted-window +self-attention processor and the multi-backend dense-attention dispatcher are +specific to SwiftVR; the same checkpoint runs bit-identically across PyTorch SDPA, +FlashAttention-2/3, SageAttention and xFormers. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from types import SimpleNamespace + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from telefuser.core.config import AttentionConfig, AttnImplType +from telefuser.ops.attention import attention +from telefuser.ops.ffn import FeedForward +from telefuser.utils.logging import logger + + +@dataclass +class Transformer2DModelOutput: + sample: torch.Tensor + + +class FP32LayerNorm(nn.LayerNorm): + """LayerNorm with upstream-compatible FP32 accumulation.""" + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + origin_dtype = inputs.dtype + return F.layer_norm( + inputs.float(), + self.normalized_shape, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ).to(origin_dtype) + + +def get_1d_rotary_pos_embed( + dim: int, + length: int, + theta: float = 10000.0, + *, + use_real: bool = True, + repeat_interleave_real: bool = True, + freqs_dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + """Local copy of the Diffusers real-valued 1D RoPE construction.""" + if dim % 2: + raise ValueError(f"RoPE dimension must be even, got {dim}") + position = torch.arange(length) + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=freqs_dtype) / dim)) + freqs = torch.outer(position, freqs) + if not use_real: + raise ValueError("SwiftVR requires real-valued rotary embeddings") + if repeat_interleave_real: + return ( + freqs.cos().repeat_interleave(2, dim=1, output_size=dim).float(), + freqs.sin().repeat_interleave(2, dim=1, output_size=dim).float(), + ) + return torch.cat([freqs.cos(), freqs.cos()], dim=-1).float(), torch.cat([freqs.sin(), freqs.sin()], dim=-1).float() + + +class Timesteps(nn.Module): + def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float) -> None: + super().__init__() + self.num_channels = num_channels + self.flip_sin_to_cos = flip_sin_to_cos + self.downscale_freq_shift = downscale_freq_shift + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + half_dim = self.num_channels // 2 + exponent = -math.log(10000) * torch.arange(half_dim, dtype=torch.float32, device=timesteps.device) + exponent = exponent / (half_dim - self.downscale_freq_shift) + embedding = timesteps[:, None].float() * torch.exp(exponent)[None, :] + embedding = torch.cat([torch.sin(embedding), torch.cos(embedding)], dim=-1) + if self.flip_sin_to_cos: + embedding = torch.cat([embedding[:, half_dim:], embedding[:, :half_dim]], dim=-1) + if self.num_channels % 2: + embedding = F.pad(embedding, (0, 1)) + return embedding + + +class TimestepEmbedding(nn.Module): + def __init__(self, in_channels: int, time_embed_dim: int) -> None: + super().__init__() + self.linear_1 = nn.Linear(in_channels, time_embed_dim) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + return self.linear_2(self.act(self.linear_1(sample))) + + +class PixArtAlphaTextProjection(nn.Module): + def __init__(self, in_features: int, hidden_size: int, act_fn: str = "gelu_tanh") -> None: + super().__init__() + if act_fn != "gelu_tanh": + raise ValueError(f"Unsupported SwiftVR text projection activation: {act_fn}") + self.linear_1 = nn.Linear(in_features, hidden_size) + self.act_1 = nn.GELU(approximate="tanh") + self.linear_2 = nn.Linear(hidden_size, hidden_size) + + def forward(self, caption: torch.Tensor) -> torch.Tensor: + return self.linear_2(self.act_1(self.linear_1(caption))) + + +# --------------------------------------------------------------------------- # +# Attention dispatch # +# --------------------------------------------------------------------------- # + + +def _dense_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attention_config: AttentionConfig, + attn_mask: torch.Tensor | None = None, +) -> torch.Tensor: + """Mask-free dense attention. Inputs/outputs are (B, N, H, D).""" + return attention( + q, + k, + v, + attention_config=attention_config, + attn_mask=attn_mask, + input_layout="BSND", + output_layout="BSND", + ) + + +# --------------------------------------------------------------------------- # +# Window partition caches # +# --------------------------------------------------------------------------- # + + +def _make_hw_starts(H, W, wh, ww, do_shift, device=None): + device = device or torch.device("cpu") + + def _axis_starts(size: int, win: int) -> torch.Tensor: + if size <= win: + return torch.zeros(1, dtype=torch.long, device=device) + shift = (win // 2) if do_shift else 0 + max_start = size - win + n = (size + win - 1) // win + 2 + k = torch.arange(n, dtype=torch.long, device=device) + starts = (k * win - shift).clamp_(0, max_start) + starts = torch.unique(starts, sorted=True) + if starts.numel() > 2: + prev, nxt = starts[:-2], starts[2:] + covered = nxt <= (prev + win) + keep = torch.ones_like(starts, dtype=torch.bool) + keep[1:-1] = ~covered + starts = starts[keep] + return starts + + return _axis_starts(H, wh), _axis_starts(W, ww) + + +def _build_hw_lin_indices(T, H, W, h_starts, w_starts, wh, ww): + device = h_starts.device + dh = torch.arange(wh, device=device) + dw = torch.arange(ww, device=device) + dt = torch.arange(T, device=device) + h_idx = h_starts[:, None] + dh[None, :] + w_idx = w_starts[:, None] + dw[None, :] + spatial_lin = (h_idx[:, None, :, None] * W + w_idx[None, :, None, :]).reshape(-1, wh * ww) + full_lin = dt[None, :, None] * (H * W) + spatial_lin[:, None, :] + return full_lin.reshape(spatial_lin.shape[0], T * wh * ww) + + +class _WindowIndexCache: + _store: dict[tuple, torch.Tensor] = {} + + @classmethod + def get(cls, T, H, W, wh, ww, do_shift, device): + key = (T, H, W, wh, ww, do_shift, device.type, device.index) + if key not in cls._store: + h_s, w_s = _make_hw_starts(H, W, wh, ww, do_shift, device) + cls._store[key] = _build_hw_lin_indices(T, H, W, h_s, w_s, wh, ww) + return cls._store[key] + + @classmethod + def clear(cls): + cls._store.clear() + + +class _WindowRuntimeMeta: + __slots__ = ("lin_flat", "owner_pos", "Nw", "Lw", "THW") + + def __init__(self, lin_flat, owner_pos, Nw, Lw, THW): + self.lin_flat = lin_flat + self.owner_pos = owner_pos + self.Nw = Nw + self.Lw = Lw + self.THW = THW + + +class _WindowRuntimeMetaCache: + _store: dict[tuple, _WindowRuntimeMeta] = {} + + @staticmethod + def _build_owner_pos_cpu(lin, prefer_front, THW): + lin_cpu = lin.detach().to("cpu") + Nw, Lw = lin_cpu.shape + owner = torch.empty(THW, dtype=torch.long) + local = torch.arange(Lw, dtype=torch.long) + order_iter = range(Nw - 1, -1, -1) if prefer_front else range(Nw) + for wi in order_iter: + owner[lin_cpu[wi]] = wi * Lw + local + return owner + + @classmethod + def get(cls, T, H, W, wh, ww, do_shift, prefer_front, device): + key = (T, H, W, wh, ww, bool(do_shift), bool(prefer_front), device.type, device.index) + if key not in cls._store: + lin = _WindowIndexCache.get(T, H, W, wh, ww, do_shift, device) + Nw, Lw = lin.shape + THW = T * H * W + owner_cpu = cls._build_owner_pos_cpu(lin, prefer_front, THW) + cls._store[key] = _WindowRuntimeMeta( + lin_flat=lin.reshape(-1).contiguous(), + owner_pos=owner_cpu.to(device=device, non_blocking=True), + Nw=int(Nw), + Lw=int(Lw), + THW=int(THW), + ) + return cls._store[key] + + @classmethod + def clear(cls): + cls._store.clear() + + +# --------------------------------------------------------------------------- # +# Rotary embedding helpers # +# --------------------------------------------------------------------------- # + + +def _apply_rotary_emb(x, freqs_cos, freqs_sin): + x1, x2 = x.unflatten(-1, (-1, 2)).unbind(-1) + cos = freqs_cos[..., 0::2] + sin = freqs_sin[..., 1::2] + if cos.dtype != x.dtype: + cos, sin = cos.to(x.dtype), sin.to(x.dtype) + out = torch.empty_like(x) + out[..., 0::2] = x1 * cos - x2 * sin + out[..., 1::2] = x1 * sin + x2 * cos + return out + + +def _apply_rotary_emb_inplace(x, freqs_cos, freqs_sin): + cos = freqs_cos[..., 0::2] + sin = freqs_sin[..., 1::2] + if cos.dtype != x.dtype: + cos, sin = cos.to(x.dtype), sin.to(x.dtype) + x_pair = x.view(*x.shape[:-1], -1, 2) + x_even, x_odd = x_pair[..., 0], x_pair[..., 1] + tmp = x_even * sin + x_even.mul_(cos) + x_even.addcmul_(x_odd, sin, value=-1) + x_odd.mul_(cos) + x_odd.add_(tmp) + del tmp + return x + + +def _release_input_storage(t: torch.Tensor) -> None: + """Free a tensor's CUDA storage to bypass Python refcount held by *args.""" + try: + if t.is_cuda and t._base is None and t.is_contiguous(): + t.untyped_storage().resize_(0) + except Exception: + pass + + +def _get_qkv_projections(attn, hidden_states, encoder_hidden_states): + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + if getattr(attn, "fused_projections", False): + if attn.cross_attention_dim_head is None: + query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1) + else: + query = attn.to_q(hidden_states) + key, value = attn.to_kv(encoder_hidden_states).chunk(2, dim=-1) + else: + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + return query, key, value + + +def _get_added_kv_projections(attn, encoder_hidden_states_img): + if getattr(attn, "fused_projections", False): + key_img, value_img = attn.to_added_kv(encoder_hidden_states_img).chunk(2, dim=-1) + else: + key_img = attn.add_k_proj(encoder_hidden_states_img) + value_img = attn.add_v_proj(encoder_hidden_states_img) + return key_img, value_img + + +def _infer_local_thw(thw_global, k_local): + Tg, Hg, Wg = thw_global + if k_local == Tg * Hg * Wg: + return Tg, Hg, Wg + if k_local % (Hg * Wg) == 0: + return (k_local // (Hg * Wg), Hg, Wg) + if k_local % (Tg * Wg) == 0: + return (Tg, k_local // (Tg * Wg), Wg) + if k_local % (Tg * Hg) == 0: + return (Tg, Hg, k_local // (Tg * Hg)) + raise RuntimeError(f"Cannot infer local THW from {thw_global} and k_local={k_local}.") + + +# --------------------------------------------------------------------------- # +# Attention modules # +# --------------------------------------------------------------------------- # + + +class WanAttnProcessor: + """Standard global attention used for cross-attention.""" + + def __call__( + self, + attn: WanAttention, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + encoder_hidden_states_img = None + if attn.add_k_proj is not None: + image_context_length = encoder_hidden_states.shape[1] - 512 + encoder_hidden_states_img = encoder_hidden_states[:, :image_context_length] + encoder_hidden_states = encoder_hidden_states[:, image_context_length:] + + query, key, value = _get_qkv_projections(attn, hidden_states, encoder_hidden_states) + del hidden_states + + query = attn.norm_q(query).unflatten(2, (attn.heads, -1)) + key = attn.norm_k(key).unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + if rotary_emb is not None: + query = _apply_rotary_emb(query, *rotary_emb) + key = _apply_rotary_emb(key, *rotary_emb) + + hidden_states_img = None + if encoder_hidden_states_img is not None: + key_img, value_img = _get_added_kv_projections(attn, encoder_hidden_states_img) + key_img = attn.norm_added_k(key_img).unflatten(2, (attn.heads, -1)) + value_img = value_img.unflatten(2, (attn.heads, -1)) + hidden_states_img = _dense_attn(query, key_img, value_img, attn.attention_config) + hidden_states_img = hidden_states_img.flatten(2, 3).type_as(query) + + hidden_states = _dense_attn(query, key, value, attn.attention_config, attention_mask) + hidden_states = hidden_states.flatten(2, 3).type_as(query) + + if hidden_states_img is not None: + hidden_states = hidden_states + hidden_states_img + + hidden_states = attn.to_out[0](hidden_states) + if attn.training and isinstance(attn.to_out[1], nn.Dropout) and attn.to_out[1].p > 0.0: + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +class WanAttention(torch.nn.Module): + _default_processor_cls = WanAttnProcessor + _available_processors = [WanAttnProcessor] + + def __init__( + self, + dim: int, + heads: int = 8, + dim_head: int = 64, + eps: float = 1e-5, + dropout: float = 0.0, + added_kv_proj_dim: int | None = None, + cross_attention_dim_head: int | None = None, + processor: object | None = None, + is_cross_attention: bool | None = None, + ) -> None: + super().__init__() + + self.inner_dim = dim_head * heads + self.heads = heads + self.added_kv_proj_dim = added_kv_proj_dim + self.cross_attention_dim_head = cross_attention_dim_head + self.kv_inner_dim = self.inner_dim if cross_attention_dim_head is None else cross_attention_dim_head * heads + + self.to_q = nn.Linear(dim, self.inner_dim, bias=True) + self.to_k = nn.Linear(dim, self.kv_inner_dim, bias=True) + self.to_v = nn.Linear(dim, self.kv_inner_dim, bias=True) + self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, dim, bias=True), nn.Dropout(dropout)]) + self.norm_q = nn.RMSNorm(dim_head * heads, eps=eps, elementwise_affine=True) + self.norm_k = nn.RMSNorm(dim_head * heads, eps=eps, elementwise_affine=True) + + self.add_k_proj = self.add_v_proj = None + if added_kv_proj_dim is not None: + self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=True) + self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=True) + self.norm_added_k = nn.RMSNorm(dim_head * heads, eps=eps) + + self.is_cross_attention = cross_attention_dim_head is not None + self.attention_config = AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA) + self.fused_projections = False + self.set_processor(processor) + + def set_processor(self, processor: object | None) -> None: + self.processor = processor + + def fuse_projections(self) -> None: + if self.fused_projections: + return + + if self.cross_attention_dim_head is None: + w = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data]) + b = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data]) + out_f, in_f = w.shape + with torch.device("meta"): + self.to_qkv = nn.Linear(in_f, out_f, bias=True) + self.to_qkv.load_state_dict({"weight": w, "bias": b}, strict=True, assign=True) + else: + w = torch.cat([self.to_k.weight.data, self.to_v.weight.data]) + b = torch.cat([self.to_k.bias.data, self.to_v.bias.data]) + out_f, in_f = w.shape + with torch.device("meta"): + self.to_kv = nn.Linear(in_f, out_f, bias=True) + self.to_kv.load_state_dict({"weight": w, "bias": b}, strict=True, assign=True) + + if self.added_kv_proj_dim is not None: + w = torch.cat([self.add_k_proj.weight.data, self.add_v_proj.weight.data]) + b = torch.cat([self.add_k_proj.bias.data, self.add_v_proj.bias.data]) + out_f, in_f = w.shape + with torch.device("meta"): + self.to_added_kv = nn.Linear(in_f, out_f, bias=True) + self.to_added_kv.load_state_dict({"weight": w, "bias": b}, strict=True, assign=True) + + self.fused_projections = True + + @torch.no_grad() + def unfuse_projections(self) -> None: + for attr in ("to_qkv", "to_kv", "to_added_kv"): + if hasattr(self, attr): + delattr(self, attr) + self.fused_projections = False + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs: object, + ) -> torch.Tensor: + return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, rotary_emb, **kwargs) + + +class WanShiftWindow2DInferProcessor: + """Mask-free 2D spatial shifted-window self-attention (full temporal view). + + Each window is densely pre-gathered (boundary-clamped) so attention reduces + to a single SDPA-style call with no mask, padding or cyclic shift. Alternate + layers use a half-window shift; the reverse step uses a priority-coherent + scatter. RoPE is applied globally before partitioning. + """ + + def __init__(self, window_hw: tuple[int, int] = (16, 16), shift_every_other_layer: bool = True) -> None: + wh, ww = window_hw + assert wh > 0 and ww > 0 + self.window_hw = window_hw + self.shift_every_other_layer = shift_every_other_layer + + def __call__( + self, + attn: WanAttention, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + if encoder_hidden_states is not None or getattr(attn, "is_cross_attention", False): + raise RuntimeError("WanShiftWindow2DInferProcessor only supports self-attention.") + if attention_mask is not None: + raise RuntimeError("External attention_mask is not supported.") + if not hasattr(attn, "_thw") or attn._thw is None: + raise RuntimeError("attn._thw=(T,H,W) must be set before forward.") + + Tg, Hg, Wg = attn._thw + B, K, _ = hidden_states.shape + T, H, W = _infer_local_thw((Tg, Hg, Wg), K) + if K != T * H * W: + raise RuntimeError(f"K mismatch: K={K}, inferred T*H*W={T * H * W}.") + + cfg_wh, cfg_ww = self.window_hw + wh, ww = min(cfg_wh, H), min(cfg_ww, W) + + # Do NOT derive the window shift from a Python integer module attribute + # such as attn._layer_id here. torch.compile treats integer attributes on + # nn.Module as static guards, so different layer ids trigger repeated + # recompilation. Instead, _do_shift is assigned once per block in + # enable_shifted_window_self_attention(). + do_shift = bool(getattr(attn, "_do_shift", False)) + prefer_front = not do_shift + + meta = _WindowRuntimeMetaCache.get( + T, H, W, wh, ww, do_shift=do_shift, prefer_front=prefer_front, device=hidden_states.device + ) + Nw, Lw = meta.Nw, meta.Lw + Hn = attn.heads + Dh = attn.inner_dim // Hn + + query, key, value = _get_qkv_projections(attn, hidden_states, None) + _release_input_storage(hidden_states) + del hidden_states + + query = attn.norm_q(query).unflatten(2, (Hn, Dh)) + key = attn.norm_k(key).unflatten(2, (Hn, Dh)) + value = value.unflatten(2, (Hn, Dh)) + + value = torch.index_select(value, 1, meta.lin_flat).view(B * Nw, Lw, Hn, Dh) + + if rotary_emb is not None: + query = _apply_rotary_emb_inplace(query, *rotary_emb) + key = _apply_rotary_emb_inplace(key, *rotary_emb) + + query = torch.index_select(query, 1, meta.lin_flat).view(B * Nw, Lw, Hn, Dh) + key = torch.index_select(key, 1, meta.lin_flat).view(B * Nw, Lw, Hn, Dh) + + o_win = _dense_attn(query, key, value, attn.attention_config) + del query, key, value + + o_flat = o_win.reshape(B, Nw * Lw, Hn, Dh) + del o_win + out = torch.index_select(o_flat, 1, meta.owner_pos) + del o_flat + + out = out.reshape(B, K, Hn * Dh) + out = attn.to_out[0](out) + if attn.training and isinstance(attn.to_out[1], nn.Dropout) and attn.to_out[1].p > 0.0: + out = attn.to_out[1](out) + return out + + +# --------------------------------------------------------------------------- # +# Embeddings and transformer block # +# --------------------------------------------------------------------------- # + + +class WanImageEmbedding(nn.Module): + def __init__(self, in_features: int, out_features: int, pos_embed_seq_len: int | None = None) -> None: + super().__init__() + self.norm1 = FP32LayerNorm(in_features) + self.ff = FeedForward(in_features, out_features, mult=1, activation_fn="gelu") + self.norm2 = FP32LayerNorm(out_features) + self.pos_embed = ( + nn.Parameter(torch.zeros(1, pos_embed_seq_len, in_features)) if pos_embed_seq_len is not None else None + ) + + def forward(self, encoder_hidden_states_image: torch.Tensor) -> torch.Tensor: + if self.pos_embed is not None: + B, S, C = encoder_hidden_states_image.shape + encoder_hidden_states_image = encoder_hidden_states_image.view(-1, 2 * S, C) + encoder_hidden_states_image = encoder_hidden_states_image + self.pos_embed + x = self.norm1(encoder_hidden_states_image) + x = self.ff(x) + return self.norm2(x) + + +class WanTimeTextImageEmbedding(nn.Module): + def __init__( + self, + dim: int, + time_freq_dim: int, + time_proj_dim: int, + text_embed_dim: int, + image_embed_dim: int | None = None, + pos_embed_seq_len: int | None = None, + ) -> None: + super().__init__() + self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0) + self.time_embedder = TimestepEmbedding(in_channels=time_freq_dim, time_embed_dim=dim) + self.act_fn = nn.SiLU() + self.time_proj = nn.Linear(dim, time_proj_dim) + self.text_embedder = PixArtAlphaTextProjection(text_embed_dim, dim, act_fn="gelu_tanh") + self.image_embedder = ( + WanImageEmbedding(image_embed_dim, dim, pos_embed_seq_len=pos_embed_seq_len) + if image_embed_dim is not None + else None + ) + + def forward( + self, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_image: torch.Tensor | None = None, + timestep_seq_len: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + timestep = self.timesteps_proj(timestep) + if timestep_seq_len is not None: + timestep = timestep.unflatten(0, (-1, timestep_seq_len)) + + dtype = next(iter(self.time_embedder.parameters())).dtype + if timestep.dtype != dtype and dtype != torch.int8: + timestep = timestep.to(dtype) + + temb = self.time_embedder(timestep).type_as(encoder_hidden_states) + timestep_proj = self.time_proj(self.act_fn(temb)) + encoder_hidden_states = self.text_embedder(encoder_hidden_states) + if encoder_hidden_states_image is not None: + encoder_hidden_states_image = self.image_embedder(encoder_hidden_states_image) + return temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image + + +class WanRotaryPosEmbed(nn.Module): + def __init__( + self, + attention_head_dim: int, + patch_size: tuple[int, int, int] | list[int], + max_seq_len: int, + theta: float = 10000.0, + ) -> None: + super().__init__() + self.attention_head_dim = attention_head_dim + self.patch_size = patch_size + self.max_seq_len = max_seq_len + + h_dim = w_dim = 2 * (attention_head_dim // 6) + t_dim = attention_head_dim - h_dim - w_dim + self.t_dim, self.h_dim, self.w_dim = t_dim, h_dim, w_dim + + freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64 + freqs_cos, freqs_sin = [], [] + for dim in [t_dim, h_dim, w_dim]: + fc, fs = get_1d_rotary_pos_embed( + dim, max_seq_len, theta, use_real=True, repeat_interleave_real=True, freqs_dtype=freqs_dtype + ) + freqs_cos.append(fc) + freqs_sin.append(fs) + self.register_buffer("freqs_cos", torch.cat(freqs_cos, dim=1), persistent=False) + self.register_buffer("freqs_sin", torch.cat(freqs_sin, dim=1), persistent=False) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + B, C, F_, H, W = hidden_states.shape + p_t, p_h, p_w = self.patch_size + ppf, pph, ppw = F_ // p_t, H // p_h, W // p_w + + split_sizes = [self.t_dim, self.h_dim, self.w_dim] + freqs_cos = self.freqs_cos.split(split_sizes, dim=1) + freqs_sin = self.freqs_sin.split(split_sizes, dim=1) + + fc_f = freqs_cos[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + fc_h = freqs_cos[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + fc_w = freqs_cos[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + fs_f = freqs_sin[0][:ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + fs_h = freqs_sin[1][:pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + fs_w = freqs_sin[2][:ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + + freqs_cos_ = torch.cat([fc_f, fc_h, fc_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1) + freqs_sin_ = torch.cat([fs_f, fs_h, fs_w], dim=-1).reshape(1, ppf * pph * ppw, 1, -1) + return freqs_cos_, freqs_sin_ + + +class WanTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + ffn_dim: int, + num_heads: int, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + ) -> None: + super().__init__() + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.attn1 = WanAttention( + dim=dim, + heads=num_heads, + dim_head=dim // num_heads, + eps=eps, + cross_attention_dim_head=None, + processor=WanAttnProcessor(), + ) + self.attn2 = WanAttention( + dim=dim, + heads=num_heads, + dim_head=dim // num_heads, + eps=eps, + added_kv_proj_dim=added_kv_proj_dim, + cross_attention_dim_head=dim // num_heads, + processor=WanAttnProcessor(), + ) + self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.ffn = FeedForward(dim, inner_dim=ffn_dim, activation_fn="gelu-approximate") + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + rotary_emb: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + h_dtype = hidden_states.dtype + + if temb.ndim == 4: + mods = (self.scale_shift_table.unsqueeze(0) + temb.float()).to(h_dtype) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = mods.chunk(6, dim=2) + shift_msa, scale_msa, gate_msa = shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2) + c_shift_msa, c_scale_msa, c_gate_msa = c_shift_msa.squeeze(2), c_scale_msa.squeeze(2), c_gate_msa.squeeze(2) + else: + mods = (self.scale_shift_table + temb.float()).to(h_dtype) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = mods.chunk(6, dim=1) + + attn_output = self.attn1( + self.norm1(hidden_states).mul_(1.0 + scale_msa).add_(shift_msa), None, None, rotary_emb + ) + hidden_states.addcmul_(attn_output, gate_msa) + del attn_output + + attn_output = self.attn2(self.norm2(hidden_states), encoder_hidden_states, None, None) + hidden_states.add_(attn_output) + del attn_output + + ff_output = self.ffn(self.norm3(hidden_states).mul_(1.0 + c_scale_msa).add_(c_shift_msa)) + hidden_states.addcmul_(ff_output, c_gate_msa) + del ff_output + + return hidden_states + + +# --------------------------------------------------------------------------- # +# Inference setup helpers # +# --------------------------------------------------------------------------- # + + +def enable_shifted_window_self_attention(model: nn.Module, window_hw: tuple[int, int] = (16, 16)) -> None: + """Fuse QKV projections and install the shifted-window self-attn processor. + + Important for torch.compile: + The shifted-window parity is stored as a boolean _do_shift on each self-attn + module. Do not use per-layer integer attributes such as _layer_id inside the + compiled attention processor, because they become static guards and cause + repeated recompilation across transformer blocks. + """ + proc = WanShiftWindow2DInferProcessor(window_hw=window_hw) + + for i, blk in enumerate(getattr(model, "blocks", [])): + underlying = getattr(blk, "_orig_mod", blk) + if hasattr(underlying, "attn1"): + underlying.attn1._do_shift = bool(i % 2 == 1) + if hasattr(underlying.attn1, "_layer_id"): + delattr(underlying.attn1, "_layer_id") + + for _, m in model.named_modules(): + if isinstance(m, WanAttention): + m.fuse_projections() + if not getattr(m, "is_cross_attention", False): + m.set_processor(proc) + + +def compile_transformer_blocks(model: nn.Module, mode: str = "default") -> None: + if not hasattr(torch, "compile"): + logger.warning("torch.compile not available (requires PyTorch 2.0+). Skipping.") + return + if mode == "reduce-overhead": + logger.warning( + "compile_mode='reduce-overhead' is incompatible with the in-place residuals; falling back to 'default'." + ) + mode = "default" + for i, blk in enumerate(getattr(model, "blocks", [])): + if isinstance(blk, WanTransformerBlock): + model.blocks[i] = torch.compile(blk, mode=mode, fullgraph=False) + + +# --------------------------------------------------------------------------- # +# Main model # +# --------------------------------------------------------------------------- # + + +class SwiftVRWanTransformer3DModel(nn.Module): + _supports_gradient_checkpointing = True + _skip_layerwise_casting_patterns = ["patch_embedding", "condition_embedder", "norm"] + _no_split_modules = ["WanTransformerBlock"] + _keep_in_fp32_modules = ["time_embedder", "scale_shift_table", "norm1", "norm2", "norm3"] + _keys_to_ignore_on_load_unexpected = ["norm_added_q"] + _repeated_blocks = ["WanTransformerBlock"] + + def __init__( + self, + patch_size: tuple[int, int, int] | list[int] = (1, 2, 2), + num_attention_heads: int = 40, + attention_head_dim: int = 128, + in_channels: int = 16, + out_channels: int | None = 16, + text_dim: int = 4096, + freq_dim: int = 256, + ffn_dim: int = 13824, + num_layers: int = 40, + cross_attn_norm: bool = True, + qk_norm: str = "rms_norm_across_heads", + eps: float = 1e-6, + image_dim: int | None = None, + added_kv_proj_dim: int | None = None, + rope_max_seq_len: int = 1024, + pos_embed_seq_len: int | None = None, + enable_swa: bool = True, + self_attn_window_hw: tuple[int, int] = (16, 16), + use_torch_compile: bool = False, + compile_mode: str = "default", + ) -> None: + super().__init__() + + self.config = SimpleNamespace( + patch_size=tuple(patch_size), + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + in_channels=in_channels, + out_channels=out_channels, + text_dim=text_dim, + freq_dim=freq_dim, + ffn_dim=ffn_dim, + num_layers=num_layers, + cross_attn_norm=cross_attn_norm, + qk_norm=qk_norm, + eps=eps, + image_dim=image_dim, + added_kv_proj_dim=added_kv_proj_dim, + rope_max_seq_len=rope_max_seq_len, + pos_embed_seq_len=pos_embed_seq_len, + ) + + inner_dim = num_attention_heads * attention_head_dim + out_channels = out_channels or in_channels + + self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len) + self.patch_embedding = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size) + self.condition_embedder = WanTimeTextImageEmbedding( + dim=inner_dim, + time_freq_dim=freq_dim, + time_proj_dim=inner_dim * 6, + text_embed_dim=text_dim, + image_embed_dim=image_dim, + pos_embed_seq_len=pos_embed_seq_len, + ) + self.blocks = nn.ModuleList( + [ + WanTransformerBlock( + inner_dim, ffn_dim, num_attention_heads, qk_norm, cross_attn_norm, eps, added_kv_proj_dim + ) + for _ in range(num_layers) + ] + ) + self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False) + self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + self.gradient_checkpointing = False + self._enable_swa = enable_swa + self._self_attn_window_hw = self_attn_window_hw + self.attention_config = AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA) + + @classmethod + def state_dict_converter(cls, config: dict | None = None) -> SwiftVRTransformerStateDictConverter: + return SwiftVRTransformerStateDictConverter(config) + + def set_attention_config(self, attention_config: AttentionConfig) -> None: + self.attention_config = attention_config + for module in self.modules(): + if isinstance(module, WanAttention): + module.attention_config = attention_config + + def prepare_for_inference( + self, + attention_config: AttentionConfig | None = None, + use_torch_compile: bool = False, + compile_mode: str = "default", + ) -> None: + self.set_attention_config(attention_config or AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA)) + enable_shifted_window_self_attention(self, window_hw=self._self_attn_window_hw) + if use_torch_compile: + compile_transformer_blocks(self, mode=compile_mode) + _WindowIndexCache.clear() + _WindowRuntimeMetaCache.clear() + self.eval() + + @torch.inference_mode() + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_image: torch.Tensor | None = None, + return_dict: bool = True, + attention_kwargs: dict[str, object] | None = None, + ) -> Transformer2DModelOutput | tuple[torch.Tensor]: + if attention_kwargs is not None: + attention_kwargs = attention_kwargs.copy() + attention_kwargs.pop("scale", None) + B, C, F_, H, W = hidden_states.shape + p_t, p_h, p_w = self.config.patch_size + ppf, pph, ppw = F_ // p_t, H // p_h, W // p_w + + rotary_emb = self.rope(hidden_states) + hidden_states = self.patch_embedding(hidden_states).flatten(2).transpose(1, 2).contiguous() + + ts_seq_len = None + if timestep.ndim == 2: + ts_seq_len = timestep.shape[1] + timestep = timestep.flatten() + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len + ) + + timestep_proj = ( + timestep_proj.unflatten(2, (6, -1)) if ts_seq_len is not None else timestep_proj.unflatten(1, (6, -1)) + ) + + if encoder_hidden_states_image is not None: + encoder_hidden_states = torch.cat([encoder_hidden_states_image, encoder_hidden_states], dim=1) + + thw_global = (ppf, pph, ppw) + cfg_wh, cfg_ww = self._self_attn_window_hw + dev = hidden_states.device + _WindowRuntimeMetaCache.get( + ppf, pph, ppw, min(cfg_wh, pph), min(cfg_ww, ppw), do_shift=False, prefer_front=True, device=dev + ) + _WindowRuntimeMetaCache.get( + ppf, pph, ppw, min(cfg_wh, pph), min(cfg_ww, ppw), do_shift=True, prefer_front=False, device=dev + ) + + for blk in self.blocks: + underlying = getattr(blk, "_orig_mod", blk) + if hasattr(underlying, "attn1"): + underlying.attn1._thw = thw_global + + for blk in self.blocks: + hidden_states = blk(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb) + + h_dtype = hidden_states.dtype + if temb.ndim == 3: + mods = (self.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).to(h_dtype) + shift, scale = mods.chunk(2, dim=2) + shift, scale = shift.squeeze(2), scale.squeeze(2) + else: + mods = (self.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).to(h_dtype) + shift, scale = mods.chunk(2, dim=1) + + normed = self.norm_out(hidden_states) + normed.mul_(1.0 + scale).add_(shift) + hidden_states = self.proj_out(normed) + del normed + + hidden_states = hidden_states.reshape(B, ppf, pph, ppw, p_t, p_h, p_w, -1) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) + + +class SwiftVRTransformerStateDictConverter: + """Keep the released Diffusers checkpoint keys unchanged.""" + + def __init__(self, config: dict | None = None) -> None: + self.config = dict(config or {}) + + def from_diffusers(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, object]]: + config = {key: value for key, value in self.config.items() if not key.startswith("_")} + return state_dict, config + + def from_official(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, object]]: + return self.from_diffusers(state_dict) diff --git a/telefuser/pipelines/swiftvr/__init__.py b/telefuser/pipelines/swiftvr/__init__.py new file mode 100644 index 00000000..ddc99ddc --- /dev/null +++ b/telefuser/pipelines/swiftvr/__init__.py @@ -0,0 +1,9 @@ +"""SwiftVR video-restoration pipeline.""" + +from .pipeline import SwiftVRPipeline, SwiftVRPipelineConfig, SwiftVRStreamSession + +__all__ = [ + "SwiftVRPipeline", + "SwiftVRPipelineConfig", + "SwiftVRStreamSession", +] diff --git a/telefuser/pipelines/swiftvr/chunk.py b/telefuser/pipelines/swiftvr/chunk.py new file mode 100644 index 00000000..b0dfd199 --- /dev/null +++ b/telefuser/pipelines/swiftvr/chunk.py @@ -0,0 +1,49 @@ +"""Fixed-size causal chunk protocol for streaming restoration. + +A clip of length ``t = 4a + 1`` is split into one FIRST chunk, zero or more +MIDDLE chunks and one LAST chunk so that the total number of input frames equals +the total number of output frames. ``clip_len`` is the MIDDLE chunk size and +must be a multiple of 4. +""" + +from dataclasses import dataclass +from enum import Enum + + +class ChunkType(Enum): + FIRST = "first" + MIDDLE = "middle" + LAST = "last" + + +@dataclass +class ChunkSpec: + ctype: ChunkType + frame_start: int + frame_count: int + b: int # LAST only: 4b + 1 input frames -> b + 1 latents + clip_idx: int + is_first_decode: bool # trim the decoder's causal-padding head frames + + +def build_chunk_specs(t: int, clip_len: int) -> list[ChunkSpec]: + assert clip_len % 4 == 0, f"clip_len must be a multiple of 4, got {clip_len}" + + if t <= clip_len + 4: + return [ChunkSpec(ChunkType.LAST, 0, t, (t - 1) // 4, 0, True)] + + specs = [ChunkSpec(ChunkType.FIRST, 0, clip_len + 4, 0, 0, True)] + + remaining = t - (clip_len + 4) + pos = clip_len + 4 + cidx = 1 + while remaining > 0: + if remaining <= clip_len: + specs.append(ChunkSpec(ChunkType.LAST, pos, remaining, (remaining - 1) // 4, cidx, False)) + break + specs.append(ChunkSpec(ChunkType.MIDDLE, pos, clip_len, 0, cidx, False)) + remaining -= clip_len + pos += clip_len + cidx += 1 + + return specs diff --git a/telefuser/pipelines/swiftvr/io.py b/telefuser/pipelines/swiftvr/io.py new file mode 100644 index 00000000..0ffbf98e --- /dev/null +++ b/telefuser/pipelines/swiftvr/io.py @@ -0,0 +1,56 @@ +"""Tensor preprocessing and output conversion for SwiftVR.""" + +from __future__ import annotations + +import numpy as np +import torch +import torch.nn.functional as F + +_INTERP_NEEDS_ALIGN = ("linear", "bilinear", "bicubic", "trilinear") + + +def preprocess_clip_uint8( + frames_uint8: torch.Tensor, + out_h: int, + out_w: int, + mode: str, + pad_h: int, + pad_w: int, + dtype: torch.dtype, +) -> torch.Tensor: + """Convert CUDA uint8 THWC frames to padded target-dtype NTCHW frames in [0, 1].""" + frames = frames_uint8.permute(0, 3, 1, 2).contiguous().to(dtype=dtype) + _, _, height, width = frames.shape + if (height, width) != (out_h, out_w): + if mode in _INTERP_NEEDS_ALIGN: + frames = F.interpolate(frames, size=(out_h, out_w), mode=mode, align_corners=False) + else: + frames = F.interpolate(frames, size=(out_h, out_w), mode=mode) + frames = frames / 255.0 + if pad_h > 0 or pad_w > 0: + frames = F.pad(frames, (0, pad_w, 0, pad_h), mode="constant", value=0) + return frames.unsqueeze(0) + + +def crop_spatial_padding_ntchw(video: torch.Tensor | None, pad_h: int = 0, pad_w: int = 0) -> torch.Tensor | None: + """Remove bottom/right spatial padding from an NTCHW tensor.""" + if video is None: + return None + if pad_h > 0: + video = video[:, :, :, :-pad_h, :] + if pad_w > 0: + video = video[:, :, :, :, :-pad_w] + return video + + +def ntchw_to_uint8_frames(video: torch.Tensor | None) -> np.ndarray | None: + """Convert [0, 1] NTCHW output to uint8 THWC frames on the host.""" + if video is None or video.numel() == 0 or video.shape[1] == 0: + return None + frames = (video[0].permute(0, 2, 3, 1).contiguous() * 255).clamp(0, 255).to(torch.uint8) + if frames.device.type != "cuda": + return frames.cpu().numpy() + cpu_frames = torch.empty_like(frames, device="cpu", pin_memory=True) + cpu_frames.copy_(frames, non_blocking=True) + torch.cuda.current_stream(frames.device).synchronize() + return cpu_frames.numpy() diff --git a/telefuser/pipelines/swiftvr/pipeline.py b/telefuser/pipelines/swiftvr/pipeline.py new file mode 100644 index 00000000..32a14b7d --- /dev/null +++ b/telefuser/pipelines/swiftvr/pipeline.py @@ -0,0 +1,369 @@ +"""Sequential SwiftVR restoration pipeline. + +The implementation follows H-oliday/SwiftVR commit +5ca168cef6ca7200f135fdfea85e5e13d12c5b53. Model execution remains sequential +to preserve the approved differential path. +""" + +from __future__ import annotations + +import json +import threading +from dataclasses import dataclass, field +from pathlib import Path +from typing import Iterable, Iterator + +import torch +from safetensors.torch import load_file + +from telefuser.core.base_pipeline import BasePipeline +from telefuser.core.config import AttentionConfig, AttnImplType, ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.models.swiftvr_reae import ReAE +from telefuser.models.swiftvr_transformer import SwiftVRWanTransformer3DModel + +from .chunk import ChunkSpec, ChunkType, build_chunk_specs +from .io import ( + crop_spatial_padding_ntchw, + preprocess_clip_uint8, +) +from .streaming_dit import StreamingDiT +from .streaming_tae import StreamingTAE + +_DTYPES = { + "float16": torch.float16, + "fp16": torch.float16, + "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, + "float32": torch.float32, + "fp32": torch.float32, +} + + +def _as_dtype(dtype: torch.dtype | str) -> torch.dtype: + if isinstance(dtype, torch.dtype): + return dtype + key = str(dtype).lower() + if key not in _DTYPES: + raise ValueError(f"Unsupported dtype {dtype!r}. Choose float16, bfloat16, or float32.") + return _DTYPES[key] + + +def aligned_pad(size: int, multiple: int = 32) -> int: + """Return right/bottom padding needed to align a spatial dimension.""" + if size <= 0: + raise ValueError(f"size must be positive, got {size}") + if multiple <= 0: + raise ValueError(f"multiple must be positive, got {multiple}") + return (multiple - size % multiple) % multiple + + +@dataclass +class SwiftVRPipelineConfig: + """Runtime configuration using existing TeleFuser model controls.""" + + dit_config: ModelRuntimeConfig = field( + default_factory=lambda: ModelRuntimeConfig( + attention_config=AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA) + ) + ) + + +class SwiftVRPipeline(BasePipeline): + """Faithful single-device SwiftVR video-restoration pipeline.""" + + upscale_mode = "bilinear" + + def __init__(self, device: str | torch.device, torch_dtype: torch.dtype = torch.bfloat16) -> None: + super().__init__(device=device, torch_dtype=torch_dtype) + self.device = torch.device(device) + self._execution_lock = threading.RLock() + + def init( + self, + module_manager: ModuleManager, + config: SwiftVRPipelineConfig, + prompt_emb: torch.Tensor, + ) -> None: + self._model_info = module_manager.get_model_info() + self.config = config + self.reae = module_manager.fetch_module("swiftvr_reae") + self.transformer = module_manager.fetch_module("swiftvr_transformer") + if self.reae is None or self.transformer is None: + raise RuntimeError("SwiftVR requires both swiftvr_reae and swiftvr_transformer") + self.prompt_emb = prompt_emb.to(dtype=self.torch_dtype) + self._prepare_for_inference() + + def _prepare_for_inference(self) -> None: + self.reae.to(device=self.device, dtype=self.torch_dtype).eval() + self.transformer.to(device=self.device, dtype=self.torch_dtype).eval() + if self.device.type == "cuda": + # Match the released SwiftVR runtime before selecting convolution + # and attention kernels; these flags are part of the frozen parity + # configuration used by the upstream pipeline. + torch.backends.cudnn.benchmark = True + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + self.transformer.prepare_for_inference(attention_config=self.config.dit_config.attention_config) + + @classmethod + def from_pretrained( + cls, + model_id_or_path: str, + *, + device: str | torch.device = "cuda", + torch_dtype: torch.dtype | str = torch.bfloat16, + attention_config: AttentionConfig | None = None, + ) -> "SwiftVRPipeline": + """Load the released local checkpoint through ModuleManager.""" + root = Path(model_id_or_path) + dtype = _as_dtype(torch_dtype) + transformer_dir = root / "transformer" + config_path = transformer_dir / "config.json" + with config_path.open(encoding="utf-8") as handle: + transformer_config = json.load(handle) + + # Upstream constructs both models on CPU, then moves ReAE followed by + # the transformer during pipeline preparation. Keep that allocation + # order because cuDNN benchmarking can otherwise select a numerically + # different convolution plan. + module_manager = ModuleManager(torch_dtype=dtype, device="cpu") + module_manager.load_model( + str(root / "reae.safetensors"), + device="cpu", + torch_dtype=dtype, + low_cpu_mem_usage=True, + name="swiftvr_reae", + model_class=ReAE, + model_resource="official", + ) + module_manager.load_model( + str(transformer_dir), + device="cpu", + torch_dtype=dtype, + low_cpu_mem_usage=True, + name="swiftvr_transformer", + model_class=SwiftVRWanTransformer3DModel, + model_resource="diffusers", + converter_kwargs={"config": transformer_config}, + ) + prompt_payload = load_file(str(root / "prompt_embedding.safetensors"), device="cpu") + prompt_emb = prompt_payload["prompt_emb"][0] + + pipeline_config = SwiftVRPipelineConfig() + if attention_config is not None: + pipeline_config.dit_config.attention_config = attention_config + pipeline = cls(device=device, torch_dtype=dtype) + pipeline.init(module_manager, pipeline_config, prompt_emb) + return pipeline + + @staticmethod + def _validate_clip_len(clip_len: int) -> None: + if clip_len <= 0 or clip_len % 4: + raise ValueError(f"clip_len must be a positive multiple of 4, got {clip_len}") + + def _target_size( + self, + lq_h: int, + lq_w: int, + resolution: tuple[int, int] | None, + upscale: int, + ) -> tuple[int, int, int, int]: + if resolution is not None: + out_w, out_h = int(resolution[0]), int(resolution[1]) + else: + if upscale <= 0: + raise ValueError(f"upscale must be positive, got {upscale}") + out_h, out_w = lq_h * upscale, lq_w * upscale + return out_h, out_w, aligned_pad(out_h), aligned_pad(out_w) + + def _restored_chunks( + self, + chunks: Iterable[tuple[ChunkSpec, torch.Tensor]], + *, + clip_len: int, + out_h: int, + out_w: int, + pad_h: int, + pad_w: int, + dit_overlap: int, + ) -> Iterator[torch.Tensor]: + tae_stream = StreamingTAE(self.reae) + dit_stream = StreamingDiT(self.transformer, overlap=dit_overlap) + n_lat = clip_len // 4 + prev_dit_out_cpu = None + + with self._execution_lock: + for spec, frames_uint8 in chunks: + gpu_frames = frames_uint8.to(device=self.device) + clip = preprocess_clip_uint8( + gpu_frames, + out_h, + out_w, + self.upscale_mode, + pad_h, + pad_w, + self.torch_dtype, + ) + encoded = tae_stream.encode_chunk_fixed(clip, spec) + if spec.ctype == ChunkType.LAST: + denoised = dit_stream.denoise_last_chunk( + encoded, + spec, + self.prompt_emb, + prev_dit_out_cpu, + n_lat, + self.device, + self.torch_dtype, + ) + else: + encoded_bcfhw = encoded.permute(0, 2, 1, 3, 4).contiguous() + denoised_bcfhw = dit_stream.denoise(encoded_bcfhw, self.prompt_emb) + denoised = denoised_bcfhw.permute(0, 2, 1, 3, 4).contiguous() + prev_dit_out_cpu = encoded_bcfhw[:, :, -n_lat:].detach().cpu().clone() + decoded = tae_stream.decode_chunk_fixed(denoised, spec) + if decoded is not None and decoded.shape[1]: + yield crop_spatial_padding_ntchw(decoded, pad_h, pad_w) + + @torch.inference_mode() + def __call__( + self, + frames_uint8: torch.Tensor, + *, + resolution: tuple[int, int] | None = None, + upscale: int = 4, + clip_len: int = 24, + dit_overlap: int = 0, + ) -> torch.Tensor: + """Restore [T,H,W,3] uint8 frames and return [1,T,3,H,W].""" + self._validate_clip_len(clip_len) + if frames_uint8.ndim != 4 or frames_uint8.shape[-1] != 3 or frames_uint8.dtype != torch.uint8: + raise ValueError("frames_uint8 must have shape [T,H,W,3] and dtype uint8") + total_frames = 4 * ((int(frames_uint8.shape[0]) - 1) // 4) + 1 + if total_frames <= 0: + raise ValueError("frames_uint8 must contain at least one frame") + frames_uint8 = frames_uint8[:total_frames] + out_h, out_w, pad_h, pad_w = self._target_size( + int(frames_uint8.shape[1]), + int(frames_uint8.shape[2]), + resolution, + upscale, + ) + chunks = ( + ( + spec, + frames_uint8[spec.frame_start : spec.frame_start + spec.frame_count], + ) + for spec in build_chunk_specs(total_frames, clip_len) + ) + outputs = list( + self._restored_chunks( + chunks, + clip_len=clip_len, + out_h=out_h, + out_w=out_w, + pad_h=pad_h, + pad_w=pad_w, + dit_overlap=dit_overlap, + ) + ) + if not outputs: + return torch.empty((1, 0, 3, out_h, out_w), dtype=self.torch_dtype, device=self.device) + return torch.cat(outputs, dim=1) + + def stream( + self, + *, + clip_len: int = 24, + resolution: tuple[int, int] | None = None, + upscale: int = 4, + dit_overlap: int = 1, + ) -> "SwiftVRStreamSession": + self._validate_clip_len(clip_len) + return SwiftVRStreamSession( + self, + clip_len=clip_len, + resolution=resolution, + upscale=upscale, + dit_overlap=dit_overlap, + ) + + +class SwiftVRStreamSession: + """Per-session ReAE, overlap, and RoPE state for causal restoration.""" + + def __init__( + self, + pipeline: SwiftVRPipeline, + *, + clip_len: int, + resolution: tuple[int, int] | None, + upscale: int, + dit_overlap: int, + ) -> None: + self.pipeline = pipeline + self.clip_len = clip_len + self.resolution = resolution + self.upscale = upscale + self._sizes: tuple[int, int, int, int] | None = None + self._tae = StreamingTAE(pipeline.reae) + self._dit = StreamingDiT(pipeline.transformer, overlap=dit_overlap) + self._closed = False + + def _ensure_open(self) -> None: + if self._closed: + raise RuntimeError("SwiftVR stream session is closed") + + def _ensure_sizes(self, lq_h: int, lq_w: int) -> tuple[int, int, int, int]: + if self._sizes is None: + self._sizes = self.pipeline._target_size(lq_h, lq_w, self.resolution, self.upscale) + return self._sizes + + def _run_latents(self, encoded: torch.Tensor) -> torch.Tensor: + encoded_bcfhw = encoded.permute(0, 2, 1, 3, 4).contiguous() + denoised = self._dit.denoise(encoded_bcfhw, self.pipeline.prompt_emb) + return denoised.permute(0, 2, 1, 3, 4).contiguous() + + @torch.inference_mode() + def step(self, frames_uint8: torch.Tensor) -> torch.Tensor | None: + self._ensure_open() + if frames_uint8.ndim != 4 or frames_uint8.shape[-1] != 3 or frames_uint8.dtype != torch.uint8: + raise ValueError("frames_uint8 must have shape [T,H,W,3] and dtype uint8") + with self.pipeline._execution_lock: + frames = frames_uint8.to(self.pipeline.device) + out_h, out_w, pad_h, pad_w = self._ensure_sizes(int(frames.shape[1]), int(frames.shape[2])) + clip = preprocess_clip_uint8( + frames, + out_h, + out_w, + self.pipeline.upscale_mode, + pad_h, + pad_w, + self.pipeline.torch_dtype, + ) + encoded = self._tae.encode_chunk(clip) + if encoded is None: + return None + decoded = self._tae.decode_chunk(self._run_latents(encoded)) + return crop_spatial_padding_ntchw(decoded, pad_h, pad_w) + + @torch.inference_mode() + def flush(self) -> torch.Tensor | None: + self._ensure_open() + with self.pipeline._execution_lock: + encoded = self._tae.flush_encoder() + if encoded is None or self._sizes is None: + return None + _, _, pad_h, pad_w = self._sizes + decoded = self._tae.decode_chunk(self._run_latents(encoded)) + return crop_spatial_padding_ntchw(decoded, pad_h, pad_w) + + def close(self) -> None: + if self._closed: + return + self._tae.reset() + self._dit.reset() + self._dit._cond_cache = None + self._dit._cond_cache_key = None + self._sizes = None + self._closed = True diff --git a/telefuser/pipelines/swiftvr/streaming_dit.py b/telefuser/pipelines/swiftvr/streaming_dit.py new file mode 100644 index 00000000..540f8934 --- /dev/null +++ b/telefuser/pipelines/swiftvr/streaming_dit.py @@ -0,0 +1,207 @@ +"""Streaming one-step DiT for SwiftVR. + +SwiftVR collapses iterative diffusion sampling to a single forward pass taken at the +fully-degraded endpoint of the flow (t = 1). No sampling scheduler is therefore +needed at inference time: the conditioning timestep is the constant +``INFERENCE_TIMESTEP`` below, equivalent to a flow-matching schedule evaluated by +``set_timesteps(1)`` (whose ``scale_model_input`` is the identity). Adjust it if +your training schedule uses a different number of train timesteps. +""" + +import torch + +from telefuser.models.swiftvr_transformer import SwiftVRWanTransformer3DModel + +from .chunk import ChunkSpec + +INFERENCE_TIMESTEP = 1000.0 +ROPE_EXTEND_MARGIN = 256 + + +def _ensure_rope_cache_len(rope, required_len): + cos, sin = rope.freqs_cos, rope.freqs_sin + old_len = cos.shape[0] + if required_len <= old_len: + return + if old_len < 2: + raise RuntimeError(f"Cannot extend RoPE cache: length={old_len}") + with torch.no_grad(): + dev, old_dtype = cos.device, cos.dtype + cos64, sin64 = cos.to(torch.float64), sin.to(torch.float64) + cos0, sin0, cos1, sin1 = cos64[0:1], sin64[0:1], cos64[1:2], sin64[1:2] + cos_delta = cos1 * cos0 + sin1 * sin0 + sin_delta = sin1 * cos0 - cos1 * sin0 + angle0 = torch.atan2(sin0, cos0) + step = torch.atan2(sin_delta, cos_delta) + pos = torch.arange(old_len, required_len, device=dev, dtype=torch.float64).view(-1, 1) + angle = angle0 + pos * step + rope.freqs_cos = torch.cat([cos, torch.cos(angle).to(old_dtype)], 0).contiguous() + rope.freqs_sin = torch.cat([sin, torch.sin(angle).to(old_dtype)], 0).contiguous() + + +def _rope_with_offset(rope, ppf, pph, ppw, t_off=0, h_off=0, w_off=0): + required_len = max(t_off + ppf, h_off + pph, w_off + ppw) + if required_len > rope.freqs_cos.shape[0]: + required_len += max(0, int(ROPE_EXTEND_MARGIN)) + _ensure_rope_cache_len(rope, required_len) + sp = [rope.t_dim, rope.h_dim, rope.w_dim] + fc = rope.freqs_cos.split(sp, dim=1) + fs = rope.freqs_sin.split(sp, dim=1) + cf = fc[0][t_off : t_off + ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + ch = fc[1][h_off : h_off + pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + cw = fc[2][w_off : w_off + ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + sf = fs[0][t_off : t_off + ppf].view(ppf, 1, 1, -1).expand(ppf, pph, ppw, -1) + sh = fs[1][h_off : h_off + pph].view(1, pph, 1, -1).expand(ppf, pph, ppw, -1) + sw = fs[2][w_off : w_off + ppw].view(1, 1, ppw, -1).expand(ppf, pph, ppw, -1) + return ( + torch.cat([cf, ch, cw], -1).reshape(1, ppf * pph * ppw, 1, -1), + torch.cat([sf, sh, sw], -1).reshape(1, ppf * pph * ppw, 1, -1), + ) + + +def _precompute_cond(transformer, B, prompt_emb, timestep): + pe = prompt_emb.clone() + if pe.ndim == 2: + pe = pe.unsqueeze(0).expand(B, -1, -1) + elif pe.shape[0] != B: + pe = pe.expand(B, -1, -1) + ts = timestep.clone() + ts_seq = None + if ts.ndim == 2: + ts_seq = ts.shape[1] + ts = ts.flatten() + temb, tp, enc_hs, enc_img = transformer.condition_embedder(ts, pe, None, timestep_seq_len=ts_seq) + tp = tp.unflatten(2 if ts_seq else 1, (6, -1)) + if enc_img is not None: + enc_hs = torch.cat([enc_img, enc_hs], dim=1) + return temb, tp, enc_hs + + +@torch.inference_mode() +def _dit_forward_chunk(transformer, chunk, temb, tp, enc_hs, t_off=0): + """One forward pass of the DiT, returning the predicted degradation velocity.""" + p_t, p_h, p_w = transformer.config.patch_size + B, C, F, H, W = chunk.shape + ppf, pph, ppw = F // p_t, H // p_h, W // p_w + rope = _rope_with_offset(transformer.rope, ppf, pph, ppw, t_off) + hs = transformer.patch_embedding(chunk).flatten(2).transpose(1, 2) + + thw_global = (ppf, pph, ppw) + for i, blk in enumerate(transformer.blocks): + underlying = getattr(blk, "_orig_mod", blk) + if hasattr(underlying, "attn1"): + underlying.attn1._thw = thw_global + underlying.attn1._layer_id = i + for blk in transformer.blocks: + hs = blk(hs, enc_hs, tp, rope) + + if temb.ndim == 3: + shift, scale = (transformer.scale_shift_table.unsqueeze(0).to(temb.device) + temb.unsqueeze(2)).chunk(2, dim=2) + shift, scale = shift.squeeze(2), scale.squeeze(2) + else: + shift, scale = (transformer.scale_shift_table.to(temb.device) + temb.unsqueeze(1)).chunk(2, dim=1) + hs = (transformer.norm_out(hs.float()) * (1 + scale.to(hs.device)) + shift.to(hs.device)).type_as(hs) + hs = transformer.proj_out(hs) + hs = hs.reshape(B, ppf, pph, ppw, p_t, p_h, p_w, -1).permute(0, 7, 1, 4, 2, 5, 3, 6) + return hs.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + +class StreamingDiT: + """One-step DiT with temporal overlap blending across chunks.""" + + def __init__(self, transformer: SwiftVRWanTransformer3DModel, overlap: int = 0) -> None: + self.transformer = transformer + self.overlap = overlap + self._prev_lq = self._prev_out = None + self._g_off = 0 + self._cond_cache_key = self._cond_cache = None + + def reset(self) -> None: + self._prev_lq = self._prev_out = None + self._g_off = 0 + + def _get_cached_condition(self, B, prompt_emb, dev, dt): + cache_key = (int(B), dev.type, dev.index, str(dt), tuple(prompt_emb.shape)) + if self._cond_cache_key != cache_key or self._cond_cache is None: + ts = torch.full((B,), INFERENCE_TIMESTEP, device=dev, dtype=torch.float32) + self._cond_cache = _precompute_cond(self.transformer, B, prompt_emb.to(dev, dt), ts) + self._cond_cache_key = cache_key + return self._cond_cache + + @torch.inference_mode() + def denoise(self, lq: torch.Tensor, prompt_emb: torch.Tensor) -> torch.Tensor: + """Restore one chunk of latents ``lq`` (shape [B, C, F, H, W]).""" + dev, dt = lq.device, lq.dtype + B, C, F_cur, H, W = lq.shape + + ol = 0 + if self._prev_lq is not None and self.overlap > 0: + ol = self._prev_lq.shape[2] + lq_ext = torch.cat([self._prev_lq.to(dev), lq], dim=2) + t_rope = self._g_off - ol + else: + lq_ext = lq + t_rope = self._g_off + + temb, tp, enc_hs = self._get_cached_condition(B, prompt_emb, dev, dt) + pred = _dit_forward_chunk(self.transformer, lq_ext, temb, tp, enc_hs, t_off=t_rope) + den_ext = lq_ext - pred + + if ol > 0 and self._prev_out is not None: + ramp = torch.linspace(0, 1, ol, device=dev, dtype=dt).view(1, 1, ol, 1, 1) + den_ext[:, :, :ol] = self._prev_out.to(dev) * (1 - ramp) + den_ext[:, :, :ol] * ramp + den_out = den_ext[:, :, ol:] + else: + den_out = den_ext + + k = min(self.overlap, F_cur) + if k > 0: + self._prev_lq = lq[:, :, -k:].detach().cpu().clone() + self._prev_out = den_out[:, :, -k:].detach().cpu().clone() + else: + self._prev_lq = self._prev_out = None + + self._g_off += F_cur + return den_out + + @torch.inference_mode() + def denoise_last_chunk( + self, + z_new_ntchw: torch.Tensor, + spec: ChunkSpec, + prompt_emb: torch.Tensor, + prev_dit_out_cpu: torch.Tensor | None, + n_lat: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + """LAST chunk: pad the (b+1) new latents up to (n_lat+1) with the previous + chunk's latents (or zeros) for a correct RoPE offset, run one pass, and + keep only the new ``b+1`` denoised latents. + """ + lat_count = spec.b + 1 + pad_count = (n_lat + 1) - lat_count + + z_bcfhw = z_new_ntchw.permute(0, 2, 1, 3, 4).contiguous() + if pad_count > 0: + if prev_dit_out_cpu is not None: + pad_z = prev_dit_out_cpu[:, :, -pad_count:].to(device=device, dtype=dtype) + else: + pad_z = torch.zeros( + z_bcfhw.shape[0], + z_bcfhw.shape[1], + pad_count, + z_bcfhw.shape[3], + z_bcfhw.shape[4], + device=device, + dtype=dtype, + ) + z_bcfhw = torch.cat([pad_z, z_bcfhw], dim=2) + + t_off = max(0, self._g_off - pad_count) + temb, tp, enc_hs = self._get_cached_condition(z_bcfhw.shape[0], prompt_emb, device, dtype) + pred = _dit_forward_chunk(self.transformer, z_bcfhw, temb, tp, enc_hs, t_off=t_off) + z_den = (z_bcfhw - pred)[:, :, -lat_count:].contiguous() + + self._g_off += lat_count + return z_den.permute(0, 2, 1, 3, 4).contiguous() diff --git a/telefuser/pipelines/swiftvr/streaming_tae.py b/telefuser/pipelines/swiftvr/streaming_tae.py new file mode 100644 index 00000000..703946df --- /dev/null +++ b/telefuser/pipelines/swiftvr/streaming_tae.py @@ -0,0 +1,165 @@ +"""Streaming wrapper around the Restoration-aware Autoencoder. + +It runs the encoder/decoder clip-by-clip while passing the MemBlock and TPool +boundary state across chunks, so the result is identical to encoding/decoding +the whole clip at once. +""" + +import torch +import torch.nn.functional as F + +from telefuser.models.swiftvr_reae import MemBlock, ReAE, TGrow, TPool + +from .chunk import ChunkSpec, ChunkType + + +def apply_parallel_with_boundary( + model: torch.nn.Sequential, + x: torch.Tensor, + state: dict[str, torch.Tensor | None] | None = None, +) -> tuple[torch.Tensor | None, dict[str, torch.Tensor | None]]: + """Run ``model`` (a Sequential of streaming blocks) over ``x``. + + ``x`` has shape ``[N, T, C, H, W]``. ``state`` carries the MemBlock/TPool + boundary buffers from the previous chunk; the updated state is returned. + """ + if state is None: + state = {} + new_state = {} + N, T, C, H, W = x.shape + x = x.reshape(N * T, C, H, W) + + for i, b in enumerate(model): + if isinstance(b, MemBlock): + NT, C, H, W = x.shape + T_ = NT // N + _x = x.reshape(N, T_, C, H, W) + key = f"mem_{i}" + if key in state: + mem = torch.cat([state[key], _x[:, :-1]], dim=1) + else: + mem = F.pad(_x, (0, 0, 0, 0, 0, 0, 1, 0), value=0)[:, :T_] + new_state[key] = _x[:, -1:].detach().clone() + x = b(x, mem.reshape(NT, C, H, W)) + + elif isinstance(b, TPool): + NT, C, H, W = x.shape + T_ = NT // N + _x = x.reshape(N, T_, C, H, W) + key = f"tpool_{i}" + if key in state and state[key] is not None: + _x = torch.cat([state[key], _x], dim=1) + T_ = _x.shape[1] + n_full = (T_ // b.stride) * b.stride + rem = T_ - n_full + new_state[key] = _x[:, n_full:].detach().clone() if rem > 0 else None + if n_full > 0: + x = b(_x[:, :n_full].reshape(N * n_full, C, H, W)) + else: + return None, new_state + + elif isinstance(b, TGrow): + x = b(x) + else: + x = b(x) + + NT, C, H, W = x.shape + return x.view(N, NT // N, C, H, W), new_state + + +class StreamingTAE: + def __init__(self, model: ReAE) -> None: + self.model = model + self._enc_st = None + self._dec_st = None + self._enc_left = None + self._first_dec = True + + def reset(self) -> None: + self._enc_st = self._dec_st = None + self._enc_left = None + self._first_dec = True + + # ----- Fixed-size chunk interface (offline, frame-count preserving) ----- # + + @torch.no_grad() + def encode_chunk_fixed(self, x: torch.Tensor, spec: ChunkSpec) -> torch.Tensor: + ps = self.model.patch_size + if ps > 1: + N, T, C, H, W = x.shape + x = F.pixel_unshuffle(x.reshape(N * T, C, H, W), ps) + x = x.reshape(N, T, *x.shape[1:]) + + if spec.ctype == ChunkType.LAST: + x = torch.cat([x, x[:, -1:].expand(-1, 3, -1, -1, -1)], dim=1) + + z, self._enc_st = apply_parallel_with_boundary(self.model.encoder, x, self._enc_st) + return z + + @torch.no_grad() + def decode_chunk_fixed(self, z: torch.Tensor, spec: ChunkSpec) -> torch.Tensor | None: + x, self._dec_st = apply_parallel_with_boundary(self.model.decoder, z, self._dec_st) + if x is None: + return None + x = self._postprocess(x) + if spec.is_first_decode: + x = x[:, self.model.frames_to_trim :] + return x + + # ----- Generic streaming interface (online, arbitrary chunk lengths) ---- # + + @torch.no_grad() + def encode_chunk(self, x: torch.Tensor) -> torch.Tensor | None: + ps = self.model.patch_size + if ps > 1: + N, T, C, H, W = x.shape + x = F.pixel_unshuffle(x.reshape(N * T, C, H, W), ps) + x = x.reshape(N, T, *x.shape[1:]) + if self._enc_left is not None: + x = torch.cat([self._enc_left, x], dim=1) + self._enc_left = None + T = x.shape[1] + rem = T % 4 + if rem: + keep = T - rem + if keep > 0: + self._enc_left = x[:, keep:].detach().clone() + x = x[:, :keep] + else: + self._enc_left = x.detach().clone() + return None + z, self._enc_st = apply_parallel_with_boundary(self.model.encoder, x, self._enc_st) + return z + + @torch.no_grad() + def flush_encoder(self) -> torch.Tensor | None: + if self._enc_left is None: + return None + x = self._enc_left + self._enc_left = None + T = x.shape[1] + if T % 4: + p = 4 - T % 4 + x = torch.cat([x, x[:, -1:].expand(-1, p, -1, -1, -1)], dim=1) + z, self._enc_st = apply_parallel_with_boundary(self.model.encoder, x, self._enc_st) + return z + + @torch.no_grad() + def decode_chunk(self, z: torch.Tensor) -> torch.Tensor | None: + x, self._dec_st = apply_parallel_with_boundary(self.model.decoder, z, self._dec_st) + if x is None: + return None + x = self._postprocess(x) + if self._first_dec: + x = x[:, self.model.frames_to_trim :] + self._first_dec = False + return x + + def _postprocess(self, x: torch.Tensor) -> torch.Tensor: + x = torch.clamp(x, 0, 1) + ps = self.model.patch_size + if ps > 1: + N, T, C, H, W = x.shape + x = F.pixel_shuffle(x.reshape(N * T, C, H, W), ps) + x = x.reshape(N, T, *x.shape[1:]) + return x diff --git a/tests/unit/models/test_swiftvr_transformer.py b/tests/unit/models/test_swiftvr_transformer.py new file mode 100644 index 00000000..c815435f --- /dev/null +++ b/tests/unit/models/test_swiftvr_transformer.py @@ -0,0 +1,72 @@ +from types import SimpleNamespace + +import torch + +from telefuser.models.swiftvr_transformer import ( + _WindowRuntimeMetaCache, + _make_hw_starts, + get_1d_rotary_pos_embed, +) +from telefuser.pipelines.swiftvr.streaming_dit import _ensure_rope_cache_len, _rope_with_offset + + +def _rope(length: int = 4) -> SimpleNamespace: + parts = [get_1d_rotary_pos_embed(4, length) for _ in range(3)] + return SimpleNamespace( + t_dim=4, + h_dim=4, + w_dim=4, + freqs_cos=torch.cat([part[0] for part in parts], dim=1), + freqs_sin=torch.cat([part[1] for part in parts], dim=1), + ) + + +def test_window_starts_cover_boundary_without_redundant_interior() -> None: + h_starts, w_starts = _make_hw_starts(11, 10, 4, 4, False) + + assert h_starts.tolist() == [0, 4, 7] + assert w_starts.tolist() == [0, 4, 6] + + shifted_h, shifted_w = _make_hw_starts(11, 10, 4, 4, True) + assert shifted_h.tolist() == [0, 2, 6, 7] + assert shifted_w.tolist() == [0, 2, 6] + + +def test_shifted_window_owner_scatter_returns_each_global_token_once() -> None: + for shifted, prefer_front in ((False, True), (True, False)): + meta = _WindowRuntimeMetaCache.get( + 2, + 5, + 6, + 4, + 4, + do_shift=shifted, + prefer_front=prefer_front, + device=torch.device("cpu"), + ) + gathered_global_indices = meta.lin_flat + restored = torch.index_select(gathered_global_indices, 0, meta.owner_pos) + assert torch.equal(restored, torch.arange(meta.THW)) + + +def test_rope_extension_preserves_existing_values() -> None: + rope = _rope() + original_cos = rope.freqs_cos.clone() + original_sin = rope.freqs_sin.clone() + + _ensure_rope_cache_len(rope, 12) + + assert rope.freqs_cos.shape == (12, 12) + assert rope.freqs_sin.shape == (12, 12) + torch.testing.assert_close(rope.freqs_cos[:4], original_cos, rtol=0, atol=0) + torch.testing.assert_close(rope.freqs_sin[:4], original_sin, rtol=0, atol=0) + + +def test_rope_offset_uses_global_temporal_position() -> None: + rope = _rope() + cos, sin = _rope_with_offset(rope, 3, 2, 2, t_off=5) + cos_grid = cos.view(3, 2, 2, 12) + sin_grid = sin.view(3, 2, 2, 12) + + torch.testing.assert_close(cos_grid[:, 0, 0, :4], rope.freqs_cos[5:8, :4]) + torch.testing.assert_close(sin_grid[:, 0, 0, :4], rope.freqs_sin[5:8, :4]) diff --git a/tests/unit/pipelines/swiftvr/test_example.py b/tests/unit/pipelines/swiftvr/test_example.py new file mode 100644 index 00000000..a9509c5f --- /dev/null +++ b/tests/unit/pipelines/swiftvr/test_example.py @@ -0,0 +1,95 @@ +from pathlib import Path + +import numpy as np +import pytest +import torch +from PIL import Image + +from examples.swiftvr.swiftvr_restore_h100 import _warmup_frame_count, main, run + + +class FakePipeline: + def __init__(self) -> None: + self.frames: torch.Tensor | None = None + self.options: dict[str, object] = {} + self.step_sizes: list[int] = [] + self.closed = False + + def stream(self, **options: object) -> "FakePipeline": + self.options = options + return self + + def step(self, frames: torch.Tensor) -> torch.Tensor: + self.frames = frames + self.step_sizes.append(len(frames)) + return frames.permute(0, 3, 1, 2).unsqueeze(0).float() / 255 + + def flush(self) -> None: + return None + + def close(self) -> None: + self.closed = True + + +def test_run_uses_flashvsr_style_pil_interface() -> None: + pipeline = FakePipeline() + inputs = [Image.fromarray(np.full((10, 18, 3), value, dtype=np.uint8)) for value in (17, 193)] + + outputs = run( + pipeline, + inputs, + scale=2, + ) + + assert pipeline.frames is not None + assert pipeline.frames.shape == (2, 8, 16, 3) + assert pipeline.frames.dtype == torch.uint8 + assert pipeline.options == {"upscale": 2} + assert pipeline.step_sizes == [2] + assert pipeline.closed is True + assert [image.size for image in outputs] == [(16, 8), (16, 8)] + assert np.asarray(outputs[0])[0, 0].tolist() == [17, 17, 17] + assert np.asarray(outputs[1])[0, 0].tolist() == [193, 193, 193] + + +def test_run_rejects_empty_video() -> None: + with pytest.raises(ValueError, match="at least one frame"): + run(FakePipeline(), []) + + +def test_run_processes_24_frame_streaming_chunks() -> None: + pipeline = FakePipeline() + inputs = [Image.fromarray(np.full((8, 8, 3), value, dtype=np.uint8)) for value in range(25)] + + outputs = run(pipeline, inputs, scale=3) + + assert pipeline.step_sizes == [24, 1] + assert pipeline.options == {"upscale": 3} + assert len(outputs) == 25 + + +def test_warmup_covers_full_and_tail_shapes() -> None: + assert _warmup_frame_count(81) == 57 + assert _warmup_frame_count(72) == 48 + assert _warmup_frame_count(25) == 25 + + +def test_default_input_is_flashvsr_example_video() -> None: + input_option = next(parameter for parameter in main.params if parameter.name == "input_video") + input_video = Path(input_option.default) + + assert input_video.name == "dag.mp4" + assert input_video.parent.name == "data" + assert input_video.is_file() + + +def test_cli_uses_flashvsr_style_options() -> None: + assert {parameter.name for parameter in main.params} == { + "input_video", + "scale", + "height", + "width", + "gpu_num", + "model_root", + "output", + } diff --git a/tests/unit/pipelines/swiftvr/test_streaming.py b/tests/unit/pipelines/swiftvr/test_streaming.py new file mode 100644 index 00000000..77abcd8a --- /dev/null +++ b/tests/unit/pipelines/swiftvr/test_streaming.py @@ -0,0 +1,157 @@ +import threading +from types import SimpleNamespace + +import torch +from torch import nn + +from telefuser.models.swiftvr_reae import MemBlock, TPool +from telefuser.pipelines.swiftvr import pipeline as pipeline_module +from telefuser.pipelines.swiftvr.chunk import ChunkType, build_chunk_specs +from telefuser.pipelines.swiftvr.pipeline import SwiftVRStreamSession, aligned_pad +from telefuser.pipelines.swiftvr.streaming_tae import apply_parallel_with_boundary + + +def test_aligned_pad() -> None: + assert aligned_pad(32) == 0 + assert aligned_pad(33) == 31 + assert aligned_pad(1080) == 8 + assert aligned_pad(1440) == 0 + + +def test_chunk_specs_preserve_frame_count_and_tail() -> None: + specs = build_chunk_specs(53, 24) + + assert [spec.ctype for spec in specs] == [ChunkType.FIRST, ChunkType.MIDDLE, ChunkType.LAST] + assert sum(spec.frame_count for spec in specs) == 53 + assert specs[-1].frame_count == 1 + assert specs[-1].b == 0 + assert specs[0].is_first_decode is True + + +def test_memblock_boundary_matches_whole_sequence() -> None: + torch.manual_seed(7) + model = nn.Sequential(MemBlock(2, 2)).eval() + inputs = torch.randn(1, 5, 2, 4, 4) + + whole, _ = apply_parallel_with_boundary(model, inputs) + first, state = apply_parallel_with_boundary(model, inputs[:, :2]) + second, _ = apply_parallel_with_boundary(model, inputs[:, 2:], state) + + torch.testing.assert_close(torch.cat([first, second], dim=1), whole) + + +def test_temporal_pool_boundary_carries_non_aligned_tail() -> None: + torch.manual_seed(11) + model = nn.Sequential(TPool(2, 2)).eval() + inputs = torch.randn(1, 6, 2, 3, 3) + + whole, _ = apply_parallel_with_boundary(model, inputs) + first, state = apply_parallel_with_boundary(model, inputs[:, :3]) + second, state = apply_parallel_with_boundary(model, inputs[:, 3:], state) + + assert state["tpool_0"] is None + torch.testing.assert_close(torch.cat([first, second], dim=1), whole) + + +def test_stream_sessions_do_not_share_boundary_or_rope_state() -> None: + pipeline = SimpleNamespace( + reae=object(), + transformer=object(), + _execution_lock=threading.RLock(), + ) + first = SwiftVRStreamSession( + pipeline, + clip_len=24, + resolution=None, + upscale=4, + dit_overlap=1, + ) + second = SwiftVRStreamSession( + pipeline, + clip_len=24, + resolution=None, + upscale=4, + dit_overlap=1, + ) + + first._tae._enc_st = {"mem_1": torch.ones(1)} + first._dit._g_off = 9 + + assert second._tae._enc_st is None + assert second._dit._g_off == 0 + first.close() + assert first._tae._enc_st is None + assert first._dit._g_off == 0 + assert second._closed is False + + +def test_interleaved_stream_sessions_do_not_exchange_frames(monkeypatch) -> None: + class FakeTAE: + def __init__(self, _model: object) -> None: + self.calls = 0 + + def encode_chunk(self, frames: torch.Tensor) -> torch.Tensor: + self.calls += 1 + return frames + self.calls / 100 + + def decode_chunk(self, latents: torch.Tensor) -> torch.Tensor: + return latents + + def flush_encoder(self) -> None: + return None + + def reset(self) -> None: + self.calls = 0 + + class FakeDiT: + def __init__(self, _model: object, overlap: int) -> None: + self.calls = 0 + self.overlap = overlap + self._cond_cache = None + self._cond_cache_key = None + + def denoise(self, latents: torch.Tensor, _prompt: torch.Tensor) -> torch.Tensor: + self.calls += 1 + return latents + self.calls / 10 + + def reset(self) -> None: + self.calls = 0 + + class FakePipeline: + reae = object() + transformer = object() + prompt_emb = torch.empty(0) + device = torch.device("cpu") + torch_dtype = torch.float32 + upscale_mode = "nearest" + _execution_lock = threading.RLock() + + @staticmethod + def _target_size( + lq_h: int, + lq_w: int, + _resolution: tuple[int, int] | None, + _upscale: int, + ) -> tuple[int, int, int, int]: + return lq_h, lq_w, 0, 0 + + monkeypatch.setattr(pipeline_module, "StreamingTAE", FakeTAE) + monkeypatch.setattr(pipeline_module, "StreamingDiT", FakeDiT) + pipeline = FakePipeline() + first = SwiftVRStreamSession(pipeline, clip_len=24, resolution=None, upscale=1, dit_overlap=1) + second = SwiftVRStreamSession(pipeline, clip_len=24, resolution=None, upscale=1, dit_overlap=1) + control = SwiftVRStreamSession(pipeline, clip_len=24, resolution=None, upscale=1, dit_overlap=1) + first_frames = torch.full((4, 2, 2, 3), 17, dtype=torch.uint8) + second_frames = torch.full((4, 2, 2, 3), 193, dtype=torch.uint8) + + first.step(first_frames) + second_first = second.step(second_frames) + control_first = control.step(second_frames) + first.step(first_frames) + first.close() + second_second = second.step(second_frames) + control_second = control.step(second_frames) + + torch.testing.assert_close(second_first, control_first, rtol=0, atol=0) + torch.testing.assert_close(second_second, control_second, rtol=0, atol=0) + assert second._closed is False From d467ceb132d9b6f867ce42be4859d34d3ba18ac4 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Mon, 10 Aug 2026 09:17:36 +0000 Subject: [PATCH 2/3] fix(attention): avoid SDPA layout copies Keep the generic TORCH_SDPA path on transpose views instead of forcing contiguous copies. Add a unit test that verifies the SDPA backend receives transpose views for BSND input. --- telefuser/ops/attention/attention_impl.py | 20 ++++++++----- tests/unit/ops/test_attention_backends.py | 34 +++++++++++++++++++++++ 2 files changed, 47 insertions(+), 7 deletions(-) diff --git a/telefuser/ops/attention/attention_impl.py b/telefuser/ops/attention/attention_impl.py index eac66305..ca322c46 100755 --- a/telefuser/ops/attention/attention_impl.py +++ b/telefuser/ops/attention/attention_impl.py @@ -263,9 +263,13 @@ def attention( current_layout = input_layout if input_layout == "BSND" and attn_impl in BNSD_IMPLS: - q = q.transpose(1, 2).contiguous() - k = k.transpose(1, 2).contiguous() - v = v.transpose(1, 2).contiguous() + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + if attn_impl != AttnImplType.TORCH_SDPA: + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() current_layout = "BNSD" elif ( input_layout == "BNSD" @@ -437,9 +441,9 @@ def attention( _warned_attn_fallback.add(msg) logger.warning(msg) if current_layout == "BSND": - q = q.transpose(1, 2).contiguous() - k = k.transpose(1, 2).contiguous() - v = v.transpose(1, 2).contiguous() + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) current_layout = "BNSD" if sequence_lengths is None: @@ -452,7 +456,9 @@ def attention( # Handle output layout conversion - output matches current_layout, may need to convert to output_layout if current_layout != output_layout: - output = output.transpose(1, 2).contiguous() + output = output.transpose(1, 2) + if attn_impl != AttnImplType.TORCH_SDPA: + output = output.contiguous() if return_lse: return output, lse diff --git a/tests/unit/ops/test_attention_backends.py b/tests/unit/ops/test_attention_backends.py index 296268e5..4599e1f7 100644 --- a/tests/unit/ops/test_attention_backends.py +++ b/tests/unit/ops/test_attention_backends.py @@ -127,6 +127,40 @@ def test_flash_attn4_packed_falls_back_to_sdpa_without_varlen_backend() -> None: torch.testing.assert_close(output, expected) +def test_torch_sdpa_bsnd_layout_uses_transpose_views(monkeypatch) -> None: + q = torch.randn(1, 7, 3, 5) + captured: dict[str, torch.Tensor] = {} + + def fake_sdpa( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + **_kwargs: object, + ) -> torch.Tensor: + captured["query"] = query + captured["key"] = key + captured["value"] = value + return torch.zeros_like(query) + + monkeypatch.setattr(attention_impl.F, "scaled_dot_product_attention", fake_sdpa) + + output = attention_impl.attention( + q, + q, + q, + attn_impl=attention_impl.AttnImplType.TORCH_SDPA, + input_layout="BSND", + output_layout="BSND", + ) + + assert captured["query"]._base is q + assert captured["key"]._base is q + assert captured["value"]._base is q + assert captured["query"].shape == (1, 3, 7, 5) + assert not captured["query"].is_contiguous() + assert output.shape == q.shape + + def test_sage_attention_prefers_tf_kernel() -> None: imported_modules: list[str] = [] tf_kernel_module = ModuleType("tf_kernel") From 43d8d404767b630e1d79ebcebc0726aaa7e642c6 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Mon, 10 Aug 2026 11:09:22 +0000 Subject: [PATCH 3/3] feat(swiftvr): add optimized streaming execution Align the default offline path with the upstream causal benchmark and return PIL frames directly from pipeline sessions.\n\nAdd opt-in torch.compile, tf-kernel FP8Linear, selectable attention backends, and three-stage tensor-channel execution. Record H100 parity and acceleration measurements, including FP8 numerical validation.\n\nVerification: pytest tests/unit/models/test_swiftvr_transformer.py tests/unit/pipelines/swiftvr/test_example.py tests/unit/pipelines/swiftvr/test_streaming.py -q (22 passed); ruff check; ruff format --check; git diff --check. --- docs/en/swiftvr.md | 73 ++- examples/swiftvr/README.md | 131 ++++- examples/swiftvr/swiftvr_restore_h100.py | 130 ++++- telefuser/models/swiftvr_transformer.py | 76 ++- telefuser/pipelines/swiftvr/io.py | 56 -- telefuser/pipelines/swiftvr/pipeline.py | 518 +++++++++++++++++- tests/unit/models/test_swiftvr_transformer.py | 118 ++++ tests/unit/pipelines/swiftvr/__init__.py | 1 + tests/unit/pipelines/swiftvr/test_example.py | 54 +- .../unit/pipelines/swiftvr/test_streaming.py | 57 +- 10 files changed, 1069 insertions(+), 145 deletions(-) delete mode 100644 telefuser/pipelines/swiftvr/io.py create mode 100644 tests/unit/pipelines/swiftvr/__init__.py diff --git a/docs/en/swiftvr.md b/docs/en/swiftvr.md index 1b00dab2..0621b3d3 100644 --- a/docs/en/swiftvr.md +++ b/docs/en/swiftvr.md @@ -4,7 +4,8 @@ TeleFuser integrates the released `H-oliday/SwiftVR` video-restoration model as a faithful, sequential, single-GPU pipeline. The default path preserves the upstream BF16 model, one-step DiT, mask-free shifted-window attention, causal ReAE state, and fixed chunk/flush behavior. Dense attention is dispatched -through `telefuser.ops.attention` with Torch SDPA. +through `telefuser.ops.attention`; Torch SDPA remains the default, while other +TeleFuser dense attention backends can be selected explicitly. ## Provenance and checkpoint @@ -48,7 +49,16 @@ encoding. state, DiT overlap state, RoPE position state, and condition cache. `step()` accepts arbitrary uint8 frame counts, buffers non-four-aligned tails, and `flush()` pads only the final encoder group while preserving output frame -count. GPU execution is serialized by the shared pipeline lock. +count. Pipeline calls and stream sessions return PIL RGB frames; partial chunks +that do not yet produce causal output return an empty list. GPU execution is +serialized by the shared pipeline lock in the default single-process path. + +The optional stage-parallel path splits ReAE encode, DiT denoise, and ReAE +decode into `ParallelWorker` stages. The encode-to-DiT and DiT-to-decode latent +handoffs use `WorkerTensorChannel`, so CUDA tensors are transported by direct +worker-to-worker IPC when profiles fit the channel pool instead of being +materialized in the parent process. This path supports one active stream +session per pipeline because each worker owns its causal state. Direct sessions are isolated and tested with interleaved inputs so ReAE boundary state, DiT overlap, RoPE offsets, and frames cannot cross sessions. @@ -59,18 +69,42 @@ track or frame payload, so a local queue adapter would not provide a usable ## H100 baseline -The target-side compute measurements exclude model loading, video encoding, -LiveKit delivery, and client playback. Each measured chunk contains 24 output -frames after two warmup chunks. - -| Resolution | Compute FPS | TTFC | P50 / P95 chunk | Peak allocated | Retained session state | +The published SwiftVR QHD result is 31.32 FPS on one H100 for 24 frames. +A local probe of the released implementation measured 31.05 FPS under the same +checkpoint and resolution. For an apples-to-apples core comparison, the +following synchronized GPU timings use BF16, 24-frame chunks, dit_overlap=0, +and exclude PIL conversion and device-to-host output transfer. + +| Output resolution | Official SwiftVR | TeleFuser default | TeleFuser opt-in compile | +| --- | ---: | ---: | ---: | +| 2560x1440 | 31.05 FPS local probe (31.32 published) | 32.42 FPS | 40.3-40.5 FPS steady | + +The TeleFuser default is therefore faster than the released implementation on +the parity path. The compile result uses torch.compile for the DiT blocks; the +first shape compile takes about 44 seconds on this host, so it is intended for +long-lived processes. The default CLI enables the parity setting dit_overlap=0 +used by the published offline benchmark. The public stream() API still +defaults to dit_overlap=1, matching upstream streaming semantics. + +The delivered List[PIL.Image.Image] path includes output conversion and +device-to-host transfer and is expected to be lower and more host-variable +than the core GPU number. The existing end-to-end example figures below report +that delivered path separately. + +The optional stage-parallel path was measured on three H100 GPUs with +WorkerTensorChannel latent handoff enabled. A 240-input-frame steady run +produced 237 output frames at 39.87-40.11 FPS, with PIL conversion removed from +the benchmark to isolate stage throughput. This is a pipeline throughput +measurement and is not directly comparable to single-GPU memory or latency. + +| Resolution | Delivered PIL FPS | TTFC | P50 / P95 chunk | Peak allocated | Retained session state | | --- | ---: | ---: | ---: | ---: | ---: | | 1920x1080 | 47.28 | 15.01 s | 0.506 / 0.511 s | 27.08 GiB | 235.5 MiB | | 2560x1440 | 26.50 | 16.31 s | 0.895 / 0.960 s | 35.93 GiB | 411.8 MiB | The complete example path was separately accepted on one H100 using all 81 -frames of `examples/data/dag.mp4`, resized to `640x360` and restored at 3x. The -measured session produced the `1920x1080` frames in 2.33 seconds (34.78 FPS), +frames of examples/data/dag.mp4, resized to 640x360 and restored at 3x. The +measured session produced the 1920x1080 frames in 2.33 seconds (34.78 FPS), and H.264 encoding took 1.32 seconds, for 22.18 end-to-end FPS. The resulting video contains 81 frames at 16 FPS. These example figures exclude one-time model loading, input decoding, and shape warmup; they include output conversion @@ -82,21 +116,16 @@ Retained memory in the table is the state of one direct causal session. No multi-session service capacity is claimed because there is no supported SwiftVR video transport in the shared streaming server. -The 1080p long-stream run completed 1176 measured chunks and 28,224 frames over -600.46 seconds at 47.01 compute FPS. P50/P95/max chunk latency was -0.508/0.523/0.654 seconds. The first and last 100 chunks had P50 values of -0.50810 and 0.50796 seconds respectively, so no latency growth was observed. -Session cleanup completed and the benchmark process returned all device memory -to the driver. - ## Current constraints -- No quantization, feature cache, sparse-attention substitution, or - `torch.compile` path is enabled. -- No stage splitting or pipeline overlap is enabled; execution remains in the - upstream order. -- The public example uses eager Torch SDPA and does not expose alternative - attention backends. +- Quantization, `torch.compile`, alternative dense attention backends, and + stage-parallel execution are opt-in. The default remains BF16 eager Torch + SDPA for parity. +- The stage-parallel path uses direct tensor channels for latent handoff, but it + does not change SwiftVR's causal execution order or enable tensor/model + parallelism inside the DiT blocks. +- Feature cache and sparse-attention substitution are still not enabled for + SwiftVR. - The supported online surface is the direct causal session API; the example does not expose a partial `stream-serve` adapter. - RTX 5090 measurements require that hardware and are not represented by H100 diff --git a/examples/swiftvr/README.md b/examples/swiftvr/README.md index 01725cdb..7d743622 100644 --- a/examples/swiftvr/README.md +++ b/examples/swiftvr/README.md @@ -4,7 +4,37 @@ This example restores videos with the released `H-oliday/SwiftVR` checkpoint on one CUDA GPU. The parity path uses BF16, dense Torch SDPA, the upstream fixed timestep, and the original 24-frame causal chunk protocol. -The checkpoint root must contain: +## Resources + +| Resource | Link | +| --- | --- | +| SwiftVR source | [H-oliday/SwiftVR](https://github.com/H-oliday/SwiftVR) | +| Pretrained checkpoint | [H-oliday/SwiftVR on Hugging Face](https://huggingface.co/H-oliday/SwiftVR) | +| Project page | [h-oliday.github.io/SwiftVR](https://h-oliday.github.io/SwiftVR) | +| Paper | [arXiv:2606.09516](https://arxiv.org/abs/2606.09516) | + +The TeleFuser integration was ported from official SwiftVR commit +`5ca168cef6ca7200f135fdfea85e5e13d12c5b53`. The checkpoint is downloaded from +Hugging Face and is compatible with the released model files below. + +## Installation + +Install PyTorch for the target CUDA version first, then install TeleFuser in +editable mode: + +```bash +pip install -e ".[dev]" +``` + +Download the checkpoint with the Hugging Face CLI: + +```bash +pip install -U huggingface_hub +huggingface-cli download H-oliday/SwiftVR --local-dir /data/SwiftVR +``` + +For gated or private repositories, authenticate first with `huggingface-cli login`. +The directory passed to `--model_root` must contain: ```text reae.safetensors @@ -35,8 +65,30 @@ actual tail shape to absorb cold cuDNN plan selection. The reported processing FPS then covers stateful inference and output transfer, while end-to-end FPS also includes H.264 encoding. -The direct causal API accepts uint8 `[T,H,W,3]` tensors and returns BF16 -`[1,T,3,H,W]` tensors in `[0,1]`: +The CLI also accepts optional acceleration controls: + +```bash +python examples/swiftvr/swiftvr_restore_h100.py \ + --model_root /data/SwiftVR \ + --attn_impl TORCH_SDPA \ + --compile_dit \ + --quantization tf-kernel-fp8 +``` + +For multi-GPU stage execution, pass three stage devices. Latents between ReAE +encode, DiT, and ReAE decode are handed off through `WorkerTensorChannel`: + +```bash +python examples/swiftvr/swiftvr_restore_h100.py \ + --model_root /data/SwiftVR \ + --gpu_num 3 \ + --enable_stage_parallel \ + --stage_devices 0,1,2 +``` + +The direct causal API accepts uint8 `[T,H,W,3]` tensors and returns PIL RGB +frames. A partial chunk can return an empty list until enough causal context is +available: ```python import torch @@ -46,25 +98,70 @@ from examples.swiftvr.swiftvr_restore_h100 import get_pipeline pipeline = get_pipeline(model_root="/data/SwiftVR") session = pipeline.stream(resolution=(1920, 1080), clip_len=24, dit_overlap=1) try: - first = session.step(torch.zeros((24, 540, 960, 3), dtype=torch.uint8)) - tail = session.step(torch.zeros((5, 540, 960, 3), dtype=torch.uint8)) - flushed = session.flush() + first_frames = session.step(torch.zeros((24, 540, 960, 3), dtype=torch.uint8)) + tail_frames = session.step(torch.zeros((5, 540, 960, 3), dtype=torch.uint8)) + flushed_frames = session.flush() finally: session.close() ``` ## H100 performance -The following baseline was measured on one NVIDIA H100 80GB HBM3 with BF16, -eager Torch SDPA, and 24-frame chunks. The values are synchronized target-side -compute metrics; model loading, video encoding, network delivery, and client -playback are excluded. - -| Output resolution | Compute FPS | TTFC | P50 / P95 chunk latency | Peak allocated memory | +The official README reports the following single-H100 results for causal +streaming with 24 frames: + +| Resolution | Official FPS | Official average time | Official peak memory | +| --- | ---: | ---: | ---: | +| 2560x1440 | 31.32 | 0.766 s | 38.01 GB | +| 3840x2160 | 14.00 | 1.714 s | Not reported | + +At 4K, the official README reports that every compared diffusion-based VR +baseline OOMs on one H100 while SwiftVR sustains 14 FPS. + +The released implementation measured 31.05 FPS in a local run under the same +checkpoint and resolution. The fair core comparison below uses BF16, eager +Torch SDPA, `dit_overlap=0`, and synchronized GPU timing that excludes PIL +conversion and device-to-host output transfer. + +| Output resolution | Official local probe | TeleFuser default | Compile | FP8Linear | Compile + FP8Linear | +| --- | ---: | ---: | ---: | ---: | ---: | +| 2560x1440 | 31.05 FPS | 32.42 FPS | 40.3-40.5 FPS | 35.8-36.1 FPS | 45.2-45.7 FPS | + +The default TeleFuser path is faster than the released implementation. +`torch.compile` has a one-time shape compilation cost of about 44 +seconds on the benchmark host, so enable it for long-lived workers. The example +defaults to `dit_overlap=0` for parity with the published offline +result; direct `stream()` keeps upstream's default +`dit_overlap=1`. + +The `tf-kernel-fp8` option uses the same public +`telefuser.ops.fp8_gemm.FP8Linear` W8A8 path as LiveAct. It wraps 360 +SwiftVR DiT Linear layers with per-token activation quantization and cached FP8 +weights. In a QHD BF16-to-FP8 output check, it produced PSNR 52.24 dB, MAE +0.00146, and maximum absolute error 0.015625. Treat it as an opt-in speed and +quality tradeoff and validate on the target content. + +SageAttention remains selectable through `--attn_impl`, but it is not +the recommended H100 default for this model: SAGE_ATTN_2_8_8_SM90 measured +30.8-31.1 FPS eager and 38.1-38.4 FPS with compile, below the corresponding +Torch SDPA paths. + +The pipeline returns `List[PIL.Image.Image]`. Its delivered throughput includes +PIL conversion and device-to-host transfer, so it is lower and more sensitive +to host scheduling than the core GPU number. The existing delivered-path +measurements are: + +| Output resolution | Delivered PIL FPS | TTFC | P50 / P95 chunk latency | Peak allocated memory | | --- | ---: | ---: | ---: | ---: | | 1920x1080 | 47.28 | 15.01 s | 0.506 / 0.511 s | 27.08 GiB | | 2560x1440 | 26.50 | 16.31 s | 0.895 / 0.960 s | 35.93 GiB | +The optional three-GPU stage-parallel path uses `WorkerTensorChannel` for latent +handoff. A 240-input-frame steady run produced 237 output frames at 39.87-40.11 +FPS with PIL conversion removed to isolate stage throughput. This is a pipeline +throughput result and is not directly comparable to single-GPU memory or +latency. + The full example command above was also verified on one H100 with all 81 frames of `dag.mp4`. The measured session produced `1920x1080` frames at 34.78 FPS; including H.264 encoding, end-to-end throughput was 22.18 FPS versus the source @@ -77,10 +174,12 @@ in 1,176 chunks at 47.01 compute FPS. P50/P95/max chunk latency was 0.508/0.523/0.654 seconds, with no latency growth between the beginning and end of the run. Session cleanup returned all device memory to the driver. -`SwiftVRPipeline.stream()` is the supported stateful interface. The example -does not expose `get_service()` because the current LiveKit protocol has no -video-input transport; presenting a partial service adapter would not make the -model usable through `telefuser stream-serve`. +`SwiftVRPipeline.stream()` is the supported stateful interface. In stage-parallel +mode, one active stream session is supported per pipeline because the worker +stages retain causal ReAE/DiT state. The example does not expose `get_service()` +because the current LiveKit protocol has no video-input transport; presenting a +partial service adapter would not make the model usable through +`telefuser stream-serve`. See [the SwiftVR integration guide](../../docs/en/swiftvr.md) for checkpoint loading, constraints, and measured results. diff --git a/examples/swiftvr/swiftvr_restore_h100.py b/examples/swiftvr/swiftvr_restore_h100.py index 09a454d8..9c7fac39 100644 --- a/examples/swiftvr/swiftvr_restore_h100.py +++ b/examples/swiftvr/swiftvr_restore_h100.py @@ -10,9 +10,15 @@ import torch from PIL import Image -from telefuser.core.config import AttentionConfig, AttnImplType +from telefuser.core.config import ( + AttentionConfig, + AttnImplType, + CompileConfig, + QuantConfig, + QuantKernelBackend, + QuantType, +) from telefuser.pipelines.swiftvr import SwiftVRPipeline -from telefuser.pipelines.swiftvr.io import ntchw_to_uint8_frames from telefuser.utils.utils import get_example_name from telefuser.utils.video import VideoData, save_video @@ -25,21 +31,72 @@ video_quality=6, chunk_size=24, warmup_chunks=2, + dit_overlap=0, + attn_impl=AttnImplType.TORCH_SDPA, + compile_mode="default", ) +def _parse_attn_impl(attn_impl: AttnImplType | str) -> AttnImplType: + if isinstance(attn_impl, AttnImplType): + return attn_impl + key = attn_impl.upper().replace("-", "_") + try: + return AttnImplType[key] + except KeyError as exc: + raise ValueError(f"Unsupported attention backend: {attn_impl}") from exc + + +def _quant_config(quantization: str | None) -> QuantConfig | None: + if quantization is None: + return None + normalized = quantization.lower().replace("_", "-") + if normalized == "torchao-fp8": + return QuantConfig(enabled=True, quant_type=QuantType.TORCHAO_FP8, kernel_backend=QuantKernelBackend.TORCHAO) + if normalized == "tf-kernel-fp8": + return QuantConfig(enabled=True, quant_type=QuantType.FP8, kernel_backend=QuantKernelBackend.TF_KERNEL) + raise ValueError("quantization must be 'torchao-fp8', 'tf-kernel-fp8', or None") + + +def _parse_stage_devices(stage_devices: str | None, parallelism: int) -> list[int] | None: + if stage_devices: + devices = [int(item.strip()) for item in stage_devices.split(",") if item.strip()] + elif parallelism > 1: + devices = list(range(parallelism)) + else: + return None + if len(devices) != 3: + raise ValueError("SwiftVR stage parallelism requires exactly three devices: encode,dit,decode") + return devices + + def get_pipeline( parallelism: int = 1, model_root: str = PPL_CONFIG["model_root"], + attn_impl: AttnImplType | str = PPL_CONFIG["attn_impl"], + compile_dit: bool = False, + compile_mode: str = PPL_CONFIG["compile_mode"], + quantization: str | None = None, + enable_stage_parallel: bool = False, + stage_devices: str | None = None, + tensor_channel_slots: int = 2, ) -> SwiftVRPipeline: """Initialize the SwiftVR pipeline.""" - if parallelism != 1: - raise ValueError("SwiftVR supports exactly one GPU") + stage_device_ids = _parse_stage_devices(stage_devices, parallelism) if enable_stage_parallel else None + if not enable_stage_parallel and parallelism != 1: + raise ValueError("SwiftVR supports multiple GPUs only through --enable_stage_parallel") return SwiftVRPipeline.from_pretrained( model_root, device="cuda", torch_dtype=torch.bfloat16, - attention_config=AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA), + attention_config=AttentionConfig.dense_attention(_parse_attn_impl(attn_impl)), + compile_config=CompileConfig(enabled=compile_dit, mode=compile_mode) if compile_dit else None, + quant_config=_quant_config(quantization), + enable_stage_parallel=enable_stage_parallel, + enable_stage_overlap=enable_stage_parallel, + stage_dit_overlap=PPL_CONFIG["dit_overlap"], + stage_device_ids=stage_device_ids, + tensor_channel_slots=tensor_channel_slots, ) @@ -59,16 +116,19 @@ def run( np.asarray(frame.convert("RGB").crop((0, 0, crop_width, crop_height)), dtype=np.uint8) for frame in input_video ] frames = torch.from_numpy(np.stack(arrays)).contiguous() + if pipeline.config.enable_stage_parallel: + return pipeline( + frames, + upscale=scale, + clip_len=PPL_CONFIG["chunk_size"], + dit_overlap=PPL_CONFIG["dit_overlap"], + ) restored_frames: list[Image.Image] = [] - session = pipeline.stream(upscale=scale) + session = pipeline.stream(upscale=scale, dit_overlap=PPL_CONFIG["dit_overlap"]) try: for start in range(0, len(frames), PPL_CONFIG["chunk_size"]): - output = ntchw_to_uint8_frames(session.step(frames[start : start + PPL_CONFIG["chunk_size"]])) - if output is not None: - restored_frames.extend(Image.fromarray(frame) for frame in output) - output = ntchw_to_uint8_frames(session.flush()) - if output is not None: - restored_frames.extend(Image.fromarray(frame) for frame in output) + restored_frames.extend(session.step(frames[start : start + PPL_CONFIG["chunk_size"]])) + restored_frames.extend(session.flush()) finally: session.close() return restored_frames @@ -97,6 +157,26 @@ def _warmup_frame_count(frame_count: int) -> int: help=f"Root directory containing model files (default: {PPL_CONFIG['model_root']})", ) @click.option("--output", "-o", default=None, help="Output video path (default: auto-generated)") +@click.option( + "--attn_impl", + default=PPL_CONFIG["attn_impl"].name, + type=click.Choice([item.name for item in AttnImplType], case_sensitive=False), + help="DiT attention backend", +) +@click.option("--compile_dit/--no_compile_dit", default=False, help="Enable torch.compile for SwiftVR DiT blocks") +@click.option( + "--compile_mode", + default=PPL_CONFIG["compile_mode"], + type=click.Choice( + ["default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"], + case_sensitive=False, + ), + help="torch.compile mode / kernel fusion preset", +) +@click.option("--quantization", type=click.Choice(["torchao-fp8", "tf-kernel-fp8"]), default=None) +@click.option("--enable_stage_parallel", is_flag=True, help="Split encode, DiT, and decode into worker stages") +@click.option("--stage_devices", default=None, help="Comma-separated encode,dit,decode GPU ids, e.g. 0,1,2") +@click.option("--tensor_channel_slots", default=2, type=int, help="CUDA IPC slots per stage tensor-channel profile") def main( input_video: str, scale: int, @@ -105,6 +185,13 @@ def main( gpu_num: int, model_root: str, output: str | None, + attn_impl: str, + compile_dit: bool, + compile_mode: str, + quantization: str | None, + enable_stage_parallel: bool, + stage_devices: str | None, + tensor_channel_slots: int, ) -> None: """SwiftVR streaming video restoration.""" if not os.path.exists(input_video): @@ -120,10 +207,27 @@ def main( click.echo(f"Scale: {scale}x") click.echo(f"GPUs: {gpu_num}") click.echo(f"Model root: {model_root}") + click.echo(f"Attention: {attn_impl}") + click.echo(f"Compile DiT: {compile_dit}") + click.echo(f"Compile mode: {compile_mode}") + click.echo(f"Quantization: {quantization or 'none'}") + click.echo(f"Stage parallel: {enable_stage_parallel}") + if stage_devices: + click.echo(f"Stage devices: {stage_devices}") click.echo(f"Output: {output}") click.echo("Loading pipeline...") - pipeline = get_pipeline(gpu_num, model_root) + pipeline = get_pipeline( + gpu_num, + model_root, + attn_impl=attn_impl, + compile_dit=compile_dit, + compile_mode=compile_mode, + quantization=quantization, + enable_stage_parallel=enable_stage_parallel, + stage_devices=stage_devices, + tensor_channel_slots=tensor_channel_slots, + ) click.echo("Loading video...") input_frames = VideoData(video_file=input_video, height=height, width=width).raw_data() diff --git a/telefuser/models/swiftvr_transformer.py b/telefuser/models/swiftvr_transformer.py index 6c7d7daa..35749252 100644 --- a/telefuser/models/swiftvr_transformer.py +++ b/telefuser/models/swiftvr_transformer.py @@ -17,7 +17,7 @@ import torch.nn as nn import torch.nn.functional as F -from telefuser.core.config import AttentionConfig, AttnImplType +from telefuser.core.config import AttentionConfig, AttnImplType, CompileConfig from telefuser.ops.attention import attention from telefuser.ops.ffn import FeedForward from telefuser.utils.logging import logger @@ -782,6 +782,23 @@ def compile_transformer_blocks(model: nn.Module, mode: str = "default") -> None: model.blocks[i] = torch.compile(blk, mode=mode, fullgraph=False) +def compile_transformer_blocks_with_config(model: nn.Module, compile_config: CompileConfig) -> None: + if not hasattr(torch, "compile"): + logger.warning("torch.compile not available (requires PyTorch 2.0+). Skipping.") + return + kwargs = compile_config.get_compile_kwargs() + if kwargs.get("mode") == "reduce-overhead": + logger.warning( + "compile mode 'reduce-overhead' is incompatible with the in-place residuals; falling back to 'default'." + ) + kwargs = dict(kwargs) + kwargs["mode"] = "default" + kwargs.pop("disable", None) + for i, blk in enumerate(getattr(model, "blocks", [])): + if isinstance(blk, WanTransformerBlock): + model.blocks[i] = torch.compile(blk, **kwargs) + + # --------------------------------------------------------------------------- # # Main model # # --------------------------------------------------------------------------- # @@ -893,6 +910,63 @@ def prepare_for_inference( _WindowRuntimeMetaCache.clear() self.eval() + def enable_quant(self, quant_type: object) -> None: + """Apply supported online quantization to SwiftVR transformer Linear layers.""" + from telefuser.core.config import QuantConfig, QuantKernelBackend, QuantType + + if not isinstance(quant_type, QuantConfig): + if quant_type in (torch.float8_e4m3fn, "float8_e4m3fn", "fp8"): + from telefuser.ops.quantized_linear import replace_linear_layers + + replace_linear_layers(self.blocks, torch.float8_e4m3fn) + self.quant_type = torch.float8_e4m3fn + return + raise ValueError(f"SwiftVR does not support quantization setting {quant_type!r}") + if not quant_type.enabled: + return + + include_names = quant_type.quantize_modules or ("blocks.",) + replaced = 0 + if quant_type.quant_type == QuantType.TORCHAO_FP8: + from telefuser.ops.torchao_fp8_linear import replace_linear_layers_with_torchao_fp8 + + replaced = replace_linear_layers_with_torchao_fp8( + self, + include_names=include_names, + exclude_names=quant_type.skip_modules, + ) + self.torchao_fp8_replaced_linear = replaced + elif quant_type.quant_type == QuantType.FP8: + if quant_type.kernel_backend not in (QuantKernelBackend.AUTO, QuantKernelBackend.TF_KERNEL): + raise ValueError( + "SwiftVR FP8 online quantization requires the tf-kernel backend; " + f"got {quant_type.kernel_backend.name}" + ) + from telefuser.ops.fp8_gemm import FP8GemmOptions, count_linear_layers, enable_fp8_gemm + + def module_filter(name: str, _module: nn.Module) -> bool: + return any(token in name for token in include_names) and not any( + token and token in name for token in quant_type.skip_modules + ) + + replaced = count_linear_layers(self, module_filter=module_filter) + enable_fp8_gemm( + self, + options=FP8GemmOptions( + fp16_weight_storage="keep" if quant_type.keep_fp16_weight else "discard", + materialize_fp8_on_wrap=True, + ), + module_filter=module_filter, + ) + self.tf_kernel_fp8_replaced_linear = replaced + else: + raise ValueError(f"SwiftVR does not support online quantization type {quant_type.quant_type.name}") + + if replaced == 0: + raise RuntimeError("SwiftVR online quantization did not select any Linear layers") + self.quant_type = quant_type.quant_type + logger.info(f"SwiftVR {quant_type.quant_type.name} converted {replaced} transformer Linear layers") + @torch.inference_mode() def forward( self, diff --git a/telefuser/pipelines/swiftvr/io.py b/telefuser/pipelines/swiftvr/io.py deleted file mode 100644 index 0ffbf98e..00000000 --- a/telefuser/pipelines/swiftvr/io.py +++ /dev/null @@ -1,56 +0,0 @@ -"""Tensor preprocessing and output conversion for SwiftVR.""" - -from __future__ import annotations - -import numpy as np -import torch -import torch.nn.functional as F - -_INTERP_NEEDS_ALIGN = ("linear", "bilinear", "bicubic", "trilinear") - - -def preprocess_clip_uint8( - frames_uint8: torch.Tensor, - out_h: int, - out_w: int, - mode: str, - pad_h: int, - pad_w: int, - dtype: torch.dtype, -) -> torch.Tensor: - """Convert CUDA uint8 THWC frames to padded target-dtype NTCHW frames in [0, 1].""" - frames = frames_uint8.permute(0, 3, 1, 2).contiguous().to(dtype=dtype) - _, _, height, width = frames.shape - if (height, width) != (out_h, out_w): - if mode in _INTERP_NEEDS_ALIGN: - frames = F.interpolate(frames, size=(out_h, out_w), mode=mode, align_corners=False) - else: - frames = F.interpolate(frames, size=(out_h, out_w), mode=mode) - frames = frames / 255.0 - if pad_h > 0 or pad_w > 0: - frames = F.pad(frames, (0, pad_w, 0, pad_h), mode="constant", value=0) - return frames.unsqueeze(0) - - -def crop_spatial_padding_ntchw(video: torch.Tensor | None, pad_h: int = 0, pad_w: int = 0) -> torch.Tensor | None: - """Remove bottom/right spatial padding from an NTCHW tensor.""" - if video is None: - return None - if pad_h > 0: - video = video[:, :, :, :-pad_h, :] - if pad_w > 0: - video = video[:, :, :, :, :-pad_w] - return video - - -def ntchw_to_uint8_frames(video: torch.Tensor | None) -> np.ndarray | None: - """Convert [0, 1] NTCHW output to uint8 THWC frames on the host.""" - if video is None or video.numel() == 0 or video.shape[1] == 0: - return None - frames = (video[0].permute(0, 2, 3, 1).contiguous() * 255).clamp(0, 255).to(torch.uint8) - if frames.device.type != "cuda": - return frames.cpu().numpy() - cpu_frames = torch.empty_like(frames, device="cpu", pin_memory=True) - cpu_frames.copy_(frames, non_blocking=True) - torch.cuda.current_stream(frames.device).synchronize() - return cpu_frames.numpy() diff --git a/telefuser/pipelines/swiftvr/pipeline.py b/telefuser/pipelines/swiftvr/pipeline.py index 32a14b7d..b11ee106 100644 --- a/telefuser/pipelines/swiftvr/pipeline.py +++ b/telefuser/pipelines/swiftvr/pipeline.py @@ -14,19 +14,31 @@ from typing import Iterable, Iterator import torch +import torch.nn.functional as F +from PIL import Image from safetensors.torch import load_file from telefuser.core.base_pipeline import BasePipeline -from telefuser.core.config import AttentionConfig, AttnImplType, ModelRuntimeConfig +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ( + AttentionConfig, + AttnImplType, + CompileConfig, + ModelRuntimeConfig, + ParallelConfig, + QuantConfig, +) from telefuser.core.module_manager import ModuleManager from telefuser.models.swiftvr_reae import ReAE -from telefuser.models.swiftvr_transformer import SwiftVRWanTransformer3DModel +from telefuser.models.swiftvr_transformer import ( + SwiftVRWanTransformer3DModel, + compile_transformer_blocks_with_config, +) +from telefuser.utils.logging import logger +from telefuser.worker.parallel_worker import ParallelWorker +from telefuser.worker.tensor_channel import WorkerTensorChannel from .chunk import ChunkSpec, ChunkType, build_chunk_specs -from .io import ( - crop_spatial_padding_ntchw, - preprocess_clip_uint8, -) from .streaming_dit import StreamingDiT from .streaming_tae import StreamingTAE @@ -58,15 +70,202 @@ def aligned_pad(size: int, multiple: int = 32) -> int: return (multiple - size % multiple) % multiple +_INTERP_NEEDS_ALIGN = ("linear", "bilinear", "bicubic", "trilinear") + + +def preprocess_clip_uint8( + frames_uint8: torch.Tensor, + out_h: int, + out_w: int, + mode: str, + pad_h: int, + pad_w: int, + dtype: torch.dtype, +) -> torch.Tensor: + """Convert uint8 THWC frames to padded target-dtype NTCHW frames in [0, 1].""" + frames = frames_uint8.permute(0, 3, 1, 2).contiguous().to(dtype=dtype) + _, _, height, width = frames.shape + if (height, width) != (out_h, out_w): + if mode in _INTERP_NEEDS_ALIGN: + frames = F.interpolate(frames, size=(out_h, out_w), mode=mode, align_corners=False) + else: + frames = F.interpolate(frames, size=(out_h, out_w), mode=mode) + frames = frames / 255.0 + if pad_h > 0 or pad_w > 0: + frames = F.pad(frames, (0, pad_w, 0, pad_h), mode="constant", value=0) + return frames.unsqueeze(0) + + +def crop_spatial_padding_ntchw(video: torch.Tensor | None, pad_h: int = 0, pad_w: int = 0) -> torch.Tensor | None: + """Remove bottom/right spatial padding from an NTCHW tensor.""" + if video is None: + return None + if pad_h > 0: + video = video[:, :, :, :-pad_h, :] + if pad_w > 0: + video = video[:, :, :, :, :-pad_w] + return video + + +def ntchw_to_pil_frames(video: torch.Tensor | None) -> list[Image.Image]: + """Convert [0, 1] NTCHW output to PIL RGB frames on the host.""" + if video is None or video.numel() == 0 or video.shape[1] == 0: + return [] + frames = (video[0].permute(0, 2, 3, 1).contiguous() * 255).clamp(0, 255).to(torch.uint8) + if frames.device.type == "cuda": + cpu_frames = torch.empty_like(frames, device="cpu", pin_memory=True) + cpu_frames.copy_(frames, non_blocking=True) + torch.cuda.current_stream(frames.device).synchronize() + frames = cpu_frames + else: + frames = frames.cpu() + return [Image.fromarray(frame.numpy()) for frame in frames] + + @dataclass class SwiftVRPipelineConfig: """Runtime configuration using existing TeleFuser model controls.""" + encode_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) dit_config: ModelRuntimeConfig = field( default_factory=lambda: ModelRuntimeConfig( attention_config=AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA) ) ) + decode_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + enable_stage_parallel: bool = False + enable_stage_overlap: bool = False + dit_overlap: int = 1 + tensor_channel_slots: int = 2 + + +def _apply_cuda_runtime_flags(device: torch.device) -> None: + if device.type != "cuda": + return + torch.backends.cudnn.benchmark = True + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.set_float32_matmul_precision("high") + + +def _prepare_swiftvr_transformer( + transformer: SwiftVRWanTransformer3DModel, + runtime_config: ModelRuntimeConfig, +) -> None: + transformer.prepare_for_inference(attention_config=runtime_config.attention_config) + if runtime_config.quant_config.enabled: + transformer.enable_quant(runtime_config.quant_config) + if runtime_config.parallel_config.world_size == 1 and runtime_config.compile_config.enabled: + compile_transformer_blocks_with_config(transformer, runtime_config.compile_config) + + +class _SwiftVRReAEEncodeStage(BaseStage): + def __init__(self, name: str, module_manager: ModuleManager, model_runtime_config: ModelRuntimeConfig) -> None: + super().__init__(name, model_runtime_config) + self.reae: ReAE = module_manager.fetch_module("swiftvr_reae") + self.model_names = ["reae"] + self._tae: StreamingTAE | None = None + + def _ensure_session(self) -> StreamingTAE: + self.reae.to(device=self.device, dtype=self.torch_dtype).eval() + _apply_cuda_runtime_flags(self.device) + if self._tae is None: + self._tae = StreamingTAE(self.reae) + return self._tae + + @with_model_offload(["reae"]) + @torch.inference_mode() + def process( + self, + frames_uint8: torch.Tensor, + out_h: int, + out_w: int, + pad_h: int, + pad_w: int, + upscale_mode: str, + ) -> torch.Tensor | None: + tae = self._ensure_session() + clip = preprocess_clip_uint8(frames_uint8, out_h, out_w, upscale_mode, pad_h, pad_w, self.torch_dtype) + return tae.encode_chunk(clip) + + @with_model_offload(["reae"]) + @torch.inference_mode() + def flush_encoder(self) -> torch.Tensor | None: + return self._ensure_session().flush_encoder() + + def reset_session(self) -> None: + if self._tae is not None: + self._tae.reset() + self._tae = None + + +class _SwiftVRDiTStage(BaseStage): + def __init__( + self, + name: str, + module_manager: ModuleManager, + model_runtime_config: ModelRuntimeConfig, + prompt_emb: torch.Tensor, + dit_overlap: int, + ) -> None: + super().__init__(name, model_runtime_config) + self.transformer: SwiftVRWanTransformer3DModel = module_manager.fetch_module("swiftvr_transformer") + self.prompt_emb = prompt_emb.to(dtype=self.torch_dtype) + self.model_names = ["transformer"] + self.dit_overlap = dit_overlap + self._dit: StreamingDiT | None = None + self._prepared = False + + def _ensure_session(self) -> StreamingDiT: + self.transformer.to(device=self.device, dtype=self.torch_dtype).eval() + _apply_cuda_runtime_flags(self.device) + if not self._prepared: + _prepare_swiftvr_transformer(self.transformer, self.model_runtime_config) + self._prepared = True + if self._dit is None: + self._dit = StreamingDiT(self.transformer, overlap=self.dit_overlap) + return self._dit + + @with_model_offload(["transformer"]) + @torch.inference_mode() + def process(self, encoded: torch.Tensor) -> torch.Tensor: + dit = self._ensure_session() + encoded_bcfhw = encoded.permute(0, 2, 1, 3, 4).contiguous() + denoised = dit.denoise(encoded_bcfhw, self.prompt_emb) + return denoised.permute(0, 2, 1, 3, 4).contiguous() + + def reset_session(self) -> None: + if self._dit is not None: + self._dit.reset() + self._dit._cond_cache = None + self._dit._cond_cache_key = None + self._dit = None + + +class _SwiftVRReAEDecodeStage(BaseStage): + def __init__(self, name: str, module_manager: ModuleManager, model_runtime_config: ModelRuntimeConfig) -> None: + super().__init__(name, model_runtime_config) + self.reae: ReAE = module_manager.fetch_module("swiftvr_reae") + self.model_names = ["reae"] + self._tae: StreamingTAE | None = None + + def _ensure_session(self) -> StreamingTAE: + self.reae.to(device=self.device, dtype=self.torch_dtype).eval() + _apply_cuda_runtime_flags(self.device) + if self._tae is None: + self._tae = StreamingTAE(self.reae) + return self._tae + + @with_model_offload(["reae"]) + @torch.inference_mode() + def process(self, denoised: torch.Tensor, pad_h: int, pad_w: int) -> torch.Tensor | None: + decoded = self._ensure_session().decode_chunk(denoised) + return crop_spatial_padding_ntchw(decoded, pad_h, pad_w) + + def reset_session(self) -> None: + if self._tae is not None: + self._tae.reset() + self._tae = None class SwiftVRPipeline(BasePipeline): @@ -92,19 +291,61 @@ def init( if self.reae is None or self.transformer is None: raise RuntimeError("SwiftVR requires both swiftvr_reae and swiftvr_transformer") self.prompt_emb = prompt_emb.to(dtype=self.torch_dtype) - self._prepare_for_inference() + self._worker_tensor_channels: list[WorkerTensorChannel] = [] + self._stage_parallel_active_session = False + self.encode_stage = None + self.dit_stage = None + self.decode_stage = None + if config.enable_stage_parallel: + self._init_stage_workers(module_manager, prompt_emb) + else: + self._prepare_for_inference() def _prepare_for_inference(self) -> None: self.reae.to(device=self.device, dtype=self.torch_dtype).eval() self.transformer.to(device=self.device, dtype=self.torch_dtype).eval() - if self.device.type == "cuda": - # Match the released SwiftVR runtime before selecting convolution - # and attention kernels; these flags are part of the frozen parity - # configuration used by the upstream pipeline. - torch.backends.cudnn.benchmark = True - torch.backends.cuda.matmul.allow_tf32 = True - torch.backends.cudnn.allow_tf32 = True - self.transformer.prepare_for_inference(attention_config=self.config.dit_config.attention_config) + _apply_cuda_runtime_flags(self.device) + _prepare_swiftvr_transformer(self.transformer, self.config.dit_config) + + def _init_stage_workers(self, module_manager: ModuleManager, prompt_emb: torch.Tensor) -> None: + timeout = max( + self.config.encode_config.parallel_config.timeout, + self.config.dit_config.parallel_config.timeout, + self.config.decode_config.parallel_config.timeout, + ) + encode_to_dit = WorkerTensorChannel( + self.config.dit_config.parallel_config.world_size, + timeout=timeout, + cuda_ipc_slots=self.config.tensor_channel_slots, + ) + dit_to_decode = WorkerTensorChannel( + self.config.decode_config.parallel_config.world_size, + timeout=timeout, + cuda_ipc_slots=self.config.tensor_channel_slots, + ) + self._worker_tensor_channels.extend((encode_to_dit, dit_to_decode)) + self.encode_stage = ParallelWorker( + _SwiftVRReAEEncodeStage("swiftvr_encode", module_manager, self.config.encode_config), + tensor_output_channel=encode_to_dit, + tensor_output_methods=("process", "flush_encoder"), + ) + self.dit_stage = ParallelWorker( + _SwiftVRDiTStage( + "swiftvr_dit", + module_manager, + self.config.dit_config, + prompt_emb, + self.config.dit_overlap, + ), + tensor_output_channel=dit_to_decode, + tensor_output_methods=("process",), + tensor_input_channels=(encode_to_dit,), + ) + self.decode_stage = ParallelWorker( + _SwiftVRReAEDecodeStage("swiftvr_decode", module_manager, self.config.decode_config), + tensor_input_channels=(dit_to_decode,), + ) + logger.info("SwiftVR stage-parallel workers initialized with direct tensor channels") @classmethod def from_pretrained( @@ -114,6 +355,13 @@ def from_pretrained( device: str | torch.device = "cuda", torch_dtype: torch.dtype | str = torch.bfloat16, attention_config: AttentionConfig | None = None, + compile_config: CompileConfig | None = None, + quant_config: QuantConfig | None = None, + enable_stage_parallel: bool = False, + enable_stage_overlap: bool = False, + stage_dit_overlap: int = 1, + stage_device_ids: tuple[int, int, int] | list[int] | None = None, + tensor_channel_slots: int = 2, ) -> "SwiftVRPipeline": """Load the released local checkpoint through ModuleManager.""" root = Path(model_id_or_path) @@ -150,9 +398,32 @@ def from_pretrained( prompt_payload = load_file(str(root / "prompt_embedding.safetensors"), device="cpu") prompt_emb = prompt_payload["prompt_emb"][0] - pipeline_config = SwiftVRPipelineConfig() + pipeline_config = SwiftVRPipelineConfig( + enable_stage_parallel=enable_stage_parallel, + enable_stage_overlap=enable_stage_overlap, + dit_overlap=stage_dit_overlap, + tensor_channel_slots=tensor_channel_slots, + ) + runtime_configs = (pipeline_config.encode_config, pipeline_config.dit_config, pipeline_config.decode_config) + for runtime_config in runtime_configs: + runtime_config.torch_dtype = dtype if attention_config is not None: pipeline_config.dit_config.attention_config = attention_config + if compile_config is not None: + pipeline_config.dit_config.compile_config = compile_config + if quant_config is not None: + pipeline_config.dit_config.quant_config = quant_config + if stage_device_ids is not None: + if len(stage_device_ids) != 3: + raise ValueError("stage_device_ids must contain encode, dit, and decode device ids") + for runtime_config, device_id in zip( + (pipeline_config.encode_config, pipeline_config.dit_config, pipeline_config.decode_config), + stage_device_ids, + strict=True, + ): + runtime_config.device_type = torch.device(device).type + runtime_config.device_id = int(device_id) + runtime_config.parallel_config = ParallelConfig(device_ids=[int(device_id)]) pipeline = cls(device=device, torch_dtype=dtype) pipeline.init(module_manager, pipeline_config, prompt_emb) return pipeline @@ -234,8 +505,8 @@ def __call__( upscale: int = 4, clip_len: int = 24, dit_overlap: int = 0, - ) -> torch.Tensor: - """Restore [T,H,W,3] uint8 frames and return [1,T,3,H,W].""" + ) -> list[Image.Image]: + """Restore [T,H,W,3] uint8 frames and return PIL RGB frames.""" self._validate_clip_len(clip_len) if frames_uint8.ndim != 4 or frames_uint8.shape[-1] != 3 or frames_uint8.dtype != torch.uint8: raise ValueError("frames_uint8 must have shape [T,H,W,3] and dtype uint8") @@ -243,6 +514,24 @@ def __call__( if total_frames <= 0: raise ValueError("frames_uint8 must contain at least one frame") frames_uint8 = frames_uint8[:total_frames] + if self.config.enable_stage_parallel: + session = self.stream( + clip_len=clip_len, + resolution=resolution, + upscale=upscale, + dit_overlap=self.config.dit_overlap, + ) + try: + if self.config.enable_stage_overlap and isinstance(session, SwiftVRStagedStreamSession): + outputs = session.restore_chunks(frames_uint8, clip_len) + else: + outputs = [] + for start in range(0, int(frames_uint8.shape[0]), clip_len): + outputs.extend(session.step(frames_uint8[start : start + clip_len])) + outputs.extend(session.flush()) + finally: + session.close() + return outputs out_h, out_w, pad_h, pad_w = self._target_size( int(frames_uint8.shape[1]), int(frames_uint8.shape[2]), @@ -268,8 +557,8 @@ def __call__( ) ) if not outputs: - return torch.empty((1, 0, 3, out_h, out_w), dtype=self.torch_dtype, device=self.device) - return torch.cat(outputs, dim=1) + return [] + return ntchw_to_pil_frames(torch.cat(outputs, dim=1)) def stream( self, @@ -278,8 +567,20 @@ def stream( resolution: tuple[int, int] | None = None, upscale: int = 4, dit_overlap: int = 1, - ) -> "SwiftVRStreamSession": + ) -> "SwiftVRStreamSession | SwiftVRStagedStreamSession": self._validate_clip_len(clip_len) + if self.config.enable_stage_parallel: + if dit_overlap != self.config.dit_overlap: + raise ValueError("stage-parallel SwiftVR requires dit_overlap to match the configured stage overlap") + if self._stage_parallel_active_session: + raise RuntimeError("stage-parallel SwiftVR supports one active stream session per pipeline") + self._stage_parallel_active_session = True + return SwiftVRStagedStreamSession( + self, + clip_len=clip_len, + resolution=resolution, + upscale=upscale, + ) return SwiftVRStreamSession( self, clip_len=clip_len, @@ -288,6 +589,169 @@ def stream( dit_overlap=dit_overlap, ) + def close(self) -> None: + for stage in (self.decode_stage, self.dit_stage, self.encode_stage): + if isinstance(stage, ParallelWorker): + stage.close() + for channel in getattr(self, "_worker_tensor_channels", ()): + channel.close() + self._worker_tensor_channels = [] + self._stage_parallel_active_session = False + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass + + +class SwiftVRStagedStreamSession: + """Stage-parallel SwiftVR session using WorkerTensorChannel between stages.""" + + def __init__( + self, + pipeline: SwiftVRPipeline, + *, + clip_len: int, + resolution: tuple[int, int] | None, + upscale: int, + ) -> None: + self.pipeline = pipeline + self.clip_len = clip_len + self.resolution = resolution + self.upscale = upscale + self._sizes: tuple[int, int, int, int] | None = None + self._closed = False + + def _ensure_open(self) -> None: + if self._closed: + raise RuntimeError("SwiftVR stream session is closed") + + def _ensure_sizes(self, lq_h: int, lq_w: int) -> tuple[int, int, int, int]: + if self._sizes is None: + self._sizes = self.pipeline._target_size(lq_h, lq_w, self.resolution, self.upscale) + return self._sizes + + def _run_encoded(self, encoded: object, pad_h: int, pad_w: int) -> list[Image.Image]: + if encoded is None: + return [] + if not isinstance(self.pipeline.dit_stage, ParallelWorker) or not isinstance( + self.pipeline.decode_stage, ParallelWorker + ): + raise RuntimeError("SwiftVR stage workers are not initialized") + denoised_wait = self.pipeline.dit_stage.process(encoded, _tensor_transport=True) + denoised = denoised_wait() + return ntchw_to_pil_frames(self.pipeline.decode_stage.process(denoised, pad_h, pad_w, sync=True)) + + def _submit_encode(self, frames_uint8: torch.Tensor, pad: tuple[int, int, int, int]): + if not isinstance(self.pipeline.encode_stage, ParallelWorker): + raise RuntimeError("SwiftVR encode stage is not initialized") + out_h, out_w, pad_h, pad_w = pad + return self.pipeline.encode_stage.process( + frames_uint8, + out_h, + out_w, + pad_h, + pad_w, + self.pipeline.upscale_mode, + _tensor_transport=True, + ) + + def restore_chunks(self, frames_uint8: torch.Tensor, clip_len: int) -> list[Image.Image]: + self._ensure_open() + if frames_uint8.ndim != 4 or frames_uint8.shape[-1] != 3 or frames_uint8.dtype != torch.uint8: + raise ValueError("frames_uint8 must have shape [T,H,W,3] and dtype uint8") + if not isinstance(self.pipeline.dit_stage, ParallelWorker) or not isinstance( + self.pipeline.decode_stage, ParallelWorker + ): + raise RuntimeError("SwiftVR stage workers are not initialized") + pad = self._ensure_sizes(int(frames_uint8.shape[1]), int(frames_uint8.shape[2])) + _, _, pad_h, pad_w = pad + chunks = [frames_uint8[start : start + clip_len] for start in range(0, int(frames_uint8.shape[0]), clip_len)] + outputs: list[Image.Image] = [] + decode_wait = None + encode_wait = self._submit_encode(chunks[0], pad) if chunks else None + + for index, _chunk in enumerate(chunks): + encoded = encode_wait() if encode_wait is not None else None + next_index = index + 1 + encode_wait = self._submit_encode(chunks[next_index], pad) if next_index < len(chunks) else None + if encoded is None: + continue + denoised_wait = self.pipeline.dit_stage.process(encoded, _tensor_transport=True) + if decode_wait is not None: + output = decode_wait() + if output is not None: + outputs.extend(ntchw_to_pil_frames(output)) + denoised = denoised_wait() + decode_wait = self.pipeline.decode_stage.process(denoised, pad_h, pad_w) + + if encode_wait is not None: + encoded = encode_wait() + if encoded is not None: + denoised = self.pipeline.dit_stage.process(encoded, _tensor_transport=True)() + if decode_wait is not None: + output = decode_wait() + if output is not None: + outputs.extend(ntchw_to_pil_frames(output)) + decode_wait = self.pipeline.decode_stage.process(denoised, pad_h, pad_w) + + if isinstance(self.pipeline.encode_stage, ParallelWorker): + encoded = self.pipeline.encode_stage.flush_encoder(_tensor_transport=True)() + if encoded is not None: + denoised = self.pipeline.dit_stage.process(encoded, _tensor_transport=True)() + if decode_wait is not None: + output = decode_wait() + if output is not None: + outputs.extend(ntchw_to_pil_frames(output)) + decode_wait = self.pipeline.decode_stage.process(denoised, pad_h, pad_w) + + if decode_wait is not None: + output = decode_wait() + if output is not None: + outputs.extend(ntchw_to_pil_frames(output)) + return outputs + + @torch.inference_mode() + def step(self, frames_uint8: torch.Tensor) -> list[Image.Image]: + self._ensure_open() + if frames_uint8.ndim != 4 or frames_uint8.shape[-1] != 3 or frames_uint8.dtype != torch.uint8: + raise ValueError("frames_uint8 must have shape [T,H,W,3] and dtype uint8") + if not isinstance(self.pipeline.encode_stage, ParallelWorker): + raise RuntimeError("SwiftVR encode stage is not initialized") + out_h, out_w, pad_h, pad_w = self._ensure_sizes(int(frames_uint8.shape[1]), int(frames_uint8.shape[2])) + encoded_wait = self.pipeline.encode_stage.process( + frames_uint8, + out_h, + out_w, + pad_h, + pad_w, + self.pipeline.upscale_mode, + _tensor_transport=True, + ) + return self._run_encoded(encoded_wait(), pad_h, pad_w) + + @torch.inference_mode() + def flush(self) -> list[Image.Image]: + self._ensure_open() + if self._sizes is None: + return [] + if not isinstance(self.pipeline.encode_stage, ParallelWorker): + raise RuntimeError("SwiftVR encode stage is not initialized") + _, _, pad_h, pad_w = self._sizes + encoded_wait = self.pipeline.encode_stage.flush_encoder(_tensor_transport=True) + return self._run_encoded(encoded_wait(), pad_h, pad_w) + + def close(self) -> None: + if self._closed: + return + for stage in (self.pipeline.decode_stage, self.pipeline.dit_stage, self.pipeline.encode_stage): + if isinstance(stage, ParallelWorker): + stage.reset_session(sync=True) + self._sizes = None + self._closed = True + self.pipeline._stage_parallel_active_session = False + class SwiftVRStreamSession: """Per-session ReAE, overlap, and RoPE state for causal restoration.""" @@ -325,7 +789,7 @@ def _run_latents(self, encoded: torch.Tensor) -> torch.Tensor: return denoised.permute(0, 2, 1, 3, 4).contiguous() @torch.inference_mode() - def step(self, frames_uint8: torch.Tensor) -> torch.Tensor | None: + def step(self, frames_uint8: torch.Tensor) -> list[Image.Image]: self._ensure_open() if frames_uint8.ndim != 4 or frames_uint8.shape[-1] != 3 or frames_uint8.dtype != torch.uint8: raise ValueError("frames_uint8 must have shape [T,H,W,3] and dtype uint8") @@ -343,20 +807,20 @@ def step(self, frames_uint8: torch.Tensor) -> torch.Tensor | None: ) encoded = self._tae.encode_chunk(clip) if encoded is None: - return None + return [] decoded = self._tae.decode_chunk(self._run_latents(encoded)) - return crop_spatial_padding_ntchw(decoded, pad_h, pad_w) + return ntchw_to_pil_frames(crop_spatial_padding_ntchw(decoded, pad_h, pad_w)) @torch.inference_mode() - def flush(self) -> torch.Tensor | None: + def flush(self) -> list[Image.Image]: self._ensure_open() with self.pipeline._execution_lock: encoded = self._tae.flush_encoder() if encoded is None or self._sizes is None: - return None + return [] _, _, pad_h, pad_w = self._sizes decoded = self._tae.decode_chunk(self._run_latents(encoded)) - return crop_spatial_padding_ntchw(decoded, pad_h, pad_w) + return ntchw_to_pil_frames(crop_spatial_padding_ntchw(decoded, pad_h, pad_w)) def close(self) -> None: if self._closed: diff --git a/tests/unit/models/test_swiftvr_transformer.py b/tests/unit/models/test_swiftvr_transformer.py index c815435f..ca6e836d 100644 --- a/tests/unit/models/test_swiftvr_transformer.py +++ b/tests/unit/models/test_swiftvr_transformer.py @@ -2,9 +2,12 @@ import torch +from telefuser.core.config import CompileConfig, QuantConfig, QuantKernelBackend, QuantType from telefuser.models.swiftvr_transformer import ( + SwiftVRWanTransformer3DModel, _WindowRuntimeMetaCache, _make_hw_starts, + compile_transformer_blocks_with_config, get_1d_rotary_pos_embed, ) from telefuser.pipelines.swiftvr.streaming_dit import _ensure_rope_cache_len, _rope_with_offset @@ -70,3 +73,118 @@ def test_rope_offset_uses_global_temporal_position() -> None: torch.testing.assert_close(cos_grid[:, 0, 0, :4], rope.freqs_cos[5:8, :4]) torch.testing.assert_close(sin_grid[:, 0, 0, :4], rope.freqs_sin[5:8, :4]) + + +def test_swiftvr_transformer_torchao_fp8_quantization_uses_block_linears(monkeypatch) -> None: + calls = [] + + def replace_linear_layers_with_torchao_fp8(module, *, include_names=None, exclude_names=()): + calls.append((module, include_names, exclude_names)) + return 2 + + monkeypatch.setattr( + "telefuser.ops.torchao_fp8_linear.replace_linear_layers_with_torchao_fp8", + replace_linear_layers_with_torchao_fp8, + ) + model = SwiftVRWanTransformer3DModel( + patch_size=(1, 1, 1), + num_attention_heads=1, + attention_head_dim=4, + in_channels=4, + out_channels=4, + text_dim=4, + freq_dim=4, + ffn_dim=8, + num_layers=1, + cross_attn_norm=False, + ) + + model.enable_quant(QuantConfig(enabled=True, quant_type=QuantType.TORCHAO_FP8)) + + assert calls == [(model, ("blocks.",), ("head", "time_embedding", "time_projection", "patch_embedding"))] + assert model.torchao_fp8_replaced_linear == 2 + assert model.quant_type is QuantType.TORCHAO_FP8 + + +def test_compile_transformer_blocks_uses_runtime_config(monkeypatch) -> None: + model = SwiftVRWanTransformer3DModel( + patch_size=(1, 1, 1), + num_attention_heads=1, + attention_head_dim=4, + in_channels=4, + out_channels=4, + text_dim=4, + freq_dim=4, + ffn_dim=8, + num_layers=1, + cross_attn_norm=False, + ) + block = model.blocks[0] + calls = [] + + def compile_block(module, **kwargs): + calls.append((module, kwargs)) + return module + + monkeypatch.setattr(torch, "compile", compile_block) + compile_transformer_blocks_with_config( + model, + CompileConfig(enabled=True, mode="max-autotune-no-cudagraphs", fullgraph=False, dynamic=False), + ) + + assert calls == [ + ( + block, + { + "backend": "inductor", + "fullgraph": False, + "dynamic": False, + "mode": "max-autotune-no-cudagraphs", + }, + ) + ] + + +def test_swiftvr_transformer_tf_kernel_fp8_quantization_uses_block_linears(monkeypatch) -> None: + calls = [] + + def count_linear_layers(module, *, module_filter=None): + calls.append(("count", module, module_filter)) + return 3 + + def enable_fp8_gemm(module, *, options, module_filter=None): + calls.append(("enable", module, options, module_filter)) + return module + + monkeypatch.setattr("telefuser.ops.fp8_gemm.count_linear_layers", count_linear_layers) + monkeypatch.setattr("telefuser.ops.fp8_gemm.enable_fp8_gemm", enable_fp8_gemm) + model = SwiftVRWanTransformer3DModel( + patch_size=(1, 1, 1), + num_attention_heads=1, + attention_head_dim=4, + in_channels=4, + out_channels=4, + text_dim=4, + freq_dim=4, + ffn_dim=8, + num_layers=1, + cross_attn_norm=False, + ) + + model.enable_quant( + QuantConfig( + enabled=True, + quant_type=QuantType.FP8, + kernel_backend=QuantKernelBackend.TF_KERNEL, + ) + ) + + assert [call[0] for call in calls] == ["count", "enable"] + assert model.tf_kernel_fp8_replaced_linear == 3 + assert model.quant_type is QuantType.FP8 + options = calls[1][2] + assert options.fp16_weight_storage == "discard" + assert options.materialize_fp8_on_wrap is True + module_filter = calls[0][2] + assert module_filter("blocks.0.ffn.net.0.proj", torch.nn.Linear(4, 4)) + assert not module_filter("patch_embedding", torch.nn.Linear(4, 4)) diff --git a/tests/unit/pipelines/swiftvr/__init__.py b/tests/unit/pipelines/swiftvr/__init__.py new file mode 100644 index 00000000..0d8675dc --- /dev/null +++ b/tests/unit/pipelines/swiftvr/__init__.py @@ -0,0 +1 @@ +"""SwiftVR pipeline tests.""" diff --git a/tests/unit/pipelines/swiftvr/test_example.py b/tests/unit/pipelines/swiftvr/test_example.py index a9509c5f..a3a46bc2 100644 --- a/tests/unit/pipelines/swiftvr/test_example.py +++ b/tests/unit/pipelines/swiftvr/test_example.py @@ -1,15 +1,17 @@ from pathlib import Path +from types import SimpleNamespace import numpy as np import pytest import torch from PIL import Image -from examples.swiftvr.swiftvr_restore_h100 import _warmup_frame_count, main, run +from examples.swiftvr.swiftvr_restore_h100 import _parse_stage_devices, _warmup_frame_count, main, run class FakePipeline: def __init__(self) -> None: + self.config = SimpleNamespace(enable_stage_parallel=False) self.frames: torch.Tensor | None = None self.options: dict[str, object] = {} self.step_sizes: list[int] = [] @@ -19,13 +21,13 @@ def stream(self, **options: object) -> "FakePipeline": self.options = options return self - def step(self, frames: torch.Tensor) -> torch.Tensor: + def step(self, frames: torch.Tensor) -> list[Image.Image]: self.frames = frames self.step_sizes.append(len(frames)) - return frames.permute(0, 3, 1, 2).unsqueeze(0).float() / 255 + return [Image.fromarray(frame.numpy()) for frame in frames] - def flush(self) -> None: - return None + def flush(self) -> list[Image.Image]: + return [] def close(self) -> None: self.closed = True @@ -44,7 +46,7 @@ def test_run_uses_flashvsr_style_pil_interface() -> None: assert pipeline.frames is not None assert pipeline.frames.shape == (2, 8, 16, 3) assert pipeline.frames.dtype == torch.uint8 - assert pipeline.options == {"upscale": 2} + assert pipeline.options == {"upscale": 2, "dit_overlap": 0} assert pipeline.step_sizes == [2] assert pipeline.closed is True assert [image.size for image in outputs] == [(16, 8), (16, 8)] @@ -64,10 +66,33 @@ def test_run_processes_24_frame_streaming_chunks() -> None: outputs = run(pipeline, inputs, scale=3) assert pipeline.step_sizes == [24, 1] - assert pipeline.options == {"upscale": 3} + assert pipeline.options == {"upscale": 3, "dit_overlap": 0} assert len(outputs) == 25 +def test_run_uses_overlapped_pipeline_call_for_stage_parallel() -> None: + class FakeStagePipeline(FakePipeline): + def __init__(self) -> None: + super().__init__() + self.config.enable_stage_parallel = True + + def __call__(self, frames: torch.Tensor, **options: object) -> list[Image.Image]: + self.frames = frames + self.options = options + return [Image.fromarray(frame.numpy()) for frame in frames] + + def stream(self, **options: object) -> "FakePipeline": + raise AssertionError("stage-parallel run must use the overlapped pipeline call") + + pipeline = FakeStagePipeline() + inputs = [Image.fromarray(np.full((8, 8, 3), value, dtype=np.uint8)) for value in range(25)] + + outputs = run(pipeline, inputs, scale=4) + + assert len(outputs) == 25 + assert pipeline.options == {"upscale": 4, "clip_len": 24, "dit_overlap": 0} + + def test_warmup_covers_full_and_tail_shapes() -> None: assert _warmup_frame_count(81) == 57 assert _warmup_frame_count(72) == 48 @@ -92,4 +117,19 @@ def test_cli_uses_flashvsr_style_options() -> None: "gpu_num", "model_root", "output", + "attn_impl", + "compile_dit", + "compile_mode", + "quantization", + "enable_stage_parallel", + "stage_devices", + "tensor_channel_slots", } + + +def test_stage_device_parser_requires_encode_dit_decode_devices() -> None: + assert _parse_stage_devices("0,1,2", 1) == [0, 1, 2] + assert _parse_stage_devices(None, 3) == [0, 1, 2] + assert _parse_stage_devices(None, 1) is None + with pytest.raises(ValueError, match="exactly three devices"): + _parse_stage_devices("0,1", 1) diff --git a/tests/unit/pipelines/swiftvr/test_streaming.py b/tests/unit/pipelines/swiftvr/test_streaming.py index 77abcd8a..b2b7278b 100644 --- a/tests/unit/pipelines/swiftvr/test_streaming.py +++ b/tests/unit/pipelines/swiftvr/test_streaming.py @@ -1,13 +1,15 @@ import threading from types import SimpleNamespace +import numpy as np import torch +from PIL import Image from torch import nn from telefuser.models.swiftvr_reae import MemBlock, TPool from telefuser.pipelines.swiftvr import pipeline as pipeline_module from telefuser.pipelines.swiftvr.chunk import ChunkType, build_chunk_specs -from telefuser.pipelines.swiftvr.pipeline import SwiftVRStreamSession, aligned_pad +from telefuser.pipelines.swiftvr.pipeline import SwiftVRStagedStreamSession, SwiftVRStreamSession, aligned_pad from telefuser.pipelines.swiftvr.streaming_tae import apply_parallel_with_boundary @@ -152,6 +154,55 @@ def _target_size( second_second = second.step(second_frames) control_second = control.step(second_frames) - torch.testing.assert_close(second_first, control_first, rtol=0, atol=0) - torch.testing.assert_close(second_second, control_second, rtol=0, atol=0) + assert [np.asarray(frame).tolist() for frame in second_first] == [ + np.asarray(frame).tolist() for frame in control_first + ] + assert [np.asarray(frame).tolist() for frame in second_second] == [ + np.asarray(frame).tolist() for frame in control_second + ] assert second._closed is False + + +def test_stage_restore_chunks_converts_decoded_tensors_to_pil(monkeypatch) -> None: + class FakeWorker: + def __init__(self, role: str) -> None: + self.role = role + self.calls = 0 + + def process(self, value: object, *args: object, **kwargs: object): + self.calls += 1 + if self.role == "encode": + result = torch.tensor(float(self.calls)) + elif self.role == "dit": + result = value + else: + result = torch.full((1, 1, 3, 2, 2), float(self.calls) / 10) + return lambda: result + + def flush_encoder(self, **kwargs: object): + return lambda: None + + def reset_session(self, **kwargs: object) -> None: + return None + + monkeypatch.setattr(pipeline_module, "ParallelWorker", FakeWorker) + pipeline = SimpleNamespace( + encode_stage=FakeWorker("encode"), + dit_stage=FakeWorker("dit"), + decode_stage=FakeWorker("decode"), + upscale_mode="nearest", + _target_size=lambda height, width, resolution, upscale: (height, width, 0, 0), + _stage_parallel_active_session=True, + ) + session = SwiftVRStagedStreamSession( + pipeline, + clip_len=24, + resolution=None, + upscale=1, + ) + + output = session.restore_chunks(torch.zeros((48, 2, 2, 3), dtype=torch.uint8), clip_len=24) + + assert len(output) == 2 + assert all(isinstance(frame, Image.Image) for frame in output) + assert [np.asarray(frame)[0, 0, 0] for frame in output] == [25, 51]