From 97b089dc95807bf1c48e99ca7f68a0e9036ad506 Mon Sep 17 00:00:00 2001 From: bymyself Date: Sun, 1 Jun 2025 00:11:00 -0700 Subject: [PATCH] [fix] Fix CI issues with CUDA dependencies and linting - Mock CUDA-dependent modules in tests to avoid CI failures on CPU-only runners - Fix ruff linting issues for code style compliance --- comfy_extras/nodes_images.py | 2 +- .../comfy_extras_test/image_stitch_test.py | 139 +++++++++--------- 2 files changed, 72 insertions(+), 69 deletions(-) diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py index 57b7ef0cf..d42b319e6 100644 --- a/comfy_extras/nodes_images.py +++ b/comfy_extras/nodes_images.py @@ -306,7 +306,7 @@ If image2 is not provided, returns image1 unchanged. if not match_image_size: h1, w1 = image1.shape[1:3] h2, w2 = image2.shape[1:3] - + if direction in ["left", "right"]: # For horizontal concat, pad heights to match if h1 != h2: diff --git a/tests-unit/comfy_extras_test/image_stitch_test.py b/tests-unit/comfy_extras_test/image_stitch_test.py index ebeb91b97..75c10aa64 100644 --- a/tests-unit/comfy_extras_test/image_stitch_test.py +++ b/tests-unit/comfy_extras_test/image_stitch_test.py @@ -1,234 +1,237 @@ -import pytest import torch import sys import os +from unittest.mock import patch # Add the project root to Python path sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..')) -from comfy_extras.nodes_images import ImageStitch +# Mock CUDA-dependent modules to avoid CI failures +with patch('comfy.model_management.get_torch_device', return_value=torch.device('cpu')): + with patch('comfy.model_management.get_total_memory', return_value=8192): + from comfy_extras.nodes_images import ImageStitch class TestImageStitch: - + def create_test_image(self, batch_size=1, height=64, width=64, channels=3): """Helper to create test images with specific dimensions""" return torch.rand(batch_size, height, width, channels) - + def test_no_image2_passthrough(self): """Test that when image2 is None, image1 is returned unchanged""" node = ImageStitch() image1 = self.create_test_image() - + result = node.stitch(image1, "right", True, 0, "white", image2=None) - + assert len(result) == 1 assert torch.equal(result[0], image1) - + def test_basic_horizontal_stitch_right(self): """Test basic horizontal stitching to the right""" node = ImageStitch() image1 = self.create_test_image(height=32, width=32) image2 = self.create_test_image(height=32, width=24) - + result = node.stitch(image1, "right", False, 0, "white", image2) - + assert result[0].shape == (1, 32, 56, 3) # 32 + 24 width - + def test_basic_horizontal_stitch_left(self): """Test basic horizontal stitching to the left""" node = ImageStitch() image1 = self.create_test_image(height=32, width=32) image2 = self.create_test_image(height=32, width=24) - + result = node.stitch(image1, "left", False, 0, "white", image2) - + assert result[0].shape == (1, 32, 56, 3) # 24 + 32 width - + def test_basic_vertical_stitch_down(self): """Test basic vertical stitching downward""" node = ImageStitch() image1 = self.create_test_image(height=32, width=32) image2 = self.create_test_image(height=24, width=32) - + result = node.stitch(image1, "down", False, 0, "white", image2) - + assert result[0].shape == (1, 56, 32, 3) # 32 + 24 height - + def test_basic_vertical_stitch_up(self): """Test basic vertical stitching upward""" node = ImageStitch() image1 = self.create_test_image(height=32, width=32) image2 = self.create_test_image(height=24, width=32) - + result = node.stitch(image1, "up", False, 0, "white", image2) - + assert result[0].shape == (1, 56, 32, 3) # 24 + 32 height - + def test_size_matching_horizontal(self): """Test size matching for horizontal concatenation""" node = ImageStitch() image1 = self.create_test_image(height=64, width=64) image2 = self.create_test_image(height=32, width=32) # Different aspect ratio - + result = node.stitch(image1, "right", True, 0, "white", image2) - + # image2 should be resized to match image1's height (64) with preserved aspect ratio expected_width = 64 + 64 # original + resized (32*64/32 = 64) assert result[0].shape == (1, 64, expected_width, 3) - + def test_size_matching_vertical(self): """Test size matching for vertical concatenation""" node = ImageStitch() image1 = self.create_test_image(height=64, width=64) image2 = self.create_test_image(height=32, width=32) - + result = node.stitch(image1, "down", True, 0, "white", image2) - + # image2 should be resized to match image1's width (64) with preserved aspect ratio expected_height = 64 + 64 # original + resized (32*64/32 = 64) assert result[0].shape == (1, expected_height, 64, 3) - + def test_padding_for_mismatched_heights_horizontal(self): """Test padding when heights don't match in horizontal concatenation""" node = ImageStitch() image1 = self.create_test_image(height=64, width=32) image2 = self.create_test_image(height=48, width=24) # Shorter height - + result = node.stitch(image1, "right", False, 0, "white", image2) - + # Both images should be padded to height 64 assert result[0].shape == (1, 64, 56, 3) # 32 + 24 width, max(64,48) height - + def test_padding_for_mismatched_widths_vertical(self): """Test padding when widths don't match in vertical concatenation""" node = ImageStitch() image1 = self.create_test_image(height=32, width=64) image2 = self.create_test_image(height=24, width=48) # Narrower width - + result = node.stitch(image1, "down", False, 0, "white", image2) - + # Both images should be padded to width 64 assert result[0].shape == (1, 56, 64, 3) # 32 + 24 height, max(64,48) width - + def test_border_horizontal(self): """Test border addition in horizontal concatenation""" node = ImageStitch() image1 = self.create_test_image(height=32, width=32) image2 = self.create_test_image(height=32, width=24) border_width = 16 - + result = node.stitch(image1, "right", False, border_width, "white", image2) - + # Expected width: 32 + 16 (border) + 24 = 72 assert result[0].shape == (1, 32, 72, 3) - + def test_border_vertical(self): """Test border addition in vertical concatenation""" node = ImageStitch() image1 = self.create_test_image(height=32, width=32) image2 = self.create_test_image(height=24, width=32) border_width = 16 - + result = node.stitch(image1, "down", False, border_width, "white", image2) - + # Expected height: 32 + 16 (border) + 24 = 72 assert result[0].shape == (1, 72, 32, 3) - + def test_border_color_values(self): """Test that border colors are applied correctly""" node = ImageStitch() image1 = self.create_test_image(height=32, width=32) image2 = self.create_test_image(height=32, width=32) - + # Test white border result_white = node.stitch(image1, "right", False, 16, "white", image2) # Check that border region contains white values (close to 1.0) border_region = result_white[0][:, :, 32:48, :] # Middle 16 pixels assert torch.all(border_region >= 0.9) # Should be close to white - + # Test black border result_black = node.stitch(image1, "right", False, 16, "black", image2) border_region = result_black[0][:, :, 32:48, :] assert torch.all(border_region <= 0.1) # Should be close to black - + def test_odd_border_width_made_even(self): """Test that odd border widths are made even""" node = ImageStitch() image1 = self.create_test_image(height=32, width=32) image2 = self.create_test_image(height=32, width=32) - + # Use odd border width result = node.stitch(image1, "right", False, 15, "white", image2) - + # Should be made even (16), so total width = 32 + 16 + 32 = 80 assert result[0].shape == (1, 32, 80, 3) - + def test_batch_size_matching(self): """Test that different batch sizes are handled correctly""" node = ImageStitch() image1 = self.create_test_image(batch_size=2, height=32, width=32) image2 = self.create_test_image(batch_size=1, height=32, width=32) - + result = node.stitch(image1, "right", False, 0, "white", image2) - + # Should match larger batch size assert result[0].shape == (2, 32, 64, 3) - + def test_channel_matching_rgb_to_rgba(self): """Test that channel differences are handled (RGB + alpha)""" node = ImageStitch() image1 = self.create_test_image(channels=3) # RGB image2 = self.create_test_image(channels=4) # RGBA - + result = node.stitch(image1, "right", False, 0, "white", image2) - + # Should have 4 channels (RGBA) assert result[0].shape[-1] == 4 - + def test_channel_matching_rgba_to_rgb(self): """Test that channel differences are handled (RGBA + RGB)""" node = ImageStitch() image1 = self.create_test_image(channels=4) # RGBA image2 = self.create_test_image(channels=3) # RGB - + result = node.stitch(image1, "right", False, 0, "white", image2) - + # Should have 4 channels (RGBA) assert result[0].shape[-1] == 4 - + def test_all_color_options(self): """Test all available color options""" node = ImageStitch() image1 = self.create_test_image(height=32, width=32) image2 = self.create_test_image(height=32, width=32) - + colors = ["white", "black", "red", "green", "blue"] - + for color in colors: result = node.stitch(image1, "right", False, 16, color, image2) assert result[0].shape == (1, 32, 80, 3) # Basic shape check - + def test_all_directions(self): """Test all direction options""" node = ImageStitch() image1 = self.create_test_image(height=32, width=32) image2 = self.create_test_image(height=32, width=32) - + directions = ["right", "left", "up", "down"] - + for direction in directions: result = node.stitch(image1, direction, False, 0, "white", image2) assert result[0].shape == (1, 32, 64, 3) if direction in ["right", "left"] else (1, 64, 32, 3) - + def test_complex_scenario(self): """Test complex scenario with all features enabled""" node = ImageStitch() image1 = self.create_test_image(batch_size=2, height=64, width=48, channels=3) image2 = self.create_test_image(batch_size=1, height=32, width=32, channels=4) - + result = node.stitch(image1, "right", True, 8, "red", image2) - + # Should handle: batch matching, size matching, channel matching, border assert result[0].shape[0] == 2 # Batch size matched assert result[0].shape[-1] == 4 # Channels matched to max @@ -237,35 +240,35 @@ class TestImageStitch: expected_image2_width = int(64 * (32/32)) # Resized to height 64 expected_total_width = 48 + 8 + expected_image2_width assert result[0].shape[2] == expected_total_width - + def test_input_types_structure(self): """Test that INPUT_TYPES returns correct structure""" input_types = ImageStitch.INPUT_TYPES() - + assert "required" in input_types assert "optional" in input_types - + required = input_types["required"] assert "image1" in required assert "direction" in required assert "match_image_size" in required assert "border_width" in required assert "border_color" in required - + optional = input_types["optional"] assert "image2" in optional - + # Check direction options directions = required["direction"][0] assert set(directions) == {"right", "down", "left", "up"} - + # Check color options colors = required["border_color"][0] assert set(colors) == {"white", "black", "red", "green", "blue"} - + def test_node_metadata(self): """Test node metadata and properties""" assert ImageStitch.RETURN_TYPES == ("IMAGE",) assert ImageStitch.FUNCTION == "stitch" assert ImageStitch.CATEGORY == "image/transform" - assert hasattr(ImageStitch, "DESCRIPTION") \ No newline at end of file + assert hasattr(ImageStitch, "DESCRIPTION")