91 lines
4.8 KiB
Python
91 lines
4.8 KiB
Python
"""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.attention import AVAILABLE_BACKENDS
|
|
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, 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=AVAILABLE_BACKENDS, 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,
|
|
_model_sigma(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
|