zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-08-25 18:32:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

Split the alpha off by image layout, not by model channel count

Browse Source 
Keying the split off upscale_model.input_channels was wrong for any model
that does not take 3 channels: a single channel model fed a plain RGB image
took only the red channel and returned 1 - green as the alpha mask, where it
used to fail loudly. Alpha is channel 3 of an RGBA image, nothing else.
This commit is contained in:
christian-byrne
2026-08-06 19:00:32 -07:00
committed by bymyself
parent 1e881c08ea
commit 76cc69a2af
2 changed files with 38 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
tests-unit/comfy_extras_test/upscale_model_alpha_test.py
View File

@@ -1,3 +1,4 @@
import pytest
import torch
from comfy.cli_args import args
@@ -19,11 +20,21 @@ class StubUpscaleModel:
class patcher:
load_device = torch.device("cpu")
def __init__(self):
self.seen_channels = None
def __call__(self, image):
self.seen_channels = image.shape[1]
assert image.shape[1] == self.input_channels, "model was handed a non-RGB tensor"
return torch.nn.functional.interpolate(image, scale_factor=self.scale, mode="nearest")
class GrayscaleUpscaleModel(StubUpscaleModel):
"""A model that wants a single channel. It must not be fed a colour channel as alpha."""
input_channels = 1
def rgba_image():
"""16x16 RGBA where the top half is transparent and the bottom half opaque."""
image = torch.zeros(1, 16, 16, 4)
@@ -32,14 +43,17 @@ def rgba_image():
return image
def upscale(image, monkeypatch):
def upscale(image, monkeypatch, model=None):
monkeypatch.setattr(comfy.model_management, "load_models_gpu", lambda *args, **kwargs: None)
return ImageUpscaleWithModel.execute(StubUpscaleModel(), image).result
model = model or StubUpscaleModel()
return model, ImageUpscaleWithModel.execute(model, image)
def test_rgba_input_does_not_reach_the_model_and_alpha_is_upscaled(monkeypatch):
image, mask = upscale(rgba_image(), monkeypatch)
model, out = upscale(rgba_image(), monkeypatch)
image, mask = out.result
assert model.seen_channels == 3
assert image.shape == (1, 32, 32, 3)
assert mask.shape == (1, 32, 32)
# inverted convention: transparent -> 1, opaque -> 0
@@ -48,15 +62,31 @@ def test_rgba_input_does_not_reach_the_model_and_alpha_is_upscaled(monkeypatch):
def test_rgb_input_reports_a_fully_opaque_mask(monkeypatch):
image, mask = upscale(torch.zeros(1, 16, 16, 3), monkeypatch)
model, out = upscale(torch.zeros(1, 16, 16, 3), monkeypatch)
image, mask = out.result
assert model.seen_channels == 3
assert image.shape == (1, 32, 32, 3)
assert mask.shape == (1, 32, 32)
assert mask.max() == 0.0
def test_rgb_input_is_untouched_by_a_model_that_wants_fewer_channels(monkeypatch):
"""The alpha split keys off the image being RGBA, never off the model's channel count.
A single channel model must still receive all three colour channels and fail loudly,
rather than having the green channel silently reinterpreted as alpha.
"""
image = torch.zeros(1, 16, 16, 3)
image[..., 1] = 0.7
with pytest.raises(AssertionError):
upscale(image, monkeypatch, model=GrayscaleUpscaleModel())
def test_mask_round_trips_through_join_image_with_alpha(monkeypatch):
image, mask = upscale(rgba_image(), monkeypatch)
_, out = upscale(rgba_image(), monkeypatch)
image, mask = out.result
rgba = JoinImageWithAlpha.execute(image, mask).result[0]