diff --git a/tools/compare_vae_temporal_assembly.py b/tools/compare_vae_temporal_assembly.py index b655cc4..f51fbb8 100644 --- a/tools/compare_vae_temporal_assembly.py +++ b/tools/compare_vae_temporal_assembly.py @@ -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) +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 @@ -70,6 +81,14 @@ with torch.inference_mode(): 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)