diff --git a/src/h3_blackwell_runtime/packing.py b/src/h3_blackwell_runtime/packing.py index 5221154..802664c 100644 --- a/src/h3_blackwell_runtime/packing.py +++ b/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)) - 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}) + 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) diff --git a/src/h3_blackwell_runtime/sampler.py b/src/h3_blackwell_runtime/sampler.py index 598656b..d07ff10 100644 --- a/src/h3_blackwell_runtime/sampler.py +++ b/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) diff --git a/tools/compare_fl2va_free_run.py b/tools/compare_fl2va_free_run.py new file mode 100644 index 0000000..bd82de5 --- /dev/null +++ b/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