From 876cee04850bbea9175b364587cc4273e79d2abb Mon Sep 17 00:00:00 2001 From: ssitu Date: Sun, 11 Jan 2026 20:32:52 -0500 Subject: [PATCH] test: use blur before measure differences --- test/tensor_utils.py | 15 +++++++++++++++ test/test_main_workflow.py | 36 +++++++++++++++++++----------------- 2 files changed, 34 insertions(+), 17 deletions(-) diff --git a/test/tensor_utils.py b/test/tensor_utils.py index 4b56ee2..ee666af 100644 --- a/test/tensor_utils.py +++ b/test/tensor_utils.py @@ -1,3 +1,6 @@ +import torchvision.transforms.functional as TF + + def img_tensor_mae(tensor1, tensor2): """Calculate the mean absolute difference between two image tensors.""" # Remove batch dimensions if present @@ -8,3 +11,15 @@ def img_tensor_mae(tensor1, tensor2): f"Tensors must have the same shape for comparison. Got {tensor1.shape=} and {tensor2.shape=}." ) return (tensor1 - tensor2).abs().mean().item() + + +def blur(tensor, kernel_size=9, sigma=None): + """Apply Gaussian blur to an image tensor.""" + # [1, H, W, C] -> [1, C, H, W] + if tensor.ndim == 4: + tensor = tensor.permute(0, 3, 1, 2) + elif tensor.ndim == 3: + tensor = tensor.permute(2, 0, 1).unsqueeze(0) + else: + raise ValueError(f"Expected a 3D or 4D tensor, got {tensor.ndim=}") + return TF.gaussian_blur(tensor, kernel_size=kernel_size, sigma=sigma).permute(0, 2, 3, 1) # type: ignore diff --git a/test/test_main_workflow.py b/test/test_main_workflow.py index d6cccd7..ae2f105 100644 --- a/test/test_main_workflow.py +++ b/test/test_main_workflow.py @@ -2,6 +2,7 @@ Tests a common workflow for UltimateSDUpscale. """ +import logging import pathlib import pytest import torch @@ -9,7 +10,7 @@ from PIL import Image import usdu_utils from setup_utils import execute -from tensor_utils import img_tensor_mae +from tensor_utils import img_tensor_mae, blur from io_utils import save_image, load_image from configs import DirectoryConfig @@ -74,8 +75,8 @@ class TestMainWorkflow: sample_dir = test_dirs.sample_images base_img1_path = sample_dir / BASE_IMAGE_1 base_img2_path = sample_dir / BASE_IMAGE_2 - save_image(image[0], base_img1_path) - save_image(image[1], base_img2_path) + save_image(image[0:1], base_img1_path) + save_image(image[1:2], base_img2_path) # Load images back as tensors to account for compression image = torch.cat([load_image(base_img1_path), load_image(base_img2_path)]) @@ -86,19 +87,21 @@ class TestMainWorkflow: Verify generated base images match reference images. This is just to check if the checkpoint and generation pipeline are as expected for the tests dependent on their behavior. """ + logger = logging.getLogger("test_base_image_matches_reference") image, _, _ = base_image test_image_dir = test_dirs.test_images - im1 = image[0] - im2 = image[1] + im1 = image[0:1] + im2 = image[1:2] test_im1 = load_image(test_image_dir / BASE_IMAGE_1) test_im2 = load_image(test_image_dir / BASE_IMAGE_2) - diff1 = img_tensor_mae(im1, test_im1) - diff2 = img_tensor_mae(im2, test_im2) - - assert diff1 < 0.02, f"Image 1 does not match test image. Diff: {diff1}" - assert diff2 < 0.02, f"Image 2 does not match test image. Diff: {diff2}" + # Reduce high-frequency noise differences with gaussian blur. Using perceptual metrics are probably overkill. + diff1 = img_tensor_mae(blur(im1), blur(test_im1)) + diff2 = img_tensor_mae(blur(im2), blur(test_im2)) + logger.info(f"Base Image Diff1: {diff1}, Diff2: {diff2}") + assert diff1 < 0.05, "Image 1 does not match its test image." + assert diff2 < 0.05, "Image 2 does not match its test image." @pytest.fixture(scope="class") def upscaled_image( @@ -164,15 +167,13 @@ class TestMainWorkflow: upscaled = torch.cat( [load_image(upscaled_img1_path), load_image(upscaled_img2_path)] ) - # Save latent - with torch.inference_mode(): - (latent, ) = execute(node_classes["VAEEncode"], vae=vae, pixels=upscaled[0:1]) return upscaled def test_upscale_with_custom_sampler( self, upscaled_image, test_dirs: DirectoryConfig ): """Test upscaling with custom sampler and sigmas.""" + logger = logging.getLogger("test_upscale_with_custom_sampler") # Verify results test_image_dir = test_dirs.test_images im1_upscaled = upscaled_image[0] @@ -181,12 +182,13 @@ class TestMainWorkflow: test_im1_upscaled = load_image(test_image_dir / UPSCALED_IMAGE_1) test_im2_upscaled = load_image(test_image_dir / UPSCALED_IMAGE_2) - diff1 = img_tensor_mae(im1_upscaled, test_im1_upscaled) - diff2 = img_tensor_mae(im2_upscaled, test_im2_upscaled) + diff1 = img_tensor_mae(blur(im1_upscaled), blur(test_im1_upscaled)) + diff2 = img_tensor_mae(blur(im2_upscaled), blur(test_im2_upscaled)) # This tolerance is enough to handle both cpu and gpu as the device, as well as jpg compression differences. - assert diff1 < 0.02, f"Upscaled Image 1 doesn't match. Diff: {diff1}" - assert diff2 < 0.02, f"Upscaled Image 2 doesn't match. Diff: {diff2}" + logger.info(f"Diff1: {diff1}, Diff2: {diff2}") + assert diff1 < 0.05, "Upscaled Image 1 doesn't match its test image." + assert diff2 < 0.05, "Upscaled Image 2 doesn't match its test image." def test_save_sample_images(self, upscaled_image, test_dirs: DirectoryConfig): """Save sample images for visual inspection (optional utility test)."""