From c38171ddb93368ee6a6bbc677b92e4b50cead865 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jukka=20Sepp=C3=A4nen?= <40791699+kijai@users.noreply.github.com> Date: Wed, 29 Jul 2026 01:25:55 +0300 Subject: [PATCH] Support Pruna LTX VAE (#15129) --- comfy/ldm/lightricks/vae/causal_conv3d.py | 10 ++++++++-- comfy/ldm/lightricks/vae/causal_video_autoencoder.py | 6 +++--- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/comfy/ldm/lightricks/vae/causal_conv3d.py b/comfy/ldm/lightricks/vae/causal_conv3d.py index 7515f0d4e..bb1803f12 100644 --- a/comfy/ldm/lightricks/vae/causal_conv3d.py +++ b/comfy/ldm/lightricks/vae/causal_conv3d.py @@ -49,6 +49,12 @@ class CausalConv3d(nn.Module): ) self.temporal_cache_state={} + def _empty_output(self, x): + # empty (0 frame) outputs must still have the conv's output channels and spatial dims + h = (x.shape[3] + 2 * self.conv.padding[1] - self.conv.kernel_size[1]) // self.conv.stride[1] + 1 + w = (x.shape[4] + 2 * self.conv.padding[2] - self.conv.kernel_size[2]) // self.conv.stride[2] + 1 + return x.new_empty((x.shape[0], self.out_channels, 0, h, w)) + def forward(self, x, causal: bool = True): tid = threading.get_ident() @@ -58,7 +64,7 @@ class CausalConv3d(nn.Module): if not causal: padding_length = padding_length // 2 if x.shape[2] == 0: - return x + return self._empty_output(x) cached = x[:, :, :1, :, :].repeat((1, 1, padding_length, 1, 1)) pieces = [ cached, x ] if is_end and not causal: @@ -83,7 +89,7 @@ class CausalConv3d(nn.Module): elif is_end: self.temporal_cache_state[tid] = (None, True) - return self.conv(x) if x.shape[2] >= self.time_kernel_size else x[:, :, :0, :, :] + return self.conv(x) if x.shape[2] >= self.time_kernel_size else self._empty_output(x) @property def weight(self): diff --git a/comfy/ldm/lightricks/vae/causal_video_autoencoder.py b/comfy/ldm/lightricks/vae/causal_video_autoencoder.py index 5975015e2..5d0eec5b8 100644 --- a/comfy/ldm/lightricks/vae/causal_video_autoencoder.py +++ b/comfy/ldm/lightricks/vae/causal_video_autoencoder.py @@ -390,10 +390,10 @@ class Decoder(nn.Module): # Compute output channel to be product of all channel-multiplier blocks output_channel = base_channels - for block_name, block_params in list(reversed(blocks)): + for block_name, block_params in blocks: block_params = block_params if isinstance(block_params, dict) else {} if block_name == "res_x_y": - output_channel = output_channel * block_params.get("multiplier", 2) + output_channel = block_params.get("in_channels", output_channel * block_params.get("multiplier", 2)) if block_name == "compress_all": output_channel = output_channel * block_params.get("multiplier", 1) if block_name == "compress_space": @@ -432,7 +432,7 @@ class Decoder(nn.Module): spatial_padding_mode=spatial_padding_mode, ) elif block_name == "res_x_y": - output_channel = output_channel // block_params.get("multiplier", 2) + output_channel = block_params.get("out_channels", output_channel // block_params.get("multiplier", 2)) block = ResnetBlock3D( dims=dims, in_channels=input_channel,