44 lines
2.9 KiB
Python
44 lines
2.9 KiB
Python
"""Compare direct FL2VA beta/RES steps with captured Comfy sampler state."""
|
|
|
|
import glob
|
|
|
|
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.qwen3vl_text import Qwen3VLPromptConditioner
|
|
from h3_blackwell_runtime.sampler import _audio_sigma, _unpack_audio, res_multistep_update
|
|
from h3_blackwell_runtime.token_refiner import H3TokenRefiner
|
|
|
|
|
|
root = "/artifacts/fl2va-sampler-reference"
|
|
initial = torch.load(f"{root}/initial.pt", map_location="cuda", weights_only=False)
|
|
steps = [torch.load(path, map_location="cuda", weights_only=False) for path in sorted(glob.glob(f"{root}/step_*.pt"))]
|
|
sigmas = initial["sigmas"].to("cuda")
|
|
checkpoint = H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
|
text = H3TokenRefiner(checkpoint)(Qwen3VLPromptConditioner("/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors", "/opt/h3-blackwell-runtime/qwen25_tokenizer")("A brass-and-paper dragon flies above a rain-washed old city at blue hour."))
|
|
model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend="sage2").eval()
|
|
packer = H3PromptPacker(checkpoint)
|
|
|
|
video, audio_carried = initial["initial_x"]
|
|
video, audio_carried = video.to("cuda"), audio_carried.to("cuda")
|
|
old_video = old_audio = old_sigma = None
|
|
for index, reference in enumerate(steps):
|
|
sigma, sigma_down = sigmas[index], sigmas[index + 1]
|
|
native_audio = audio_carried * (_audio_sigma(sigma) / 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)
|
|
velocity_video = -unpatchify_video(raw_video, video.shape[2], video.shape[-2], video.shape[-1])
|
|
velocity_audio = -_unpack_audio(raw_audio)
|
|
carry = _audio_sigma(sigma) / sigma
|
|
velocity_audio = (1.0 - 4.0) * (audio_carried * carry) + (1.0 + 3.0 * _audio_sigma(sigma)) * velocity_audio
|
|
denoised = (video - sigma * velocity_video, audio_carried - sigma * velocity_audio)
|
|
reference_denoised = reference["denoised"]
|
|
denoised_delta = (denoised[0].float() - reference_denoised[0].float()).abs()
|
|
previous_sigma = sigmas[index - 1] if index else None
|
|
video = res_multistep_update(video, denoised[0], sigma, sigma_down, old_video, old_sigma, previous_sigma)
|
|
audio_carried = res_multistep_update(audio_carried, denoised[1], sigma, sigma_down, old_audio, old_sigma, previous_sigma)
|
|
latent_delta = (video.float() - reference["x"][0].float()).abs()
|
|
print(f"step={index:02d} x0_video_mean={denoised_delta.mean().item():.6g} x0_video_max={denoised_delta.max().item():.6g} latent_video_mean={latent_delta.mean().item():.6g} latent_video_max={latent_delta.max().item():.6g}")
|
|
old_video, old_audio, old_sigma = denoised[0], denoised[1], sigma_down
|