mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-18 23:38:29 +08:00
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:
committed by
bymyself
parent
a464ac3358
commit
3bd20204f5
@@ -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:
|
||||
|
||||
42
tests-unit/comfy_test/vae_channel_trim_test.py
Normal file
42
tests-unit/comfy_test/vae_channel_trim_test.py
Normal 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()]
|
||||
Reference in New Issue
Block a user