Add H3 free-run parity comparator

src/h3_blackwell_runtime/packing.py

@@ -65,7 +65,14 @@ class H3PromptPacker:
         self.text_weight = checkpoint.tensor("condition_proj.weight", dtype=torch.bfloat16)
         self.text_bias = checkpoint.tensor("condition_proj.bias", dtype=torch.bfloat16)
 
-    def __call__(self, text: torch.Tensor, video: torch.Tensor, audio: torch.Tensor, sigma: float) -> tuple[torch.Tensor, torch.Tensor, list[tuple[int, int, int]], tuple[int, int, int], tuple[int, int, int]]:
+    def __call__(
+        self,
+        text: torch.Tensor,
+        video: torch.Tensor,
+        audio: torch.Tensor,
+        sigma: float,
+        model_timesteps: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, list[tuple[int, int, int]], tuple[int, int, int], tuple[int, int, int]]:
         if text.shape[-1] == 5120:
             text_rows = functional.linear(text[0].to(self.text_weight.dtype), self.text_weight, self.text_bias).to(torch.bfloat16)
         elif text.shape[-1] == 5376:
@@ -76,11 +83,18 @@ class H3PromptPacker:
         audio_rows = functional.linear(pack_audio(audio.to(torch.bfloat16)).float(), self.audio_weight, self.audio_bias).to(torch.bfloat16)
         text_length, audio_length = text_rows.shape[0], audio_rows.shape[0]
         hidden = torch.cat((text_rows, audio_rows, video_rows))
+        if model_timesteps is None:
             video_sigma = torch.tensor(float(sigma), device=hidden.device).clamp(min=1e-6)
             base = video_sigma / (12.0 + video_sigma * (1.0 - 12.0))
             audio_sigma = 3.0 * base / (1.0 + (3.0 - 1.0) * base)
             video_time, audio_time = (1.0 - video_sigma).item(), (1.0 - audio_sigma).item()
             unique_times = sorted({video_time, audio_time})
+        else:
+            times_override = model_timesteps.to(device=hidden.device, dtype=torch.float32).flatten()
+            if times_override.numel() not in (1, 2):
+                raise ValueError("Prompt-only H3 expects one or two model timesteps.")
+            unique_times = times_override.tolist()
+            video_time, audio_time = unique_times[0], unique_times[-1]
         row = {value: index for index, value in enumerate(unique_times)}
         video_row, audio_row = row[video_time] * 3, row[audio_time] * 3
         times = torch.tensor(unique_times, device=hidden.device, dtype=torch.float32)

src/h3_blackwell_runtime/sampler.py

@@ -48,7 +48,16 @@ def _audio_sigma(video_sigma: torch.Tensor) -> torch.Tensor:
 
 
 @torch.inference_mode()
-def sample_video_res_multistep(model, packer: H3PromptPacker, text: torch.Tensor, video: torch.Tensor, audio: torch.Tensor, *, steps: int = 12) -> torch.Tensor:
+def sample_video_res_multistep(
+    model,
+    packer: H3PromptPacker,
+    text: torch.Tensor,
+    video: torch.Tensor,
+    audio: torch.Tensor,
+    *,
+    steps: int = 12,
+    model_timesteps: list[torch.Tensor] | tuple[torch.Tensor, ...] | None = None,
+) -> torch.Tensor:
     """Direct H3 beta/RES sampling with Comfy-equivalent joint AV carry semantics."""
     sigmas = beta_sigmas(steps, device=video.device)
     audio_carried = audio
@@ -63,7 +72,8 @@ def sample_video_res_multistep(model, packer: H3PromptPacker, text: torch.Tensor
         sigma_audio = _audio_sigma(sigma)
         carry = sigma_audio / sigma
         native_audio = audio_carried.to(torch.bfloat16) * carry
-        hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, float(sigma))
+        step_timesteps = None if model_timesteps is None else model_timesteps[previous_index]
+        hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, float(sigma), step_timesteps)
         raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment)
         raw_video = raw_video.to(torch.bfloat16).float()
         raw_audio = raw_audio.to(torch.bfloat16)

tools/compare_fl2va_free_run.py

