diff --git a/node_mappings.py b/node_mappings.py index 36e82ee..01ae330 100644 --- a/node_mappings.py +++ b/node_mappings.py @@ -147,6 +147,7 @@ NODE_CLASS_MAPPINGS = { "TT_LatentFactory": TT_LatentFactory, # Latent V3 Nodes "TT_LatentFactoryNode": TT_LatentFactoryNode, + "TT_LatentMultiTransformNode": TT_LatentMultiTransformNode, # Loaders "TT_SdxlModelsLoader": TT_SdxlModelsLoader, "TT_SdxlModelsLoaderAdvanced": TT_SdxlModelsLoaderAdvanced, @@ -228,6 +229,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "TT_LatentFactory": "TT Latent Factory (Deprecated)", # Latent V3 Nodes "TT_LatentFactoryNode": "TT Latent Factory", + "TT_LatentMultiTransformNode": "TT Latent Multi Transform", # Loaders "TT_SdxlModelsLoader": "TT SDXL Models Loader", "TT_SdxlModelsLoaderAdvanced": "TT SDXL Models Loader (Advanced)", diff --git a/nodes_latent.py b/nodes_latent.py index 61a68cd..43c047f 100644 --- a/nodes_latent.py +++ b/nodes_latent.py @@ -5,16 +5,19 @@ from typing import Any, override import torch -from comfy import sample +from comfy import sample, utils from comfy_api.latest import IO, ComfyExtension CATEGORY = "TenserTensor/Latent" ASPECT_RATIOS = ["1:1", "4:3", "3:2", "16:9", "21:9"] -MEGAPIXELS = ["0.25 MP", "0.5 MP", "1 MP", "2 MP", "4 MP", "8 MP"] -ORIENTATIONS = ["landscape", "portrait"] -MODEL_TYPES = ["FLUX1.D", "FLUX2.D", "SDXL"] 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°"] +SCALE_METHODS = ["nearest-exact", "bilinear", "area", "bicubic", "bislerp"] +SCALE_FACTORS = ["0.25x", "0.5x", "1x", "2x", "4x", "8x"] class RandomNoise: @@ -131,6 +134,86 @@ class TT_LatentFactoryNode(IO.ComfyNode): return IO.NodeOutput(*args.values()) +def set_latent_mask(latent, mask): + samples = latent + samples["noise_mask"] = mask.unsqueeze(0) if mask.ndim == 2 else mask + + return samples + + +def scale_latent(**kwargs): + samples = kwargs.get("latent")["samples"] + orig_height, orig_width = samples.shape[2:] + scale_factor = float(kwargs.get("scale_factor").replace("x", "")) + final_width, final_height = round(orig_width * scale_factor), round(orig_height * scale_factor) + scaled = utils.common_upscale( + samples, + final_width, + final_height, kwargs.get("scale_method"), + "disabled" + ) + + return {"samples": scaled} + + +def rotate_latent(latent, turns): + samples = latent["samples"] + + return {"samples": torch.rot90(samples, k=turns, dims=[3, 2])} + + +def flip_latent(latent, axis): + samples = latent["samples"] + + return {"samples": torch.flip(samples, dims=[2 if axis == "x" else 3])} + + +class TT_LatentMultiTransformNode(IO.ComfyNode): + @classmethod + def define_schema(cls) -> IO.Schema: + return IO.Schema( + node_id="TT_LatentMultiTransformNode", + display_name="TT Latent MultiTransform", + category=CATEGORY, + description="", + inputs=[ + 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() + + mask = kwargs.get("mask") + if mask is not None: + latent = set_latent_mask(latent, mask) + + if kwargs.get("scale_latent"): + latent = scale_latent(**kwargs) + + if kwargs.get("rotate_latent"): + turns = int(kwargs.get("rotate_angle").replace("°", "")) // 90 + latent = rotate_latent(latent, turns) + + if kwargs.get("flip_latent"): + axis = "x" if kwargs.get("flip_direction") == "vertical" else "y" + latent = flip_latent(latent, axis) + + return IO.NodeOutput(latent) + + # ============================================================================== # V3 entrypoint — registers context nodes with ComfyUI # ============================================================================== @@ -140,6 +223,7 @@ class LatentNodesExtension(ComfyExtension): async def get_node_list(self) -> list[type[IO.ComfyNode]]: return [ TT_LatentFactoryNode, + TT_LatentMultiTransformNode, ] @@ -153,4 +237,5 @@ async def comfy_entrypoint() -> LatentNodesExtension: __all__ = [ "TT_LatentFactoryNode", + "TT_LatentMultiTransformNode", ] diff --git a/nodes_sampling.py b/nodes_sampling.py index bde06bc..b9f16bd 100644 --- a/nodes_sampling.py +++ b/nodes_sampling.py @@ -5,11 +5,9 @@ from typing import override import latent_preview from comfy import sample, utils, model_management from comfy_api.latest import ComfyExtension, IO +from .nodes_latent import SCALE_FACTORS, SCALE_METHODS, scale_latent CATEGORY = "TenserTensor/Sampling" -SCALE_METHODS = ["nearest-exact", "bilinear", "area", "bicubic", "bislerp"] -MEMORY_POLICIES = ["default", "balanced", "aggressive"] -SCALE_FACTORS = ["0.25x", "0.5x", "1x", "2x", "4x", "8x"] def execute_guided_sampling(**kwargs): @@ -79,21 +77,6 @@ class TT_GuidedKSamplerNode(IO.ComfyNode): return IO.NodeOutput(samples) -def scale_latent(latent, **kwargs): - samples = latent.copy()["samples"] - orig_height, orig_width = samples.shape[2:] - scale_factor = float(kwargs.get("scale_factor").replace("x", "")) - final_width, final_height = round(orig_width * scale_factor), round(orig_height * scale_factor) - scaled = utils.common_upscale( - samples, - final_width, - final_height, kwargs.get("scale_method"), - "disabled" - ) - - return {"samples": scaled} - - class TT_GuidedUpscaleKSamplerNode(IO.ComfyNode): @classmethod def define_schema(cls) -> IO.Schema: @@ -108,7 +91,7 @@ class TT_GuidedUpscaleKSamplerNode(IO.ComfyNode): IO.Sigmas.Input("sigmas"), IO.Sampler.Input("sampler"), IO.Noise.Input("random_noise"), - IO.Boolean.Input("scale_latent", default=True, label_on="Scale latent", label_off="Keep size"), + IO.Boolean.Input("scale_latent", default=False, label_on="Scale", label_off="Skip"), IO.Combo.Input("scale_factor", options=SCALE_FACTORS, default="1x", advanced=True), IO.Combo.Input("scale_method", options=SCALE_METHODS, default="nearest-exact", advanced=True), ], @@ -122,12 +105,10 @@ class TT_GuidedUpscaleKSamplerNode(IO.ComfyNode): samples = execute_guided_sampling(**kwargs) if kwargs.get("scale_latent"): - samples = scale_latent(samples, **kwargs) + samples = scale_latent(**kwargs) return IO.NodeOutput(samples) -class TT_LatentMultiTransformNode(IO.ComfyNode): - pass # ============================================================================== # V3 entrypoint — registers context nodes with ComfyUI diff --git a/web/docs/TT_LatentMultiTransformNode.md b/web/docs/TT_LatentMultiTransformNode.md new file mode 100644 index 0000000..fc99ca8 --- /dev/null +++ b/web/docs/TT_LatentMultiTransformNode.md @@ -0,0 +1,35 @@ +## TT Latent MultiTransform *TT_LatentMultiTransformNode* + +*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 in a single node: scaling, rotation, and flipping. Each transformation can be enabled or +disabled independently. + +### Inputs + +| Parameter | Type | Required | Description | +|------------------|---------|----------|--------------------------------------------------------| +| `latent` | LATENT | Yes | Input latent to transform | +| `mask` | MASK | No | Optional noise mask to attach to the latent | +| `scale_latent` | BOOLEAN | Yes | Scale latent (Scale) or keep original size (Keep size) | +| `scale_factor` | COMBO | No | Scale multiplier (advanced) | +| `scale_method` | COMBO | No | Interpolation method for scaling (advanced) | +| `rotate_latent` | BOOLEAN | Yes | Rotate latent (Rotate) or keep orientation (Keep) | +| `rotate_angle` | COMBO | Yes | Rotation angle in degrees | +| `flip_latent` | BOOLEAN | Yes | Flip latent (Flip) or keep orientation (Keep) | +| `flip_direction` | COMBO | Yes | Flip axis: horizontal or vertical | + +### Outputs + +| Parameter | Type | Description | +|-----------|--------|--------------------| +| `LATENT` | LATENT | Transformed latent | + +### Usage + +Apply scale, rotation, and flip to a latent before or after sampling. Transformations are applied in order: scale → rotate → flip. Use the boolean +toggles to enable only the transformations you need. + +--- + +[View on GitHub](https://github.com/tenser-tensor/ComfyUI-TenserTensor)