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..0621b3d3 --- /dev/null +++ b/docs/en/swiftvr.md @@ -0,0 +1,132 @@ +# 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`; Torch SDPA remains the default, while other +TeleFuser dense attention backends can be selected explicitly. + +## 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. 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. +`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 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), +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. + +## Current constraints + +- 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 + results. diff --git a/examples/swiftvr/README.md b/examples/swiftvr/README.md new file mode 100644 index 00000000..7d743622 --- /dev/null +++ b/examples/swiftvr/README.md @@ -0,0 +1,185 @@ +# 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. + +## 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 +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 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 + +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_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 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 +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. 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 new file mode 100644 index 00000000..9c7fac39 --- /dev/null +++ b/examples/swiftvr/swiftvr_restore_h100.py @@ -0,0 +1,258 @@ +"""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, + CompileConfig, + QuantConfig, + QuantKernelBackend, + QuantType, +) +from telefuser.pipelines.swiftvr import SwiftVRPipeline +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, + 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.""" + 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(_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, + ) + + +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() + 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, dit_overlap=PPL_CONFIG["dit_overlap"]) + try: + for start in range(0, len(frames), PPL_CONFIG["chunk_size"]): + restored_frames.extend(session.step(frames[start : start + PPL_CONFIG["chunk_size"]])) + restored_frames.extend(session.flush()) + 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)") +@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, + height: int | None, + width: int | None, + 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): + 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"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, + 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() + 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..35749252 --- /dev/null +++ b/telefuser/models/swiftvr_transformer.py @@ -0,0 +1,1058 @@ +"""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, CompileConfig +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) + + +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 # +# --------------------------------------------------------------------------- # + + +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() + + 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, + 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/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/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/pipeline.py b/telefuser/pipelines/swiftvr/pipeline.py new file mode 100644 index 00000000..b11ee106 --- /dev/null +++ b/telefuser/pipelines/swiftvr/pipeline.py @@ -0,0 +1,833 @@ +"""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 +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.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, + 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 .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 + + +_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): + """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._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() + _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( + cls, + model_id_or_path: str, + *, + 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) + 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( + 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 + + @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, + ) -> 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") + 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] + 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]), + 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 [] + return ntchw_to_pil_frames(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 | 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, + resolution=resolution, + upscale=upscale, + 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.""" + + 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) -> 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") + 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 [] + decoded = self._tae.decode_chunk(self._run_latents(encoded)) + return ntchw_to_pil_frames(crop_spatial_padding_ntchw(decoded, pad_h, pad_w)) + + @torch.inference_mode() + 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 [] + _, _, pad_h, pad_w = self._sizes + decoded = self._tae.decode_chunk(self._run_latents(encoded)) + return ntchw_to_pil_frames(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..ca6e836d --- /dev/null +++ b/tests/unit/models/test_swiftvr_transformer.py @@ -0,0 +1,190 @@ +from types import SimpleNamespace + +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 + + +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]) + + +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/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") 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 new file mode 100644 index 00000000..a3a46bc2 --- /dev/null +++ b/tests/unit/pipelines/swiftvr/test_example.py @@ -0,0 +1,135 @@ +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 _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] = [] + self.closed = False + + def stream(self, **options: object) -> "FakePipeline": + self.options = options + return self + + def step(self, frames: torch.Tensor) -> list[Image.Image]: + self.frames = frames + self.step_sizes.append(len(frames)) + return [Image.fromarray(frame.numpy()) for frame in frames] + + def flush(self) -> list[Image.Image]: + return [] + + 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, "dit_overlap": 0} + 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, "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 + 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", + "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 new file mode 100644 index 00000000..b2b7278b --- /dev/null +++ b/tests/unit/pipelines/swiftvr/test_streaming.py @@ -0,0 +1,208 @@ +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 SwiftVRStagedStreamSession, 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) + + 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]