Files
jaidaken f09734b0ee
Some checks failed
Python Linting / Run Ruff (push) Has been cancelled
Python Linting / Run Pylint (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Has been cancelled
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Has been cancelled
Execution Tests / test (macos-latest) (push) Has been cancelled
Execution Tests / test (ubuntu-latest) (push) Has been cancelled
Execution Tests / test (windows-latest) (push) Has been cancelled
Test server launches without errors / test (push) Has been cancelled
Unit Tests / test (macos-latest) (push) Has been cancelled
Unit Tests / test (ubuntu-latest) (push) Has been cancelled
Unit Tests / test (windows-2022) (push) Has been cancelled
Add custom nodes, Civitai loras (LFS), and vast.ai setup script
Includes 30 custom nodes committed directly, 7 Civitai-exclusive
loras stored via Git LFS, and a setup script that installs all
dependencies and downloads HuggingFace-hosted models on vast.ai.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-09 00:56:42 +00:00

97 lines
3.2 KiB
Python

"""
Test using controlnet in the upscaling workflow.
"""
import logging
import pathlib
import pytest
import torch
from setup_utils import execute
from tensor_utils import img_tensor_mae, blur
from io_utils import save_image, load_image
from configs import DirectoryConfig
from fixtures_images import EXT
CATEGORY = pathlib.Path(pathlib.Path(__file__).stem.removeprefix("test_"))
CONTROLNET_TILE_OUTPUT_IMAGE = "controlnet_tile" + EXT
TEST_CONTROLNET_TILE_MODEL = "control_v11f1e_sd15_tile.pth"
class TestControlNet:
"""Integration tests for the upscaling workflow with ControlNet."""
@pytest.fixture(scope="class")
def controlnet_upscaled_image(
self,
base_image,
loaded_checkpoint,
upscale_model,
node_classes,
seed,
test_dirs,
):
"""Generate upscaled images using ControlNet."""
image, positive, negative = base_image
model, clip, vae = loaded_checkpoint
image = image[0:1]
(controlnet_tile_model,) = execute(
node_classes["ControlNetLoader"], TEST_CONTROLNET_TILE_MODEL
)
(positive,) = execute(
node_classes["ControlNetApply"], positive, controlnet_tile_model, image, 1.0
)
with torch.inference_mode():
# Run upscale with ControlNet
usdu = node_classes["UltimateSDUpscale"]
(upscaled,) = usdu().upscale(
image=image,
model=model,
positive=positive,
negative=negative,
vae=vae,
upscale_by=2.0,
seed=seed,
steps=5,
cfg=8,
sampler_name="euler",
scheduler="normal",
denoise=1.0,
upscale_model=None,
mode_type="Chess",
tile_width=512,
tile_height=512,
mask_blur=8,
tile_padding=32,
seam_fix_mode="None",
seam_fix_denoise=1.0,
seam_fix_width=64,
seam_fix_mask_blur=8,
seam_fix_padding=16,
force_uniform_tiles=True,
tiled_decode=False,
)
# Save and reload sample image
sample_dir = test_dirs.sample_images
filename = CATEGORY / CONTROLNET_TILE_OUTPUT_IMAGE
save_image(upscaled[0], sample_dir / filename)
upscaled = load_image(sample_dir / filename)
return upscaled
def test_controlnet_upscaled_image_matches_reference(
self, controlnet_upscaled_image, test_dirs: DirectoryConfig
):
"""
Verify ControlNet upscaled images match reference images.
"""
logger = logging.getLogger("test_controlnet_upscaled_image_matches_reference")
test_img_dir = test_dirs.test_images
test_img = load_image(test_img_dir / CATEGORY / CONTROLNET_TILE_OUTPUT_IMAGE)
# Reduce high-frequency noise differences with gaussian blur
diff = img_tensor_mae(blur(controlnet_upscaled_image), blur(test_img))
logger.info(f"ControlNet Upscaled Image Diff: {diff}")
assert diff < 0.05, "ControlNet upscaled image does not match its test image."