daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Trace VAE temporal blend metadata

Browse source
This commit is contained in:
Daniel Maddern 2026-08-14 12:56:36 +07:00
parent 26fe7be966
commit 81714ed50f
1 changed files with 19 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
tools/compare_vae_temporal_assembly.py
View file

@ -24,6 +24,17 @@ def diff_stats(label: str, a: torch.Tensor, b: torch.Tensor) -> None:
print({"label": label, "shape": tuple(a.shape), "max": float(diff.max()), "mean": float(diff.mean())}, flush=True) 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): def plan(vae, z_len: int):
pseudo = z_len + vae.token_drop pseudo = z_len + vae.token_drop
pad = (-pseudo) % vae.tokens_chunk_size pad = (-pseudo) % vae.tokens_chunk_size
@ -70,6 +81,14 @@ with torch.inference_mode():
part_u = clip_u[:, :, frame_start:frame_end][:, :, upstream.frame_pre_padding:] part_u = clip_u[:, :, frame_start:frame_end][:, :, upstream.frame_pre_padding:]
if split == 0 and prev_d is not None: if split == 0 and prev_d is not None:
diff_stats(f"chunk_{chunk_index}_preblend_current", part_d, part_u) 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}_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)) 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_d = direct.blend(prev_d, part_d, direct.frame_overlap, -3)