Comfy-Org / Comfy-Org/ComfyUI

VAE.decode chunked loop shape mismatch with TAEHV on multi-frame Wan latents

Open
#14,114 2 comments 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 7h
Merged PRs (30d)
158

Description

### Custom Node Testing

Not minimally reproducible without a custom node, but the failing assignment is in core code (`comfy/sd.py`, the batched `VAE.decode` loop). The trigger is sending a multi-frame latent through a TAEHV `VAE` instance, which is what custom previewers (e.g. VHS) do during sampling. Vanilla `TAEHVPreviewerImpl.decode_latent_to_preview` slices `x0[:1, :, :1]` first so it never exercises this path.

### Expected Behavior

`VAE.decode(samples)` with a TAEHV decoder on a multi-frame Wan 2.x latent returns a tensor without crashing on the internal `pixel_samples[x:x+batch_number] = out` assignment.

### Actual Behavior

```
RuntimeError: The expanded size of the tensor (73) must match the existing size (5) at non-singleton dimension 2.
Target sizes: [2, 3, 73, 1024, 2048]. Tensor sizes: [3, 5, 1024, 2048]
```

The chunked-decode loop preallocates `pixel_samples` as `(samples_in.shape[0],) + tuple(out.shape[1:])` from the first iteration's `out`. On a later iteration the next chunk's `out` has a different rank/shape (5-D `[B, C, T_pix, H, W]` first time, 4-D `[C, T, H, W]` shown in trace), so the slice assignment fails.

### Steps to Reproduce

1. Use a Wan 2.1 / Wan 2.2 model with a multi-frame latent (e.g. `[B, 16, T_lat, H_lat, W_lat]` with T_lat > 1).
2. Trigger `VAE.decode` on that latent through a TAEHV-class previewer (one of `taehv`, `lighttaew2_1`, `lighttaew2_2`, `lighttaehy1_5`, `taeltx_2` installed in `models/vae_approx/`).
3. The chunked loop's second-or-later iteration assignment crashes with the shape mismatch above.

Easiest user-facing trigger today is ComfyUI-VideoHelperSuite's animated previewer ([Kosinkadink/ComfyUI-VideoHelperSuite#674](https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite/pull/674)), which calls `self.taesd.decode(x0)` on the full multi-frame latent. A direct Python repro would be `VAE(load_torch_file('models/vae_approx/taehv.pth')).decode(latent)` on a 5-D Wan latent.

Workaround for users: switch preview method to `latent2rgb` (Wan21/Wan22 both define `latent_rgb_factors`, so Latent2RGB fallback works without TAEHV).

### Debug Logs

User-reported trace from Wan 2.1 + TAEHV + RTX PRO 6000 (CUDA 13.1):

```powershell
got prompt
VAE load device: cuda:0, offload device: cpu, dtype: torch.bfloat16
Model WAN21 prepared for dynamic VRAM loading. 27252MB Staged. 400 patches attached.
0%| | 0/2 [00:00

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.