@@ -0,0 +1,90 @@
+"""Compare direct free-run FL2VA sampling against captured Comfy sampler state."""
+
+import argparse
+
+import torch
+
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
+from h3_blackwell_runtime.denoiser import H3PackedDenoiser
+from h3_blackwell_runtime.packing import H3PromptPacker, unpatchify_video
+from h3_blackwell_runtime.sampler import _audio_sigma, _unpack_audio, res_multistep_update
+from h3_blackwell_runtime.t2v import random_av_latents
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--root", default="/artifacts/fl2va-sampler-reference")
+parser.add_argument("--capture", default="/artifacts/capture")
+parser.add_argument("--width", type=int, default=320)
+parser.add_argument("--height", type=int, default=192)
+parser.add_argument("--frames", type=int, default=22)
+parser.add_argument("--seed", type=int, default=440204)
+parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default="sage2")
+parser.add_argument("--oracle-timesteps", action="store_true")
+args = parser.parse_args()
+
+checkpoint = H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
+model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend=args.attention).eval()
+packer = H3PromptPacker(checkpoint)
+text = torch.load(f"{args.capture}/input_00.pt", map_location="cuda", weights_only=False)["hidden"][:17].unsqueeze(0).to("cuda")
+video, audio_carried, _ = random_av_latents(args.width, args.height, args.frames, args.seed)
+sigmas = torch.load(f"{args.root}/initial.pt", map_location="cuda", weights_only=False)["sigmas"].to("cuda")
+
+video_shape = video.shape
+audio_shape = audio_carried.shape
+video_count = video.numel()
+audio_count = audio_carried.numel()
+video_history = audio_history = history_sigma = None
+
+for index, sigma in enumerate(sigmas[:-1]):
+    reference = torch.load(f"{args.root}/step_{index:02d}.pt", map_location="cuda", weights_only=False)
+    h3_input = torch.load(f"{args.capture}/input_{index:02d}.pt", map_location="cuda", weights_only=False)
+    reference_x = reference["x"].to("cuda").reshape(-1)
+    reference_video = reference_x[:video_count].reshape(video_shape)
+    reference_audio = reference_x[video_count:video_count + audio_count].reshape(audio_shape)
+    pre_video = (video.float() - reference_video.float()).abs()
+    pre_audio = (audio_carried.float() - reference_audio.float()).abs()
+
+    sigma_audio = _audio_sigma(sigma)
+    carry = sigma_audio / sigma
+    model_timesteps = h3_input["timesteps"] if args.oracle_timesteps else None
+    hidden, times, segments, positions, video_segment, audio_segment = packer(
+        text,
+        video,
+        audio_carried.to(torch.bfloat16) * carry,
+        float(sigma),
+        model_timesteps,
+    )
+    hidden_delta = (hidden.float() - h3_input["hidden"].to("cuda").float()).abs()
+    time_delta = (times.float() - h3_input["timesteps"].to("cuda").float()).abs()
+
+    with torch.inference_mode():
+        raw_video, raw_audio = model(hidden, times, positions, segments, video_segment, audio_segment)
+    raw_video = raw_video.to(torch.bfloat16).float()
+    raw_audio = raw_audio.to(torch.bfloat16)
+    video_denoised = video + sigma * unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1])
+    audio_model_output = (
+        (1.0 - 4.0) * (audio_carried.to(torch.bfloat16) * carry.to(torch.bfloat16))
+        + (1.0 + 3.0 * sigma_audio).to(torch.bfloat16) * (-_unpack_audio(raw_audio))
+    ).float()
+    audio_denoised = audio_carried - sigma * audio_model_output
+
+    reference_denoised = reference["denoised"].to("cuda").reshape(-1)
+    reference_video_denoised = reference_denoised[:video_count].reshape(video_shape)
+    reference_audio_denoised = reference_denoised[video_count:video_count + audio_count].reshape(audio_shape)
+    denoised_video = (video_denoised.float() - reference_video_denoised.float()).abs()
+    denoised_audio = (audio_denoised.float() - reference_audio_denoised.float()).abs()
+
+    print(
+        f"step={index:02d} "
+        f"pre_max={max(pre_video.max().item(), pre_audio.max().item()):.6g} "
+        f"hidden_max={hidden_delta.max().item():.6g} "
+        f"time_max={time_delta.max().item():.6g} "
+        f"denoised_max={max(denoised_video.max().item(), denoised_audio.max().item()):.6g} "
+        f"denoised_mean={max(denoised_video.mean().item(), denoised_audio.mean().item()):.6g}"
+    )
+
+    previous_sigma = sigmas[index - 1] if index else None
+    sigma_down = sigmas[index + 1]
+    video = res_multistep_update(video, video_denoised, sigma, sigma_down, video_history, history_sigma, previous_sigma)
+    audio_carried = res_multistep_update(audio_carried, audio_denoised, sigma, sigma_down, audio_history, history_sigma, previous_sigma)
+    video_history, audio_history, history_sigma = video_denoised, audio_denoised, sigma_down