From a26ecd410e9ca75ab98c32ed69c7b875e7c9d6b9 Mon Sep 17 00:00:00 2001 From: Tasha Upchurch Date: Sat, 21 Jun 2025 21:17:25 -0600 Subject: [PATCH 1/2] Add default transform list --- tests/test_torchvision_nodes.py | 18 +++++-- torchvision_nodes.py | 87 ++++++++++++++++++++++++++++----- 2 files changed, 90 insertions(+), 15 deletions(-) diff --git a/tests/test_torchvision_nodes.py b/tests/test_torchvision_nodes.py index 08833c5..bedca23 100644 --- a/tests/test_torchvision_nodes.py +++ b/tests/test_torchvision_nodes.py @@ -1,7 +1,10 @@ import pytest -from ..torchvision_nodes import TorchVisionTransformCompositionList -from ..torchvision_nodes import TorchVisionTransformNode -from ..torchvision_nodes import TorchVisionCallComposition +from ..torchvision_nodes import ( + TorchVisionTransformCompositionList, + TorchVisionTransformStringParser, + TorchVisionTransformNode, + TorchVisionCallComposition, +) # when testing set "UNIT_TEST" to true in the environment to avoid loading all of ComfyUI @@ -13,6 +16,15 @@ def test_TorchVisionTransformCompositionList(): ret = ret[0] assert ret is not None + +def test_TorchVisionTransformStringParser(): + node = TorchVisionTransformStringParser() + ret = node.compose("Resize(64,64)\nGaussianBlur(1)", "test") + transforms, name = ret + assert name == "test" + assert transforms[0][0] == "Resize" + assert transforms[1][0] == "GaussianBlur" + def test_TorchVisionTransformNode(): node1 = TorchVisionTransformCompositionList() node2 = TorchVisionTransformNode() diff --git a/torchvision_nodes.py b/torchvision_nodes.py index 71c4323..0591983 100644 --- a/torchvision_nodes.py +++ b/torchvision_nodes.py @@ -59,6 +59,64 @@ class TorchVisionTransformCompositionList: return (transforms, name,) +@ETK_torchvision_base +class TorchVisionTransformStringParser: + """Parses a newline separated list of torchvision transform calls. + + Each line should be either ``TransformName`` or ``TransformName(arg1, arg2)``. + This node helps users quickly define a composition pipeline without having + to add a separate node for each transform. + """ + + @classmethod + def INPUT_TYPES(cls): + import inspect + import torchvision.transforms as transforms + import torch.nn as nn + + # build default definition text listing all available transform + # classes and their constructor parameters for quick reference + lines = [] + for name, obj in inspect.getmembers(transforms, inspect.isclass): + if obj.__module__.startswith("torchvision.transforms") and issubclass(obj, nn.Module): + sig = inspect.signature(obj.__init__) + params = [p for p in sig.parameters.values() if p.name != "self"] + p_names = ", ".join(p.name for p in params) + lines.append(f"{name}({p_names})") + + default_def = "\n".join(lines) + + return { + "required": { + "definition": ("STRING", {"default": default_def, "multiline": True}), + "name": ("STRING", {"default": "t_comp_00"}), + }, + } + + RETURN_TYPES = ("TV_COMPOSITION_LIST", "STRING") + RETURN_NAMES = ("transforms_list", "name") + FUNCTION = "compose" + + def compose(self, definition, name): + import re + + lines = [l.strip() for l in definition.splitlines() if l.strip()] + transforms = [] + for line in lines: + m = re.match(r"(\w+)(?:\((.*)\))?", line) + if not m: + continue + t_name = m.group(1) + args_str = m.group(2) + if args_str: + args = [a.strip() for a in args_str.split(',') if a.strip()] + else: + args = None + transforms.append([t_name, args]) + + return (transforms, name,) + + @ETK_torchvision_base class TorchVisionTransformNode: """Provides the base class for torchvision transforms""" @@ -67,14 +125,15 @@ class TorchVisionTransformNode: def INPUT_TYPES(cls): import torchvision.transforms as transforms import inspect - import torch + import torch.nn as nn + + # gather torchvision transform classes for dropdown selection t_list = [] c_list = [] - for i in transforms.__dict__.values(): - if "torchvision.transforms.transforms." in str(i): - nm = str(i).split(".")[-1][:-2] - t_list.append([nm, i]) - c_list.append(nm) + for name, obj in inspect.getmembers(transforms, inspect.isclass): + if obj.__module__.startswith("torchvision.transforms") and issubclass(obj, nn.Module): + t_list.append([name, obj]) + c_list.append(name) return { "required": { @@ -156,17 +215,21 @@ class TorchVisionCallComposition: # create the transform composition list that will be used with transforms.Compose transform_composition = [] for l_c_name, l_c_args in composition: - if c_name: + if l_c_name: + if l_c_args is None: + l_c_args = [] transform_composition.append(getattr(torchvision.transforms, l_c_name)(*l_c_args)) # now that the list is composed use it to create the compostion object transform_composition = torchvision.transforms.Compose(transform_composition) - # apply the transform to the tensor - # these transforms require b,c,h,w - # so then permute - tensor = tensor.permute(0, 3, 1, 2) - tensor = transform_composition(tensor) + # apply the transform to each item in the batch + # many torchvision transforms expect a single image tensor of shape C,H,W + tensor = tensor.permute(0, 3, 1, 2) # B,C,H,W + out_list = [] + for img in tensor: + out_list.append(transform_composition(img)) + tensor = torch.stack(out_list, dim=0) tensor = tensor.permute(0, 2, 3, 1) return (tensor, tensor, transform_composition,) From fe60b1462416d7f25d67094ff3afeabd0338f96d Mon Sep 17 00:00:00 2001 From: Tasha Upchurch Date: Sat, 21 Jun 2025 21:19:59 -0600 Subject: [PATCH 2/2] Update tests/test_torchvision_nodes.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- tests/test_torchvision_nodes.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_torchvision_nodes.py b/tests/test_torchvision_nodes.py index bedca23..12c7408 100644 --- a/tests/test_torchvision_nodes.py +++ b/tests/test_torchvision_nodes.py @@ -23,7 +23,9 @@ def test_TorchVisionTransformStringParser(): transforms, name = ret assert name == "test" assert transforms[0][0] == "Resize" + assert transforms[0][1] == ["64", "64"] assert transforms[1][0] == "GaussianBlur" + assert transforms[1][1] == ["1"] def test_TorchVisionTransformNode(): node1 = TorchVisionTransformCompositionList()