Support Pruna LTX VAE (#15129)

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):

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,