Tiled VAE encode/decode reserve VRAM for the full tensor, not the tile, causing unnecessary model eviction
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
### Custom Node Testing
- [x] I have tried disabling custom nodes and the issue persists
### Expected Behavior
`VAEDecodeTiled` / `VAEEncodeTiled` exist to bound the VAE activation
workspace, so the VRAM ComfyUI reserves before running them should scale with
`tile_size`. Switching to the tiled node should therefore let other resident
models (e.g. the diffusion model) stay in VRAM.
### Actual Behavior
`VAE.decode_tiled` and `VAE.encode_tiled` size their reservation from the
**full** latent/pixel shape and ignore the tile size entirely:
```python
# comfy/sd.py, decode_tiled
memory_used = self.memory_used_decode(samples.shape, self.vae_dtype) #TODO: calculate mem required for tile
# comfy/sd.py, encode_tiled
memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype) # TODO: calculate mem required for tile
```
Both `#TODO`s are still present on master (`f06a187`).
`load_models_gpu(memory_required=...)` is therefore asked for several GB more
than a tiled run can ever allocate, so it evicts other models even when the
tiled path would have fit comfortably. On a memory-tight setup that eviction
turns into a full model reload on the next prompt — and reaching for the tiled
node, which is the usual remedy for exactly this situation, changes nothing.
### Measurements
RX 9060 XT 16GB (gfx1200), Windows 11, torch 2.10.0+rocm, Qwen-image VAE,
1024x1024, bf16. VAE loaded alone, peak taken with
`torch.cuda.max_memory_allocated()` around the call:
| | plain | tile 512 | tile 256 |
|---|---:|---:|---:|
| **decode** workspace | 4140.5 MB | 1035.1 MB | 258.8 MB |
| **encode** workspace | 2310.2 MB | 578.0 MB | 144.5 MB |
Reserved vs actually used at tile 512:
| | reserved | actually used | over-reserved |
|---|---:|---:|---:|
| decode | 4400.0 MB | 1035.1 MB | 4.3x |
| encode | 3000.0 MB | 578.0 MB | 5.2x |
With a 12.5GB diffusion model resident the server logs
`14472.36 MB usable, 12532.86 MB loaded`, i.e. **1939.5 MB free**.
- A tile-512 decode really needs 1035 + 242 (VAE weights) = 1277 MB, which
fits in that 1939.5 MB.
- But because the reservation is computed from the full shape, it still evicted:
```
Requested to load WanVAE
Unloaded partially: 2848.78 MB freed, 9684.11 MB remains loaded
```
So the tiled node cut real usage by 75% and changed nothing about the eviction.
### Steps to Reproduce
1. Load a model large enough that it plus the VAE decode workspace exceeds VRAM
(12.5GB diffusion model on a 16GB card here).
2. Run a 1024x1024 workflow using `VAEDecodeTiled` with `tile_size=512`.
3. Watch the log: an `Unloaded partially: ...` line still appears immediately
before `Requested to load `, and the next prompt re-loads the diffusion
model from scratch.
### Suggested fix
Size the reservation from the per-tile shape.
`comfy.utils.tiled_scale_multidim` narrows each tiled axis to `min(tile, dim)`
before invoking the model, so the estimate can mirror that.
One wrinkle worth flagging: `decode_tiled_` and `encode_tiled_` blend three
passes with different geometries — `(tx, ty)`, `(tx * 2, ty // 2)` and
`(tx // 2, ty * 2)` — so clamping to `(tx, ty)` alone under-estimates when one
axis is smaller than its tile and the other is not. All three geometries need
to be evaluated and the largest taken.
I have a PR ready that does this. Resulting reservations stay above the
measured peaks, so it does not trade an eviction for an OOM:
| | reserved (patched) | measured | margin |
|---|---:|---:|---:|
| decode tile 512 | 1100.0 MB | 1035.1 MB | +64.9 MB |
| decode tile 256 | 275.0 MB | 258.8 MB | +16.2 MB |
| encode tile 512 | 750.0 MB | 578.0 MB | +172.0 MB |
| encode tile 256 | 187.5 MB | 144.5 MB | +43.0 MB |
Model-owned tilers (`handles_tiling`) and the 1d path keep the previous
full-shape estimate, since their chunking is not described by these tile args.
After the change, the VAE loads without touching the diffusion model and the
next prompt starts sampling with no reload:
```
Requested to load WanVAE
loaded completely; 1077.39 MB usable, 242.03 MB loaded, full load: True
Prompt executed in 153.72 seconds
got prompt
0%| | 0/8 [00:00
Contributor guide
Assessment
This issue has not been assessed yet.