From 45b000223ced73e17e1d29d5d411317cbef0ba93 Mon Sep 17 00:00:00 2001 From: TenserTensor Date: Wed, 4 Mar 2026 02:41:46 +0300 Subject: [PATCH] added TT_LatentMultiTransformOnPixelSpaceNode w/docs, minor bugfixes --- node_mappings.py | 2 + nodes_image.py | 165 ++++++++---------- nodes_latent.py | 63 ++++++- nodes_vae.py | 29 +++ ...TT_LatentMultiTransformOnPixelSpaceNode.md | 36 ++++ 5 files changed, 206 insertions(+), 89 deletions(-) create mode 100644 web/docs/TT_LatentMultiTransformOnPixelSpaceNode.md diff --git a/node_mappings.py b/node_mappings.py index 01ae330..a7a89b3 100644 --- a/node_mappings.py +++ b/node_mappings.py @@ -148,6 +148,7 @@ NODE_CLASS_MAPPINGS = { # Latent V3 Nodes "TT_LatentFactoryNode": TT_LatentFactoryNode, "TT_LatentMultiTransformNode": TT_LatentMultiTransformNode, + "TT_LatentMultiTransformOnPixelSpaceNode": TT_LatentMultiTransformOnPixelSpaceNode, # Loaders "TT_SdxlModelsLoader": TT_SdxlModelsLoader, "TT_SdxlModelsLoaderAdvanced": TT_SdxlModelsLoaderAdvanced, @@ -230,6 +231,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { # Latent V3 Nodes "TT_LatentFactoryNode": "TT Latent Factory", "TT_LatentMultiTransformNode": "TT Latent Multi Transform", + "TT_LatentMultiTransformOnPixelSpaceNode": "TT Latent MultiTransform On Pixel Space", # Loaders "TT_SdxlModelsLoader": "TT SDXL Models Loader", "TT_SdxlModelsLoaderAdvanced": "TT SDXL Models Loader (Advanced)", diff --git a/nodes_image.py b/nodes_image.py index b211a32..462b9cd 100644 --- a/nodes_image.py +++ b/nodes_image.py @@ -14,16 +14,15 @@ from spandrel import ModelLoader, ImageModelDescriptor import folder_paths from comfy import model_management, utils, samplers -from comfy_api.latest import ComfyExtension, IO, ui -from comfy_api.latest._io import UploadType -from node_helpers import conditioning_set_values -from node_helpers import pillow -from .nodes_latent import MEGAPIXELS +from comfy_api.latest import ComfyExtension, io, ui +from node_helpers import conditioning_set_values, pillow +from .nodes_vae import vae_encode CATEGORY = "TenserTensor/Image" RESIZE_METHODS = ["nearest-exact", "bilinear", "area", "bicubic", "bislerp"] FILENAME_FORMATS = ["name-###", "date-name-###", "name-datetime"] FORMAT_EXT = {"PNG": ".png", "JPEG": ".jpg", "WEBP": ".webp", } +MEGAPIXELS = ["0.25 MP", "0.5 MP", "1 MP", "2 MP", "4 MP", "8 MP"] class SingleCondCFGGuider(samplers.CFGGuider): @@ -88,6 +87,14 @@ def resize_image_to_megapixels(timage, resize_method, megapixels, dimension_step return resized.squeeze(-1) if is_mask else resized +def rotate_image(pixels, turns): + return torch.rot90(pixels, k=turns, dims=[1, 2]) + + +def flip_image(pixels, axis): + return torch.flip(pixels, dims=[2 if axis == "x" else 1]) + + def get_torch_device(device=None): return ( model_management.get_torch_device() @@ -156,40 +163,40 @@ def get_image_files(): ]) -class TT_ImageLoaderResizerNode(IO.ComfyNode): +class TT_ImageLoaderResizerNode(io.ComfyNode): @classmethod - def define_schema(cls) -> IO.Schema: - return IO.Schema( + def define_schema(cls) -> io.Schema: + return io.Schema( node_id="TT_ImageLoaderResizerNode", display_name="TT Image Loader / Resizer", category=CATEGORY, description="", inputs=[ - IO.Boolean.Input("resize_image", default=False, label_on="Scale", label_off="Keep"), - IO.Combo.Input("resize_method", options=RESIZE_METHODS, advanced=True), - IO.Int.Input("dimension_step", default=1, min=1, max=256, advanced=True), - IO.Combo.Input("megapixels", options=MEGAPIXELS, advanced=True), - IO.Combo.Input("image", options=get_image_files(), upload=UploadType.image) + io.Boolean.Input("resize_image", default=False, label_on="Scale", label_off="Skip"), + io.Combo.Input("resize_method", options=RESIZE_METHODS, advanced=True), + io.Int.Input("dimension_step", default=1, min=1, max=256, advanced=True), + io.Combo.Input("megapixels", options=MEGAPIXELS, advanced=True), + io.Combo.Input("image", options=get_image_files(), upload=io.UploadType.image) ], outputs=[ - IO.Image.Output(display_name="IMAGE"), - IO.Mask.Output(display_name="MASK"), + io.Image.Output(display_name="IMAGE"), + io.Mask.Output(display_name="MASK"), ], ) @classmethod def fingerprint_inputs(cls, **kwargs) -> dict[str, Any]: path = folder_paths.get_annotated_filepath(kwargs.get("image")) - resize = kwargs.get("resize_image") + resize_image = kwargs.get("resize_image") result = { "image": path + str(os.path.getmtime(path)), - "resize_image": resize, - "resize_method": kwargs.get("resize_method"), - "dimension_step": kwargs.get("dimension_step"), + "resize_image": resize_image, } - if resize: + if resize_image: result["megapixels"] = kwargs.get("megapixels") + result["resize_method"] = kwargs.get("resize_method") + result["dimension_step"] = kwargs.get("dimension_step") return result @@ -202,7 +209,7 @@ class TT_ImageLoaderResizerNode(IO.ComfyNode): return True @classmethod - def execute(cls, **kwargs) -> IO.NodeOutput: + def execute(cls, **kwargs) -> io.NodeOutput: timage, tmask = load_image(kwargs.get("image")) if kwargs.get("resize_image"): resize_method, megapixels, dimension_step = ( @@ -210,20 +217,10 @@ class TT_ImageLoaderResizerNode(IO.ComfyNode): kwargs.get("megapixels"), kwargs.get("dimension_step") ) - timage = resize_image_to_megapixels( - timage, - resize_method, - megapixels, - dimension_step - ) - tmask = resize_image_to_megapixels( - tmask, - resize_method, - megapixels, - dimension_step - ) + timage = resize_image_to_megapixels(timage, resize_method, megapixels, dimension_step) + tmask = resize_image_to_megapixels(tmask, resize_method, megapixels, dimension_step) - return IO.NodeOutput(timage, tmask, ui=ui.PreviewImage(timage, cls=cls)) + return io.NodeOutput(timage, tmask, ui=ui.PreviewImage(timage, cls=cls)) class InvalidSavePathError(Exception): @@ -283,36 +280,36 @@ def save_image(**kwargs) -> None: tframe.save(os.path.join(save_path, filename), format=image_format, **params) -class TT_ImagePreviewSaveNode(IO.ComfyNode): +class TT_ImagePreviewSaveNode(io.ComfyNode): @classmethod - def define_schema(cls) -> IO.Schema: - return IO.Schema( + def define_schema(cls) -> io.Schema: + return io.Schema( node_id="TT_ImagePreviewSaveNode", display_name="TT Image Preview / Save", category=CATEGORY, description="", is_output_node=True, inputs=[ - IO.Image.Input("image"), - IO.Boolean.Input("save_image", default=True, label_on="Save image", label_off="Only preview"), - IO.String.Input("filename_prefix", default="TT"), - IO.Combo.Input("filename_format", options=FILENAME_FORMATS, default="name-datetime", advanced=True), - IO.String.Input("subfolder", default="", advanced=True), - IO.Combo.Input("image_format", options=list(FORMAT_EXT.keys()), default="PNG", advanced=True), - IO.Int.Input("image_quality", default=75, min=1, max=100, advanced=True), - IO.Int.Input("compression_level", default=6, min=0, max=9, advanced=True), + io.Image.Input("image"), + io.Boolean.Input("save_image", default=True, label_on="Save image", label_off="Only preview"), + io.String.Input("filename_prefix", default="TT"), + io.Combo.Input("filename_format", options=FILENAME_FORMATS, default="name-datetime", advanced=True), + io.String.Input("subfolder", default="", advanced=True), + io.Combo.Input("image_format", options=list(FORMAT_EXT.keys()), default="PNG", advanced=True), + io.Int.Input("image_quality", default=75, min=1, max=100, advanced=True), + io.Int.Input("compression_level", default=6, min=0, max=9, advanced=True), ], outputs=[] ) @classmethod - def execute(cls, **kwargs) -> IO.NodeOutput: + def execute(cls, **kwargs) -> io.NodeOutput: timage = kwargs.get("image") if kwargs.get("save_image"): save_image(**kwargs) - return IO.NodeOutput(ui=ui.PreviewImage(timage, cls=cls)) + return io.NodeOutput(ui=ui.PreviewImage(timage, cls=cls)) DEVICES = ["default", "cpu"] @@ -391,35 +388,35 @@ def upscale(timage, **kwargs): return torch.clamp(scaled.movedim(1, -1), min=0.0, max=1.0) -class TT_ImagePreviewUpscaleSaveNode(IO.ComfyNode): +class TT_ImagePreviewUpscaleSaveNode(io.ComfyNode): @classmethod - def define_schema(cls) -> IO.Schema: - return IO.Schema( + def define_schema(cls) -> io.Schema: + return io.Schema( node_id="TT_ImagePreviewUpscaleSaveNode", display_name="TT Image Preview / Upscale / Save", category=CATEGORY, description="", is_output_node=True, inputs=[ - IO.Image.Input("image"), - IO.Boolean.Input("save_image", default=True, label_on="Save image", label_off="Only preview"), - IO.Boolean.Input("upscale_image", default=True, label_on="Upscale image", label_off="Keep size"), - IO.String.Input("filename_prefix", default="TT"), - IO.Combo.Input("filename_format", options=FILENAME_FORMATS, default="name-datetime", advanced=True), - IO.String.Input("subfolder", default="", advanced=True), - IO.Combo.Input("image_format", options=list(FORMAT_EXT.keys()), default="PNG", advanced=True), - IO.Int.Input("image_quality", default=75, min=1, max=100, advanced=True), - IO.Int.Input("compression_level", default=6, min=0, max=9, advanced=True), - IO.Combo.Input("upscaler_device", options=DEVICES, advanced=True), - IO.Combo.Input("upscale_model", options=get_upscale_models(), advanced=True), - IO.Int.Input("upscale_tile", default=512, min=128, max=4096, step=64, advanced=True), - IO.Int.Input("upscale_overlap", default=64, min=8, max=256, step=8, advanced=True), + io.Image.Input("image"), + io.Boolean.Input("save_image", default=True, label_on="Save image", label_off="Only preview"), + io.Boolean.Input("upscale_image", default=True, label_on="Upscale image", label_off="Keep size"), + io.String.Input("filename_prefix", default="TT"), + io.Combo.Input("filename_format", options=FILENAME_FORMATS, default="name-datetime", advanced=True), + io.String.Input("subfolder", default="", advanced=True), + io.Combo.Input("image_format", options=list(FORMAT_EXT.keys()), default="PNG", advanced=True), + io.Int.Input("image_quality", default=75, min=1, max=100, advanced=True), + io.Int.Input("compression_level", default=6, min=0, max=9, advanced=True), + io.Combo.Input("upscaler_device", options=DEVICES, advanced=True), + io.Combo.Input("upscale_model", options=get_upscale_models(), advanced=True), + io.Int.Input("upscale_tile", default=512, min=128, max=4096, step=64, advanced=True), + io.Int.Input("upscale_overlap", default=64, min=8, max=256, step=8, advanced=True), ], outputs=[] ) @classmethod - def execute(cls, **kwargs) -> IO.NodeOutput: + def execute(cls, **kwargs) -> io.NodeOutput: timage = kwargs.get("image") if kwargs.get("upscale_image"): @@ -429,36 +426,33 @@ class TT_ImagePreviewUpscaleSaveNode(IO.ComfyNode): kwargs["image"] = timage save_image(**kwargs) - return IO.NodeOutput(ui=ui.PreviewImage(timage, cls=cls)) + return io.NodeOutput(ui=ui.PreviewImage(timage, cls=cls)) -TILE_SIZE, OVERLAP = 512, 64 - - -class TT_GuiderImageReferenceNode(IO.ComfyNode): +class TT_GuiderImageReferenceNode(io.ComfyNode): @classmethod - def define_schema(cls) -> IO.Schema: - return IO.Schema( + def define_schema(cls) -> io.Schema: + return io.Schema( node_id="TT_GuiderImageReferenceNode", display_name="TT Guider Image Reference", category=CATEGORY, description="", inputs=[ - IO.Vae.Input("vae"), - IO.Guider.Input("guider"), - IO.Combo.Input("megapixels", options=MEGAPIXELS), - IO.Combo.Input("resize_method", options=RESIZE_METHODS, advanced=True), - IO.Int.Input("dimension_step", default=1, min=1, max=256, advanced=True), - IO.Combo.Input("image", options=get_image_files(), upload=UploadType.image) + io.Vae.Input("vae"), + io.Guider.Input("guider"), + io.Combo.Input("megapixels", options=MEGAPIXELS), + io.Combo.Input("resize_method", options=RESIZE_METHODS, advanced=True), + io.Int.Input("dimension_step", default=1, min=1, max=256, advanced=True), + io.Combo.Input("image", options=get_image_files(), upload=io.UploadType.image) ], outputs=[ - IO.Vae.Output("VAE"), - IO.Guider.Output("GUIDER"), + io.Vae.Output("VAE"), + io.Guider.Output("GUIDER"), ] ) @classmethod - def execute(cls, **kwargs) -> IO.NodeOutput: + def execute(cls, **kwargs) -> io.NodeOutput: timage, _ = load_image(kwargs.get("image")) vae, guider, resize_method, megapixels, dimension_step = ( kwargs.get("vae"), @@ -469,12 +463,7 @@ class TT_GuiderImageReferenceNode(IO.ComfyNode): ) resized = resize_image_to_megapixels(timage, resize_method, megapixels, dimension_step) - samples = vae.encode_tiled( - resized, - tile_x=TILE_SIZE, - tile_y=TILE_SIZE, - overlap=OVERLAP - ) + samples = vae_encode(resized, vae) conditioning = guider.get_conds() @@ -484,7 +473,7 @@ class TT_GuiderImageReferenceNode(IO.ComfyNode): else: raise ValueError("ERROR: Guider has no conditioning") - return IO.NodeOutput(vae, guider, ui=ui.PreviewImage(timage, cls=cls)) + return io.NodeOutput(vae, guider, ui=ui.PreviewImage(timage, cls=cls)) # ============================================================================== @@ -493,7 +482,7 @@ class TT_GuiderImageReferenceNode(IO.ComfyNode): class ImageNodesExtension(ComfyExtension): @override - async def get_node_list(self) -> list[type[IO.ComfyNode]]: + async def get_node_list(self) -> list[type[io.ComfyNode]]: return [ TT_ImageLoaderResizerNode, TT_ImagePreviewSaveNode, diff --git a/nodes_latent.py b/nodes_latent.py index 43c047f..016e3ce 100644 --- a/nodes_latent.py +++ b/nodes_latent.py @@ -7,12 +7,13 @@ import torch from comfy import sample, utils from comfy_api.latest import IO, ComfyExtension +from .nodes_image import MEGAPIXELS, resize_image, rotate_image, flip_image +from .nodes_vae import vae_decode, vae_encode CATEGORY = "TenserTensor/Latent" ASPECT_RATIOS = ["1:1", "4:3", "3:2", "16:9", "21:9"] CLIP_MULTIPLIERS = ["1x", "2x", "4x"] -MEGAPIXELS = ["0.25 MP", "0.5 MP", "1 MP", "2 MP", "4 MP", "8 MP"] MODEL_TYPES = ["FLUX1.D", "FLUX2.D", "SDXL"] ORIENTATIONS = ["landscape", "portrait"] ROTATE_ANGLES = ["90°", "180°", "270°"] @@ -214,6 +215,64 @@ class TT_LatentMultiTransformNode(IO.ComfyNode): return IO.NodeOutput(latent) +class TT_LatentMultiTransformOnPixelSpaceNode(IO.ComfyNode): + @classmethod + def define_schema(cls) -> IO.Schema: + return IO.Schema( + node_id="TT_LatentMultiTransformOnPixelSpaceNode", + display_name="TT Latent MultiTransform On Pixel Space", + category=CATEGORY, + description="", + inputs=[ + IO.Vae.Input("vae"), + IO.Latent.Input("latent"), + IO.Mask.Input("mask", optional=True), + IO.Boolean.Input("scale_latent", default=True, label_on="Scale", label_off="Skip"), + IO.Combo.Input("scale_factor", options=SCALE_FACTORS, default="1x"), + IO.Combo.Input("scale_method", options=SCALE_METHODS, default="nearest-exact"), + IO.Boolean.Input("rotate_latent", default=True, label_on="Rotate", label_off="Skip"), + IO.Combo.Input("rotate_angle", options=ROTATE_ANGLES), + IO.Boolean.Input("flip_latent", default=True, label_on="Flip", label_off="Skip"), + IO.Combo.Input("flip_direction", options=["horizontal", "vertical"]), + ], + outputs=[ + IO.Latent.Output("LATENT"), + ] + ) + + @classmethod + def execute(cls, **kwargs) -> IO.NodeOutput: + latent = kwargs.get("latent").copy() + vae = kwargs.get("vae") + pixels = vae_decode(latent, vae) + orig_height, orig_width = pixels.shape[1:3] + + if kwargs.get("scale_latent"): + scale_factor = float(kwargs.get("scale_factor").replace("x", "")) + pixels = resize_image( + pixels, + int(orig_width * scale_factor), + int(orig_height * scale_factor), + kwargs.get("scale_method"), + ) + + if kwargs.get("rotate_latent"): + turns = int(kwargs.get("rotate_angle").replace("°", "")) // 90 + pixels = rotate_image(pixels, turns) + + if kwargs.get("flip_latent"): + axis = "x" if kwargs.get("flip_direction") == "vertical" else "y" + pixels = flip_image(pixels, axis) + + latent = vae_encode(pixels, vae) + + mask = kwargs.get("mask") + if mask is not None: + latent = set_latent_mask(latent, mask) + + return IO.NodeOutput(latent) + + # ============================================================================== # V3 entrypoint — registers context nodes with ComfyUI # ============================================================================== @@ -224,6 +283,7 @@ class LatentNodesExtension(ComfyExtension): return [ TT_LatentFactoryNode, TT_LatentMultiTransformNode, + TT_LatentMultiTransformOnPixelSpaceNode, ] @@ -238,4 +298,5 @@ async def comfy_entrypoint() -> LatentNodesExtension: __all__ = [ "TT_LatentFactoryNode", "TT_LatentMultiTransformNode", + "TT_LatentMultiTransformOnPixelSpaceNode", ] diff --git a/nodes_vae.py b/nodes_vae.py index 6c563da..9f46d5b 100644 --- a/nodes_vae.py +++ b/nodes_vae.py @@ -2,8 +2,37 @@ from typing import override +import torch + from comfy_api.latest import IO, ComfyExtension +TILE_SIZE, OVERLAP = 512, 64 + + +def vae_encode(image, vae, tile_width=TILE_SIZE, tile_height=TILE_SIZE, overlap=OVERLAP): + samples = vae.encode_tiled( + image, + tile_x=tile_width, + tile_y=tile_height, + overlap=overlap + ) + + return {"samples": samples} + + +def vae_decode(latent, vae, tile_width=TILE_SIZE, tile_height=TILE_SIZE, overlap=OVERLAP): + samples = latent["samples"] + + compression = vae.spacial_compression_decode() + images = vae.decode_tiled( + samples, + tile_x=tile_width // compression, + tile_y=tile_height // compression, + overlap=overlap // compression + ) + + return images + # ============================================================================== # V3 entrypoint — registers context nodes with ComfyUI diff --git a/web/docs/TT_LatentMultiTransformOnPixelSpaceNode.md b/web/docs/TT_LatentMultiTransformOnPixelSpaceNode.md new file mode 100644 index 0000000..3be2c29 --- /dev/null +++ b/web/docs/TT_LatentMultiTransformOnPixelSpaceNode.md @@ -0,0 +1,36 @@ +## TT Latent MultiTransform On Pixel Space *TT_LatentMultiTransformOnPixelSpaceNode* + +*This description was generated with AI assistance. If you spot any errors, please create an issue on GitHub.* + +Applies multiple geometric transformations to a latent by first decoding it to pixel space, transforming, and re-encoding. Produces better quality +results than direct latent manipulation. + +### Inputs + +| Parameter | Type | Required | Description | +|------------------|---------|----------|----------------------------------------------------| +| `vae` | VAE | Yes | VAE model for decoding and re-encoding the latent | +| `latent` | LATENT | Yes | Input latent to transform | +| `mask` | MASK | No | Optional noise mask to attach to the output latent | +| `scale_latent` | BOOLEAN | Yes | Scale image (Scale) or skip (Skip) | +| `scale_factor` | COMBO | Yes | Scale multiplier | +| `scale_method` | COMBO | Yes | Interpolation method for scaling | +| `rotate_latent` | BOOLEAN | Yes | Rotate image (Rotate) or skip (Skip) | +| `rotate_angle` | COMBO | Yes | Rotation angle in degrees, clockwise | +| `flip_latent` | BOOLEAN | Yes | Flip image (Flip) or skip (Skip) | +| `flip_direction` | COMBO | Yes | Flip axis: horizontal or vertical | + +### Outputs + +| Parameter | Type | Description | +|-----------|--------|-----------------------------------| +| `LATENT` | LATENT | Transformed and re-encoded latent | + +### Usage + +Decode the latent to pixel space, apply scale, rotation, and flip transformations, then re-encode back to latent space. Transformations are applied in +order: scale → rotate → flip. Requires a VAE — use the same VAE that was used to encode the original latent. + +--- + +[View on GitHub](https://github.com/tenser-tensor/ComfyUI-TenserTensor)