From a2c49ad93b8717999f67ce0d8dd8b50605685f0f Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Thu, 13 Aug 2026 16:28:51 +0700 Subject: [PATCH] Match Comfy sampler BF16 output boundary --- src/h3_blackwell_runtime/sampler.py | 2 ++ tools/compare_fl2va_steps.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/h3_blackwell_runtime/sampler.py b/src/h3_blackwell_runtime/sampler.py index 1a3755e..1079be5 100644 --- a/src/h3_blackwell_runtime/sampler.py +++ b/src/h3_blackwell_runtime/sampler.py @@ -64,6 +64,8 @@ def sample_video_res_multistep(model, packer: H3PromptPacker, text: torch.Tensor native_audio = audio_carried * (sigma_audio / sigma) 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() 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 diff --git a/tools/compare_fl2va_steps.py b/tools/compare_fl2va_steps.py index cfd461d..ae5914a 100644 --- a/tools/compare_fl2va_steps.py +++ b/tools/compare_fl2va_steps.py @@ -39,7 +39,7 @@ for index, reference in enumerate(steps): state_video = reference["x"].to("cuda").reshape(-1)[:video_count].reshape(video_shape) reference_denoised = reference["denoised"].to("cuda").reshape(-1)[:video_count].reshape(video_shape) - converted_denoised = state_video + sigmas[index] * h3_output["video"].to("cuda") + converted_denoised = state_video + sigmas[index] * h3_output["video"].to("cuda").to(torch.bfloat16).float() denoised_delta = (converted_denoised.float() - reference_denoised.float()).abs() if index + 1 < len(steps):