103 lines
5.2 KiB
Python
103 lines
5.2 KiB
Python
"""Trace direct/upstream temporal VAE assembly from matching tiled clips."""
|
|
|
|
import argparse
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
from safetensors.torch import load_file
|
|
|
|
from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE as DirectVAE
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--latent", type=Path, required=True)
|
|
parser.add_argument("--comfy-path", default="/opt/ComfyUI")
|
|
args = parser.parse_args()
|
|
|
|
sys.path.insert(0, args.comfy_path)
|
|
from h3_blackwell_runtime.upstream_vae import MiniMaxH3VideoVAE as UpstreamVAE # noqa: E402
|
|
|
|
|
|
def diff_stats(label: str, a: torch.Tensor, b: torch.Tensor) -> None:
|
|
diff = (a.float() - b.float()).abs()
|
|
print({"label": label, "shape": tuple(a.shape), "max": float(diff.max()), "mean": float(diff.mean())}, flush=True)
|
|
|
|
|
|
def tensor_meta(label: str, tensor: torch.Tensor) -> None:
|
|
print({
|
|
"label": label,
|
|
"shape": tuple(tensor.shape),
|
|
"stride": tuple(tensor.stride()),
|
|
"dtype": str(tensor.dtype),
|
|
"contiguous": tensor.is_contiguous(),
|
|
"storage_offset": tensor.storage_offset(),
|
|
}, flush=True)
|
|
|
|
|
|
def plan(vae, z_len: int):
|
|
pseudo = z_len + vae.token_drop
|
|
pad = (-pseudo) % vae.tokens_chunk_size
|
|
pseudo += pad
|
|
chunks = pseudo // vae.tokens_chunk_size - int(vae.token_drop > 0)
|
|
if chunks < 1:
|
|
pad += vae.tokens_chunk_size
|
|
chunks += 1
|
|
return chunks, pad, vae._decode_temporal_frame_plan(z_len + pad, chunks, pad)
|
|
|
|
|
|
state = torch.load(args.latent, map_location="cuda", weights_only=False)
|
|
latent = state["latent"].to("cuda") if isinstance(state, dict) else state.to("cuda")
|
|
|
|
direct = DirectVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda", tiling=True).eval()
|
|
upstream = UpstreamVAE(tiling=True).to("cuda").eval()
|
|
upstream.load_state_dict(load_file("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda"), strict=True)
|
|
|
|
with torch.inference_mode():
|
|
zd = latent.to(next(direct.parameters()).dtype)
|
|
zu = latent.to(next(upstream.parameters()).dtype)
|
|
zd = zd * direct.latents_std.view(1, -1, 1, 1, 1).to(zd) + direct.latents_mean.view(1, -1, 1, 1, 1).to(zd)
|
|
zu = zu * upstream.latents_std.view(1, -1, 1, 1, 1).to(zu) + upstream.latents_mean.view(1, -1, 1, 1, 1).to(zu)
|
|
chunks, pad, frames = plan(direct, zd.shape[2])
|
|
if pad:
|
|
zd = torch.cat((zd, zd[:, :, -1:].repeat(1, 1, pad, 1, 1)), dim=2)
|
|
zu = torch.cat((zu, zu[:, :, -1:].repeat(1, 1, pad, 1, 1)), dim=2)
|
|
print({"chunks": chunks, "pad": pad, "frames": frames}, flush=True)
|
|
|
|
chunk_dec = direct.tokens_chunk_size * direct.vae_ratio_t
|
|
split_count = int(direct.token_drop > 0) + 1
|
|
prev_d = prev_u = None
|
|
write_pos = 0
|
|
for chunk_index in range(chunks):
|
|
start = chunk_index * direct.tokens_chunk_size
|
|
end = start + direct.tokens_chunk_size + direct.token_overlap
|
|
clip_d = direct.tiled_decode(zd[:, :, start:end])
|
|
clip_u = upstream.tiled_decode(zu[:, :, start:end])
|
|
diff_stats(f"clip_{chunk_index}", clip_d, clip_u)
|
|
for split in range(split_count):
|
|
frame_start = split * chunk_dec
|
|
frame_end = min(frame_start + chunk_dec, clip_d.shape[2])
|
|
part_d = clip_d[:, :, frame_start:frame_end][:, :, direct.frame_pre_padding:]
|
|
part_u = clip_u[:, :, frame_start:frame_end][:, :, upstream.frame_pre_padding:]
|
|
if split == 0 and prev_d is not None:
|
|
diff_stats(f"chunk_{chunk_index}_preblend_current", part_d, part_u)
|
|
if chunk_index == 1:
|
|
tensor_meta("prev_d", prev_d)
|
|
tensor_meta("prev_u", prev_u)
|
|
tensor_meta("part_d", part_d)
|
|
tensor_meta("part_u", part_u)
|
|
diff_stats("chunk_1_prev", prev_d, prev_u)
|
|
diff_stats("chunk_1_direct_blend_contiguous", direct.blend(prev_d.contiguous(), part_d.contiguous(), direct.frame_overlap, -3), upstream.blend(prev_d.contiguous(), part_d.contiguous(), direct.frame_overlap, -3))
|
|
diff_stats("chunk_1_cross_respective_contiguous", direct.blend(prev_d.contiguous(), part_d.contiguous(), direct.frame_overlap, -3), upstream.blend(prev_u.contiguous(), part_u.contiguous(), upstream.frame_overlap, -3))
|
|
diff_stats(f"chunk_{chunk_index}_direct_blend_same_inputs", direct.blend(prev_d, part_d, direct.frame_overlap, -3), upstream.blend(prev_d, part_d, direct.frame_overlap, -3))
|
|
diff_stats(f"chunk_{chunk_index}_upstream_blend_same_inputs", direct.blend(prev_u, part_u, upstream.frame_overlap, -3), upstream.blend(prev_u, part_u, upstream.frame_overlap, -3))
|
|
part_d = direct.blend(prev_d, part_d, direct.frame_overlap, -3)
|
|
part_u = upstream.blend(prev_u, part_u, upstream.frame_overlap, -3)
|
|
diff_stats(f"chunk_{chunk_index}_blended_at_{write_pos}", part_d, part_u)
|
|
elif split == 0:
|
|
diff_stats(f"chunk_{chunk_index}_write_at_{write_pos}", part_d, part_u)
|
|
else:
|
|
prev_d, prev_u = part_d.contiguous(), part_u.contiguous()
|
|
diff_stats(f"chunk_{chunk_index}_overlap", prev_d, prev_u)
|
|
continue
|
|
write_pos += part_d.shape[2]
|