diff --git a/nodes.py b/nodes.py index 585fa80b..a55e8c94 100644 --- a/nodes.py +++ b/nodes.py @@ -198,6 +198,29 @@ class LatentUpscale: s = torch.nn.functional.interpolate(s, size=(height // 8, width // 8), mode=upscale_method) return (s,) +class LatentRotate: + @classmethod + def INPUT_TYPES(s): + return {"required": { "samples": ("LATENT",), + "rotation": (["none", "90 degrees", "180 degrees", "270 degrees"],), + }} + RETURN_TYPES = ("LATENT",) + FUNCTION = "rotate" + + CATEGORY = "latent" + + def rotate(self, samples, rotation): + s = samples.clone() + rotate_by = 0 + if rotation.startswith("90"): + rotate_by = 1 + elif rotation.startswith("180"): + rotate_by = 2 + elif rotation.startswith("270"): + rotate_by = 3 + + s = torch.rot90(samples, k=rotate_by, dims=[3, 2]) + return (s,) class KSampler: def __init__(self, device="cuda"): self.device = device @@ -342,6 +365,7 @@ NODE_CLASS_MAPPINGS = { "LoadImage": LoadImage, "ConditioningCombine": ConditioningCombine, "ConditioningSetArea": ConditioningSetArea, + "LatentRotate": LatentRotate, }