From 2aac3f4e42d33471b711e81914a654d46b3c8996 Mon Sep 17 00:00:00 2001 From: Musisoul Date: Thu, 22 Oct 2026 12:13:49 +0000 Subject: [PATCH 1/6] Support relayvsr --- configs/relayvsr/relayvsr.json | 10 + configs/relayvsr/relayvsr_block_offload.json | 11 + configs/relayvsr/relayvsr_sp.json | 14 + lightx2v/common/kvcache/relayvsr.py | 26 ++ lightx2v/models/networks/relayvsr/__init__.py | 0 .../networks/relayvsr/flash_attention.py | 120 ++++++++ .../models/networks/relayvsr/flash_infer.py | 80 +++++ lightx2v/models/networks/relayvsr/model.py | 76 +++++ lightx2v/models/networks/relayvsr/parallel.py | 35 +++ lightx2v/models/networks/relayvsr/rope.py | 60 ++++ .../models/networks/relayvsr/wan_attention.py | 193 ++++++++++++ .../models/networks/relayvsr/wan_infer.py | 175 +++++++++++ lightx2v/models/networks/relayvsr/weights.py | 185 +++++++++++ lightx2v/models/runners/relayvsr/__init__.py | 0 .../runners/relayvsr/relayvsr_runner.py | 287 ++++++++++++++++++ lightx2v/models/runners/relayvsr/video_io.py | 119 ++++++++ lightx2v/models/runners/runner_factory.py | 1 + .../models/schedulers/relayvsr/__init__.py | 0 .../models/schedulers/relayvsr/scheduler.py | 23 ++ scripts/relayvsr/run_relayvsr_sr.sh | 22 ++ 20 files changed, 1437 insertions(+) create mode 100644 configs/relayvsr/relayvsr.json create mode 100644 configs/relayvsr/relayvsr_block_offload.json create mode 100644 configs/relayvsr/relayvsr_sp.json create mode 100644 lightx2v/common/kvcache/relayvsr.py create mode 100644 lightx2v/models/networks/relayvsr/__init__.py create mode 100644 lightx2v/models/networks/relayvsr/flash_attention.py create mode 100644 lightx2v/models/networks/relayvsr/flash_infer.py create mode 100644 lightx2v/models/networks/relayvsr/model.py create mode 100644 lightx2v/models/networks/relayvsr/parallel.py create mode 100644 lightx2v/models/networks/relayvsr/rope.py create mode 100644 lightx2v/models/networks/relayvsr/wan_attention.py create mode 100644 lightx2v/models/networks/relayvsr/wan_infer.py create mode 100644 lightx2v/models/networks/relayvsr/weights.py create mode 100644 lightx2v/models/runners/relayvsr/__init__.py create mode 100644 lightx2v/models/runners/relayvsr/relayvsr_runner.py create mode 100644 lightx2v/models/runners/relayvsr/video_io.py create mode 100644 lightx2v/models/schedulers/relayvsr/__init__.py create mode 100644 lightx2v/models/schedulers/relayvsr/scheduler.py create mode 100644 scripts/relayvsr/run_relayvsr_sr.sh diff --git a/configs/relayvsr/relayvsr.json b/configs/relayvsr/relayvsr.json new file mode 100644 index 000000000..0b3028007 --- /dev/null +++ b/configs/relayvsr/relayvsr.json @@ -0,0 +1,10 @@ +{ + "cross_attention_backend": "torch_sdpa", + "cpu_offload": false, + "clip_len": 16, + "video_codec": "libx264", + "video_crf": 18, + "ffmpeg_preset": "fast", + "encoder_threads": 8, + "queue_size": 3 +} diff --git a/configs/relayvsr/relayvsr_block_offload.json b/configs/relayvsr/relayvsr_block_offload.json new file mode 100644 index 000000000..778d6e30a --- /dev/null +++ b/configs/relayvsr/relayvsr_block_offload.json @@ -0,0 +1,11 @@ +{ + "cross_attention_backend": "torch_sdpa", + "cpu_offload": true, + "offload_granularity": "block", + "clip_len": 16, + "video_codec": "libx264", + "video_crf": 18, + "ffmpeg_preset": "fast", + "encoder_threads": 8, + "queue_size": 3 +} diff --git a/configs/relayvsr/relayvsr_sp.json b/configs/relayvsr/relayvsr_sp.json new file mode 100644 index 000000000..d1a6c1b93 --- /dev/null +++ b/configs/relayvsr/relayvsr_sp.json @@ -0,0 +1,14 @@ +{ + "cross_attention_backend": "torch_sdpa", + "cpu_offload": false, + "clip_len": 16, + "video_codec": "libx264", + "video_crf": 18, + "ffmpeg_preset": "fast", + "encoder_threads": 8, + "queue_size": 3, + "parallel": { + "seq_p_size": 2, + "cfg_p_size": 1 + } +} diff --git a/lightx2v/common/kvcache/relayvsr.py b/lightx2v/common/kvcache/relayvsr.py new file mode 100644 index 000000000..7d4d1e6f7 --- /dev/null +++ b/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) diff --git a/lightx2v/models/networks/relayvsr/__init__.py b/lightx2v/models/networks/relayvsr/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/lightx2v/models/networks/relayvsr/flash_attention.py b/lightx2v/models/networks/relayvsr/flash_attention.py new file mode 100644 index 000000000..75dab2d6f --- /dev/null +++ b/lightx2v/models/networks/relayvsr/flash_attention.py @@ -0,0 +1,120 @@ +# 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 + + +@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): + 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 diff --git a/lightx2v/models/networks/relayvsr/flash_infer.py b/lightx2v/models/networks/relayvsr/flash_infer.py new file mode 100644 index 000000000..8559d0f17 --- /dev/null +++ b/lightx2v/models/networks/relayvsr/flash_infer.py @@ -0,0 +1,80 @@ +import torch +import torch.nn.functional as F + +from .flash_attention import triton_local_attention +from .rope import apply_rope + + +class FlashPreInfer: + def frame(self, weights, lr): + x = lr + for i in (0, 2, 4): + x = F.silu(getattr(weights, f"conv{i}").apply(x)) + x = weights.conv6.apply(x).flatten(2).transpose(1, 2)[0] + return weights.lr_norm.apply(x) + weights.video_type.tensor[0] + + def condition(self, weights, latent): + x = weights.condition.apply(latent).flatten(2).transpose(1, 2)[0] + return weights.condition_norm.apply(x) + weights.condition_type.tensor[0] + + +class FlashTransformerInfer: + def __init__(self, config, parallel): + self.config, self.parallel = config, parallel + + def infer_block(self, block, x, condition, past_key, past_value, q_rope, k_rope, windows, height, width): + tokens = height * width + normalized = block.attention_norm.apply(x) + query = block.query.apply(normalized).view(-1, 8, 64) + key_op = block.condition_key if condition else block.key + value_op = block.condition_value if condition else block.value + key = block.key_norm.apply(key_op.apply(normalized).view(-1, 2, 64)) + value = block.value_norm.apply(value_op.apply(normalized).view(-1, 2, 64)) + query = self.parallel.to_heads(query, tokens) + key = self.parallel.to_heads(key, tokens) + value = self.parallel.to_heads(value, tokens) + all_key = key if condition else torch.cat((past_key, key), dim=0) + all_value = value if condition else torch.cat((past_value, value), dim=0) + query = apply_rope(query, q_rope) + all_key = apply_rope(all_key, k_rope) + query, all_key, all_value = query.transpose(0, 1)[None], all_key.transpose(0, 1)[None], all_value.transpose(0, 1)[None] + attended = triton_local_attention(query, all_key, all_value, height, width, windows) + attended = attended[0].transpose(0, 1) + x = x + block.output.apply(self.parallel.to_sequence(attended)) + gate, up = block.gate_up.apply(block.feed_forward_norm.apply(x)).chunk(2, -1) + x = x + block.down.apply(F.silu(gate) * up) + return x, key, value + + def prepare_condition(self, blocks, x, rope, height, width): + frequencies = rope.for_positions(height, width, (0,)) + window = tuple(self.config["spatial_attention"]["condition_self_window_size"]) + caches = [] + for block in blocks: + x, key, value = self.infer_block(block, x, True, None, None, frequencies, frequencies, (window,), height, width) + caches.append((key, value)) + return caches + + def infer(self, blocks, x, conditions, history, rope, frame, height, width): + spatial = self.config["spatial_attention"] + previous = min(frame, self.config["window_size"] - 1) + q_rope = rope.for_positions(height, width, (0,)) + k_rope = rope.for_positions(height, width, (0, 0, *range(-previous, 1))) + windows = (tuple(spatial["start_condition_window_size"]), tuple(spatial["end_condition_window_size"])) + (tuple(spatial["video_window_size"]),) * (previous + 1) + for index, block in enumerate(blocks): + ck, cv = conditions[index] + if previous: + past_key = torch.cat((ck, history.k_cache(index)), 0) + past_value = torch.cat((cv, history.v_cache(index)), 0) + else: + past_key, past_value = ck, cv + x, key, value = self.infer_block(block, x, False, past_key, past_value, q_rope, k_rope, windows, height, width) + history.store_kv(key, value, index) + return x + + +class FlashPostInfer: + def infer(self, weights, hidden, lr): + height, width = lr.shape[-2] // 4, lr.shape[-1] // 4 + pixels = weights.projection.apply(F.silu(weights.head.apply(hidden))) + residual = F.pixel_shuffle(pixels.t().reshape(1, 768, height, width), 16) + return F.interpolate(lr, scale_factor=4, mode="bicubic", align_corners=False) + residual diff --git a/lightx2v/models/networks/relayvsr/model.py b/lightx2v/models/networks/relayvsr/model.py new file mode 100644 index 000000000..28a4b193f --- /dev/null +++ b/lightx2v/models/networks/relayvsr/model.py @@ -0,0 +1,76 @@ +import torch + +from lightx2v.utils.envs import GET_DTYPE + +from .flash_infer import FlashPostInfer, FlashPreInfer, FlashTransformerInfer +from .parallel import SequenceParallel +from .rope import RotaryEmbedding +from .wan_infer import RelayWanPostInfer, RelayWanPreInfer, RelayWanTransformerInfer +from .weights import FlashWeights, WanWeights, checkpoint_config, resolve_checkpoints + + +def validate_checkpoint_configs(wan, flash): + expected_wan = { + "patch_size": [1, 2, 2], + "in_channels": 48, + "out_channels": 48, + "qk_norm": "rms_norm_across_heads", + "cross_attn_norm": True, + "num_frames_per_block": 1, + "spatial_attention_mode": "local", + } + expected_flash = {"variant": "S", "upscale_factor": 4, "window_size": 2, "condition_channels": 48, "condition_mode": "endpoints", "temporal_rope_policy": "current_relative"} + for actual, expected in ((wan, expected_wan), (flash, expected_flash)): + for key, value in expected.items(): + if actual.get(key) != value: + raise ValueError(f"Unsupported RelayVSR checkpoint: {key}={actual.get(key)!r}, expected {value!r}") + if wan.get("temporal_rope", {}).get("policy") != "fixed_end": + raise ValueError("RelayVSR requires fixed_end Wan temporal RoPE") + + +class RelayVSRModel: + def __init__(self, config, group=None): + self.config = config + wan_path, flash_path = resolve_checkpoints(config) + self.wan_config = checkpoint_config(wan_path, "wan") + self.flash_config = checkpoint_config(flash_path, "flash") + validate_checkpoint_configs(self.wan_config, self.flash_config) + self.parallel = SequenceParallel(group) + self.wan_weights = WanWeights(wan_path, self.wan_config, config) + self.flash_weights = FlashWeights(flash_path, self.flash_config) + self.wan_pre = RelayWanPreInfer(self.wan_config) + self.wan_transformer = RelayWanTransformerInfer(self.wan_config, config, self.parallel, self.wan_weights) + self.wan_post = RelayWanPostInfer({"out_dim": self.wan_config["out_channels"]}) + self.flash_pre = FlashPreInfer() + self.flash_transformer = FlashTransformerInfer(self.flash_config, self.parallel) + self.flash_post = FlashPostInfer() + self.wan_rope = RotaryEmbedding(wan=True, device=torch.device("cuda"), dtype=GET_DTYPE(), head_dim=self.wan_config["attention_head_dim"], max_length=self.wan_config["rope_max_seq_len"]) + self.flash_rope = RotaryEmbedding(wan=False, device=torch.device("cuda"), dtype=GET_DTYPE(), theta=self.flash_config["rope_theta"]) + self.temb, self.modulation, self.context = self.wan_pre.conditioning(self.wan_weights.pre) + self.wan_transformer.prepare_cross_attention(self.wan_weights.blocks, self.context) + + def set_scheduler(self, scheduler): + self.wan_transformer.set_scheduler(scheduler) + self.wan_post.set_scheduler(scheduler) + + def infer(self, scheduler, lr, cache, frame): + height, width = lr.shape[-2] // 8, lr.shape[-1] // 8 + hidden = self.parallel.split(self.wan_pre.infer(self.wan_weights.pre, scheduler.latents.to(GET_DTYPE()), lr)) + pre = self.wan_transformer.prepare_input(hidden, self.temb, self.modulation, self.context, self.wan_rope, cache, frame, height, width) + hidden = self.wan_transformer.infer(self.wan_weights, pre) + scheduler.noise_pred = self.wan_post.infer(hidden, pre) + + def prepare_condition(self, latent): + height, width = latent.shape[-2:] + hidden = self.parallel.split(self.flash_pre.condition(self.flash_weights.pre, latent)) + return self.flash_transformer.prepare_condition(self.flash_weights.blocks, hidden, self.flash_rope, height, width) + + def decode(self, lr, conditions, history, frame): + height, width = lr.shape[-2] // 4, lr.shape[-1] // 4 + hidden = self.parallel.split(self.flash_pre.frame(self.flash_weights.pre, lr)) + hidden = self.flash_transformer.infer(self.flash_weights.blocks, hidden, conditions, history, self.flash_rope, frame, height, width) + return self.flash_post.infer(self.flash_weights.post, self.parallel.gather(hidden, height * width), lr) + + def clear(self): + self.wan_rope.clear() + self.flash_rope.clear() diff --git a/lightx2v/models/networks/relayvsr/parallel.py b/lightx2v/models/networks/relayvsr/parallel.py new file mode 100644 index 000000000..a7d6949ff --- /dev/null +++ b/lightx2v/models/networks/relayvsr/parallel.py @@ -0,0 +1,35 @@ +"""Reuse LightX2V Ulysses collectives around full-grid local attention.""" + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from lightx2v.common.ops.attn.utils.all2all import all2all_head2seq, all2all_seq2head + + +class SequenceParallel: + def __init__(self, group=None): + self.group = group + self.size = dist.get_world_size(group) if group is not None else 1 + self.rank = dist.get_rank(group) if group is not None else 0 + + def split(self, x): + if self.size == 1: + return x + return F.pad(x, (0, 0, 0, -x.shape[0] % self.size)).chunk(self.size)[self.rank].contiguous() + + def gather(self, x, tokens): + if self.size == 1: + return x + output = torch.empty((x.shape[0] * self.size, *x.shape[1:]), device=x.device, dtype=x.dtype) + dist.all_gather_into_tensor(output, x.contiguous(), group=self.group) + return output[:tokens] + + def to_heads(self, x, tokens): + return all2all_seq2head(x, group=self.group)[:tokens] if self.size > 1 else x + + def to_sequence(self, x): + if self.size == 1: + return x.flatten(1) + x = F.pad(x, (0, 0, 0, 0, 0, -x.shape[0] % self.size)) + return all2all_head2seq(x, group=self.group).flatten(1) diff --git a/lightx2v/models/networks/relayvsr/rope.py b/lightx2v/models/networks/relayvsr/rope.py new file mode 100644 index 000000000..e4013ee38 --- /dev/null +++ b/lightx2v/models/networks/relayvsr/rope.py @@ -0,0 +1,60 @@ +import torch + + +def apply_rope(x, rope): + cosine, sine = rope + even, odd = x[..., 0::2], x[..., 1::2] + out = torch.empty_like(x) + out[..., 0::2] = even * cosine[:, None, 0::2] - odd * sine[:, None, 1::2] + out[..., 1::2] = even * sine[:, None, 1::2] + odd * cosine[:, None, 0::2] + return out + + +class RotaryEmbedding: + def __init__(self, *, wan, device, dtype, head_dim=128, max_length=1024, theta=10000.0): + self.wan = wan + self.device, self.dtype = device, dtype + self.max_length = max_length + self.dims = (head_dim - 4 * (head_dim // 6), 2 * (head_dim // 6), 2 * (head_dim // 6)) if wan else (16, 24, 24) + self.cache = {} + self.tables = [] + for dim in self.dims: + if wan: + # Upstream builds these tables on CPU in float64, then casts to float32. + inverse = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float64) / dim)) + angles = torch.outer(torch.arange(max_length, dtype=torch.float64), inverse) + self.tables.append((angles.cos().repeat_interleave(2, -1).float().to(device), angles.sin().repeat_interleave(2, -1).float().to(device))) + else: + self.tables.append(1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))) + + def for_positions(self, height, width, positions): + key = (height, width, positions) + if key in self.cache: + return self.cache[key] + if self.wan and (max(height, width) > self.max_length or min(positions) < 0 or max(positions) >= self.max_length): + raise ValueError("RelayVSR spatial or temporal position exceeds the checkpoint's RoPE range") + frames = len(positions) + coordinates = ( + torch.tensor(positions, device=self.device).view(frames, 1, 1).expand(frames, height, width).reshape(-1), + torch.arange(height, device=self.device).view(1, height, 1).expand(frames, height, width).reshape(-1), + torch.arange(width, device=self.device).view(1, 1, width).expand(frames, height, width).reshape(-1), + ) + cosines, sines = [], [] + for coordinate, table in zip(coordinates, self.tables): + if self.wan: + cosine, sine = table[0][coordinate], table[1][coordinate] + else: + angles = coordinate.float().unsqueeze(-1) * table + cosine, sine = angles.cos().repeat_interleave(2, -1), angles.sin().repeat_interleave(2, -1) + cosines.append(cosine) + sines.append(sine) + result = (torch.cat(cosines, -1), torch.cat(sines, -1)) + if not self.wan: + result = tuple(value.to(self.dtype) for value in result) + if len(self.cache) >= 32: + self.cache.clear() + self.cache[key] = result + return result + + def clear(self): + self.cache.clear() diff --git a/lightx2v/models/networks/relayvsr/wan_attention.py b/lightx2v/models/networks/relayvsr/wan_attention.py new file mode 100644 index 000000000..17791b46b --- /dev/null +++ b/lightx2v/models/networks/relayvsr/wan_attention.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from kopperx/RelayVSR. +# Copyright RelayVSR contributors. + +"""Exact boundary-preserving spatial attention for inference on CUDA. + +Q/K already carry spatial RoPE and their resolved temporal RoPE. Each program handles a 2-D query +tile and directly loads its rectangular KV neighborhood. No replicated +KV windows or quadratic mask are materialized. +""" + +from dataclasses import dataclass + +import torch +import triton +import triton.language as tl + + +@dataclass(frozen=True) +class SpatialAttentionSpec: + frames: int + height: int + width: int + window_height: int + window_width: int + temporal_window: int + key_frames: int + query_frame_offset: int = 0 + key_frame_offset: int = 0 + + +@triton.jit +def _spatial_forward( + Q, + K, + V, + OUT, + qs0: tl.constexpr, + qs1: tl.constexpr, + qs2: tl.constexpr, + ks0: tl.constexpr, + ks1: tl.constexpr, + ks2: tl.constexpr, + vs0: tl.constexpr, + vs1: tl.constexpr, + vs2: tl.constexpr, + FRAMES: tl.constexpr, + KEY_FRAMES: tl.constexpr, + QUERY_OFFSET: tl.constexpr, + KEY_OFFSET: tl.constexpr, + HEIGHT: tl.constexpr, + WIDTH: tl.constexpr, + WH: tl.constexpr, + WW: tl.constexpr, + TW: tl.constexpr, + HEADS: tl.constexpr, + DIM: tl.constexpr, + SCALE: tl.constexpr, + BH: tl.constexpr, + BW: tl.constexpr, + BM: tl.constexpr, + BN: tl.constexpr, + COLS: tl.constexpr, + ROWS: tl.constexpr, +): + row = tl.program_id(0) + segment = tl.program_id(1) + bh = tl.program_id(2) + batch, head = bh // HEADS, bh % HEADS + qt, row_tile = row // triton.cdiv(HEIGHT, BH), row % triton.cdiv(HEIGHT, BH) + qy = row_tile * BH + tl.arange(0, BM) // BW + qx = segment * BW + tl.arange(0, BM) % BW + qi = (qt * HEIGHT + qy) * WIDTH + qx + valid_q = (qx < WIDTH) & (qy < HEIGHT) + d = tl.arange(0, DIM) + top = tl.minimum(tl.maximum(qy - WH // 2, 0), HEIGHT - WH) + left = tl.minimum(tl.maximum(qx - WW // 2, 0), WIDTH - WW) + base_left = tl.minimum(tl.maximum(segment * BW - WW // 2, 0), WIDTH - WW) + base_top = tl.minimum(tl.maximum(row_tile * BH - WH // 2, 0), HEIGHT - WH) + q = tl.load( + Q + batch * qs0 + qi[:, None] * qs1 + head * qs2 + d[None, :], + valid_q[:, None], + 0, + ) + acc = tl.full((BM, DIM), 0, tl.float32) + maximum = tl.full((BM,), -float("inf"), tl.float32) + denominator = tl.full((BM,), 0, tl.float32) + logical_qt = qt + QUERY_OFFSET - KEY_OFFSET + first_t = tl.maximum(0, logical_qt - TW + 1) + last_t = tl.minimum(KEY_FRAMES, logical_qt + 1) + for kt in range(first_t, last_t): + for offset in range(triton.cdiv(ROWS * COLS, BN)): + p = offset * BN + tl.arange(0, BN) + ky = base_top + p // COLS + kx = base_left + p % COLS + valid_k = (p < ROWS * COLS) & (kx < WIDTH) & (ky < HEIGHT) + ki = kt * HEIGHT * WIDTH + ky * WIDTH + kx + k = tl.load( + K + batch * ks0 + ki[None, :] * ks1 + head * ks2 + d[:, None], + valid_k[None, :], + 0, + ) + scores = tl.dot(q, k).to(tl.float32) * SCALE + allowed = valid_k[None, :] & (ky[None, :] >= top[:, None]) & (ky[None, :] < top[:, None] + WH) & (kx[None, :] >= left[:, None]) & (kx[None, :] < left[:, None] + WW) + 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 * vs0 + ki[:, None] * vs1 + head * vs2 + d[None, :], + valid_k[:, None], + 0, + ) + acc = acc * correction[:, None] + acc += tl.dot(probabilities.to(v.dtype), v) + denominator = denominator * correction + tl.sum(probabilities, 1) + maximum = new_max + output = acc / denominator[:, None] + tl.store( + OUT + ((batch * FRAMES * HEIGHT * WIDTH + qi[:, None]) * HEADS + head) * DIM + d[None, :], + output, + valid_q[:, None], + ) + + +def spatial_attention(query, key, value, spec: SpatialAttentionSpec): + """[B,T*H*W,heads,dim] -> same layout, exact shifted rectangular mask.""" + if ( + key.shape != value.shape + or query.shape[0] != key.shape[0] + or query.shape[2:] != key.shape[2:] + or query.dtype != key.dtype + or key.dtype != value.dtype + or query.device != key.device + or key.device != value.device + ): + raise ValueError("Spatial attention requires matching batch, heads, dtype and device") + key_frames = spec.key_frames + if key.shape[1] != key_frames * spec.height * spec.width: + raise ValueError("Spatial attention K/V length disagrees with key_frames") + first_q = spec.query_frame_offset - spec.key_frame_offset + if first_q < 0 or first_q + spec.frames > key_frames: + raise ValueError("Every query must have a corresponding key frame") + b, n, heads, dim = query.shape + if n != spec.frames * spec.height * spec.width or dim not in (32, 64, 128): + raise ValueError("Unsupported spatial attention shape") + if any(t.stride(-1) != 1 for t in (query, key, value)): + raise ValueError("Spatial attention requires contiguous head dimensions") + wh, ww = min(spec.window_height, spec.height), min(spec.window_width, spec.width) + if min(wh, ww, spec.temporal_window) <= 0: + raise ValueError("Spatial and temporal windows must be positive") + output = torch.empty_like(query, memory_format=torch.contiguous_format) + bh, bw, bn = 4, 16, 64 + bm = bh * bw + cols = triton.next_power_of_2(ww + bw - 1) + rows = min(spec.height, wh + bh - 1) + _spatial_forward[ + ( + spec.frames * triton.cdiv(spec.height, bh), + triton.cdiv(spec.width, bw), + b * heads, + ) + ]( + query, + key, + value, + output, + *query.stride()[:3], + *key.stride()[:3], + *value.stride()[:3], + spec.frames, + key_frames, + spec.query_frame_offset, + spec.key_frame_offset, + spec.height, + spec.width, + wh, + ww, + spec.temporal_window, + heads, + dim, + dim**-0.5 * 1.4426950408889634, + bh, + bw, + bm, + bn, + cols, + rows, + num_warps=4, + num_stages=2, + ) + return output diff --git a/lightx2v/models/networks/relayvsr/wan_infer.py b/lightx2v/models/networks/relayvsr/wan_infer.py new file mode 100644 index 000000000..5a0d509c9 --- /dev/null +++ b/lightx2v/models/networks/relayvsr/wan_infer.py @@ -0,0 +1,175 @@ +import math + +import torch +import torch.nn.functional as F + +from lightx2v.models.networks.wan.infer.module_io import GridOutput, WanPreInferModuleOutput +from lightx2v.models.networks.wan.infer.offload.transformer_infer import WanOffloadTransformerInfer +from lightx2v.models.networks.wan.infer.post_infer import WanPostInfer +from lightx2v.models.networks.wan.infer.transformer_infer import modulate + +from .rope import apply_rope +from .wan_attention import SpatialAttentionSpec, spatial_attention + + +class RelayWanPreInfer: + def __init__(self, config): + self.config = config + + def conditioning(self, weights): + dim = self.config["freq_dim"] + exponent = -math.log(10000) * torch.arange(dim // 2, device=weights.prompt.tensor.device, dtype=torch.float32) / (dim // 2) + angles = 1000 * torch.exp(exponent) + timestep = torch.cat((angles.cos(), angles.sin())).unsqueeze(0).to(weights.prompt.tensor.dtype) + temb = weights.time2.apply(F.silu(weights.time1.apply(timestep))) + modulation = weights.time_proj.apply(F.silu(temb)).view(6, -1) + context = weights.text2.apply(F.gelu(weights.text1.apply(weights.prompt.tensor[0]), approximate="tanh")) + return temb, modulation, context + + def infer(self, weights, latents, lr): + hidden = weights.patch.apply(latents).flatten(2).transpose(1, 2)[0] + upscaled = F.interpolate(lr, scale_factor=4, mode="bilinear", align_corners=False) + lq = F.pixel_unshuffle(upscaled, 16) + lq = F.silu(weights.lq_norm1.apply(weights.lq_conv1.apply(lq))) + lq = F.silu(weights.lq_norm2.apply(weights.lq_conv2.apply(lq))) + # Flattening also preserves nn.Linear's leading stride for one token. + lq = lq.flatten(2).transpose(1, 2).flatten(0, 1) + lq = weights.lq_out.apply(lq) if lq.is_contiguous() else weights.lq_out_unbiased.apply(lq) + weights.lq_out.bias + return hidden + lq + + +class RelayWanTransformerInfer(WanOffloadTransformerInfer): + """Wan execution/offload/FFN with RelayVSR's causal reference attention. + + Keep the shared block pipeline, offload, normal-size FFN and final residuals. + Relay-specific math is limited to input preparation, local self attention, + fixed-context KV, and the checkpoint's FP32 normalization boundaries. + """ + + def __init__(self, config, deployment, parallel, weights): + family_config = { + **deployment, + "task": "sr", + "model_cls": "relayvsr", + "seq_parallel": parallel.size > 1, + "num_layers": config["num_layers"], + "num_heads": config["num_attention_heads"], + "dim": config["num_attention_heads"] * config["attention_head_dim"], + "modulate_type": "torch", + "mxfp8_fuse_enable": False, + "clean_cuda_cache": False, + } + super().__init__(family_config) + self.relay_config, self.parallel = config, parallel + # Norms and modulation run in FP32, but residuals round to the inference + # dtype after every sublayer, matching the official distilled checkpoint. + self.sensitive_layer_dtype = self.infer_dtype + self.modulate_func = self.modulate_and_cast + self.cross_cache = [] + if weights.offload: + self.offload_manager.init_cuda_buffer(blocks_cuda_buffer=weights.offload_block_cuda_buffers) + + def modulate_and_cast(self, x, scale, shift): + return modulate(x, scale, shift).to(self.infer_dtype) + + def prepare_cross_attention(self, blocks, context): + self.cross_cache = [] + pre = WanPreInferModuleOutput(embed=None, grid_sizes=None, x=None, embed0=None, context=context, adapter_args={"prepare_cross_kv": True}) + # Use the same family traversal/offload path for initialization as for + # inference. Each staged block's KV remains resident after it is evicted. + self.infer_func(blocks, None, pre) + return self.cross_cache + + def infer_block(self, block, x, pre_infer_out): + if pre_infer_out.adapter_args.get("prepare_cross_kv", False): + phase, context = block.compute_phases[1], pre_infer_out.context + key = phase.cross_attn_norm_k.apply(phase.cross_attn_k.apply(context)).view(-1, self.num_heads, self.head_dim) + value = phase.cross_attn_v.apply(context).view(-1, self.num_heads, self.head_dim) + self.cross_cache.append((key, value)) + return x + return super().infer_block(block, x, pre_infer_out, self_attn_kwargs=pre_infer_out.adapter_args) + + def infer_cross_attn(self, phase, x, context, y_out, gate_msa): + # RelayVSR has one fixed prompt and no image/CFG branch. Keep its + # persistent context KV here without changing the Wan family contract. + x.add_(y_out * gate_msa.squeeze()) + query = phase.cross_attn_norm_q.apply(phase.cross_attn_q.apply(phase.norm3.apply(x))).view(-1, self.num_heads, self.head_dim) + key, value = self.cross_cache[self.block_idx] + attended = phase.cross_attn_1.apply( + q=query, + k=key, + v=value, + cu_seqlens_q=self.cross_attn_cu_seqlens_q, + cu_seqlens_kv=self.cross_attn_cu_seqlens_kv, + max_seqlen_q=query.size(0), + max_seqlen_kv=key.size(0), + ) + return x, phase.cross_attn_o.apply(attended) + + def infer_ffn(self, phase, x, attn_out, c_shift_msa, c_scale_msa, c_gate_msa=None): + if x.shape[0] != 1: + return super().infer_ffn(phase, x, attn_out, c_shift_msa, c_scale_msa, c_gate_msa) + # A padded 8x8 LR image has a single Wan token. The base FFN squeezes + # away that axis; handle only this boundary case in the specialization. + x.add_(attn_out) + normalized = self.modulate_func(phase.norm2.apply(x), scale=c_scale_msa, shift=c_shift_msa).reshape(1, -1) + return phase.ffn_2.apply(F.gelu(phase.ffn_0.apply(normalized), approximate="tanh")) + + def infer_self_attn(self, phase, x, shift_msa, scale_msa, grid_sizes=None, *, q_rope, k_rope, cache, frame): + _, height, width = grid_sizes + normalized = self.modulate_func(phase.norm1.apply(x), scale_msa, shift_msa) + tokens = height * width + query = self.parallel.to_heads(phase.self_attn_norm_q.apply(phase.self_attn_q.apply(normalized)).view(-1, self.num_heads, self.head_dim), tokens) + key = self.parallel.to_heads(phase.self_attn_norm_k.apply(phase.self_attn_k.apply(normalized)).view(-1, self.num_heads, self.head_dim), tokens) + value = self.parallel.to_heads(phase.self_attn_v.apply(normalized).view(-1, self.num_heads, self.head_dim), tokens) + key, value = cache.update(key, value, self.block_idx, frame) + query = apply_rope(query, q_rope) + key = apply_rope(key, k_rope) + visible = key.shape[0] // tokens + wh, ww = self.relay_config["spatial_window_size"] + if height > wh or width > ww: + spec = SpatialAttentionSpec(1, height, width, wh, ww, self.relay_config["temporal_context_frames"], key_frames=visible, query_frame_offset=visible - 1) + attended = spatial_attention(query[None], key[None], value[None], spec)[0] + else: + attended = phase.self_attn_1.apply(query, key, value).view(tokens, -1, self.head_dim) + return phase.self_attn_o.apply(self.parallel.to_sequence(attended)) + + def prepare_input(self, x, temb, modulation, context, rope, cache, frame, height, width): + window = self.relay_config["temporal_context_frames"] + visible = min(frame + 1, window) + origin = max(0, frame - self.relay_config["temporal_rope"]["training_frames"] + 1) + positions = tuple(frame - ((frame - slot) % window) - origin for slot in range(visible)) + grid = (1, height, width) + return WanPreInferModuleOutput( + embed=temb, + grid_sizes=GridOutput(tensor=torch.tensor(grid), tuple=grid), + x=x, + embed0=modulation.float().unsqueeze(0), + context=context, + adapter_args={ + "q_rope": rope.for_positions(height, width, (frame - origin,)), + "k_rope": rope.for_positions(height, width, positions), + "cache": cache, + "frame": frame, + }, + ) + + def infer_non_blocks(self, weights, x, e): + # The official head modulates FP32-normalized activations with BF16 + # coefficients, then casts before projection. Gather SP tokens first. + x = self.parallel.gather(x, self._output_tokens) + shift, scale = (weights.post.modulation.tensor[0] + e).unbind(0) + x = (weights.post.norm.apply(x) * (1 + scale) + shift).to(x.dtype) + return weights.post.projection.apply(x) + + def infer_main_blocks(self, blocks, pre_infer_out): + _, height, width = pre_infer_out.grid_sizes.tuple + self._output_tokens = height * width + return super().infer_main_blocks(blocks, pre_infer_out) + + +class RelayWanPostInfer(WanPostInfer): + def infer(self, x, pre_infer_out): + # Reuse Wan's unpatchify; the one-step scheduler expects a batched BF16 + # tensor instead of the base post-infer's list of FP32 videos. + return self.unpatchify(x, pre_infer_out.grid_sizes.tuple)[0].unsqueeze(0) diff --git a/lightx2v/models/networks/relayvsr/weights.py b/lightx2v/models/networks/relayvsr/weights.py new file mode 100644 index 000000000..a0cd35e52 --- /dev/null +++ b/lightx2v/models/networks/relayvsr/weights.py @@ -0,0 +1,185 @@ +import json +from pathlib import Path + +import torch +import torch.nn.functional as F +from safetensors import safe_open + +from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList +from lightx2v.common.ops.norm.layer_norm_weight import LNWeight +from lightx2v.common.ops.norm.rms_norm_weight import RMSWeight +from lightx2v.utils.envs import GET_DTYPE +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, CONV2D_WEIGHT_REGISTER, CONV3D_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, TENSOR_REGISTER + + +class FP32LayerNorm(LNWeight): + def __init__(self, *args, cast_output=False, **kwargs): + super().__init__(*args, **kwargs) + self.cast_output = cast_output + + def apply(self, x): + output = F.layer_norm(x.float(), (x.shape[-1],), None if self.weight is None else self.weight.float(), None if self.bias is None else self.bias.float(), self.eps) + return output.to(x.dtype) if self.cast_output else output + + +class ChannelRMSNorm(RMSWeight): + def apply(self, x): + value = x.float() + value = value * torch.rsqrt(value.square().mean(dim=1, keepdim=True) + self.eps) + return (value * self.weight).to(x.dtype) + + +class RelayWeights(WeightModule): + def linear(self, name, key, *, bias=True, buffer=False): + self.add_module(name, MM_WEIGHT_REGISTER["Default"](key + ".weight", key + ".bias" if bias else None, create_cuda_buffer=buffer)) + + def tensor(self, name, key, *, buffer=False): + self.add_module(name, TENSOR_REGISTER["Default"](key, create_cuda_buffer=buffer)) + + def rms(self, name, key, *, eps, backend="torch_native", buffer=False): + self.add_module(name, RMS_WEIGHT_REGISTER[backend](key + ".weight", eps=eps, create_cuda_buffer=buffer)) + + def conv(self, name, key, *, stride=1, padding=0): + self.add_module(name, CONV2D_WEIGHT_REGISTER["Default"](key + ".weight", key + ".bias", stride=stride, padding=padding)) + + +class WanPreWeights(RelayWeights): + def __init__(self, config): + super().__init__() + self.add_module("patch", CONV3D_WEIGHT_REGISTER["Default"]("patch_embedding.weight", "patch_embedding.bias", stride=tuple(config["patch_size"]))) + self.conv("lq_conv1", "lq_proj.conv1") + self.conv("lq_conv2", "lq_proj.conv2", stride=2, padding=1) + for i in (1, 2): + self.add_module(f"lq_norm{i}", ChannelRMSNorm(f"lq_proj.norm{i}.weight", eps=1e-6)) + # Upstream's transposed tokens normally require separate matmul/bias + # rounding. A one-token grid is contiguous and uses fused addmm instead. + self.linear("lq_out", "lq_proj.out_proj") + self.linear("lq_out_unbiased", "lq_proj.out_proj", bias=False) + for name, key in ( + ("text1", "text_embedder.linear_1"), + ("text2", "text_embedder.linear_2"), + ("time1", "time_embedder.linear_1"), + ("time2", "time_embedder.linear_2"), + ("time_proj", "time_proj"), + ): + self.linear(name, "condition_embedder." + key) + self.tensor("prompt", "prompt_embeds") + + +class WanBlockWeights(RelayWeights): + """Official checkpoint names mapped onto the Wan family's three phases.""" + + def __init__(self, index, config, *, attention_backend="torch_sdpa", buffer=False): + super().__init__() + prefix = f"blocks.{index}." + phases = WeightModuleList([RelayWeights() for _ in range(3)]) + self.add_module("compute_phases", phases) + for phase, attention, family_name in ((phases[0], "attn1", "self_attn"), (phases[1], "attn2", "cross_attn")): + for projection in ("q", "k", "v"): + phase.linear(family_name + "_" + projection, prefix + attention + ".to_" + projection, buffer=buffer) + phase.linear(family_name + "_o", prefix + attention + ".to_out.0", buffer=buffer) + for projection in ("q", "k"): + phase.rms(family_name + "_norm_" + projection, prefix + attention + ".norm_" + projection, eps=config["eps"], buffer=buffer) + phase.add_module(family_name + "_1", ATTN_WEIGHT_REGISTER["torch_sdpa" if attention == "attn1" else attention_backend]()) + phases[0].add_module("norm1", FP32LayerNorm(eps=config["eps"])) + phases[0].tensor("modulation", prefix + "scale_shift_table", buffer=buffer) + phases[1].add_module("norm3", FP32LayerNorm(prefix + "norm2.weight", prefix + "norm2.bias", eps=config["eps"], create_cuda_buffer=buffer, cast_output=True)) + phases[2].add_module("norm2", FP32LayerNorm(eps=config["eps"])) + phases[2].linear("ffn_0", prefix + "ffn.net.0.proj", buffer=buffer) + phases[2].linear("ffn_2", prefix + "ffn.net.2", buffer=buffer) + + +class WanPostWeights(RelayWeights): + def __init__(self, config): + super().__init__() + self.add_module("norm", FP32LayerNorm(eps=config["eps"])) + self.tensor("modulation", "scale_shift_table") + self.linear("projection", "proj_out") + + +class FlashPreWeights(RelayWeights): + def __init__(self, config): + super().__init__() + for i, stride, padding in ((0, 2, 1), (2, 2, 1), (4, 1, 1), (6, 1, 0)): + self.conv(f"conv{i}", f"lr_encoder.layers.{i}", stride=stride, padding=padding) + self.conv("condition", "condition_encoder.projection") + self.rms("lr_norm", "lr_encoder.norm", eps=config["norm_eps"]) + self.rms("condition_norm", "condition_encoder.norm", eps=config["norm_eps"]) + self.tensor("video_type", "video_type_embedding") + self.tensor("condition_type", "condition_type_embedding") + + +class FlashBlockWeights(RelayWeights): + def __init__(self, index, config): + super().__init__() + prefix = f"backbone.blocks.{index}." + for name in ("query", "key", "value", "condition_key", "condition_value", "output"): + self.linear(name, prefix + "attention." + name) + for name in ("key_norm", "value_norm"): + self.rms(name, prefix + "attention." + name, eps=config["norm_eps"]) + for name in ("attention_norm", "feed_forward_norm"): + self.rms(name, prefix + name, eps=config["norm_eps"]) + self.linear("gate_up", prefix + "feed_forward.gate_up_proj", bias=False) + self.linear("down", prefix + "feed_forward.down", bias=False) + + +class FlashPostWeights(RelayWeights): + def __init__(self): + super().__init__() + self.linear("head", "spatial_head.0") + self.linear("projection", "residual_projection") + + +def checkpoint_config(path, component): + with safe_open(path, framework="pt", device="cpu") as handle: + metadata = handle.metadata() or {} + if metadata.get("format") != f"refvsr-{component}-v1": + raise ValueError(f"Expected official RelayVSR {component} checkpoint with embedded config: {path}") + return json.loads(metadata["config"]) + + +def resolve_checkpoints(config): + root = Path(config["model_path"]) + return Path(config.get("dit_original_ckpt") or root / "wan.safetensors"), Path(config.get("vae_original_ckpt") or root / "flash.safetensors") + + +def load_container(container, handle, keys, device): + # Mapping and device placement happen once here, never in the inference path. + state = {key: handle.get_tensor(key).to(device=device, dtype=GET_DTYPE()) for key in keys} + container.load(state) + + +class WanWeights(WeightModule): + def __init__(self, path, model_config, deployment): + super().__init__() + self.add_module("pre", WanPreWeights(model_config)) + self.add_module("post", WanPostWeights(model_config)) + attention_backend = deployment.get("cross_attention_backend", "torch_sdpa") + blocks = WeightModuleList([WanBlockWeights(i, model_config, attention_backend=attention_backend) for i in range(model_config["num_layers"])]) + self.add_module("blocks", blocks) + self.offload = deployment.get("cpu_offload", False) + self.offload_block_cuda_buffers = [] + with safe_open(path, framework="pt", device="cpu") as handle: + keys = list(handle.keys()) + non_blocks = [key for key in keys if not key.startswith("blocks.")] + load_container(WeightModuleList([self.pre, self.post]), handle, non_blocks, "cuda") + for i, block in enumerate(blocks): + block_keys = [key for key in keys if key.startswith(f"blocks.{i}.")] + load_container(block, handle, block_keys, "cpu" if self.offload else "cuda") + if i == 0 and self.offload: + for _ in range(2): + slot = WanBlockWeights(0, model_config, attention_backend=attention_backend, buffer=True) + load_container(slot, handle, block_keys, "cuda") + self.offload_block_cuda_buffers.append(slot) + + +class FlashWeights(WeightModule): + def __init__(self, path, config): + super().__init__() + self.add_module("pre", FlashPreWeights(config)) + self.add_module("blocks", WeightModuleList([FlashBlockWeights(i, config) for i in range(10)])) + self.add_module("post", FlashPostWeights()) + with safe_open(path, framework="pt", device="cpu") as handle: + # condition_rgb_projection is an auxiliary training head, unused by decode. + keys = [key for key in handle.keys() if not key.startswith("condition_rgb_projection.")] # noqa: SIM118 - safe_open is not iterable + load_container(self, handle, keys, "cuda") diff --git a/lightx2v/models/runners/relayvsr/__init__.py b/lightx2v/models/runners/relayvsr/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/lightx2v/models/runners/relayvsr/relayvsr_runner.py b/lightx2v/models/runners/relayvsr/relayvsr_runner.py new file mode 100644 index 000000000..a5fdec7df --- /dev/null +++ b/lightx2v/models/runners/relayvsr/relayvsr_runner.py @@ -0,0 +1,287 @@ +import json +import os +import tempfile +import time +from contextlib import closing +from fractions import Fraction +from itertools import islice +from pathlib import Path +from typing import ClassVar + +import av +import numpy as np +import torch +import torch.distributed as dist +import torch.nn.functional as F +from PIL import Image +from loguru import logger + +import lightx2v.common.kvcache.relayvsr # noqa: F401 - register frame-slot cache +from lightx2v.common.kvcache.manager import KVCacheManager +from lightx2v.models.networks.relayvsr.model import RelayVSRModel +from lightx2v.models.runners.default_runner import DefaultRunner +from lightx2v.models.runners.request_fields import COMMON_REQUEST_FIELDS +from lightx2v.models.schedulers.relayvsr.scheduler import RelayVSRScheduler +from lightx2v.utils.envs import GET_DTYPE +from lightx2v.utils.registry_factory import RUNNER_REGISTER +from lightx2v.utils.utils import mux_audio_from_video +from lightx2v_platform.base.global_var import AI_DEVICE + +from .video_io import AsyncVideoWriter, DeviceTimings, prefetch_frames + + +def validate_config(config): + if config.get("cpu_offload") and config.get("offload_granularity", "block") != "block": + raise ValueError("RelayVSR supports block weight offload only") + parallel = config.get("parallel") or {} + if parallel.get("seq_p_size", 1) not in (1, 2): + raise ValueError("RelayVSR Flash has two KV groups, so sequence parallelism supports 1 or 2 GPUs") + for key, default, minimum in (("clip_len", 16, 2), ("queue_size", 3, 1), ("encoder_threads", 8, 1)): + if config.get(key, default) < minimum: + raise ValueError(f"RelayVSR {key} must be >= {minimum}") + + +@RUNNER_REGISTER("relayvsr") +class RelayVSRRunner(DefaultRunner): + """Official big RelayVSR: sparse Wan references + streaming Flash decode.""" + + supported_request_fields_by_task: ClassVar[dict] = {"sr": COMMON_REQUEST_FIELDS | {"video_path", "image_path", "sr_ratio", "size"}} + + def __init__(self, config): + for key, value in (("infer_steps", 1), ("guidance_scale", 1.0), ("enable_cfg", False), ("fps", None)): + config.setdefault(key, value) + validate_config(config) + super().__init__(config) + if AI_DEVICE != "cuda" or not torch.cuda.is_available() or GET_DTYPE() not in (torch.bfloat16, torch.float16): + raise ValueError("RelayVSR local attention requires CUDA with BF16 or FP16") + self.device = torch.device("cuda", torch.cuda.current_device()) + self.main_rank = not dist.is_initialized() or dist.get_rank() == 0 + self.seq_p_group = config["device_mesh"].get_group(mesh_dim="seq_p") if config.get("seq_parallel") else None + world_size = dist.get_world_size() if dist.is_initialized() else 1 + if world_size != (config.get("parallel") or {}).get("seq_p_size", 1): + raise ValueError("RelayVSR world size must equal parallel.seq_p_size") + + def init_scheduler(self): + self.scheduler = RelayVSRScheduler(self.config) + + @torch.no_grad() + def init_modules(self): + started = time.perf_counter() + self.model = RelayVSRModel(self.config, self.seq_p_group) + self.model.set_scheduler(self.scheduler) + torch.cuda.synchronize() + self.load_seconds = time.perf_counter() - started + self.config.lock() + logger.info("RelayVSR native 4x model loaded in {:.2f}s", self.load_seconds) + + @staticmethod + def output_size(input_info, height, width): + if input_info.size: + size = input_info.size + if len(size) != 2 or any(isinstance(v, bool) or not isinstance(v, int) for v in size): + raise ValueError("RelayVSR size must be [height, width] with integer dimensions") + target = tuple(size) + else: + if input_info.sr_ratio not in (2, 4): + raise ValueError("RelayVSR sr_ratio must be 2 or 4; 2x resizes the native 4x result") + target = (int(height * input_info.sr_ratio), int(width * input_info.sr_ratio)) + if min(target) <= 0 or target[0] > height * 4 or target[1] > width * 4: + raise ValueError("RelayVSR output dimensions must be positive and no larger than the native 4x result") + return target + + def preprocess(self, frame): + height, width = frame.shape[-2:] + x = frame[None].to(self.device, non_blocking=frame.is_pinned()).to(GET_DTYPE()) + x = x.div(127.5).sub(1.0) + return F.pad(x, (0, -width % 8, 0, -height % 8), mode="replicate") + + @staticmethod + def postprocess(prediction, target): + if prediction.shape[-2:] != target: + prediction = F.interpolate(prediction.float().clamp(-1, 1), size=target, mode="bicubic", align_corners=False, antialias=True) + return prediction[0].add(1).mul_(127.5).clamp_(0, 255).to(torch.uint8).permute(1, 2, 0).contiguous() + + def init_run(self, height, width, seed): + self.scheduler.prepare(seed, self.device) + wan = self.model.wan_config + self.wan_cache_manager = KVCacheManager( + { + "num_layers": wan["num_layers"], + "num_heads": wan["num_attention_heads"], + "dim": wan["num_attention_heads"] * wan["attention_head_dim"], + }, + device=self.device, + ) + self.wan_cache = self.wan_cache_manager.create_self_attn_kv_cache( + "wan", height // 8 * (width // 8) * wan["temporal_context_frames"], kv_cache_scheme="frame_ring", num_heads=wan["num_attention_heads"] // self.model.parallel.size + ) + self.flash_cache_manager = KVCacheManager({"num_layers": 10, "num_heads": 2, "dim": 128}, device=self.device) + self.flash_history = self.flash_cache_manager.create_self_attn_kv_cache("flash_video", height // 4 * (width // 4), kv_cache_scheme="static", num_heads=2 // self.model.parallel.size) + self.reference_count = self.decode_count = 0 + + def end_run(self): + self.scheduler.clear() + self.model.clear() + self.wan_cache = self.flash_history = None + self.wan_cache_manager = self.flash_cache_manager = None + + def reference(self, frame): + lr = self.timings.call("preprocess", self.preprocess, frame) + self.scheduler.prepare_latents(*lr.shape[-2:], self.device) + self.scheduler.step_pre(0) + self.timings.call("wan", self.model.infer, self.scheduler, lr, self.wan_cache, self.reference_count) + self.scheduler.step_post() + self.reference_count += 1 + latent = self.scheduler.latents[:, :, 0].to(GET_DTYPE()) + return self.timings.call("flash_condition", self.model.prepare_condition, latent) + + def restore_frames(self, source, seed): + """Yield cropped native 4x predictions; overlapping endpoint is emitted once.""" + iterator = iter(source) + first = next(iterator, None) + if first is None: + raise ValueError("RelayVSR cannot restore an empty stream") + shape = first.shape + height, width = shape[-2:] + self.timings = DeviceTimings() + try: + self.init_run(height + (-height) % 8, width + (-width) % 8, seed) + start_condition = self.reference(first) + first_window = True + start = first + while True: + window = [start, *islice(iterator, self.config.get("clip_len", 16) - 1)] + if any(frame.shape != shape for frame in window): + raise ValueError("All RelayVSR input frames must have the same dimensions") + if len(window) == 1 and not first_window: + return + self.check_stop() + end_condition = self.reference(window[-1]) if len(window) > 1 else start_condition + conditions = [(torch.cat((a[0], b[0]), 0), torch.cat((a[1], b[1]), 0)) for a, b in zip(start_condition, end_condition)] + self.flash_history.reset() + for offset, frame in enumerate(window): + self.check_stop() + lr = self.timings.call("preprocess", self.preprocess, frame) + output = self.timings.call("flash_decode", self.model.decode, lr, conditions, self.flash_history, offset) + self.decode_count += 1 + if first_window or offset > 0: + yield output[..., : height * 4, : width * 4] + first_window = False + start, start_condition = window[-1], end_condition + if len(window) < self.config.get("clip_len", 16): + return + finally: + self.end_run() + + @torch.no_grad() + def run_pipeline(self, input_info): + if bool(input_info.image_path) == bool(input_info.video_path): + raise ValueError("RelayVSR requires exactly one image_path or video_path") + if input_info.return_result_tensor and input_info.video_path: + raise ValueError("RelayVSR streams videos to save_result_path; returning a whole video tensor is unsupported") + if not input_info.save_result_path and not input_info.return_result_tensor: + raise ValueError("RelayVSR requires save_result_path") + self.input_info = input_info + if input_info.video_path: + return self.run_video(input_info) + return self.run_image(input_info) + + def run_image(self, input_info): + started = time.perf_counter() + with Image.open(input_info.image_path) as image: + frame = torch.from_numpy(np.array(image.convert("RGB"))).permute(2, 0, 1).contiguous() + target = self.output_size(input_info, *frame.shape[-2:]) + with closing(self.restore_frames([frame], input_info.seed)) as predictions: + pixels = self.postprocess(next(predictions), target).cpu() + if self.main_rank and input_info.save_result_path: + path = Path(input_info.save_result_path) + path.parent.mkdir(parents=True, exist_ok=True) + Image.fromarray(pixels.numpy()).save(path) + stats = self.timings.collect() + stats["total_seconds"] = time.perf_counter() - started + if self.progress_callback: + self.progress_callback(100, 100) + return {"images": pixels[None].float().div_(255) if input_info.return_result_tensor else None, "stats": stats} + + def run_video(self, input_info): + started = time.perf_counter() + source, destination = Path(input_info.video_path), Path(input_info.save_result_path) + if source.resolve() == destination.resolve(): + raise ValueError("RelayVSR output must differ from its source video") + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = None + writer = None + output_frames = 0 + io_timing = {"decode_seconds": 0.0} + torch.cuda.reset_peak_memory_stats() + try: + if self.main_rank: + fd, temporary = tempfile.mkstemp(prefix=".relayvsr-", suffix=destination.suffix or ".mp4", dir=destination.parent) + os.close(fd) + with av.open(str(source)) as container: + if not container.streams.video: + raise ValueError("RelayVSR input contains no video stream") + stream = container.streams.video[0] + stream.thread_type = "AUTO" + fps = Fraction(str(self.config["fps"])) if self.config.get("fps") is not None else stream.average_rate + if fps is None or fps <= 0: + raise ValueError("Cannot determine source FPS; set startup fps explicitly") + target = self.output_size(input_info, stream.height, stream.width) + if target[0] % 2 or target[1] % 2: + raise ValueError("YUV420 video output requires even height and width") + input_info.output_fps = float(fps) + if self.main_rank: + writer = AsyncVideoWriter(temporary, fps, self.config) + with closing(prefetch_frames(container, self.config.get("clip_len", 16), io_timing)) as frames, closing(self.restore_frames(frames, input_info.seed)) as predictions: + for prediction in predictions: + if self.main_rank: + pixels = self.timings.call("postprocess", self.postprocess, prediction, target) + writer.submit(pixels) + output_frames += 1 + if self.main_rank and output_frames % 30 == 0: + logger.info("RelayVSR: {} frames written to encoder queue", output_frames) + if writer is not None: + writer.close() + stages = self.timings.collect() + mux_seconds = 0.0 + if self.main_rank: + mux_start = time.perf_counter() + if mux_audio_from_video(str(source), temporary, prefer_copy=True, trim_to_shortest=False) is None: + raise RuntimeError("RelayVSR could not mux the source audio") + mux_seconds = time.perf_counter() - mux_start + os.replace(temporary, destination) + temporary = None + stats = { + "frames": output_frames, + "fps": str(fps), + "height": target[0], + "width": target[1], + "native_scale": 4, + "wan_references": self.reference_count, + "flash_decode_calls": self.decode_count, + "load_seconds": self.load_seconds, + "cuda_stage_seconds": stages, + "decode_seconds": io_timing["decode_seconds"], + "output_copy_seconds": writer.copy_seconds if writer else 0, + "encoder_queue_wait_seconds": writer.queue_wait_seconds if writer else 0, + "encode_seconds": writer.encode_seconds if writer else 0, + "encoder_flush_seconds": writer.flush_seconds if writer else 0, + "audio_mux_seconds": mux_seconds, + "total_seconds": time.perf_counter() - started, + "peak_allocated_gib": torch.cuda.max_memory_allocated() / 2**30, + "peak_reserved_gib": torch.cuda.max_memory_reserved() / 2**30, + } + if self.main_rank: + logger.info("RelayVSR video saved to: {}", destination.resolve()) + logger.info("RelayVSR completed: {}", json.dumps(stats, ensure_ascii=False)) + if self.progress_callback: + self.progress_callback(100, 100) + return {"video": None, "stats": stats} + finally: + if writer is not None and temporary is not None: + try: + writer.close() + except (RuntimeError, ValueError, OSError): + logger.exception("RelayVSR encoder cleanup failed") + if temporary is not None: + Path(temporary).unlink(missing_ok=True) diff --git a/lightx2v/models/runners/relayvsr/video_io.py b/lightx2v/models/runners/relayvsr/video_io.py new file mode 100644 index 000000000..9c703e9f5 --- /dev/null +++ b/lightx2v/models/runners/relayvsr/video_io.py @@ -0,0 +1,119 @@ +import time +from collections import deque +from concurrent.futures import ThreadPoolExecutor +from itertools import islice + +import av +import torch + + +def prefetch_frames(container, batch_size, timing): + def frames(): + for frame in container.decode(video=0): + yield torch.from_numpy(frame.to_ndarray(format="rgb24")).permute(2, 0, 1).contiguous().pin_memory() + + iterator = iter(frames()) + + def read_batch(): + started = time.perf_counter() + batch = list(islice(iterator, batch_size)) + timing["decode_seconds"] += time.perf_counter() - started + return batch + + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="relayvsr-decode") as executor: + pending = executor.submit(read_batch) + while True: + batch = pending.result() + if not batch: + return + pending = executor.submit(read_batch) + yield from batch + + +class AsyncVideoWriter: + def __init__(self, path, fps, config): + self.path, self.fps, self.config = path, fps, config + self.executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="relayvsr-encode") + self.pending = deque() + self.copy_stream = torch.cuda.Stream() + self.container = self.stream = None + self.frames = 0 + self.encode_seconds = self.flush_seconds = 0.0 + self.copy_seconds = self.queue_wait_seconds = 0.0 + self.queue_size = config.get("queue_size", 3) + self.closed = False + + def _write(self, pixels, copy_start, event): + event.synchronize() + self.copy_seconds += copy_start.elapsed_time(event) / 1000 + start = time.perf_counter() + if self.container is None: + self.container = av.open(str(self.path), mode="w") + self.stream = self.container.add_stream(self.config.get("video_codec", "libx264"), rate=self.fps) + self.stream.width, self.stream.height = pixels.shape[1], pixels.shape[0] + self.stream.pix_fmt = "yuv420p" + self.stream.codec_context.thread_count = self.config.get("encoder_threads", 8) + self.stream.options = {"crf": str(self.config.get("video_crf", 18)), "preset": self.config.get("ffmpeg_preset", "fast")} + frame = av.VideoFrame.from_ndarray(pixels.numpy(), format="rgb24") + frame.pts = self.frames + frame.time_base = 1 / self.fps + for packet in self.stream.encode(frame): + self.container.mux(packet) + self.frames += 1 + self.encode_seconds += time.perf_counter() - start + + def submit(self, pixels): + if len(self.pending) >= self.queue_size: + started = time.perf_counter() + self.pending.popleft().result() + self.queue_wait_seconds += time.perf_counter() - started + host = torch.empty(pixels.shape, device="cpu", dtype=torch.uint8, pin_memory=True) + self.copy_stream.wait_stream(torch.cuda.current_stream(pixels.device)) + with torch.cuda.stream(self.copy_stream): + copy_start = torch.cuda.Event(enable_timing=True) + copy_start.record() + host.copy_(pixels, non_blocking=True) + pixels.record_stream(self.copy_stream) + event = torch.cuda.Event(enable_timing=True) + event.record() + self.pending.append(self.executor.submit(self._write, host, copy_start, event)) + + def close(self): + if self.closed: + return + self.closed = True + try: + for future in self.pending: + future.result() + start = time.perf_counter() + if self.stream is not None: + for packet in self.stream.encode(): + self.container.mux(packet) + self.flush_seconds = time.perf_counter() - start + finally: + self.executor.shutdown(wait=True) + if self.container is not None: + self.container.close() + + +class DeviceTimings: + """CUDA elapsed spans overlap CPU I/O; they must not be summed into wall time.""" + + def __init__(self): + self.events = [] + + def call(self, name, function, *args): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + result = function(*args) + end.record() + self.events.append((name, start, end)) + return result + + def collect(self): + torch.cuda.synchronize() + totals = {} + for name, start, end in self.events: + totals[name] = totals.get(name, 0.0) + start.elapsed_time(end) / 1000 + self.events.clear() + return totals diff --git a/lightx2v/models/runners/runner_factory.py b/lightx2v/models/runners/runner_factory.py index e2acc3254..6ed9fea4e 100644 --- a/lightx2v/models/runners/runner_factory.py +++ b/lightx2v/models/runners/runner_factory.py @@ -31,6 +31,7 @@ "qwen_image": "lightx2v.models.runners.qwen_image.qwen_image_runner", "qwen_image_21": "lightx2v.models.runners.qwen_image_21.qwen_image_21_runner", "realtimewam": "lightx2v.models.runners.wan.realtimewam_runner", + "relayvsr": "lightx2v.models.runners.relayvsr.relayvsr_runner", "seedvr2": "lightx2v.models.runners.seedvr.seedvr_runner", "seko_talk": "lightx2v.models.runners.wan.wan_audio_runner", "seko_talk_ar": "lightx2v.models.runners.wan.wan_audio_runner", diff --git a/lightx2v/models/schedulers/relayvsr/__init__.py b/lightx2v/models/schedulers/relayvsr/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/lightx2v/models/schedulers/relayvsr/scheduler.py b/lightx2v/models/schedulers/relayvsr/scheduler.py new file mode 100644 index 000000000..80b2d3085 --- /dev/null +++ b/lightx2v/models/schedulers/relayvsr/scheduler.py @@ -0,0 +1,23 @@ +import torch + +from lightx2v.models.schedulers.scheduler import BaseScheduler + + +class RelayVSRScheduler(BaseScheduler): + """Official distilled flow step: sigma 1 -> 0, no classifier-free guidance.""" + + def __init__(self, config): + super().__init__(config) + self.keep_latents_dtype_in_scheduler = True + + def prepare(self, seed, device): + self.generator = torch.Generator(device=device).manual_seed(seed) + + def prepare_latents(self, height, width, device): + self.latents = torch.randn((1, 48, 1, height // 4, width // 4), generator=self.generator, device=device, dtype=torch.float32) + + def step_post(self): + self.latents = self.latents - self.noise_pred.float() + + def clear(self): + self.latents = self.noise_pred = self.generator = None diff --git a/scripts/relayvsr/run_relayvsr_sr.sh b/scripts/relayvsr/run_relayvsr_sr.sh new file mode 100644 index 000000000..11014bcca --- /dev/null +++ b/scripts/relayvsr/run_relayvsr_sr.sh @@ -0,0 +1,22 @@ +#!/bin/bash + +# set path firstly +lightx2v_path=/path/to/LightX2V +# https://huggingface.co/kopper/RelayVSR +model_path=/path/to/RelayVSR +video_path=/path/to/input.mp4 + +export CUDA_VISIBLE_DEVICES=0 + +# set environment variables +source "${lightx2v_path}/scripts/base/base.sh" + +python -m lightx2v.infer \ + --model_cls relayvsr \ + --task sr \ + --model_path "${model_path}" \ + --config_json "${lightx2v_path}/configs/relayvsr/relayvsr.json" \ + --video_path "${video_path}" \ + --sr_ratio 2 \ + --save_result_path "${lightx2v_path}/save_results/relayvsr_x2.mp4" \ + --seed 42 From c32ce431a0b2dbf7b75947da28af2fd42e18bcd9 Mon Sep 17 00:00:00 2001 From: Musisoul Date: Fri, 23 Oct 2026 09:58:38 +0000 Subject: [PATCH 2/6] optimize cpu video decoding --- configs/relayvsr/relayvsr.json | 4 ++-- configs/relayvsr/relayvsr_block_offload.json | 4 ++-- configs/relayvsr/relayvsr_sp.json | 4 ++-- lightx2v/models/runners/relayvsr/relayvsr_runner.py | 2 +- lightx2v/models/runners/relayvsr/video_io.py | 7 ++++--- 5 files changed, 11 insertions(+), 10 deletions(-) diff --git a/configs/relayvsr/relayvsr.json b/configs/relayvsr/relayvsr.json index 0b3028007..65493e350 100644 --- a/configs/relayvsr/relayvsr.json +++ b/configs/relayvsr/relayvsr.json @@ -5,6 +5,6 @@ "video_codec": "libx264", "video_crf": 18, "ffmpeg_preset": "fast", - "encoder_threads": 8, - "queue_size": 3 + "encoder_threads": 16, + "queue_size": 16 } diff --git a/configs/relayvsr/relayvsr_block_offload.json b/configs/relayvsr/relayvsr_block_offload.json index 778d6e30a..87b6727b5 100644 --- a/configs/relayvsr/relayvsr_block_offload.json +++ b/configs/relayvsr/relayvsr_block_offload.json @@ -6,6 +6,6 @@ "video_codec": "libx264", "video_crf": 18, "ffmpeg_preset": "fast", - "encoder_threads": 8, - "queue_size": 3 + "encoder_threads": 16, + "queue_size": 16 } diff --git a/configs/relayvsr/relayvsr_sp.json b/configs/relayvsr/relayvsr_sp.json index d1a6c1b93..7ae17c20a 100644 --- a/configs/relayvsr/relayvsr_sp.json +++ b/configs/relayvsr/relayvsr_sp.json @@ -5,8 +5,8 @@ "video_codec": "libx264", "video_crf": 18, "ffmpeg_preset": "fast", - "encoder_threads": 8, - "queue_size": 3, + "encoder_threads": 16, + "queue_size": 16, "parallel": { "seq_p_size": 2, "cfg_p_size": 1 diff --git a/lightx2v/models/runners/relayvsr/relayvsr_runner.py b/lightx2v/models/runners/relayvsr/relayvsr_runner.py index a5fdec7df..5b528270d 100644 --- a/lightx2v/models/runners/relayvsr/relayvsr_runner.py +++ b/lightx2v/models/runners/relayvsr/relayvsr_runner.py @@ -36,7 +36,7 @@ def validate_config(config): parallel = config.get("parallel") or {} if parallel.get("seq_p_size", 1) not in (1, 2): raise ValueError("RelayVSR Flash has two KV groups, so sequence parallelism supports 1 or 2 GPUs") - for key, default, minimum in (("clip_len", 16, 2), ("queue_size", 3, 1), ("encoder_threads", 8, 1)): + for key, default, minimum in (("clip_len", 16, 2), ("queue_size", 16, 1), ("encoder_threads", 16, 1)): if config.get(key, default) < minimum: raise ValueError(f"RelayVSR {key} must be >= {minimum}") diff --git a/lightx2v/models/runners/relayvsr/video_io.py b/lightx2v/models/runners/relayvsr/video_io.py index 9c703e9f5..4bf60b09d 100644 --- a/lightx2v/models/runners/relayvsr/video_io.py +++ b/lightx2v/models/runners/relayvsr/video_io.py @@ -40,7 +40,7 @@ def __init__(self, path, fps, config): self.frames = 0 self.encode_seconds = self.flush_seconds = 0.0 self.copy_seconds = self.queue_wait_seconds = 0.0 - self.queue_size = config.get("queue_size", 3) + self.queue_size = config.get("queue_size", 16) self.closed = False def _write(self, pixels, copy_start, event): @@ -52,9 +52,10 @@ def _write(self, pixels, copy_start, event): self.stream = self.container.add_stream(self.config.get("video_codec", "libx264"), rate=self.fps) self.stream.width, self.stream.height = pixels.shape[1], pixels.shape[0] self.stream.pix_fmt = "yuv420p" - self.stream.codec_context.thread_count = self.config.get("encoder_threads", 8) + self.stream.codec_context.thread_count = self.config.get("encoder_threads", 16) self.stream.options = {"crf": str(self.config.get("video_crf", 18)), "preset": self.config.get("ffmpeg_preset", "fast")} - frame = av.VideoFrame.from_ndarray(pixels.numpy(), format="rgb24") + # Share the pinned RGB buffer with PyAV. + frame = av.VideoFrame.from_numpy_buffer(pixels.numpy(), format="rgb24") frame.pts = self.frames frame.time_base = 1 / self.fps for packet in self.stream.encode(frame): From 0f2020bb17e95f30a10061b5589f00305e47b62b Mon Sep 17 00:00:00 2001 From: Musisoul Date: Fri, 23 Oct 2026 12:54:05 +0000 Subject: [PATCH 3/6] Add compile and warmup --- configs/relayvsr/relayvsr.json | 4 +- configs/relayvsr/relayvsr_block_offload.json | 4 +- configs/relayvsr/relayvsr_compile.json | 12 +++++ configs/relayvsr/relayvsr_sp.json | 2 + lightx2v/models/networks/relayvsr/compile.py | 53 +++++++++++++++++++ .../networks/relayvsr/flash_attention.py | 16 ++++++ .../models/networks/relayvsr/flash_infer.py | 33 +++++++++--- lightx2v/models/networks/relayvsr/model.py | 7 ++- .../models/networks/relayvsr/wan_attention.py | 13 +++++ .../models/networks/relayvsr/wan_infer.py | 32 ++++++++--- lightx2v/models/networks/relayvsr/weights.py | 22 ++++++-- .../runners/relayvsr/relayvsr_runner.py | 33 +++++++++++- 12 files changed, 208 insertions(+), 23 deletions(-) create mode 100644 configs/relayvsr/relayvsr_compile.json create mode 100644 lightx2v/models/networks/relayvsr/compile.py diff --git a/configs/relayvsr/relayvsr.json b/configs/relayvsr/relayvsr.json index 65493e350..81d8606e6 100644 --- a/configs/relayvsr/relayvsr.json +++ b/configs/relayvsr/relayvsr.json @@ -6,5 +6,7 @@ "video_crf": 18, "ffmpeg_preset": "fast", "encoder_threads": 16, - "queue_size": 16 + "queue_size": 16, + "use_compile": false, + "warmup": false } diff --git a/configs/relayvsr/relayvsr_block_offload.json b/configs/relayvsr/relayvsr_block_offload.json index 87b6727b5..011db0435 100644 --- a/configs/relayvsr/relayvsr_block_offload.json +++ b/configs/relayvsr/relayvsr_block_offload.json @@ -7,5 +7,7 @@ "video_crf": 18, "ffmpeg_preset": "fast", "encoder_threads": 16, - "queue_size": 16 + "queue_size": 16, + "use_compile": false, + "warmup": false } diff --git a/configs/relayvsr/relayvsr_compile.json b/configs/relayvsr/relayvsr_compile.json new file mode 100644 index 000000000..8f7cf4628 --- /dev/null +++ b/configs/relayvsr/relayvsr_compile.json @@ -0,0 +1,12 @@ +{ + "cross_attention_backend": "torch_sdpa", + "cpu_offload": false, + "clip_len": 16, + "video_codec": "libx264", + "video_crf": 18, + "ffmpeg_preset": "fast", + "encoder_threads": 16, + "queue_size": 16, + "use_compile": true, + "warmup": true +} diff --git a/configs/relayvsr/relayvsr_sp.json b/configs/relayvsr/relayvsr_sp.json index 7ae17c20a..5a3ad32c0 100644 --- a/configs/relayvsr/relayvsr_sp.json +++ b/configs/relayvsr/relayvsr_sp.json @@ -7,6 +7,8 @@ "ffmpeg_preset": "fast", "encoder_threads": 16, "queue_size": 16, + "use_compile": true, + "warmup": true, "parallel": { "seq_p_size": 2, "cfg_p_size": 1 diff --git a/lightx2v/models/networks/relayvsr/compile.py b/lightx2v/models/networks/relayvsr/compile.py new file mode 100644 index 000000000..530bd0213 --- /dev/null +++ b/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) diff --git a/lightx2v/models/networks/relayvsr/flash_attention.py b/lightx2v/models/networks/relayvsr/flash_attention.py index 75dab2d6f..177f7b25e 100644 --- a/lightx2v/models/networks/relayvsr/flash_attention.py +++ b/lightx2v/models/networks/relayvsr/flash_attention.py @@ -9,6 +9,20 @@ 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, @@ -87,6 +101,8 @@ def _forward( @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 diff --git a/lightx2v/models/networks/relayvsr/flash_infer.py b/lightx2v/models/networks/relayvsr/flash_infer.py index 8559d0f17..038dd4954 100644 --- a/lightx2v/models/networks/relayvsr/flash_infer.py +++ b/lightx2v/models/networks/relayvsr/flash_infer.py @@ -1,6 +1,9 @@ import torch import torch.nn.functional as F +from lightx2v.common.transformer_infer.transformer_infer import BaseTransformerInfer + +from .compile import RelayCompileMixin, mark_sequence_dynamic from .flash_attention import triton_local_attention from .rope import apply_rope @@ -18,12 +21,20 @@ def condition(self, weights, latent): return weights.condition_norm.apply(x) + weights.condition_type.tensor[0] -class FlashTransformerInfer: - def __init__(self, config, parallel): +class FlashTransformerInfer(RelayCompileMixin, BaseTransformerInfer): + def __init__(self, config, parallel, deployment): self.config, self.parallel = config, parallel + self.init_compile(deployment) + + def run_block(self, block_idx, block, x, condition, past_key, past_value, q_rope, k_rope, windows, height, width): + if self.use_compile: + mark_sequence_dynamic(x, *q_rope, *k_rope) + if not condition: + mark_sequence_dynamic(past_key, past_value) + return super().run_block(block_idx, block, x, condition, past_key, past_value, q_rope, k_rope, windows, height, width) def infer_block(self, block, x, condition, past_key, past_value, q_rope, k_rope, windows, height, width): - tokens = height * width + tokens = q_rope[0].shape[0] normalized = block.attention_norm.apply(x) query = block.query.apply(normalized).view(-1, 8, 64) key_op = block.condition_key if condition else block.key @@ -38,6 +49,12 @@ def infer_block(self, block, x, condition, past_key, past_value, q_rope, k_rope, query = apply_rope(query, q_rope) all_key = apply_rope(all_key, k_rope) query, all_key, all_value = query.transpose(0, 1)[None], all_key.transpose(0, 1)[None], all_value.transpose(0, 1)[None] + if torch.compiler.is_compiling(): + spatial = self.config["spatial_attention"] + names = ("condition_self_window_size",) if condition else ("start_condition_window_size", "end_condition_window_size", "video_window_size") + # The attention leaf repeats the video window for the dynamic KV + # length; Python history length does not specialize the block. + windows = tuple(tuple(spatial[name]) for name in names) attended = triton_local_attention(query, all_key, all_value, height, width, windows) attended = attended[0].transpose(0, 1) x = x + block.output.apply(self.parallel.to_sequence(attended)) @@ -46,15 +63,19 @@ def infer_block(self, block, x, condition, past_key, past_value, q_rope, k_rope, return x, key, value def prepare_condition(self, blocks, x, rope, height, width): + if self.use_compile and self.parallel.size > 1: + x = x.clone() # Normalize the first SP shard's storage offset. frequencies = rope.for_positions(height, width, (0,)) window = tuple(self.config["spatial_attention"]["condition_self_window_size"]) caches = [] - for block in blocks: - x, key, value = self.infer_block(block, x, True, None, None, frequencies, frequencies, (window,), height, width) + for index, block in enumerate(blocks): + x, key, value = self.run_block(index, block, x, True, None, None, frequencies, frequencies, (window,), height, width) caches.append((key, value)) return caches def infer(self, blocks, x, conditions, history, rope, frame, height, width): + if self.use_compile and self.parallel.size > 1: + x = x.clone() spatial = self.config["spatial_attention"] previous = min(frame, self.config["window_size"] - 1) q_rope = rope.for_positions(height, width, (0,)) @@ -67,7 +88,7 @@ def infer(self, blocks, x, conditions, history, rope, frame, height, width): past_value = torch.cat((cv, history.v_cache(index)), 0) else: past_key, past_value = ck, cv - x, key, value = self.infer_block(block, x, False, past_key, past_value, q_rope, k_rope, windows, height, width) + x, key, value = self.run_block(index, block, x, False, past_key, past_value, q_rope, k_rope, windows, height, width) history.store_kv(key, value, index) return x diff --git a/lightx2v/models/networks/relayvsr/model.py b/lightx2v/models/networks/relayvsr/model.py index 28a4b193f..cd49085bd 100644 --- a/lightx2v/models/networks/relayvsr/model.py +++ b/lightx2v/models/networks/relayvsr/model.py @@ -35,14 +35,17 @@ def __init__(self, config, group=None): self.wan_config = checkpoint_config(wan_path, "wan") self.flash_config = checkpoint_config(flash_path, "flash") validate_checkpoint_configs(self.wan_config, self.flash_config) + # The distilled checkpoint relies on BF16/FP16 rounding between ops. + # Keep those casts and avoid rewriting mm + residual into addmm. + compile_config = {**config, "compile_options": {"emulate_precision_casts": True, "pattern_matcher": False, **config.get("compile_options", {})}} self.parallel = SequenceParallel(group) self.wan_weights = WanWeights(wan_path, self.wan_config, config) self.flash_weights = FlashWeights(flash_path, self.flash_config) self.wan_pre = RelayWanPreInfer(self.wan_config) - self.wan_transformer = RelayWanTransformerInfer(self.wan_config, config, self.parallel, self.wan_weights) + self.wan_transformer = RelayWanTransformerInfer(self.wan_config, compile_config, self.parallel, self.wan_weights) self.wan_post = RelayWanPostInfer({"out_dim": self.wan_config["out_channels"]}) self.flash_pre = FlashPreInfer() - self.flash_transformer = FlashTransformerInfer(self.flash_config, self.parallel) + self.flash_transformer = FlashTransformerInfer(self.flash_config, self.parallel, compile_config) self.flash_post = FlashPostInfer() self.wan_rope = RotaryEmbedding(wan=True, device=torch.device("cuda"), dtype=GET_DTYPE(), head_dim=self.wan_config["attention_head_dim"], max_length=self.wan_config["rope_max_seq_len"]) self.flash_rope = RotaryEmbedding(wan=False, device=torch.device("cuda"), dtype=GET_DTYPE(), theta=self.flash_config["rope_theta"]) diff --git a/lightx2v/models/networks/relayvsr/wan_attention.py b/lightx2v/models/networks/relayvsr/wan_attention.py index 17791b46b..36e1edcd6 100644 --- a/lightx2v/models/networks/relayvsr/wan_attention.py +++ b/lightx2v/models/networks/relayvsr/wan_attention.py @@ -29,6 +29,16 @@ class SpatialAttentionSpec: key_frame_offset: int = 0 +@torch.library.custom_op("lightx2v::relayvsr_wan_attention", mutates_args=(), device_types="cuda") +def _compiled_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, geometry: list[int]) -> torch.Tensor: + return spatial_attention(query, key, value, SpatialAttentionSpec(*geometry)) + + +@_compiled_attention.register_fake +def _compiled_attention_fake(query, key, value, geometry): + return torch.empty_like(query, memory_format=torch.contiguous_format) + + @triton.jit def _spatial_forward( Q, @@ -126,6 +136,9 @@ def _spatial_forward( def spatial_attention(query, key, value, spec: SpatialAttentionSpec): """[B,T*H*W,heads,dim] -> same layout, exact shifted rectangular mask.""" + if torch.compiler.is_compiling(): + geometry = [spec.frames, spec.height, spec.width, spec.window_height, spec.window_width, spec.temporal_window, spec.key_frames, spec.query_frame_offset, spec.key_frame_offset] + return _compiled_attention(query, key, value, geometry) if ( key.shape != value.shape or query.shape[0] != key.shape[0] diff --git a/lightx2v/models/networks/relayvsr/wan_infer.py b/lightx2v/models/networks/relayvsr/wan_infer.py index 5a0d509c9..6c95a8707 100644 --- a/lightx2v/models/networks/relayvsr/wan_infer.py +++ b/lightx2v/models/networks/relayvsr/wan_infer.py @@ -1,4 +1,5 @@ import math +from dataclasses import replace import torch import torch.nn.functional as F @@ -8,6 +9,7 @@ from lightx2v.models.networks.wan.infer.post_infer import WanPostInfer from lightx2v.models.networks.wan.infer.transformer_infer import modulate +from .compile import RelayCompileMixin, mark_sequence_dynamic from .rope import apply_rope from .wan_attention import SpatialAttentionSpec, spatial_attention @@ -38,7 +40,7 @@ def infer(self, weights, latents, lr): return hidden + lq -class RelayWanTransformerInfer(WanOffloadTransformerInfer): +class RelayWanTransformerInfer(RelayCompileMixin, WanOffloadTransformerInfer): """Wan execution/offload/FFN with RelayVSR's causal reference attention. Keep the shared block pipeline, offload, normal-size FFN and final residuals. @@ -80,21 +82,32 @@ def prepare_cross_attention(self, blocks, context): self.infer_func(blocks, None, pre) return self.cross_cache - def infer_block(self, block, x, pre_infer_out): + def run_block(self, block_idx, block, x, pre_infer_out): + if pre_infer_out.adapter_args.get("prepare_cross_kv", False): + return self.infer_block(block, x, pre_infer_out) + inputs = pre_infer_out.adapter_args + cache = (inputs["cache"].k_cache(block_idx), inputs["cache"].v_cache(block_idx)) + attention = {"q_rope": inputs["q_rope"], "k_rope": inputs["k_rope"], "cache": cache, "cache_slot": inputs["frame"] % self.relay_config["temporal_context_frames"]} + pre_infer_out = replace(pre_infer_out, context=self.cross_cache[block_idx]) + if self.use_compile: + mark_sequence_dynamic(x, *cache, *inputs["q_rope"], *inputs["k_rope"]) + return super().run_block(block_idx, block, x, pre_infer_out, attention) + + def infer_block(self, block, x, pre_infer_out, self_attn_kwargs=None): if pre_infer_out.adapter_args.get("prepare_cross_kv", False): phase, context = block.compute_phases[1], pre_infer_out.context key = phase.cross_attn_norm_k.apply(phase.cross_attn_k.apply(context)).view(-1, self.num_heads, self.head_dim) value = phase.cross_attn_v.apply(context).view(-1, self.num_heads, self.head_dim) self.cross_cache.append((key, value)) return x - return super().infer_block(block, x, pre_infer_out, self_attn_kwargs=pre_infer_out.adapter_args) + return super().infer_block(block, x, pre_infer_out, self_attn_kwargs=self_attn_kwargs) def infer_cross_attn(self, phase, x, context, y_out, gate_msa): # RelayVSR has one fixed prompt and no image/CFG branch. Keep its # persistent context KV here without changing the Wan family contract. x.add_(y_out * gate_msa.squeeze()) query = phase.cross_attn_norm_q.apply(phase.cross_attn_q.apply(phase.norm3.apply(x))).view(-1, self.num_heads, self.head_dim) - key, value = self.cross_cache[self.block_idx] + key, value = context attended = phase.cross_attn_1.apply( q=query, k=key, @@ -115,14 +128,19 @@ def infer_ffn(self, phase, x, attn_out, c_shift_msa, c_scale_msa, c_gate_msa=Non normalized = self.modulate_func(phase.norm2.apply(x), scale=c_scale_msa, shift=c_shift_msa).reshape(1, -1) return phase.ffn_2.apply(F.gelu(phase.ffn_0.apply(normalized), approximate="tanh")) - def infer_self_attn(self, phase, x, shift_msa, scale_msa, grid_sizes=None, *, q_rope, k_rope, cache, frame): + def infer_self_attn(self, phase, x, shift_msa, scale_msa, grid_sizes=None, *, q_rope, k_rope, cache, cache_slot): _, height, width = grid_sizes normalized = self.modulate_func(phase.norm1.apply(x), scale_msa, shift_msa) - tokens = height * width + tokens = q_rope[0].shape[0] query = self.parallel.to_heads(phase.self_attn_norm_q.apply(phase.self_attn_q.apply(normalized)).view(-1, self.num_heads, self.head_dim), tokens) key = self.parallel.to_heads(phase.self_attn_norm_k.apply(phase.self_attn_k.apply(normalized)).view(-1, self.num_heads, self.head_dim), tokens) value = self.parallel.to_heads(phase.self_attn_v.apply(normalized).view(-1, self.num_heads, self.head_dim), tokens) - key, value = cache.update(key, value, self.block_idx, frame) + start = cache_slot * tokens + cache[0][start : start + tokens].copy_(key) + cache[1][start : start + tokens].copy_(value) + # Derive both views inside the graph. Independent overlapping inputs + # can lose the write-before-read dependency during functionalization. + key, value = cache[0][: k_rope[0].shape[0]], cache[1][: k_rope[0].shape[0]] query = apply_rope(query, q_rope) key = apply_rope(key, k_rope) visible = key.shape[0] // tokens diff --git a/lightx2v/models/networks/relayvsr/weights.py b/lightx2v/models/networks/relayvsr/weights.py index a0cd35e52..aca005e48 100644 --- a/lightx2v/models/networks/relayvsr/weights.py +++ b/lightx2v/models/networks/relayvsr/weights.py @@ -7,9 +7,18 @@ from lightx2v.common.modules.weight_module import WeightModule, WeightModuleList from lightx2v.common.ops.norm.layer_norm_weight import LNWeight -from lightx2v.common.ops.norm.rms_norm_weight import RMSWeight +from lightx2v.common.ops.norm.rms_norm_weight import RMSWeight, RMSWeightNative from lightx2v.utils.envs import GET_DTYPE -from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, CONV2D_WEIGHT_REGISTER, CONV3D_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, RMS_WEIGHT_REGISTER, TENSOR_REGISTER +from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, CONV2D_WEIGHT_REGISTER, CONV3D_WEIGHT_REGISTER, MM_WEIGHT_REGISTER, TENSOR_REGISTER + +from .compile import native_layer_norm, native_rms_norm + + +class RelayRMSNorm(RMSWeightNative): + def apply(self, x): + if torch.compiler.is_compiling(): + return native_rms_norm(x, self._get_actual_weight(), self.eps) + return super().apply(x) class FP32LayerNorm(LNWeight): @@ -18,7 +27,10 @@ def __init__(self, *args, cast_output=False, **kwargs): self.cast_output = cast_output def apply(self, x): - output = F.layer_norm(x.float(), (x.shape[-1],), None if self.weight is None else self.weight.float(), None if self.bias is None else self.bias.float(), self.eps) + if torch.compiler.is_compiling(): + output = native_layer_norm(x, self.weight, self.bias, self.eps) + else: + output = F.layer_norm(x.float(), (x.shape[-1],), None if self.weight is None else self.weight.float(), None if self.bias is None else self.bias.float(), self.eps) return output.to(x.dtype) if self.cast_output else output @@ -36,8 +48,8 @@ def linear(self, name, key, *, bias=True, buffer=False): def tensor(self, name, key, *, buffer=False): self.add_module(name, TENSOR_REGISTER["Default"](key, create_cuda_buffer=buffer)) - def rms(self, name, key, *, eps, backend="torch_native", buffer=False): - self.add_module(name, RMS_WEIGHT_REGISTER[backend](key + ".weight", eps=eps, create_cuda_buffer=buffer)) + def rms(self, name, key, *, eps, buffer=False): + self.add_module(name, RelayRMSNorm(key + ".weight", eps=eps, create_cuda_buffer=buffer)) def conv(self, name, key, *, stride=1, padding=0): self.add_module(name, CONV2D_WEIGHT_REGISTER["Default"](key + ".weight", key + ".bias", stride=stride, padding=padding)) diff --git a/lightx2v/models/runners/relayvsr/relayvsr_runner.py b/lightx2v/models/runners/relayvsr/relayvsr_runner.py index 5b528270d..e34e714c1 100644 --- a/lightx2v/models/runners/relayvsr/relayvsr_runner.py +++ b/lightx2v/models/runners/relayvsr/relayvsr_runner.py @@ -4,7 +4,7 @@ import time from contextlib import closing from fractions import Fraction -from itertools import islice +from itertools import islice, repeat from pathlib import Path from typing import ClassVar @@ -47,6 +47,8 @@ class RelayVSRRunner(DefaultRunner): supported_request_fields_by_task: ClassVar[dict] = {"sr": COMMON_REQUEST_FIELDS | {"video_path", "image_path", "sr_ratio", "size"}} + WARMUP_RESOLUTIONS = ((768, 1024), (768, 1344)) + def __init__(self, config): for key, value in (("infer_steps", 1), ("guidance_scale", 1.0), ("enable_cfg", False), ("fps", None)): config.setdefault(key, value) @@ -74,6 +76,35 @@ def init_modules(self): self.config.lock() logger.info("RelayVSR native 4x model loaded in {:.2f}s", self.load_seconds) + @torch.no_grad() + def run_warmup(self): + if type(self) is not RelayVSRRunner or self.config["task"] != "sr": + raise NotImplementedError("RelayVSR warmup supports the native SR runner only") + if self.config.get("lazy_load", False): + raise NotImplementedError("RelayVSR warmup does not support lazy loading") + started = time.perf_counter() + # Fill the causal Wan cache and overwrite every slot once. Flash's + # condition, first-frame and history branches follow the normal stream. + references = 2 * self.model.wan_config["temporal_context_frames"] + frames = (references - 1) * (self.config.get("clip_len", 16) - 1) + 1 + for height, width in self.WARMUP_RESOLUTIONS: + frame = prediction = pixels = None + logger.info("RelayVSR warmup: {}x{}, {} reference frames", height, width, references) + try: + frame = torch.zeros((3, height, width), dtype=torch.uint8, pin_memory=True) + with closing(self.restore_frames(repeat(frame, frames), 42)) as predictions: + for prediction in predictions: + for scale in (2, 4): + pixels = self.postprocess(prediction, (height * scale, width * scale)) + torch.cuda.synchronize() + finally: + frame = prediction = pixels = None + self.timings = None + self.input_info = None + self.warmup_seconds = time.perf_counter() - started + logger.info("RelayVSR warmup completed in {:.2f}s", self.warmup_seconds) + self._maybe_freeze_gc() + @staticmethod def output_size(input_info, height, width): if input_info.size: From d7bb77139365d818a614f65d7e48b87994343396 Mon Sep 17 00:00:00 2001 From: Musisoul Date: Fri, 23 Oct 2026 13:07:55 +0000 Subject: [PATCH 4/6] libx265 ultrafast --- configs/relayvsr/relayvsr.json | 4 +-- configs/relayvsr/relayvsr_block_offload.json | 4 +-- configs/relayvsr/relayvsr_compile.json | 4 +-- configs/relayvsr/relayvsr_sp.json | 4 +-- lightx2v/models/runners/relayvsr/video_io.py | 27 +++++++++++++++++--- 5 files changed, 31 insertions(+), 12 deletions(-) diff --git a/configs/relayvsr/relayvsr.json b/configs/relayvsr/relayvsr.json index 81d8606e6..370391679 100644 --- a/configs/relayvsr/relayvsr.json +++ b/configs/relayvsr/relayvsr.json @@ -2,9 +2,9 @@ "cross_attention_backend": "torch_sdpa", "cpu_offload": false, "clip_len": 16, - "video_codec": "libx264", + "video_codec": "libx265", "video_crf": 18, - "ffmpeg_preset": "fast", + "ffmpeg_preset": "ultrafast", "encoder_threads": 16, "queue_size": 16, "use_compile": false, diff --git a/configs/relayvsr/relayvsr_block_offload.json b/configs/relayvsr/relayvsr_block_offload.json index 011db0435..1f9457e5a 100644 --- a/configs/relayvsr/relayvsr_block_offload.json +++ b/configs/relayvsr/relayvsr_block_offload.json @@ -3,9 +3,9 @@ "cpu_offload": true, "offload_granularity": "block", "clip_len": 16, - "video_codec": "libx264", + "video_codec": "libx265", "video_crf": 18, - "ffmpeg_preset": "fast", + "ffmpeg_preset": "ultrafast", "encoder_threads": 16, "queue_size": 16, "use_compile": false, diff --git a/configs/relayvsr/relayvsr_compile.json b/configs/relayvsr/relayvsr_compile.json index 8f7cf4628..5d9e132f5 100644 --- a/configs/relayvsr/relayvsr_compile.json +++ b/configs/relayvsr/relayvsr_compile.json @@ -2,9 +2,9 @@ "cross_attention_backend": "torch_sdpa", "cpu_offload": false, "clip_len": 16, - "video_codec": "libx264", + "video_codec": "libx265", "video_crf": 18, - "ffmpeg_preset": "fast", + "ffmpeg_preset": "ultrafast", "encoder_threads": 16, "queue_size": 16, "use_compile": true, diff --git a/configs/relayvsr/relayvsr_sp.json b/configs/relayvsr/relayvsr_sp.json index 5a3ad32c0..cbac682e1 100644 --- a/configs/relayvsr/relayvsr_sp.json +++ b/configs/relayvsr/relayvsr_sp.json @@ -2,9 +2,9 @@ "cross_attention_backend": "torch_sdpa", "cpu_offload": false, "clip_len": 16, - "video_codec": "libx264", + "video_codec": "libx265", "video_crf": 18, - "ffmpeg_preset": "fast", + "ffmpeg_preset": "ultrafast", "encoder_threads": 16, "queue_size": 16, "use_compile": true, diff --git a/lightx2v/models/runners/relayvsr/video_io.py b/lightx2v/models/runners/relayvsr/video_io.py index 4bf60b09d..dfdff8e92 100644 --- a/lightx2v/models/runners/relayvsr/video_io.py +++ b/lightx2v/models/runners/relayvsr/video_io.py @@ -49,11 +49,21 @@ def _write(self, pixels, copy_start, event): start = time.perf_counter() if self.container is None: self.container = av.open(str(self.path), mode="w") - self.stream = self.container.add_stream(self.config.get("video_codec", "libx264"), rate=self.fps) + codec = self.config.get("video_codec", "libx265") + self.stream = self.container.add_stream(codec, rate=self.fps) self.stream.width, self.stream.height = pixels.shape[1], pixels.shape[0] self.stream.pix_fmt = "yuv420p" - self.stream.codec_context.thread_count = self.config.get("encoder_threads", 16) - self.stream.options = {"crf": str(self.config.get("video_crf", 18)), "preset": self.config.get("ffmpeg_preset", "fast")} + threads = self.config.get("encoder_threads", 16) + options = {"crf": str(self.config.get("video_crf", 18)), "preset": self.config.get("ffmpeg_preset", "ultrafast")} + if codec == "libx265": + options["x265-params"] = f"log-level=warning:pools={threads}" + if self.stream.width < 64: + # x265's SIMD path can corrupt memory on very narrow frames. + options["x265-params"] += ":asm=0" + self.stream.codec_context.codec_tag = "hvc1" + else: + self.stream.codec_context.thread_count = threads + self.stream.options = options # Share the pinned RGB buffer with PyAV. frame = av.VideoFrame.from_numpy_buffer(pixels.numpy(), format="rgb24") frame.pts = self.frames @@ -88,7 +98,16 @@ def close(self): future.result() start = time.perf_counter() if self.stream is not None: - for packet in self.stream.encode(): + packets = self.stream.encode() + if self.stream.codec_context.name == "libx265" and self.frames < 3: + # Match SwiftVR's DTS correction for very short HEVC videos. + decoder = av.CodecContext.create("hevc", "r") + decoder.thread_count = 1 + decoder.extradata = self.stream.codec_context.extradata + decoder.open() + for packet in packets: + packet.dts = packet.pts - decoder.reorder_depth + for packet in packets: self.container.mux(packet) self.flush_seconds = time.perf_counter() - start finally: From 8f0e13863a9af6671dbd7f77641a5c0c0fa2a6e5 Mon Sep 17 00:00:00 2001 From: Musisoul Date: Fri, 23 Oct 2026 13:12:01 +0000 Subject: [PATCH 5/6] update readme --- README.md | 3 +- README_zh.md | 3 +- scripts/relayvsr/README.md | 61 +++++++++++++++++++++++++++++++++++ scripts/relayvsr/README_zh.md | 59 +++++++++++++++++++++++++++++++++ 4 files changed, 124 insertions(+), 2 deletions(-) create mode 100644 scripts/relayvsr/README.md create mode 100644 scripts/relayvsr/README_zh.md diff --git a/README.md b/README.md index 882afc9c8..496b1db42 100644 --- a/README.md +++ b/README.md @@ -239,13 +239,14 @@ LightX2V supports image, video, and world models and provides quantized weights,
🔍 Restoration & Super-Resolution

