From 1f8d8e6c772fe462f697486626c6c1720b16a37d Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 4 Apr 2024 15:06:17 -0400
Subject: [PATCH] Add InstructPixToPixConditioning node.

---
 comfy_extras/nodes_ip2p.py | 45 ++++++++++++++++++++++++++++++++++++++
 nodes.py                   |  1 +
 2 files changed, 46 insertions(+)
 create mode 100644 comfy_extras/nodes_ip2p.py

diff --git a/comfy_extras/nodes_ip2p.py b/comfy_extras/nodes_ip2p.py
new file mode 100644
index 00000000..c2e70a84
--- /dev/null
+++ b/comfy_extras/nodes_ip2p.py
@@ -0,0 +1,45 @@
+import torch
+
+class InstructPixToPixConditioning:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"positive": ("CONDITIONING", ),
+                             "negative": ("CONDITIONING", ),
+                             "vae": ("VAE", ),
+                             "pixels": ("IMAGE", ),
+                             }}
+
+    RETURN_TYPES = ("CONDITIONING","CONDITIONING","LATENT")
+    RETURN_NAMES = ("positive", "negative", "latent")
+    FUNCTION = "encode"
+
+    CATEGORY = "conditioning/instructpix2pix"
+
+    def encode(self, positive, negative, pixels, vae):
+        x = (pixels.shape[1] // 8) * 8
+        y = (pixels.shape[2] // 8) * 8
+
+        if pixels.shape[1] != x or pixels.shape[2] != y:
+            x_offset = (pixels.shape[1] % 8) // 2
+            y_offset = (pixels.shape[2] % 8) // 2
+            pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:]
+
+        concat_latent = vae.encode(pixels)
+
+        out_latent = {}
+        out_latent["samples"] = torch.zeros_like(concat_latent)
+
+        out = []
+        for conditioning in [positive, negative]:
+            c = []
+            for t in conditioning:
+                d = t[1].copy()
+                d["concat_latent_image"] = concat_latent
+                n = [t[0], d]
+                c.append(n)
+            out.append(c)
+        return (out[0], out[1], out_latent)
+
+NODE_CLASS_MAPPINGS = {
+    "InstructPixToPixConditioning": InstructPixToPixConditioning,
+}
diff --git a/nodes.py b/nodes.py
index 5cf020df..cdb82f80 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1965,6 +1965,7 @@ def init_custom_nodes():
         "nodes_morphology.py",
         "nodes_stable_cascade.py",
         "nodes_differential_diffusion.py",
+        "nodes_ip2p.py",
     ]
 
     import_failed = []