Warn when VAE.encode drops input channels

VAE.encode silently trims anything past output_channels, so an RGBA image
loses its alpha with no log line. Warn once per VAE and point at
SplitImageWithAlpha. No behavior change.
This commit is contained in:
christian-byrne
2026-08-06 17:36:55 -07:00
committed by bymyself
parent a464ac3358
commit 3bd20204f5
2 changed files with 46 additions and 0 deletions

View File

@@ -497,6 +497,7 @@ class VAE:
self.latent_dim = 2
self.output_channels = 3
self.pad_channel_value = None
self.channel_trim_warned = False
self.process_input = lambda image: image * 2.0 - 1.0
self.process_output = lambda image: image.add_(1.0).div_(2.0).clamp_(0.0, 1.0)
self.working_dtypes = [torch.bfloat16, torch.float32]
@@ -1053,6 +1054,9 @@ class VAE:
pixels = pixels.narrow(d + 1, x_offset, x)
if pixels.shape[-1] > self.output_channels:
if not self.channel_trim_warned:
self.channel_trim_warned = True
logging.warning("VAE encode: input has {} channels, this VAE encodes {}. The extra channels are dropped; use SplitImageWithAlpha to keep an image's alpha as a MASK.".format(pixels.shape[-1], self.output_channels))
pixels = pixels[..., :self.output_channels]
elif pixels.shape[-1] < self.output_channels:
if self.pad_channel_value is not None:

View File

@@ -0,0 +1,42 @@
import logging
import torch
from comfy.cli_args import args
if not torch.cuda.is_available():
args.cpu = True
from comfy.sd import VAE # noqa: E402
def make_vae():
vae = VAE(sd={})
vae.crop_input = False
return vae
def test_extra_channels_are_trimmed_and_warned_once(caplog):
vae = make_vae()
pixels = torch.zeros(1, 16, 16, 4)
with caplog.at_level(logging.WARNING):
first = vae.vae_encode_crop_pixels(pixels)
second = vae.vae_encode_crop_pixels(pixels)
assert first.shape == (1, 16, 16, 3)
assert second.shape == (1, 16, 16, 3)
trim_warnings = [r for r in caplog.records if "VAE encode" in r.getMessage()]
assert len(trim_warnings) == 1
assert "SplitImageWithAlpha" in trim_warnings[0].getMessage()
def test_matching_channels_do_not_warn(caplog):
vae = make_vae()
with caplog.at_level(logging.WARNING):
out = vae.vae_encode_crop_pixels(torch.zeros(1, 16, 16, 3))
assert out.shape == (1, 16, 16, 3)
assert not [r for r in caplog.records if "VAE encode" in r.getMessage()]