-SeedVR2 · SwiftVR +SeedVR2 · SwiftVR · RelayVSR | 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) |
diff --git a/README_zh.md b/README_zh.md index 62d937be5..0fb303628 100644 --- a/README_zh.md +++ b/README_zh.md @@ -239,13 +239,14 @@ LightX2V 支持多种图像、视频与世界模型，并提供量化权重、
🔍 修复与超分辨率

-SeedVR2 · SwiftVR +SeedVR2 · SwiftVR · RelayVSR | 模型 / 变体与权重 | 使用指南 / 脚本 | | --- | --- | | [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) |
diff --git a/scripts/relayvsr/README.md b/scripts/relayvsr/README.md new file mode 100644 index 000000000..f31622f29 --- /dev/null +++ b/scripts/relayvsr/README.md @@ -0,0 +1,61 @@ +# RelayVSR video super-resolution + +[中文说明](README_zh.md) + +Native integration of the official [RelayVSR](https://github.com/kopperx/RelayVSR) big model: Wan, with approximately 5.03B parameters, generates sparse reference frames, while Flash-S, with approximately 31.9M parameters, decodes all video frames. + +**The model always computes at 4× internally.** `--sr_ratio 4` outputs the result directly; `--sr_ratio 2` removes padding from the 4× result and downsizes it using FP32 antialiased bicubic interpolation before encoding. 2× still requires the full 4× model computation, but reduces output transfer, encoding costs and file size. + +## Setup and running + +Use the LightX2V CUDA environment with PyTorch, Triton, PyAV and FFmpeg. BF16 is the default. Download the two official checkpoints; no weight conversion is required: + +```bash +hf download kopper/RelayVSR wan.safetensors flash.safetensors --local-dir /models/RelayVSR +``` + +Edit `lightx2v_path`, `model_path` and `video_path` at the top of `scripts/relayvsr/run_relayvsr_sr.sh`, then run: + +```bash +bash scripts/relayvsr/run_relayvsr_sr.sh +``` + +GPU selection and inference arguments are set directly in the script, which initializes the environment through `scripts/base/base.sh`. Change `--sr_ratio` to `4` for 4× output and set the output path with `--save_result_path`. You can also use the unified CLI directly: + +```bash +python -m lightx2v.infer \ + --model_cls relayvsr --task sr \ + --model_path /models/RelayVSR \ + --config_json configs/relayvsr/relayvsr.json \ + --video_path /videos/input.mp4 --sr_ratio 2 --seed 42 \ + --save_result_path /videos/output_x2.mp4 +``` + +For images, replace `--video_path` with `--image_path input.png` and save the output as PNG. + +## Execution modes + +| Config | Behavior | +| --- | --- | +| `relayvsr.json` | Weights remain on GPU, with eager execution; compilation and warmup are disabled | +| `relayvsr_compile.json` | Enables `torch.compile` for both Wan and Flash Transformer blocks, with startup warmup | +| `relayvsr_block_offload.json` | Wan weights stay on CPU and are loaded onto the GPU block by block using LightX2V double buffering and prefetch; Flash remains on GPU | +| `relayvsr_sp.json` | Two-GPU sequence parallelism; both Wan and Flash reuse LightX2V Ulysses all-to-all | + +Select a configuration by changing `--config_json` in the script. For example, set it to `"${lightx2v_path}/configs/relayvsr/relayvsr_block_offload.json"` to enable block weight offload. + +The example script uses `relayvsr_compile.json` by default, enabling compilation and startup warmup. `use_compile` and `warmup` are startup JSON settings: setting both to `true` triggers compilation during warmup after model loading; enabling only `warmup` warms up eager execution; enabling only `use_compile` triggers compilation during the first inference. To disable both, change the script's `--config_json` back to `"${lightx2v_path}/configs/relayvsr/relayvsr.json"`. Compilation covers all 30 Wan blocks and all 10 Flash blocks, including Flash condition and video paths. Pre/postprocessing, weight transfers and video I/O stay outside the compiled blocks. + +Two-GPU example: + +```bash +CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 -m lightx2v.infer \ + --model_cls relayvsr --task sr --model_path /models/RelayVSR \ + --config_json configs/relayvsr/relayvsr_sp.json \ + --video_path /videos/input.mp4 --sr_ratio 2 --seed 42 \ + --save_result_path /videos/output_x2.mp4 +``` + +Flash has only two KV groups, so sequence parallelism supports at most two GPUs. Partitioned linear operations may introduce BF16 numerical differences; use the default single-GPU eager mode when elementwise parity is required. + +Each configuration explicitly lists the cross-attention backend, offload switch, clip length and encoding parameters for easy editing. Defaults are `torch_sdpa`, 16-frame clips and source FPS. CPU encoding uses `libx265`, `ultrafast` and CRF 18, with the `hvc1` tag; `encoder_threads` controls the x265 worker pool. diff --git a/scripts/relayvsr/README_zh.md b/scripts/relayvsr/README_zh.md new file mode 100644 index 000000000..ab1f6232e --- /dev/null +++ b/scripts/relayvsr/README_zh.md @@ -0,0 +1,59 @@ +# RelayVSR 视频超分 + +原生接入 [RelayVSR](https://github.com/kopperx/RelayVSR) 官方大模型：约 5.03B 参数的 Wan 生成稀疏参考帧，约 31.9M 参数的 Flash-S 解码全部视频帧。 + +**模型内部始终计算 4×。** `--sr_ratio 4` 直接输出；`--sr_ratio 2` 将裁去 padding 的 4× 结果用 FP32、抗锯齿 bicubic 缩小，再编码。2× 不会省去 4× 模型计算，但可以降低输出搬运、编码和文件大小。 + +## 准备与运行 + +使用 LightX2V 的 CUDA 环境，包含 PyTorch、Triton、PyAV 和 FFmpeg。默认 BF16。下载官方两个文件即可，不需要转换权重： + +```bash +hf download kopper/RelayVSR wan.safetensors flash.safetensors --local-dir /models/RelayVSR +``` + +先修改 `scripts/relayvsr/run_relayvsr_sr.sh` 顶部的 `lightx2v_path`、`model_path` 和 `video_path`，然后运行： + +```bash +bash scripts/relayvsr/run_relayvsr_sr.sh +``` + +显卡和推理参数直接在脚本中设置，通过 `scripts/base/base.sh` 初始化环境。将 `--sr_ratio` 改为 `4` 可直接输出 4×，输出位置通过 `--save_result_path` 设置。也可以直接调用统一 CLI： + +```bash +python -m lightx2v.infer \ + --model_cls relayvsr --task sr \ + --model_path /models/RelayVSR \ + --config_json configs/relayvsr/relayvsr.json \ + --video_path /videos/input.mp4 --sr_ratio 2 --seed 42 \ + --save_result_path /videos/output_x2.mp4 +``` + +图片用 `--image_path input.png` 替换 `--video_path`，输出设为 PNG。 + +## 执行模式 + +| 配置 | 行为 | +| --- | --- | +| `relayvsr.json` | 常驻显存，eager，关闭编译和预热 | +| `relayvsr_compile.json` | Wan 和 Flash 的 Transformer block 都启用 `torch.compile`，启动时预热 | +| `relayvsr_block_offload.json` | Wan 权重放在 CPU，用 LightX2V 双缓冲和预取逐块上卡；Flash 常驻 | +| `relayvsr_sp.json` | 两卡序列并行；Wan 和 Flash 都复用 LightX2V 的 Ulysses all-to-all | + +选择配置时修改脚本中的 `--config_json`，例如将其改为 `"${lightx2v_path}/configs/relayvsr/relayvsr_block_offload.json"`，即可启用 block 权重卸载。 + +示例脚本默认使用 `relayvsr_compile.json`，同时开启编译和启动预热。`use_compile` 和 `warmup` 是启动 JSON 参数：两者都为 `true` 时，模型加载后自动在预热过程中触发编译；只开启 `warmup` 可预热 eager；只开启 `use_compile` 则在首次推理时触发编译。需要关闭两者时，将脚本的 `--config_json` 改回 `"${lightx2v_path}/configs/relayvsr/relayvsr.json"`。编译覆盖 Wan 的 30 个 block 和 Flash 的 10 个 block，包括 Flash 条件帧与视频帧两条路径；前后处理、权重搬运和视频读写不在编译图内。 + +双卡示例： + +```bash +CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 -m lightx2v.infer \ + --model_cls relayvsr --task sr --model_path /models/RelayVSR \ + --config_json configs/relayvsr/relayvsr_sp.json \ + --video_path /videos/input.mp4 --sr_ratio 2 --seed 42 \ + --save_result_path /videos/output_x2.mp4 +``` + +Flash 只有两个 KV groups，序列并行最多支持两卡。分块线性计算可能带来 BF16 数值差异；需要逐元素对齐时用默认单卡 eager。 + +每份配置显式列出交叉注意力后端、卸载开关、窗口长度和编码参数，方便修改。默认使用 `torch_sdpa`、16 帧窗口和源帧率。CPU 编码使用 `libx265`、`ultrafast`、CRF 18，并设置 `hvc1` 标记；`encoder_threads` 控制 x265 工作线程池。 From 19b04861fc0e633a462519fff0d1c5430d38a9c2 Mon Sep 17 00:00:00 2001 From: Musisoul Date: Fri, 23 Oct 2026 15:48:17 +0000 Subject: [PATCH 6/6] update server example --- scripts/relayvsr/README.md | 42 +++++++++++++++++++++++++++++++++++ scripts/relayvsr/README_zh.md | 42 +++++++++++++++++++++++++++++++++++ 2 files changed, 84 insertions(+) diff --git a/scripts/relayvsr/README.md b/scripts/relayvsr/README.md index f31622f29..cd5d6df39 100644 --- a/scripts/relayvsr/README.md +++ b/scripts/relayvsr/README.md @@ -59,3 +59,45 @@ CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 -m lightx2v.in Flash has only two KV groups, so sequence parallelism supports at most two GPUs. Partitioned linear operations may introduce BF16 numerical differences; use the default single-GPU eager mode when elementwise parity is required. Each configuration explicitly lists the cross-attention backend, offload switch, clip length and encoding parameters for easy editing. Defaults are `torch_sdpa`, 16-frame clips and source FPS. CPU encoding uses `libx265`, `ultrafast` and CRF 18, with the `hvc1` tag; `encoder_threads` controls the x265 worker pool. + +## Serving and requests + +Set the repository and model paths, then start a single-GPU server: + +```bash +lightx2v_path=/path/to/LightX2V +model_path=/models/RelayVSR + +export CUDA_VISIBLE_DEVICES=0 +source "${lightx2v_path}/scripts/base/base.sh" + +python -m lightx2v.server \ + --model_cls relayvsr --task sr \ + --model_path "${model_path}" \ + --config_json "${lightx2v_path}/configs/relayvsr/relayvsr_compile.json" \ + --host 0.0.0.0 --port 8000 +``` + +The startup configuration enables compilation and warmup. After model loading and warmup finish, check the service status and submit a 4× request from another terminal. Both input and output paths must be absolute paths accessible on the server: + +```bash +curl -fsS http://127.0.0.1:8000/v1/service/status + +curl -fsS -X POST http://127.0.0.1:8000/v1/tasks/video/ \ + -H 'Content-Type: application/json' \ + -d '{ + "video_path": "/videos/input.mp4", + "sr_ratio": 4, + "save_result_path": "/videos/output_x4.mp4", + "seed": 42 + }' +``` + +The submission endpoint returns a `task_id` asynchronously. Set the variable below to that value to query the task status: + +```bash +task_id=REPLACE_WITH_RETURNED_TASK_ID +curl -fsS "http://127.0.0.1:8000/v1/tasks/${task_id}/status" +``` + +When the status endpoint reports `status` as `completed`, the video has been saved to the server-side `save_result_path`. If it reports `failed`, check the returned `error`. Submit further requests to the resident server to reuse the loaded model, setting the input path, output path and `seed` for each request. Set `sr_ratio` to `2` for 2× output. Compilation, warmup and encoding options are startup JSON settings; restart the server after changing them. diff --git a/scripts/relayvsr/README_zh.md b/scripts/relayvsr/README_zh.md index ab1f6232e..019cdd908 100644 --- a/scripts/relayvsr/README_zh.md +++ b/scripts/relayvsr/README_zh.md @@ -57,3 +57,45 @@ CUDA_VISIBLE_DEVICES=0,1 torchrun --standalone --nproc_per_node=2 -m lightx2v.in Flash 只有两个 KV groups，序列并行最多支持两卡。分块线性计算可能带来 BF16 数值差异；需要逐元素对齐时用默认单卡 eager。 每份配置显式列出交叉注意力后端、卸载开关、窗口长度和编码参数，方便修改。默认使用 `torch_sdpa`、16 帧窗口和源帧率。CPU 编码使用 `libx265`、`ultrafast`、CRF 18，并设置 `hvc1` 标记；`encoder_threads` 控制 x265 工作线程池。 + +## 服务启动与请求 + +修改仓库和模型路径后，启动单卡服务： + +```bash +lightx2v_path=/path/to/LightX2V +model_path=/models/RelayVSR + +export CUDA_VISIBLE_DEVICES=0 +source "${lightx2v_path}/scripts/base/base.sh" + +python -m lightx2v.server \ + --model_cls relayvsr --task sr \ + --model_path "${model_path}" \ + --config_json "${lightx2v_path}/configs/relayvsr/relayvsr_compile.json" \ + --host 0.0.0.0 --port 8000 +``` + +示例通过启动配置开启编译和预热。等待模型加载、预热完成后，在另一终端检查服务状态并提交 x4 请求。输入和输出路径都应填写服务端可访问的绝对路径： + +```bash +curl -fsS http://127.0.0.1:8000/v1/service/status + +curl -fsS -X POST http://127.0.0.1:8000/v1/tasks/video/ \ + -H 'Content-Type: application/json' \ + -d '{ + "video_path": "/videos/input.mp4", + "sr_ratio": 4, + "save_result_path": "/videos/output_x4.mp4", + "seed": 42 + }' +``` + +提交接口异步返回 `task_id`。将其填入下面的变量，查询任务状态： + +```bash +task_id=REPLACE_WITH_RETURNED_TASK_ID +curl -fsS "http://127.0.0.1:8000/v1/tasks/${task_id}/status" +``` + +状态接口的 `status` 为 `completed` 时，视频已保存到服务端的 `save_result_path`；若为 `failed`，查看返回的 `error`。常驻服务可连续提交请求，复用已加载的模型；每个请求分别设置输入路径、输出路径和 `seed`。将 `sr_ratio` 改为 `2` 即可输出 x2。编译、预热和编码参数通过启动 JSON 配置，修改后需重启服务。