diff --git a/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py index 6dba2d58f..4dd2ca445 100644 --- a/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py +++ b/fastvideo/pipelines/basic/open_h3_regenerate_2k/stages/denoising.py @@ -14,6 +14,10 @@ from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.forward_context import set_forward_context from fastvideo.logger import init_logger from fastvideo.pipelines.basic.minimax_h3.packing import MiniMaxH3PackedLayout, build_row_timesteps +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import ( + _h3_vsa_metadata_builder, + _h3_vsa_single_region_segments, +) from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY from fastvideo.pipelines.basic.open_h3_regenerate_2k.sampling import refine_sigmas from fastvideo.pipelines.basic.open_h3_regenerate_2k.sharp_curve import SharpCurveBranches @@ -62,18 +66,38 @@ class Regenerate2KDenoisingStage(PipelineStage): device = get_local_torch_device() sigmas = refine_sigmas(batch.num_inference_steps, self.start_sigma) state = batch.latents.to(device) + vsa = self._vsa_plan(batch, layout, transformer, fastvideo_args, device) with self._attention_kernel(device): - state = self._euler_loop(batch, layout, state, sigmas, device, transformer) + state = self._euler_loop(batch, layout, state, sigmas, device, transformer, vsa) batch.latents = state return batch + @staticmethod + def _vsa_plan(batch: ForwardBatch, layout: MiniMaxH3PackedLayout, transformer: Any, fastvideo_args: FastVideoArgs, + device: torch.device) -> tuple[Any, ...] | None: + """Per-step VSA metadata inputs, or None for a dense backend. + + The builder and the packed segments come from MiniMax-H3's own denoising stage so the refine sees the same + sparse layout its t2v runs do; the refine layout is a single generated-video region (no reference rows). + """ + builder = _h3_vsa_metadata_builder(transformer, fastvideo_args) + if builder is None: + return None + patch_size = fastvideo_args.pipeline_config.dit_config.patch_size + sparsity = fastvideo_args.VSA_sparsity + if sparsity is None: + sparsity = getattr(batch, "VSA_sparsity", 0.0) + tile_size = fastvideo_args.VSA_tile_size + return (builder, patch_size, _h3_vsa_single_region_segments(layout, + patch_size), float(sparsity), tile_size, device) + @staticmethod def _attention_kernel(device: torch.device) -> contextlib.AbstractContextManager[None]: """cuDNN attention: at 133k tokens PyTorch's default SDPA dispatch picks its flash kernel, 3x slower on B200.""" return sdpa_kernel(SDPBackend.CUDNN_ATTENTION) if device.type == "cuda" else contextlib.nullcontext() def _euler_loop(self, batch: ForwardBatch, layout: MiniMaxH3PackedLayout, state: torch.Tensor, sigmas: list[float], - device: torch.device, transformer: Any) -> torch.Tensor: + device: torch.device, transformer: Any, vsa: Any) -> torch.Tensor: position_ids = layout.position_ids.to(device) token_tags = layout.token_tags.to(device) video_indices = layout.video_indices.to(device) @@ -89,8 +113,19 @@ class Regenerate2KDenoisingStage(PipelineStage): condition_audio_timestep=1.0) unique, inverse = unique.to(device), inverse.to(device) self.curve.set_timesteps(unique) + attn_metadata = None + if vsa is not None: + builder, patch_size, segments, sparsity, tile_size, vsa_device = vsa + attn_metadata = builder.build(current_timestep=index, + patch_size=patch_size, + VSA_sparsity=sparsity, + packed_segments=segments, + device=vsa_device, + exempt=True, + dense_layers=(), + tile_size=tile_size) torch.compiler.cudagraph_mark_step_begin() - with set_forward_context(current_timestep=index, attn_metadata=None, forward_batch=batch): + with set_forward_context(current_timestep=index, attn_metadata=attn_metadata, forward_batch=batch): velocity, _ = transformer(hidden_states=state.to(torch.bfloat16)[None], audio_hidden_states=audio, encoder_hidden_states=prompt_embeds, diff --git a/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py b/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py index c079b30a0..98b574890 100644 --- a/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py +++ b/fastvideo/tests/pipelines/test_open_h3_regenerate_2k.py @@ -186,9 +186,13 @@ class _ConstantVelocityTransformer: def __init__(self, velocity): self.velocity = velocity self.calls = [] + self.attn_metadata = [] def __call__(self, **kwargs): + from fastvideo.forward_context import get_forward_context + self.calls.append(kwargs) + self.attn_metadata.append(get_forward_context().attn_metadata) rows = kwargs["hidden_states"].shape[1] audio_rows = kwargs["audio_hidden_states"].shape[1] return torch.full((1, rows, 8), self.velocity), torch.zeros(1, audio_rows, 4) @@ -207,13 +211,24 @@ class _RecordingCurve: self.timesteps.append(unique.clone()) -def _denoise(monkeypatch, steps, velocity=0.5): +def _refine_args(vsa_builder=None): + from types import SimpleNamespace + + return SimpleNamespace(pipeline_config=SimpleNamespace(dit_config=SimpleNamespace(patch_size=(1, 2, 2))), + VSA_sparsity=0.8, + VSA_tile_size=64, + vsa_builder=vsa_builder) + + +def _denoise(monkeypatch, steps, velocity=0.5, vsa_builder=None): from fastvideo.pipelines.basic.minimax_h3.packing import build_packed_sequence from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MINIMAX_H3_LAYOUT_KEY from fastvideo.pipelines.basic.open_h3_regenerate_2k.stages import denoising from fastvideo.pipelines.pipeline_batch_info import ForwardBatch monkeypatch.setattr(denoising, "get_local_torch_device", lambda: torch.device("cpu")) + monkeypatch.setattr(denoising, "_h3_vsa_metadata_builder", lambda *args: vsa_builder) + monkeypatch.setattr(denoising, "_h3_vsa_single_region_segments", lambda layout, patch: ("segments", )) layout = build_packed_sequence(torch.ones(3, dtype=torch.long), 2, 2, 2, 2, (1, 2, 2)) batch = ForwardBatch(data_type="video") batch.num_inference_steps = steps @@ -223,7 +238,7 @@ def _denoise(monkeypatch, steps, velocity=0.5): batch.extra[MINIMAX_H3_LAYOUT_KEY] = layout transformer, curve = _ConstantVelocityTransformer(velocity), _RecordingCurve() stage = denoising.Regenerate2KDenoisingStage(transformer, curve, None) - return stage.forward(batch, None), transformer, curve + return stage.forward(batch, _refine_args(vsa_builder)), transformer, curve def test_euler_steps_integrate_the_constant_velocity_over_sigma_zero_to_start(monkeypatch): @@ -323,3 +338,27 @@ def test_saved_draft_latent_rejects_the_wrong_rank(tmp_path): torch.save(torch.zeros(24, 31, 48), path) with pytest.raises(ValueError, match=r"\[C, T, h, w\]"): _preparation_stage(tmp_path, str(path))._load_draft_latent(str(path), torch.device("cpu")) + + +class _RecordingVSABuilder: + + def __init__(self): + self.builds = [] + + def build(self, **kwargs): + self.builds.append(kwargs) + return f"metadata-{len(self.builds)}" + + +def test_vsa_backend_builds_one_sparse_metadata_per_refine_step(monkeypatch): + builder = _RecordingVSABuilder() + _, transformer, _ = _denoise(monkeypatch, steps=3, vsa_builder=builder) + assert [build["VSA_sparsity"] for build in builder.builds] == [0.8, 0.8, 0.8] + assert all(build["tile_size"] == 64 and build["exempt"] for build in builder.builds) + assert [build["current_timestep"] for build in builder.builds] == [0, 1, 2] + assert transformer.attn_metadata == ["metadata-1", "metadata-2", "metadata-3"] + + +def test_dense_backend_passes_no_attention_metadata(monkeypatch): + _, transformer, _ = _denoise(monkeypatch, steps=3) + assert transformer.attn_metadata == [None, None, None] diff --git a/scripts/regen2k_fasth3/score_all.sh b/scripts/regen2k_fasth3/score_all.sh new file mode 100755 index 000000000..47b5f857a --- /dev/null +++ b/scripts/regen2k_fasth3/score_all.sh @@ -0,0 +1,8 @@ +#!/bin/bash +# usage: score_all.sh DIR... +REF=/mnt/nfs/vlm-aryan/regen2k/out/parity_fox/ref_frames_u8.pt +source /mnt/nfs/vlm-aryan/regen2k-fast/scripts/common.sh +for d in "$@"; do + echo "== $d" + $PY $W/scripts/score_run.py --ref $REF --dir $d +done diff --git a/scripts/regen2k_fasth3/score_run.py b/scripts/regen2k_fasth3/score_run.py index 009f1c83c..01a8c10cb 100644 --- a/scripts/regen2k_fasth3/score_run.py +++ b/scripts/regen2k_fasth3/score_run.py @@ -16,7 +16,7 @@ import torch import torch.nn.functional as F -def gaussian_window(size: int = 11, sigma: float = 1.5, device: str = "cuda") -> torch.Tensor: +def gaussian_window(size: int = 11, sigma: float = 1.5, device: str = "cpu") -> torch.Tensor: coords = torch.arange(size, dtype=torch.float32, device=device) - size // 2 kernel = torch.exp(-(coords**2) / (2 * sigma**2)) kernel = kernel / kernel.sum() @@ -29,8 +29,8 @@ def psnr(x: torch.Tensor, y: torch.Tensor) -> float: def ssim(x: torch.Tensor, y: torch.Tensor, window: torch.Tensor) -> float: - x = x.float().mean(-1)[None, None].cuda() - y = y.float().mean(-1)[None, None].cuda() + x = x.float().mean(-1)[None, None] + y = y.float().mean(-1)[None, None] c1, c2 = (0.01 * 255)**2, (0.03 * 255)**2 mu_x, mu_y = F.conv2d(x, window), F.conv2d(y, window) sxx = F.conv2d(x * x, window) - mu_x * mu_x