From f5fe966a7e3d37121145ed41375013644c64ac96 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Thu, 13 Aug 2026 21:11:41 +0700 Subject: [PATCH] Match Comfy audio sampler boundary --- src/h3_blackwell_runtime/sampler.py | 8 +++++--- tools/compare_fl2va_steps.py | 26 +++++++++++++++++++++----- 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/src/h3_blackwell_runtime/sampler.py b/src/h3_blackwell_runtime/sampler.py index 1079be5..6a73dfc 100644 --- a/src/h3_blackwell_runtime/sampler.py +++ b/src/h3_blackwell_runtime/sampler.py @@ -65,11 +65,13 @@ def sample_video_res_multistep(model, packer: H3PromptPacker, text: torch.Tensor hidden, times, segments, positions, video_segment, audio_segment = packer(text, video, native_audio, float(sigma)) 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).float() + raw_audio = raw_audio.to(torch.bfloat16) velocity_video = -unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1]) - velocity_audio_native = -_unpack_audio(raw_audio) carry = sigma_audio / sigma - velocity_audio = (1.0 - 4.0) * (audio_carried * carry) + (1.0 + 3.0 * sigma_audio) * velocity_audio_native + velocity_audio = ( + (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() video_denoised = video - sigma * velocity_video audio_denoised = audio_carried - sigma * velocity_audio previous_sigma = sigmas[previous_index - 1] if previous_index else None diff --git a/tools/compare_fl2va_steps.py b/tools/compare_fl2va_steps.py index ae5914a..4c62f28 100644 --- a/tools/compare_fl2va_steps.py +++ b/tools/compare_fl2va_steps.py @@ -7,7 +7,7 @@ import torch from h3_blackwell_runtime.checkpoint import H3Checkpoint from h3_blackwell_runtime.denoiser import H3PackedDenoiser from h3_blackwell_runtime.packing import unpatchify_video -from h3_blackwell_runtime.sampler import res_multistep_update +from h3_blackwell_runtime.sampler import _audio_sigma, res_multistep_update root = "/artifacts/fl2va-sampler-reference" @@ -19,8 +19,10 @@ model = H3PackedDenoiser.from_checkpoint( H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors"), attention_backend="sage2" ).eval() video_shape = (1, 24, 7, 12, 20) +audio_shape = (1, 32, 2, 37) video_count = torch.tensor(video_shape).prod().item() -old_video = old_sigma = None +audio_count = torch.tensor(audio_shape).prod().item() +old_video = old_audio = old_sigma = None for index, reference in enumerate(steps): h3_input = torch.load(f"{capture}/input_{index:02d}.pt", map_location="cuda", weights_only=False) @@ -38,23 +40,37 @@ for index, reference in enumerate(steps): h3_delta = (video_out.float() - h3_output["video"].to("cuda").float()).abs() state_video = reference["x"].to("cuda").reshape(-1)[:video_count].reshape(video_shape) + state_audio = reference["x"].to("cuda").reshape(-1)[video_count:video_count + audio_count].reshape(audio_shape) reference_denoised = reference["denoised"].to("cuda").reshape(-1)[:video_count].reshape(video_shape) + reference_audio_denoised = reference["denoised"].to("cuda").reshape(-1)[video_count:video_count + audio_count].reshape(audio_shape) converted_denoised = state_video + sigmas[index] * h3_output["video"].to("cuda").to(torch.bfloat16).float() denoised_delta = (converted_denoised.float() - reference_denoised.float()).abs() + sigma_audio = _audio_sigma(sigmas[index]) + carry = sigma_audio / sigmas[index] + audio_model_output = ( + (1.0 - 4.0) * (state_audio.to(torch.bfloat16) * carry.to(torch.bfloat16)) + + (1.0 + 3.0 * sigma_audio).to(torch.bfloat16) * (-h3_output["audio"].to("cuda").to(torch.bfloat16)) + ).float() + converted_audio_denoised = state_audio - sigmas[index] * audio_model_output + audio_denoised_delta = (converted_audio_denoised.float() - reference_audio_denoised.float()).abs() if index + 1 < len(steps): previous_sigma = sigmas[index - 1] if index else None updated = res_multistep_update(state_video, reference_denoised, sigmas[index], sigmas[index + 1], old_video, old_sigma, previous_sigma) + updated_audio = res_multistep_update(state_audio, reference_audio_denoised, sigmas[index], sigmas[index + 1], old_audio, old_sigma, previous_sigma) next_video = steps[index + 1]["x"].to("cuda").reshape(-1)[:video_count].reshape(video_shape) + next_audio = steps[index + 1]["x"].to("cuda").reshape(-1)[video_count:video_count + audio_count].reshape(audio_shape) update_delta = (updated.float() - next_video.float()).abs() - update_text = f"update_mean={update_delta.mean().item():.6g} update_max={update_delta.max().item():.6g}" + update_audio_delta = (updated_audio.float() - next_audio.float()).abs() + update_text = f"update_mean={update_delta.mean().item():.6g} update_max={update_delta.max().item():.6g} audio_update_mean={update_audio_delta.mean().item():.6g} audio_update_max={update_audio_delta.max().item():.6g}" else: - update_text = "update_mean=final-unobserved update_max=final-unobserved" + update_text = "update_mean=final-unobserved update_max=final-unobserved audio_update_mean=final-unobserved audio_update_max=final-unobserved" print( f"step={index:02d} " f"h3_mean={h3_delta.mean().item():.6g} h3_max={h3_delta.max().item():.6g} " f"denoised_mean={denoised_delta.mean().item():.6g} denoised_max={denoised_delta.max().item():.6g} " + f"audio_denoised_mean={audio_denoised_delta.mean().item():.6g} audio_denoised_max={audio_denoised_delta.max().item():.6g} " f"{update_text}" ) - old_video, old_sigma = reference_denoised, sigmas[index + 1] + old_video, old_audio, old_sigma = reference_denoised, reference_audio_denoised, sigmas[index + 1]