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

Match upstream VAE blend implementation

Browse source
This commit is contained in:
Daniel Maddern 2026-08-14 01:41:53 +07:00
parent 4e17e866a0
commit 28a4121df9
1 changed files with 18 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
src/h3_blackwell_runtime/vae_decoder.py
View file

@ -236,12 +236,27 @@ class MiniMaxH3VideoVAE(nn.Module):
@staticmethod
def blend(a: torch.Tensor, b: torch.Tensor, extent: int, dim: int) -> torch.Tensor:
extent = min(a.shape[dim], b.shape[dim], extent)
positions = torch.arange(extent, device=b.device, dtype=b.dtype)
weight_a = 1 - positions / extent
weight_b = positions / extent
shape = [1] * a.ndim
shape[dim] = extent
position = torch.arange(extent, device=b.device, dtype=b.dtype).view(shape)
blended = a.narrow(dim, a.shape[dim] - extent, extent) * (1 - position / extent) + b.narrow(dim, 0, extent) * (position / extent)
weight_a = weight_a.view(shape)
weight_b = weight_b.view(shape)
slice_a = [slice(None)] * a.ndim
slice_a[dim] = slice(-extent, None)
slice_b = [slice(None)] * b.ndim
slice_b[dim] = slice(0, extent)
blended = a[tuple(slice_a)] * weight_a + b[tuple(slice_b)] * weight_b
if extent < b.shape[dim]:
return torch.cat((blended, b.narrow(dim, extent, b.shape[dim] - extent)), dim=dim)
slice_b_rest = [slice(None)] * b.ndim
slice_b_rest[dim] = slice(extent, None)
return torch.cat((blended, b[tuple(slice_b_rest)]), dim=dim)
return blended
def tiled_decode(self, z: torch.Tensor) -> torch.Tensor: