From 039e7daa8da48db2d1dcdfd1ce9ab25293a53f90 Mon Sep 17 00:00:00 2001
From: nofunstudio <jimmy@nofun.io>
Date: Fri, 11 Jul 2025 21:44:00 -0600
Subject: [PATCH] kontext inpaint

---
 __init__.py                          |  7 ++-
 flux_kontext_inpaint_conditioning.py | 91 ++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+), 2 deletions(-)
 create mode 100644 flux_kontext_inpaint_conditioning.py

diff --git a/__init__.py b/__init__.py
index 27b70b0..832026a 100644
--- a/__init__.py
+++ b/__init__.py
@@ -10,6 +10,7 @@ from .fal_recraft_upscale import FalAPI_recraft_upscale
 from .fal_kling_video import FalAPI_kling_video
 from .string_lower import StringLower
 from .multi_alpha_composite import MultiAlphaComposite
+from .flux_kontext_inpaint_conditioning import FluxKontextInpaintingConditioning
 WEB_DIRECTORY = "./js"
 
 NODE_CLASS_MAPPINGS = {
@@ -23,7 +24,8 @@ NODE_CLASS_MAPPINGS = {
     "FalAPI_recraft_upscale": FalAPI_recraft_upscale,
     "FalAPI_kling_video": FalAPI_kling_video,
     "StringLower": StringLower,
-    "MultiAlphaComposite": MultiAlphaComposite
+    "MultiAlphaComposite": MultiAlphaComposite,
+    "FluxKontextInpaintingConditioning": FluxKontextInpaintingConditioning
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "LayeredInfiniteZoom": "Layered Infinite Zoom",
@@ -36,7 +38,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "FalAPI_recraft_upscale": "Fal API Recraft Upscale",
     "FalAPI_kling_video": "Fal API Kling Video",
     "StringLower": "String to Lowercase",
-    "MultiAlphaComposite": "Multi Alpha Composite"
+    "MultiAlphaComposite": "Multi Alpha Composite",
+    "FluxKontextInpaintingConditioning": "No Fun Flux Kontext Inpaint Conditioning"
 }
 
 __all__ = [
diff --git a/flux_kontext_inpaint_conditioning.py b/flux_kontext_inpaint_conditioning.py
new file mode 100644
index 0000000..8567a35
--- /dev/null
+++ b/flux_kontext_inpaint_conditioning.py
@@ -0,0 +1,91 @@
+# Copyright 2025 ZenAI, Inc.
+# Author: @Trgtuan_10, @vuongminh1907 (modified)
+import torch
+from comfy.comfy_types import IO
+import node_helpers
+
+class FluxKontextInpaintingConditioning:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "conditioning":      ("CONDITIONING", ),
+                "vae":               ("VAE", ),
+                "pixels":            ("IMAGE", ),
+                "mask":              ("MASK", ),
+                "reference_latents": ("LATENT", ),  # ← new face‐latent port
+                "noise_mask":        ("BOOLEAN", {
+                                         "default": True,
+                                         "tooltip": "Restrict noise to the mask region."
+                                     }),
+            }
+        }
+
+    RETURN_TYPES = ("CONDITIONING","LATENT")
+    RETURN_NAMES = ("conditioning", "latent")
+    FUNCTION = "encode"
+    CATEGORY = "conditioning/inpaint"
+
+    def _encode_latent(self, vae, pixels):
+        """Encode raw pixels into a latent dict."""
+        t = vae.encode(pixels[:,:,:,:3])
+        return {"samples": t}
+
+    def _append_reference_latent(self, conditioning, ref_latent):
+        """Append the face-reference latent to conditioning."""
+        if ref_latent is not None:
+            conditioning = node_helpers.conditioning_set_values(
+                conditioning,
+                {"reference_latents": [ref_latent["samples"]]},
+                append=True
+            )
+        return conditioning
+
+    def encode(self, conditioning, pixels, vae, mask, reference_latents, noise_mask=True):
+        # 1) Align dims to multiples of 8
+        h8 = (pixels.shape[1] // 8) * 8
+        w8 = (pixels.shape[2] // 8) * 8
+
+        # 2) Upsample mask to image size
+        mask = torch.nn.functional.interpolate(
+            mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])),
+            size=(pixels.shape[1], pixels.shape[2]),
+            mode="bilinear"
+        )
+
+        # 3) Crop or pad if needed
+        orig_pixels = pixels
+        pixels = orig_pixels.clone()
+        if pixels.shape[1] != h8 or pixels.shape[2] != w8:
+            y_off = (pixels.shape[1] - h8) // 2
+            x_off = (pixels.shape[2] - w8) // 2
+            pixels = pixels[:, y_off:y_off+h8, x_off:x_off+w8, :]
+            mask   = mask[:, :, y_off:y_off+h8, x_off:x_off+w8]
+
+        # 4) Prepare masked pixels for concat encoding
+        m = (1.0 - mask.round()).squeeze(1)  # 1=keep, 0=inpaint
+        for c in range(3):
+            pixels[:,:,:,c] = (pixels[:,:,:,c] - 0.5) * m + 0.5
+
+        # 5) Encode latents
+        concat_latent = vae.encode(pixels)
+        orig_latent   = vae.encode(orig_pixels)
+        ref_latent    = self._encode_latent(vae, reference_latents) if reference_latents is not None else None
+
+        # 6) Build conditioning dict
+        c = node_helpers.conditioning_set_values(
+            conditioning,
+            {
+                "concat_latent_image": concat_latent,
+                "concat_mask": mask
+            }
+        )
+        # 7) Append the face‐reference latent so it's only used in the mask
+        conditioning = self._append_reference_latent(c, ref_latent)
+
+        # 8) Prepare output latent
+        out_latent = {"samples": orig_latent}
+        if noise_mask:
+            out_latent["noise_mask"] = mask
+
+        return (conditioning, out_latent) 
\ No newline at end of file