From 3bd20204f5cdf734ba5e2a515f59b2c0328581b2 Mon Sep 17 00:00:00 2001
From: christian-byrne <abolkonsky.rem@gmail.com>
Date: Thu, 6 Aug 2026 17:36:55 -0700
Subject: [PATCH] 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.
---
 comfy/sd.py                                   |  4 ++
 .../comfy_test/vae_channel_trim_test.py       | 42 +++++++++++++++++++
 2 files changed, 46 insertions(+)
 create mode 100644 tests-unit/comfy_test/vae_channel_trim_test.py

diff --git a/comfy/sd.py b/comfy/sd.py
index 9ccd561bc..013a4ede0 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -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:
diff --git a/tests-unit/comfy_test/vae_channel_trim_test.py b/tests-unit/comfy_test/vae_channel_trim_test.py
new file mode 100644
index 000000000..eb0fbbdea
--- /dev/null
+++ b/tests-unit/comfy_test/vae_channel_trim_test.py
@@ -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()]