mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-24 18:10:27 +08:00
76cc69a2af
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.
96 lines
3.0 KiB
Python
96 lines
3.0 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from comfy.cli_args import args
|
|
|
|
if not torch.cuda.is_available():
|
|
args.cpu = True
|
|
|
|
import comfy.model_management # noqa: E402
|
|
from comfy_extras.nodes_compositing import JoinImageWithAlpha # noqa: E402
|
|
from comfy_extras.nodes_upscale_model import ImageUpscaleWithModel # noqa: E402
|
|
|
|
|
|
class StubUpscaleModel:
|
|
"""Stands in for a spandrel ImageModelDescriptor that only accepts RGB."""
|
|
|
|
scale = 2
|
|
input_channels = 3
|
|
|
|
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)
|
|
image[..., :3] = 0.5
|
|
image[0, 8:, :, 3] = 1.0
|
|
return image
|
|
|
|
|
|
def upscale(image, monkeypatch, model=None):
|
|
monkeypatch.setattr(comfy.model_management, "load_models_gpu", lambda *args, **kwargs: None)
|
|
model = model or StubUpscaleModel()
|
|
return model, ImageUpscaleWithModel.execute(model, image)
|
|
|
|
|
|
def test_rgba_input_does_not_reach_the_model_and_alpha_is_upscaled(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
|
|
assert mask[0, :14].min() > 0.9
|
|
assert mask[0, 18:].max() < 0.1
|
|
|
|
|
|
def test_rgb_input_reports_a_fully_opaque_mask(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):
|
|
_, out = upscale(rgba_image(), monkeypatch)
|
|
image, mask = out.result
|
|
|
|
rgba = JoinImageWithAlpha.execute(image, mask).result[0]
|
|
|
|
assert rgba.shape == (1, 32, 32, 4)
|
|
assert rgba[0, :14, :, 3].max() < 0.1
|
|
assert rgba[0, 18:, :, 3].min() > 0.9
|