Support Relayvsr

README.md

@@ -239,13 +239,14 @@ LightX2V supports image, video, and world models and provides quantized weights,
 
 <details>
 <summary><strong>🔍 Restoration &amp; Super-Resolution</strong><br><br>
-<a href="scripts/seedvr2/README.md">SeedVR2</a> · <a href="scripts/swiftvr/README.md">SwiftVR</a>
+<a href="scripts/seedvr2/README.md">SeedVR2</a> · <a href="scripts/swiftvr/README.md">SwiftVR</a> · <a href="scripts/relayvsr/README.md">RelayVSR</a>
 </summary>
 
 | Models / Variants and Weights | Guides / Scripts |
 | --- | --- |
 | [SeedVR2](https://huggingface.co/ByteDance-Seed/SeedVR2-3B) | [Guide](scripts/seedvr2/README.md) |
 | [SwiftVR](https://huggingface.co/H-oliday/SwiftVR) | [Guide](scripts/swiftvr/README.md) |
+| [RelayVSR](https://huggingface.co/kopper/RelayVSR) | [Guide](scripts/relayvsr/README.md) |
 
 </details>
 

README_zh.md

@@ -239,13 +239,14 @@ LightX2V 支持多种图像、视频与世界模型，并提供量化权重、
 
 <details>
 <summary><strong>🔍 修复与超分辨率</strong><br><br>
-<a href="scripts/seedvr2/README_zh.md">SeedVR2</a> · <a href="scripts/swiftvr/README_zh.md">SwiftVR</a>
+<a href="scripts/seedvr2/README_zh.md">SeedVR2</a> · <a href="scripts/swiftvr/README_zh.md">SwiftVR</a> · <a href="scripts/relayvsr/README_zh.md">RelayVSR</a>
 </summary>
 
 | 模型 / 变体与权重 | 使用指南 / 脚本 |
 | --- | --- |
 | [SeedVR2](https://huggingface.co/ByteDance-Seed/SeedVR2-3B) | [使用指南](scripts/seedvr2/README_zh.md) |
 | [SwiftVR](https://huggingface.co/H-oliday/SwiftVR) | [使用指南](scripts/swiftvr/README_zh.md) |
+| [RelayVSR](https://huggingface.co/kopper/RelayVSR) | [使用指南](scripts/relayvsr/README_zh.md) |
 
 </details>
 

configs/relayvsr/relayvsr.json

@@ -0,0 +1,12 @@
+{
+  "cross_attention_backend": "torch_sdpa",
+  "cpu_offload": false,
+  "clip_len": 16,
+  "video_codec": "libx265",
+  "video_crf": 18,
+  "ffmpeg_preset": "ultrafast",
+  "encoder_threads": 16,
+  "queue_size": 16,
+  "use_compile": false,
+  "warmup": false
+}

configs/relayvsr/relayvsr_block_offload.json

@@ -0,0 +1,13 @@
+{
+  "cross_attention_backend": "torch_sdpa",
+  "cpu_offload": true,
+  "offload_granularity": "block",
+  "clip_len": 16,
+  "video_codec": "libx265",
+  "video_crf": 18,
+  "ffmpeg_preset": "ultrafast",
+  "encoder_threads": 16,
+  "queue_size": 16,
+  "use_compile": false,
+  "warmup": false
+}

configs/relayvsr/relayvsr_compile.json

@@ -0,0 +1,12 @@
+{
+  "cross_attention_backend": "torch_sdpa",
+  "cpu_offload": false,
+  "clip_len": 16,
+  "video_codec": "libx265",
+  "video_crf": 18,
+  "ffmpeg_preset": "ultrafast",
+  "encoder_threads": 16,
+  "queue_size": 16,
+  "use_compile": true,
+  "warmup": true
+}

configs/relayvsr/relayvsr_sp.json

@@ -0,0 +1,16 @@
+{
+  "cross_attention_backend": "torch_sdpa",
+  "cpu_offload": false,
+  "clip_len": 16,
+  "video_codec": "libx265",
+  "video_crf": 18,
+  "ffmpeg_preset": "ultrafast",
+  "encoder_threads": 16,
+  "queue_size": 16,
+  "use_compile": true,
+  "warmup": true,
+  "parallel": {
+    "seq_p_size": 2,
+    "cfg_p_size": 1
+  }
+}

lightx2v/common/kvcache/relayvsr.py

@@ -0,0 +1,26 @@
+"""Frame-slot caches for RelayVSR's unrotated Wan and Flash keys/values."""
+
+from .base import BaseKVCachePool
+from .manager import register_self_attn_kv_cache
+
+
+class FrameRingKVCachePool(BaseKVCachePool):
+    """Resident physical frame slots, overwritten in place as frames advance."""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._init_kv_buffer()
+
+    def update(self, k, v, layer_id, frame_index):
+        span = k.shape[0]
+        if self.cache_size % span:
+            raise ValueError("Frame cache capacity must be a multiple of the frame token count")
+        frames = self.cache_size // span
+        start = (frame_index % frames) * span
+        self._k_buffer[layer_id, start : start + span].copy_(k)
+        self._v_buffer[layer_id, start : start + span].copy_(v)
+        visible = min(frame_index + 1, frames) * span
+        return self.k_cache(layer_id, 0, visible), self.v_cache(layer_id, 0, visible)
+
+
+register_self_attn_kv_cache("frame_ring", FrameRingKVCachePool)

lightx2v/models/networks/relayvsr/compile.py

@@ -0,0 +1,53 @@
+"""Compile dispatch, sequence-only dynamic shapes and native RelayVSR norms."""
+
+from functools import partial
+
+import torch
+import torch.nn.functional as F
+
+
+class RelayCompileMixin:
+    def get_compiled_block(self, block_idx, block):
+        key = self.get_compile_block_key(block_idx, block)
+        cached = self.compiled_blocks.get(key)
+        if cached is not None and cached[0] is block:
+            return cached[1]
+
+        # Compile the family method before binding the block so Wan and Flash
+        # do not share a wrapper's Dynamo cache and recompile limit.
+        if self.compile_backend == "default":
+            options = {"options": self.compile_options} if self.compile_options else {}
+            compiled = torch.compile(self.infer_block, dynamic=None, **options)
+        else:
+            compiled = torch.compile(self.infer_block, dynamic=None, backend=self._compile_backend_obj)
+        bound = partial(compiled, block)
+        self.compiled_blocks[key] = (block, bound)
+        return bound
+
+
+def mark_sequence_dynamic(*tensors):
+    for tensor in tensors:
+        torch._dynamo.mark_dynamic(tensor, 0)
+        torch._dynamo.mark_static(tensor, list(range(1, tensor.ndim)))
+
+
+# Preserve the checkpoint's native reductions and low-precision rounding while
+# Inductor fuses the surrounding projections, modulation and residuals.
+@torch.library.custom_op("lightx2v::relayvsr_rms_norm", mutates_args=(), device_types="cuda")
+def native_rms_norm(x: torch.Tensor, weight: torch.Tensor | None, eps: float) -> torch.Tensor:
+    return F.rms_norm(x, (x.shape[-1],), weight, eps)
+
+
+@native_rms_norm.register_fake
+def _rms_norm_fake(x, weight, eps):
+    return torch.empty_like(x, memory_format=torch.contiguous_format)
+
+
+@torch.library.custom_op("lightx2v::relayvsr_layer_norm", mutates_args=(), device_types="cuda")
+def native_layer_norm(x: torch.Tensor, weight: torch.Tensor | None, bias: torch.Tensor | None, eps: float) -> torch.Tensor:
+    return F.layer_norm(x.float(), (x.shape[-1],), None if weight is None else weight.float(), None if bias is None else bias.float(), eps)
+
+
+@native_layer_norm.register_fake
+def _layer_norm_fake(x, weight, bias, eps):
+    return torch.empty_like(x, dtype=torch.float32, memory_format=torch.contiguous_format)

lightx2v/models/networks/relayvsr/flash_attention.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+# Adapted from kopperx/RelayVSR.
+# Copyright RelayVSR contributors.
+
+"""Single-query-frame local attention over condition and video KV planes."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@torch.library.custom_op("lightx2v::relayvsr_flash_attention", mutates_args=(), device_types="cuda")
+def _compiled_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, height: int, width: int, windows: list[int]) -> torch.Tensor:
+    windows = tuple(zip(windows[::2], windows[1::2]))
+    if len(windows) > 1:
+        frames = key.shape[2] // (height * width)
+        windows = windows[:2] + (windows[2],) * (frames - 2)
+    return triton_local_attention(query, key, value, height, width, windows)
+
+
+@_compiled_attention.register_fake
+def _compiled_attention_fake(query, key, value, height, width, windows):
+    return torch.empty_like(query, memory_format=torch.contiguous_format)
+
+
+@triton.jit
+def _forward(
+    Q,
+    K,
+    V,
+    OUT,
+    HEIGHT: tl.constexpr,
+    WIDTH: tl.constexpr,
+    HQ: tl.constexpr,
+    HK: tl.constexpr,
+    D: tl.constexpr,
+    WHS: tl.constexpr,
+    WWS: tl.constexpr,
+    COLS: tl.constexpr,
+    NF: tl.constexpr,
+    BH: tl.constexpr,
+    BW: tl.constexpr,
+    BM: tl.constexpr,
+    BN: tl.constexpr,
+):
+    n: tl.constexpr = HEIGHT * WIDTH
+    tile_y, tile_x, bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    batch, head = bh // HQ, bh % HQ
+    kv_head = head // (HQ // HK)
+    row = tl.arange(0, BM)
+    qy, qx = tile_y * BH + row // BW, tile_x * BW + row % BW
+    qi = qy * WIDTH + qx
+    valid_q = (qy < HEIGHT) & (qx < WIDTH)
+    d = tl.arange(0, D)
+    q = tl.load(
+        Q + ((batch * HQ + head) * n + qi[:, None]) * D + d[None, :],
+        valid_q[:, None],
+        0,
+    )
+    acc = tl.full((BM, D), 0, tl.float32)
+    maximum = tl.full((BM,), -float("inf"), tl.float32)
+    denominator = tl.full((BM,), 0, tl.float32)
+    scale: tl.constexpr = D**-0.5 * 1.4426950408889634
+    for source in tl.static_range(len(WHS)):
+        wh = WHS[source]
+        ww = WWS[source]
+        rows = wh + BH - 1
+        top, left = tile_y * BH - wh // 2, tile_x * BW - ww // 2
+        for offset in range(tl.cdiv(rows * COLS, BN)):
+            p = offset * BN + tl.arange(0, BN)
+            ky, kx = top + p // COLS, left + p % COLS
+            valid_k = (p < rows * COLS) & (ky >= 0) & (ky < HEIGHT) & (kx >= 0) & (kx < WIDTH)
+            ki = source * n + ky * WIDTH + kx
+            k = tl.load(
+                K + ((batch * HK + kv_head) * NF * n + ki[None, :]) * D + d[:, None],
+                valid_k[None, :],
+                0,
+            )
+            scores = tl.dot(q, k).to(tl.float32) * scale
+            allowed = valid_k[None, :] & (ky[None, :] >= qy[:, None] - wh // 2) & (ky[None, :] < qy[:, None] + wh // 2) & (kx[None, :] >= qx[:, None] - ww // 2) & (kx[None, :] < qx[:, None] + ww // 2)
+            scores = tl.where(allowed, scores, -float("inf"))
+            new_max = tl.maximum(maximum, tl.max(scores, 1))
+            safe_max = tl.where(new_max == -float("inf"), 0.0, new_max)
+            probabilities = tl.exp2(scores - safe_max[:, None])
+            correction = tl.exp2(maximum - safe_max)
+            v = tl.load(
+                V + ((batch * HK + kv_head) * NF * n + ki[:, None]) * D + d[None, :],
+                valid_k[:, None],
+                0,
+            )
+            acc = acc * correction[:, None] + tl.dot(probabilities.to(v.dtype), v)
+            denominator = denominator * correction + tl.sum(probabilities, 1)
+            maximum = new_max
+    result = acc / denominator[:, None]
+    tl.store(
+        OUT + ((batch * HQ + head) * n + qi[:, None]) * D + d[None, :],
+        result,
+        valid_q[:, None],
+    )
+
+
+@torch.no_grad()
+def triton_local_attention(query, key, value, height, width, windows):
+    if torch.compiler.is_compiling():
+        return _compiled_attention(query, key, value, height, width, [size for window in windows for size in window])
+    with torch.cuda.device(query.device):
+        q, k, v = query.contiguous(), key.contiguous(), value.contiguous()
+        b, hq, _, d = q.shape
+        hk = k.shape[1]
+        output = torch.empty_like(q)
+        whs, wws = tuple(w[0] for w in windows), tuple(w[1] for w in windows)
+        bh, bw, bn, warps, stages = 4, 16, 64, 4, 2
+        grid = (triton.cdiv(height, bh), triton.cdiv(width, bw), b * hq)
+        _forward[grid](
+            q,
+            k,
+            v,
+            output,
+            height,
+            width,
+            hq,
+            hk,
+            d,
+            whs,
+            wws,
+            triton.next_power_of_2(max(wws) + bw - 1),
+            len(windows),
+            bh,
+            bw,
+            bh * bw,
+            bn,
+            num_warps=warps,
+            num_stages=stages,
+            enable_fp_fusion=True,
+        )
+        return output