kontext inpaint

This commit is contained in:
nofunstudio
2025-07-11 21:44:00 -06:00
parent 4ace62b84d
commit 039e7daa8d
2 changed files with 96 additions and 2 deletions
+5 -2
View File
@@ -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__ = [
+91
View File
@@ -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)