128 lines
7.6 KiB
Python
128 lines
7.6 KiB
Python
"""Replay a matched Comfy keyframe sampler capture with exact static conditioning."""
|
|
|
|
import argparse
|
|
from pathlib import Path
|
|
|
|
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, _model_sigma, _unpack_audio, beta_sigmas, res_multistep_update
|
|
from h3_blackwell_runtime.t2v import random_av_latents
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--sampler", type=Path, required=True)
|
|
parser.add_argument("--dit", type=Path, required=True)
|
|
parser.add_argument("--model", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
|
parser.add_argument("--result-latent", type=Path)
|
|
parser.add_argument("--production-trace", type=Path)
|
|
args = parser.parse_args()
|
|
|
|
checkpoint = H3Checkpoint(args.model)
|
|
model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend="sage2").eval()
|
|
packer = H3PromptPacker(checkpoint)
|
|
captured_input = torch.load(args.dit / "input.pt", map_location="cuda", weights_only=False)
|
|
text_length = next(start for start, _, code in captured_input["segments"] if code == 3)
|
|
prefix_stop = next(start for start, _, code in captured_input["segments"] if code % 3 == 2)
|
|
text = captured_input["hidden"][:text_length].unsqueeze(0).cuda()
|
|
prefix = captured_input["hidden"][:prefix_stop].cuda()
|
|
text_tags = torch.ones(text_length, dtype=torch.long, device="cuda")
|
|
for start, stop, code in captured_input["segments"]:
|
|
if stop <= text_length:
|
|
text_tags[start:stop] = code % 3
|
|
|
|
video, audio, frame_count = random_av_latents(384, 384, 22, 440207)
|
|
video_shape, audio_shape = video.shape, audio.shape
|
|
video_count, audio_count = video.numel(), audio.numel()
|
|
zero_cond = [torch.zeros(1, 24, 1, 24, 24, device="cuda") for _ in range(2)]
|
|
sigmas = torch.load(args.sampler / "initial.pt", map_location="cuda", weights_only=False)["sigmas"].cuda()
|
|
direct_sigmas = beta_sigmas(12, device="cuda")
|
|
sigma_delta = (direct_sigmas.float() - sigmas.float()).abs()
|
|
print({"stage": "sigmas", "direct": direct_sigmas.tolist(), "comfy": sigmas.tolist(), "mean_delta": float(sigma_delta.mean()), "max_delta": float(sigma_delta.max())}, flush=True)
|
|
video_history = audio_history = history_sigma = None
|
|
|
|
for index, sigma in enumerate(sigmas[:-1]):
|
|
reference = torch.load(args.sampler / f"step_{index:02d}.pt", map_location="cuda", weights_only=False)
|
|
reference_x = reference["x"].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.float() - reference_audio.float()).abs()
|
|
|
|
sigma_audio = _audio_sigma(sigma)
|
|
carry = sigma_audio / sigma
|
|
hidden, times, segments, positions, video_segment, audio_segment = packer(
|
|
text,
|
|
video,
|
|
audio.to(torch.bfloat16) * carry,
|
|
_model_sigma(sigma),
|
|
text_token_tags=text_tags,
|
|
cond_latents=zero_cond,
|
|
cond_frame_indices=[0, frame_count - 1],
|
|
frame_count=frame_count,
|
|
seed=440207,
|
|
)
|
|
hidden[:prefix_stop] = prefix.to(hidden)
|
|
input_hidden = hidden.detach().clone()
|
|
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.to(torch.bfloat16) * carry.to(torch.bfloat16))
|
|
+ (1.0 + 3.0 * sigma_audio).to(torch.bfloat16) * (-_unpack_audio(raw_audio))
|
|
).float()
|
|
audio_denoised = audio - sigma * audio_model_output
|
|
|
|
reference_denoised = reference["denoised"].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()
|
|
production = torch.load(args.production_trace / f"step_{index:02d}.pt", map_location="cuda", weights_only=False) if args.production_trace else None
|
|
production_video = (production["video"].float() - reference_video.float()).abs() if production else None
|
|
production_audio = (production["audio"].float() - reference_audio.float()).abs() if production else None
|
|
production_denoised = (production["video_denoised"].float() - reference_video_denoised.float()).abs() if production else None
|
|
production_hidden = (production["hidden"].float() - input_hidden.float()).abs() if production and "hidden" in production else None
|
|
production_raw = (production["raw_video"].float() - raw_video.float()).abs() if production and "raw_video" in production else None
|
|
production_segments = [
|
|
(start, stop, code, float((production["hidden"][start:stop].float() - input_hidden[start:stop].float()).abs().mean()))
|
|
for start, stop, code in segments
|
|
] if production and "hidden" in production else None
|
|
print({
|
|
"step": index,
|
|
"pre_video_mean": float(pre_video.mean()),
|
|
"pre_video_max": float(pre_video.max()),
|
|
"pre_audio_mean": float(pre_audio.mean()),
|
|
"pre_audio_max": float(pre_audio.max()),
|
|
"denoised_video_mean": float(denoised_video.mean()),
|
|
"denoised_video_max": float(denoised_video.max()),
|
|
"denoised_audio_mean": float(denoised_audio.mean()),
|
|
"denoised_audio_max": float(denoised_audio.max()),
|
|
"production_video_mean": float(production_video.mean()) if production_video is not None else None,
|
|
"production_video_max": float(production_video.max()) if production_video is not None else None,
|
|
"production_audio_mean": float(production_audio.mean()) if production_audio is not None else None,
|
|
"production_audio_max": float(production_audio.max()) if production_audio is not None else None,
|
|
"production_denoised_mean": float(production_denoised.mean()) if production_denoised is not None else None,
|
|
"production_denoised_max": float(production_denoised.max()) if production_denoised is not None else None,
|
|
"production_hidden_mean": float(production_hidden.mean()) if production_hidden is not None else None,
|
|
"production_hidden_max": float(production_hidden.max()) if production_hidden is not None else None,
|
|
"production_raw_mean": float(production_raw.mean()) if production_raw is not None else None,
|
|
"production_raw_max": float(production_raw.max()) if production_raw is not None else None,
|
|
"production_segments": production_segments,
|
|
}, flush=True)
|
|
|
|
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 = res_multistep_update(audio, audio_denoised, sigma, sigma_down, audio_history, history_sigma, previous_sigma)
|
|
video_history, audio_history, history_sigma = video_denoised, audio_denoised, sigma_down
|
|
|
|
if args.result_latent:
|
|
result = torch.load(args.result_latent, map_location="cuda", weights_only=False)
|
|
result_video = result["latent"] if isinstance(result, dict) else result
|
|
delta = (video.float() - result_video.cuda().float()).abs()
|
|
print({"stage": "final_video", "mean_delta": float(delta.mean()), "max_delta": float(delta.max())}, flush=True)
|