From 03d3488ae0e75db249f4d1a7a69e94e31dc3a54c Mon Sep 17 00:00:00 2001
From: Sean Robertson <sprobertson@gmail.com>
Date: Sat, 26 Aug 2023 23:07:01 -0700
Subject: [PATCH] initial commit

---
 __init__.py |   9 ++++
 nodes.py    | 118 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 127 insertions(+)
 create mode 100644 __init__.py
 create mode 100644 nodes.py

diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..2e6a6b4
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,9 @@
+from .nodes import LatentMirror
+
+NODE_CLASS_MAPPINGS = {
+    "LatentMirror": LatentMirror,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "LatentMirror": "Latent Mirror",
+}
diff --git a/nodes.py b/nodes.py
new file mode 100644
index 0000000..b4c4648
--- /dev/null
+++ b/nodes.py
@@ -0,0 +1,118 @@
+import torch
+
+Y_FLIP = "y_axis (left to right)"
+X_FLIP = "x_axis (top to bottom)"
+
+
+class LatentMirror:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "samples": ("LATENT",),
+                "symmetry_axis": ([Y_FLIP, X_FLIP],),
+                "flip_side": (["first", "last"],),
+                # TODO: Implement feathering to blend the seam line
+            }
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "symmetry"
+
+    CATEGORY = "latent/transform"
+
+    def symmetry(self, samples, symmetry_axis, flip_side):
+        original_s = samples["samples"].shape
+
+        cropped = self._crop_half(samples, symmetry_axis, flip_side)
+        flipped = self._flip(cropped, symmetry_axis)
+
+        x_pad = 0
+        y_pad = 0
+
+        cropped_s = cropped["samples"].shape
+
+        if symmetry_axis.startswith("x"):
+            x_pad = 0
+            if flip_side == "first":
+                y_pad = cropped["samples"].shape[2] * 8
+            else:
+                y_pad = 0
+        if symmetry_axis.startswith("y"):
+            if flip_side == "first":
+                x_pad = cropped["samples"].shape[3] * 8
+            else:
+                x_pad = 0
+            y_pad = 0
+
+        composited = self._composite(samples, flipped, x_pad, y_pad)
+
+        return (composited,)
+
+    def _crop_half(self, samples, symmetry_axis, flip_side):
+        samples_copy = samples.copy()
+        samples_data = samples["samples"]
+
+        width = samples_data.shape[3] * 8
+        height = samples_data.shape[2] * 8
+
+        x = 0
+        y = 0
+
+        if symmetry_axis.startswith("x"):
+            if flip_side == "first":
+                y = 0
+            else:
+                y = height // 2
+            return self._crop(samples_copy, width, height // 2, x, y)
+        elif symmetry_axis.startswith("y"):
+            if flip_side == "first":
+                x = 0
+                print("we shoud lool here")
+            else:
+                x = width // 2
+            return self._crop(samples_copy, width // 2, height, x, y)
+        else:
+            raise Exception("Invalid symmetry_axis")
+
+    def _crop(self, samples, width, height, x, y):
+        s = samples.copy()
+        samples = samples["samples"]
+        x = x // 8
+        y = y // 8
+
+        # enfonce minimum size of 64
+        if x > (samples.shape[3] - 8):
+            x = samples.shape[3] - 8
+        if y > (samples.shape[2] - 8):
+            y = samples.shape[2] - 8
+
+        new_height = height // 8
+        new_width = width // 8
+        to_x = new_width + x
+        to_y = new_height + y
+        s["samples"] = samples[:, :, y:to_y, x:to_x]
+        return s
+
+    def _flip(self, samples, symmetry_axis):
+        s = samples.copy()
+        if symmetry_axis.startswith("x"):
+            s["samples"] = torch.flip(samples["samples"], dims=[2])
+        elif symmetry_axis.startswith("y"):
+            s["samples"] = torch.flip(samples["samples"], dims=[3])
+        return s
+
+    def _composite(self, samples_to, samples_from, x, y, composite_method="normal"):
+        x = x // 8
+        y = y // 8
+
+        samples_out = samples_to.copy()
+        s = samples_to["samples"].clone()
+        samples_to = samples_to["samples"]
+        samples_from = samples_from["samples"]
+        s[
+            :, :, y : y + samples_from.shape[2], x : x + samples_from.shape[3]
+        ] = samples_from[:, :, : samples_to.shape[2] - y, : samples_to.shape[3] - x]
+        samples_out["samples"] = s
+        return samples_out
+