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

Compare VAE decoder projection output

Browse source
This commit is contained in:
Daniel Maddern 2026-08-14 00:34:36 +07:00
parent b55887808d
commit a262d44632
1 changed files with 15 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
tools/compare_vae_decoder_clip.py
View file

@ -62,5 +62,18 @@ with torch.inference_mode():
h_d = block_d(h_d, rope_d)
h_u = block_u(h_u, rope_u)
stats(f"block_{index:02d}", h_d, h_u)
if index >= 5 and (h_d.float() - h_u.float()).abs().mean() > 1e-3:
break
patches = h_d.shape[1] - 1 - dd.num_register_tokens
out_d = dd.proj_out(dd.norm_out(h_d))[:, :patches]
out_u = du.proj_out(du.norm_out(h_u))[:, :patches]
stats("proj_out", out_d, out_u)
out_d = out_d.view(b, latent_t, latent_h, latent_w, dd.out_channels, dd.patch_size_t, dd.patch_size, dd.patch_size)
out_u = out_u.view(b, latent_t, latent_h, latent_w, du.out_channels, du.patch_size_t, du.patch_size, du.patch_size)
out_d = out_d.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape(b, dd.out_channels, latent_t * dd.patch_size_t, latent_h * dd.patch_size, latent_w * dd.patch_size)
out_u = out_u.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape(b, du.out_channels, latent_t * du.patch_size_t, latent_h * du.patch_size, latent_w * du.patch_size)
stats("decoded", out_d, out_u)
pix_d = out_d.float().mul_(direct.pixel_std.to(out_d)).add_(direct.pixel_mean.to(out_d)).clamp_(0, 1).mul_(2).sub_(1)
pix_u = out_u.float().mul_(upstream.pixel_std.to(out_u)).add_(upstream.pixel_mean.to(out_u)).clamp_(0, 1).mul_(2).sub_(1)
stats("pixels", pix_d, pix_u)