diff --git a/README.md b/README.md index 2d4c7ae..397ecdb 100644 --- a/README.md +++ b/README.md @@ -71,9 +71,27 @@ preview actually needs it. | `Flux Empty Latent Image MXD` / `ZIT Empty Latent Image MXD` / `SDXL Empty Latent Image MXD` | Resolution presets plus vertical toggle to avoid retyping the same sizes repeatedly. | | `Save Image MXD` | Simple save modes (`Save + Preview`, `Save Only`, `Preview Only`). | | `WAN 2.2 MXD` nodes | Helpers for WAN 2.2 latent/video prep, frame tools, and I2V-focused workflows. | -| `Krea 2 Edit MXD` / `Krea2 Edit MXD` | Krea 2 source-preservation and grounded prompt encoding. No separate custom-node pack is required; model weights are still downloaded separately. | +| `Krea2 Img Edit MXD` / `Krea2 Img Edit Adv MXD` | Single-image editing with one boost. Adv adds a Base/Ref choice and a boost mask. | +| `Krea2 2 Img Edit MXD` / `Krea2 2 Img Edit Adv MXD` | Two-image editing. Simple uses fixed Base/Ref roles; Adv allows independent roles, boosts, and boost masks. | +| `Krea2 Img Edit Prompt MXD` / `Krea2 2 Img Edit Prompt MXD` | Grounded prompt encoding with one image or Base/Ref images, matching the edit nodes. Both use a fixed 768px grounding limit. | +| `Krea 2 Edit MXD` | The original model patch node is retained for published workflows. The separate grounded prompt encoder is named `Krea2 Edit Prompt MXD`. No separate custom-node pack is required; download model weights separately. | | Prompt spellcheck | Right-click a misspelled word in any prompt box for suggestions. Works offline, no node to add. | +Krea image editors output a model and `source_latent`; prompt encoding, sampling, +and VAE decoding remain separate. `source_latent` always comes from the first +image (the base in the simple two-image editor). Advanced Base/Ref choices affect +reference layout influence, not which image supplies that latent. To put two +characters into a new scene, use the two-image Adv editor, set both images to +Ref, and describe the new scene in the prompt. Boost masks restrict attention +boosts to white areas; they are not inpainting masks. + +The original `Krea2EditModelPatchMXD` key retains its published inputs, widget +order, defaults, and aligned placement for both images so existing V1.2 workflows +can load unchanged. The new editors also accept V1.2 and V2 edit LoRAs; the simple +two-image editor uses freer placement for its Ref input. The prompt encoder keeps its class key and image/prompt inputs. Its displayed +title is `Krea2 Edit Prompt MXD`, and its legacy grounding control is retained for existing workflows. The new +one/two-image prompt nodes use a fixed 768px limit; this does not set output size. + ## Companion Packs Some things that used to live in here now ship separately, so you can install @@ -118,3 +136,11 @@ and license details are recorded in [THIRD_PARTY_LICENSES.md](THIRD_PARTY_LICENSES.md) and alongside vendored code. This pack has no pip dependencies — everything it needs ships with ComfyUI. + +### Get / Set connections + +Maxed Out includes compatible `GetNode` / `SetNode` frontend nodes under +**MXD/Utilities**. They connect parts of a graph by a shared name. If KJNodes' +Set/Get extension is enabled, Maxed Out uses that provider and registers no +second copy, menus or settings. Existing workflows use the same node keys +with either pack. Refresh the browser after adding or removing a provider. diff --git a/THIRD_PARTY_LICENSES.md b/THIRD_PARTY_LICENSES.md index 0c698bd..b39345c 100644 --- a/THIRD_PARTY_LICENSES.md +++ b/THIRD_PARTY_LICENSES.md @@ -22,6 +22,10 @@ Locke, revision `86f886dac23013d88996e3a2e99093ba44d322fb`. It was modified to extract only the implementation used by the Maxed Out nodes. The upstream project is licensed under the Apache License, Version 2.0. +V2 reference placement and reference pre-encoding were adapted from revision +`13ef60f20828949becada815d3cbc371eb5ccaab`, retaining the MXD per-image controls +and latent output. + The complete license text is included at `third_party/ComfyUI-Krea2Edit/LICENSE`. @@ -36,3 +40,18 @@ distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. + +## Set/Get frontend nodes + +`web/nodes/mxd_set_get.js` adapts `web/js/setgetnodes.js` from +[ComfyUI-KJNodes](https://github.com/kijai/ComfyUI-KJNodes), licensed under +GPL-3.0 (full text in the repository-root `LICENSE`). Upstream credits +[diffus3's SetGet](https://github.com/diffus3/ComfyUI-extensions) as its origin. +The installed source snapshot used on 2026-11-08 has SHA-256 +`82d4dbcf522c23679e28132a59d907a33bc4c580e1c787afa713e3cae1e0c59d`. + +MXD adds provider detection and a singleton guard, places fallback nodes in +MXD/Utilities, updates package metadata, and makes the add-node commands +independent of KJ's other modules. KJ's enabled Set/Get script takes precedence; +the fallback does not register nodes, settings, menus or hooks alongside it. +The original SetNode/GetNode keys and KJ settings IDs are retained for compatibility. diff --git a/nodes/__init__.py b/nodes/__init__.py index 21ca9ef..26cdaaa 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -24,6 +24,7 @@ for _name in ( "krea", "masks", "media_io", + "gifs", "comparers", "checkpoints", "wan22", diff --git a/nodes/gifs.py b/nodes/gifs.py new file mode 100644 index 0000000..b3857f7 --- /dev/null +++ b/nodes/gifs.py @@ -0,0 +1,191 @@ +"""Timed GIF loading, spatial stitching, and saving. + +Registered nodes: Load GIF MXD, GIF Stitch MXD, Save GIF MXD. +MXD_GIF carries RGB IMAGE frames and their individual durations in milliseconds. +""" +from __future__ import annotations + +import bisect +import hashlib +import os + +import numpy as np +import torch +from PIL import Image, ImageSequence + +import comfy.utils +import folder_paths + + +def _animation(frames, durations): + if frames.ndim != 4 or frames.shape[0] == 0 or frames.shape[-1] != 3: + raise ValueError("GIF frames must be a nonempty RGB image batch.") + durations = [int(d) for d in durations] + if len(durations) != frames.shape[0] or any(d < 10 or d % 10 for d in durations): + raise ValueError("Each GIF frame needs a positive duration in multiples of 10 ms.") + return {"frames": frames, "durations": durations} + + +def _timeline(animation, total): + ends = np.cumsum(animation["durations"]).tolist() + period = ends[-1] + boundaries = {0, total} + for offset in range(0, total, period): + boundaries.update(offset + end for end in ends if offset + end < total) + return ends, period, boundaries + + +def _stitch_frames(image1, image2, direction, match_image_size, spacing_width, spacing_color): + """Use the same layout rules as core Image Stitch, on synchronized batches.""" + horizontal = direction in ("left", "right") + if match_image_size: + h1, w1 = image1.shape[1:3] + h2, w2 = image2.shape[1:3] + width, height = (max(1, int(h1 * w2 / h2)), h1) if horizontal else (w1, max(1, int(w1 * h2 / w2))) + image2 = comfy.utils.common_upscale( + image2.movedim(-1, 1), width, height, "lanczos", "disabled" + ).movedim(1, -1) + + colors = {"white": (1, 1, 1), "black": (0, 0, 0), "red": (1, 0, 0), + "green": (0, 1, 0), "blue": (0, 0, 1)} + color = colors[spacing_color] + # Core pads unmatched sizes with white/black, or black for RGB spacing colors. + pad_value = 1.0 if spacing_color == "white" else 0.0 + dim = 1 if horizontal else 2 + target = max(image1.shape[dim], image2.shape[dim]) + + def pad(image): + delta = target - image.shape[dim] + before, after = delta // 2, delta - delta // 2 + padding = (0, 0, 0, 0, before, after) if horizontal else (0, 0, before, after) + return torch.nn.functional.pad(image, padding, value=pad_value) if delta else image + + image1, image2 = pad(image1), pad(image2) + pieces = [image2, image1] if direction in ("left", "up") else [image1, image2] + if spacing_width: + spacing_width += spacing_width % 2 + shape = list(image1.shape) + shape[2 if horizontal else 1] = spacing_width + spacing = image1.new_empty(shape) + spacing[:] = image1.new_tensor(color) + pieces.insert(1, spacing) + return torch.cat(pieces, dim=2 if horizontal else 1) + + +class LoadGifMXD: + @classmethod + def INPUT_TYPES(cls): + root = folder_paths.get_input_directory() + files = sorted(f for f in os.listdir(root) if f.lower().endswith(".gif") + and os.path.isfile(os.path.join(root, f))) + return {"required": {"gif": (files, {"image_upload": True, "animated_image_upload": True})}} + + RETURN_TYPES = ("MXD_GIF", "IMAGE") + RETURN_NAMES = ("gif", "images") + FUNCTION = "load" + CATEGORY = "Maxed Out/GIF" + DESCRIPTION = "Load all GIF frames with their original timing. Transparency is composited onto black." + + def load(self, gif): + path = folder_paths.get_annotated_filepath(gif) + frames, durations = [], [] + with Image.open(path) as source: + if source.format != "GIF": + raise ValueError("Load GIF MXD requires a GIF file.") + for frame in ImageSequence.Iterator(source): + rgba = frame.convert("RGBA") + background = Image.new("RGBA", rgba.size, (0, 0, 0, 255)) + rgb = Image.alpha_composite(background, rgba).convert("RGB") + frames.append(torch.from_numpy(np.array(rgb).astype(np.float32) / 255.0)) + duration = int(frame.info.get("duration", 100)) + durations.append(max(10, duration) if duration > 0 else 100) + frames = torch.stack(frames) + return (_animation(frames, durations), frames) + + @classmethod + def IS_CHANGED(cls, gif): + digest = hashlib.sha256() + with open(folder_paths.get_annotated_filepath(gif), "rb") as source: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + @classmethod + def VALIDATE_INPUTS(cls, gif): + if not folder_paths.exists_annotated_filepath(gif): + return f"Invalid GIF file: {gif}" + return True + + +class GifStitchMXD: + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "gif1": ("MXD_GIF",), + "direction": (["right", "down", "left", "up"], {"default": "right"}), + "match_image_size": ("BOOLEAN", {"default": True}), + "spacing_width": ("INT", {"default": 0, "min": 0, "max": 1024, "step": 2}), + "spacing_color": (["white", "black", "red", "green", "blue"], {"default": "white"}), + }, "optional": {"gif2": ("MXD_GIF",)}} + + RETURN_TYPES = ("MXD_GIF", "IMAGE") + RETURN_NAMES = ("gif", "images") + FUNCTION = "stitch" + CATEGORY = "Maxed Out/GIF" + DESCRIPTION = "Stitch GIFs spatially with original timing. The shorter GIF loops to the longer GIF's duration." + + def stitch(self, gif1, direction, match_image_size, spacing_width, spacing_color, gif2=None): + gif1 = _animation(gif1["frames"], gif1["durations"]) + if gif2 is None: + return (gif1, gif1["frames"]) + gif2 = _animation(gif2["frames"], gif2["durations"]) + total = max(sum(gif1["durations"]), sum(gif2["durations"])) + ends1, period1, bounds1 = _timeline(gif1, total) + ends2, period2, bounds2 = _timeline(gif2, total) + bounds = sorted(bounds1 | bounds2) + # Merge consecutive identical frame pairs to avoid needless image copies. + pairs, durations = [], [] + for start, end in zip(bounds, bounds[1:]): + pair = (bisect.bisect_right(ends1, start % period1), + bisect.bisect_right(ends2, start % period2)) + if pairs and pairs[-1] == pair: + durations[-1] += end - start + else: + pairs.append(pair) + durations.append(end - start) + image1 = gif1["frames"][[p[0] for p in pairs]] + image2 = gif2["frames"][[p[1] for p in pairs]].to(image1) + frames = _stitch_frames(image1, image2, direction, match_image_size, spacing_width, spacing_color) + return (_animation(frames, durations), frames) + + +class SaveGifMXD: + @classmethod + def INPUT_TYPES(cls): + return {"required": {"gif": ("MXD_GIF",), + "filename_prefix": ("STRING", {"default": "MaxedOut/GIF"})}} + + RETURN_TYPES = ("MXD_GIF",) + RETURN_NAMES = ("gif",) + FUNCTION = "save" + OUTPUT_NODE = True + CATEGORY = "Maxed Out/GIF" + DESCRIPTION = "Save an infinitely looping GIF with its frame timing and an animated preview." + + def save(self, gif, filename_prefix="MaxedOut/GIF"): + gif = _animation(gif["frames"], gif["durations"]) + images = gif["frames"] + root, filename, counter, subfolder, _ = folder_paths.get_save_image_path( + filename_prefix, folder_paths.get_output_directory(), images.shape[2], images.shape[1]) + filename = filename.replace("%batch_num%", "0") + filename = f"{filename}_{counter:05}_.gif" + frames = [Image.fromarray((frame.detach().cpu().numpy().clip(0, 1) * 255).astype(np.uint8)) + for frame in images] + frames[0].save(os.path.join(root, filename), save_all=True, append_images=frames[1:], + duration=gif["durations"], loop=0, disposal=2, optimize=False) + return {"ui": {"images": [{"filename": filename, "subfolder": subfolder, "type": "output"}], + "animated": (True,)}, "result": (gif,)} + + +NODE_CLASS_MAPPINGS = {"Load GIF MXD": LoadGifMXD, "GIF Stitch MXD": GifStitchMXD, "Save GIF MXD": SaveGifMXD} +NODE_DISPLAY_NAME_MAPPINGS = {key: key for key in NODE_CLASS_MAPPINGS} diff --git a/nodes/krea.py b/nodes/krea.py index 00cab90..a00bc24 100644 --- a/nodes/krea.py +++ b/nodes/krea.py @@ -1,14 +1,28 @@ -"""Krea 2 editing nodes. +"""Modular Krea 2 editing nodes. -Registered nodes: +Published compatibility nodes (original mapping keys and input schemas): Krea2EditModelPatchMXD Krea 2 Edit MXD - Krea2EditGroundedEncodeMXD Krea2 Edit MXD + Krea2EditGroundedEncodeMXD Krea2 Edit Prompt MXD + +Focused editing nodes: + Krea2ImgEditMXD Krea2 Img Edit MXD + Krea2ImgEditAdvMXD Krea2 Img Edit Adv MXD + Krea2TwoImgEditMXD Krea2 2 Img Edit MXD + Krea2EditAdvancedMXD Krea2 2 Img Edit Adv MXD + Krea2ImgPromptMXD Krea2 Img Edit Prompt MXD + Krea2TwoImgPromptMXD Krea2 2 Img Edit Prompt MXD + +The advanced two-image key retains the advanced node introduced locally. +All new editors share the same edit path; conditioning and sampling stay separate. """ import comfy.patcher_extension import comfy.utils -from .krea2_edit_core import fit_encode_image, krea2_edit_forward +from .krea2_edit_core import REF_POSITIONS, fit_encode_image, fit_boost_mask, krea2_edit_forward + + +IMAGE_ROLES = {"Base": REF_POSITIONS[0], "Ref": REF_POSITIONS[1]} class Krea2EditModelPatchMXD: @@ -166,10 +180,288 @@ class Krea2EditModelPatchMXD: return (patched, source_latent) +class _KreaEditImplementation: + TITLE = "Krea 2 Edit Advanced MXD" + CATEGORY = "MXD/Krea" + DESCRIPTION = ( + "Krea 2 editing with independent Base/Ref roles, attention boosts, and " + "boost masks for each image. Masks limit attention boosts; they are not edit masks." + ) + RETURN_TYPES = ("MODEL", "LATENT") + RETURN_NAMES = ("model", "source_latent") + FUNCTION = "patch" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "image_1": ( + "IMAGE", + {"tooltip": "Primary reference image, usually the scene or image to edit."}, + ), + }, + "optional": { + "image_2": ( + "IMAGE", + { + "tooltip": ( + "Optional second reference, usually a subject to place into Image 1." + ) + }, + ), + "image_1_boost": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 1000.0, + "step": 0.01, + "round": 0.001, + "tooltip": "Attention strength for Image 1. 1.0 = unchanged.", + }, + ), + "image_1_boost_mask": ( + "MASK", + { + "tooltip": ( + "Optional mask limiting Image 1 Boost to a region such as a face. " + "White areas are boosted." + ) + }, + ), + "image_2_boost": ( + "FLOAT", + { + "default": 1.0, + "min": 0.0, + "max": 1000.0, + "step": 0.01, + "round": 0.001, + "tooltip": ( + "Attention strength for Image 2. 1.0 = unchanged; no effect when " + "Image 2 is disconnected." + ), + }, + ), + "image_2_boost_mask": ( + "MASK", + { + "tooltip": ( + "Optional mask limiting Image 2 Boost to a region such as a face. " + "White areas are boosted; no effect when Image 2 is disconnected." + ) + }, + ), + "vae": ( + "VAE", + {"tooltip": "Required. VAE used to encode the reference images."}, + ), + # Keep each image role independent. + "image_1_position": ( + list(IMAGE_ROLES), + {"default": "Base", "tooltip": "Base: follow this image's layout and framing. Ref: use this image's appearance with more freedom to reposition and reframe its contents. Base keeps existing V1.2 behavior. This does not change source_latent, which always uses Image 1."}, + ), + "image_2_position": ( + list(IMAGE_ROLES), + {"default": "Base", "tooltip": "Base: follow this image's layout and framing. Ref: use this image's appearance with more freedom to reposition and reframe its contents. Base keeps existing V1.2 behavior. No effect when Image 2 is disconnected. Set both images to Ref to restage two subjects in a new scene."}, + ), + }, + } + + def patch( + self, + model, + image_1, + image_2=None, + image_1_boost=1.0, + image_1_boost_mask=None, + image_2_boost=1.0, + image_2_boost_mask=None, + vae=None, + image_1_position="Base", + image_2_position="Base", + ): + if vae is None: + raise ValueError("Connect a VAE to Krea 2 Edit MXD.") + + images = [image_1] + boosts = [image_1_boost] + boost_masks = [image_1_boost_mask] + positions = [IMAGE_ROLES.get(image_1_position, image_1_position)] + if image_2 is not None: + images.append(image_2) + boosts.append(image_2_boost) + boost_masks.append(image_2_boost_mask) + positions.append(IMAGE_ROLES.get(image_2_position, image_2_position)) + cache = {} + + # Cache each source at the actual sampled resolution. + def encode_sources(height, width): + return [ + model.model.process_latent_in( + fit_encode_image( + image, + vae, + height, + width, + cache, + (index, height, width), + "fit", + ) + ) + for index, image in enumerate(images) + ] + + def wrapper(executor, x, timesteps, context, *args, **kwargs): + transformer_options = kwargs.get("transformer_options") + if transformer_options is None: + transformer_options = next( + (arg for arg in reversed(args) if isinstance(arg, dict)), {} + ) + height, width = x.shape[-2:] + refs = encode_sources(height, width) + return krea2_edit_forward( + executor.class_obj, + x, + timesteps, + context, + refs, + transformer_options, + image_boosts=boosts, + image_boost_masks=[fit_boost_mask(mask, image, height * 8, width * 8) + for mask, image in zip(boost_masks, images)], + ref_native=True, + pos_mode="stride1", + image_positions=positions, + ) + + patched = model.clone() + options = patched.model_options.setdefault("transformer_options", {}) + comfy.patcher_extension.add_wrapper_with_key( + comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, + "krea2_edit", + wrapper, + options, + ) + # This is deliberately only the first image's ordinary VAE latent. + # The optional second reference remains internal to the edit wrapper. + source_latent = {"samples": vae.encode(image_1[..., :3])} + return (patched, source_latent) + + +def _edit_inputs(*, two=False, advanced=False): + """Build explicit one/two-image schemas without exposing unused image controls.""" + source = _KreaEditImplementation.INPUT_TYPES() + options = source["optional"] + required = {"model": ("MODEL",)} + if two: + if advanced: + required.update({"image_1": source["required"]["image_1"], + "image_2": options["image_2"]}) + else: + required.update({ + "img_base": ("IMAGE", {"tooltip": "Image to edit. Guides layout and framing; supplies source_latent."}), + "img_ref": ("IMAGE", {"tooltip": "Appearance reference with freer placement than the base image."}), + }) + else: + required["image"] = ("IMAGE", {"tooltip": "The image to edit; supplies source_latent."}) + required["vae"] = options["vae"] + controls = {} + for i in range(1, 3 if two else 2): + if not two: + prefix, label = "", "the image" + elif advanced: + prefix, label = f"image_{i}_", f"Image {i}" + elif i == 1: + prefix, label = "img_base_", "the base image" + else: + prefix, label = "img_ref_", "the reference image" + boost = options[f"image_{i}_boost"] + if two and i == 2: + boost[1]["default"] = 2.0 + boost[1]["tooltip"] = f"Attention strength for {label}. 1.0 = unchanged." + controls[prefix + "boost"] = boost + if advanced: + controls[prefix + "boost_mask"] = ( + "MASK", {"tooltip": f"Limit the boost for {label} to white areas. This is not an edit or inpainting mask."} + ) + controls[prefix + "position"] = ( + list(IMAGE_ROLES), { + "default": "Ref" if two and i == 2 else "Base", + "tooltip": ( + "Base: follow this image's layout and framing. Ref: use its appearance " + "with freer positioning. " + "This does not change which image supplies source_latent." + ), + } + ) + return {"required": required, "optional": controls} + + +class Krea2ImgEditMXD(_KreaEditImplementation): + TITLE = "Krea2 Img Edit MXD" + DESCRIPTION = "Edit one image with one attention boost. Connect conditioning and a sampler separately." + + @classmethod + def INPUT_TYPES(cls): + return _edit_inputs() + + def patch(self, model, image, vae, boost=1.0): + return super().patch(model, image, image_1_boost=boost, vae=vae) + + +class Krea2ImgEditAdvMXD(_KreaEditImplementation): + TITLE = "Krea2 Img Edit Adv MXD" + DESCRIPTION = "Edit or restage one image with a Base/Ref role, attention boost, and optional boost mask." + + @classmethod + def INPUT_TYPES(cls): + return _edit_inputs(advanced=True) + + def patch(self, model, image, vae, boost=1.0, boost_mask=None, position="Base"): + return super().patch(model, image, image_1_boost=boost, + image_1_boost_mask=boost_mask, vae=vae, + image_1_position=position) + + +class Krea2TwoImgEditMXD(_KreaEditImplementation): + TITLE = "Krea2 2 Img Edit MXD" + DESCRIPTION = "Edit a base image using a second appearance reference. Base follows layout; Ref allows freer placement." + + @classmethod + def INPUT_TYPES(cls): + return _edit_inputs(two=True) + + def patch(self, model, img_base, img_ref, vae, img_base_boost=1.0, img_ref_boost=2.0): + return super().patch(model, img_base, image_2=img_ref, + image_1_boost=img_base_boost, image_2_boost=img_ref_boost, + vae=vae, image_1_position="Base", image_2_position="Ref") + + +class Krea2EditAdvancedMXD(_KreaEditImplementation): + TITLE = "Krea2 2 Img Edit Adv MXD" + DESCRIPTION = "Two images with independent Base/Ref roles, attention boosts, and boost masks. Use Ref for both to restage two subjects." + + @classmethod + def INPUT_TYPES(cls): + return _edit_inputs(two=True, advanced=True) + + def patch(self, model, image_1, image_2=None, image_1_boost=1.0, + image_1_boost_mask=None, image_2_boost=2.0, image_2_boost_mask=None, + vae=None, image_1_position="Base", image_2_position="Ref"): + return super().patch( + model, image_1, image_2=image_2, + image_1_boost=image_1_boost, image_1_boost_mask=image_1_boost_mask, + image_2_boost=image_2_boost, image_2_boost_mask=image_2_boost_mask, + vae=vae, image_1_position=image_1_position, image_2_position=image_2_position, + ) + + class Krea2EditGroundedEncodeMXD: """Encode a Krea 2 edit instruction together with its reference image.""" - TITLE = "Krea2 Edit MXD" + TITLE = "Krea2 Edit Prompt MXD" CATEGORY = "MXD/Krea" DESCRIPTION = ( "Encodes the edit instruction grounded on the source image using the " @@ -240,6 +532,7 @@ class Krea2EditGroundedEncodeMXD: return samples.movedim(1, -1)[:, :, :, :3] def encode(self, clip, prompt, image=None, image_b=None, grounding_px=768): + # Legacy workflows retain their grounding control. New prompt nodes use 768px. if image is None: tokens = clip.tokenize(prompt) return (clip.encode_from_tokens_scheduled(tokens),) @@ -255,12 +548,48 @@ class Krea2EditGroundedEncodeMXD: return (clip.encode_from_tokens_scheduled(tokens),) +class Krea2ImgPromptMXD(Krea2EditGroundedEncodeMXD): + TITLE = "Krea2 Img Edit Prompt MXD" + DESCRIPTION = "Encode an edit instruction with the same image used by the single-image editor. Image grounding is fixed at 768px." + + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "clip": ("CLIP",), + "image": ("IMAGE", {"tooltip": "Connect the same image as the single-image edit node."}), + "prompt": ("STRING", {"multiline": True, "default": ""}), + }} + + def encode(self, clip, image, prompt): + return super().encode(clip, prompt, image=image, grounding_px=768) + + +class Krea2TwoImgPromptMXD(Krea2EditGroundedEncodeMXD): + TITLE = "Krea2 2 Img Edit Prompt MXD" + DESCRIPTION = "Encode an edit instruction with both images in the same order as the two-image editor. Image grounding is fixed at 768px." + + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "clip": ("CLIP",), + "img_base": ("IMAGE", {"tooltip": "The base image, or Image 1 when using the advanced editor."}), + "img_ref": ("IMAGE", {"tooltip": "The reference image, or Image 2 when using the advanced editor."}), + "prompt": ("STRING", {"multiline": True, "default": ""}), + }} + + def encode(self, clip, img_base, img_ref, prompt): + return super().encode(clip, prompt, image=img_base, image_b=img_ref, grounding_px=768) + + NODE_CLASS_MAPPINGS = { "Krea2EditModelPatchMXD": Krea2EditModelPatchMXD, "Krea2EditGroundedEncodeMXD": Krea2EditGroundedEncodeMXD, + "Krea2ImgEditMXD": Krea2ImgEditMXD, + "Krea2ImgEditAdvMXD": Krea2ImgEditAdvMXD, + "Krea2TwoImgEditMXD": Krea2TwoImgEditMXD, + "Krea2EditAdvancedMXD": Krea2EditAdvancedMXD, + "Krea2ImgPromptMXD": Krea2ImgPromptMXD, + "Krea2TwoImgPromptMXD": Krea2TwoImgPromptMXD, } -NODE_DISPLAY_NAME_MAPPINGS = { - "Krea2EditModelPatchMXD": "Krea 2 Edit MXD", - "Krea2EditGroundedEncodeMXD": "Krea2 Edit MXD", -} +NODE_DISPLAY_NAME_MAPPINGS = {key: cls.TITLE for key, cls in NODE_CLASS_MAPPINGS.items()} diff --git a/nodes/krea2_edit_core.py b/nodes/krea2_edit_core.py index 5b6c30a..452d657 100644 --- a/nodes/krea2_edit_core.py +++ b/nodes/krea2_edit_core.py @@ -4,10 +4,12 @@ Adapted from ComfyUI-Krea2Edit by Conrad Locke: https://github.com/lbouaraba/comfyui-krea2edit Upstream revision: 86f886dac23013d88996e3a2e99093ba44d322fb +V2 reference placement: 13ef60f20828949becada815d3cbc371eb5ccaab Upstream license: Apache License 2.0 (see THIRD_PARTY_LICENSES.md). This file was modified for ComfyUI-MaxedOut by extracting only the image-fit -and diffusion-forward helpers needed by the MXD nodes. The upstream public +and diffusion-forward helpers needed by the MXD nodes, with per-image V2 +reference placement. The upstream public nodes, workflow, and packaging code are intentionally not duplicated. """ @@ -21,6 +23,30 @@ import comfy.ldm.common_dit from comfy.ldm.flux.layers import timestep_embedding +REF_POSITIONS = ["aligned (trained)", "beside (right of target)", "bag of features"] + + +def _place_refs(reference_ids, modes, height, width): + """Place references in RoPE space, preserving their frame indices.""" + placed = [] + next_column = float(width) + for ids, mode in zip(reference_ids, modes): + if mode == REF_POSITIONS[0]: + placed.append(ids) + continue + ids = ids.clone() + if mode == REF_POSITIONS[1]: + ids[..., 2] = ids[..., 2] - ids[..., 2].min() + next_column + next_column = float(ids[..., 2].max()) + 1.0 + elif mode == REF_POSITIONS[2]: + ids[..., 1] = (height - 1) / 2.0 + ids[..., 2] = (width - 1) / 2.0 + else: + raise ValueError(f"Unknown reference placement: {mode!r}") + placed.append(ids) + return placed + + def _imgids(bs, frame, height, width, device): ids = torch.zeros(height, width, 3, device=device, dtype=torch.float32) ids[..., 0] = frame @@ -131,6 +157,26 @@ def fit_encode_image(image, vae, height, width, cache, key, fit_mode="crop"): return latent +def fit_boost_mask(mask, image, height, width): + """Apply the same crop as fit_encode_image before its token-grid resize.""" + if mask is None: + return None + if mask.ndim == 2: + mask = mask.unsqueeze(0) + ih, iw = image.shape[1:3] + mask = F.interpolate(mask[:1, None].float(), size=(ih, iw), mode="nearest") + scale = min(height / ih, width / iw) + if ih * scale >= height * 0.92 and iw * scale >= width * 0.92: + fill = max(height / ih, width / iw) + ch, cw = min(ih, round(height / fill)), min(iw, round(width / fill)) + else: + nh = min(max(16, int(ih * scale) // 16 * 16), max(16, height // 16 * 16)) + nw = min(max(16, int(iw * scale) // 16 * 16), max(16, width // 16 * 16)) + ch, cw = min(ih, max(1, round(nh / scale))), min(iw, max(1, round(nw / scale))) + top, left = (ih - ch) // 2, (iw - cw) // 2 + return mask[:, 0, top:top + ch, left:left + cw] + + def _ref_attn_bias( boosts, boost_masks, @@ -187,6 +233,7 @@ def krea2_edit_forward( image_boost_masks=None, ref_native=False, pos_mode="anchor", + image_positions=None, ): """Run Krea 2 with clean source blocks prepended to the noisy target.""" patch = model.patch @@ -286,6 +333,12 @@ def krea2_edit_forward( ) for index, (source_h, source_w) in enumerate(source_grids) ] + reference_ids = _place_refs( + reference_ids, + image_positions or [REF_POSITIONS[0]] * len(reference_ids), + grid_h, + grid_w, + ) positions = torch.cat( [ torch.zeros( diff --git a/nodes/resolution.py b/nodes/resolution.py index fd57ec9..cd8d01c 100644 --- a/nodes/resolution.py +++ b/nodes/resolution.py @@ -1,6 +1,7 @@ """Image scaling to model-safe megapixel targets + resolution matchers. Registered nodes: + DownscaleImageToTotalPixelsMXD Downscale Image MXD Image Scale To Total Pixels (SDXL Safe) Scale SDXL Image MXD Flux Image Scale To Total Pixels (Flux Safe) Scale Flux Image MXD FluxResolutionMatcher Flux Resolution Matcher MXD @@ -11,6 +12,50 @@ from __future__ import annotations import math, comfy, comfy.utils, torch from .latents import SdxlEmptyLatentImage, ResolutionSelectorEmptyLatentImage +######################################################################################################################## +class DownscaleImageToTotalPixelsMXD: + DESCRIPTION = """Reduce images toward a megapixel target without ever increasing either dimension. Aligns width and height to resolution_steps, including images already below the target. Images with either dimension smaller than one step pass through unchanged.""" + upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"] + + @classmethod + def INPUT_TYPES(cls): + return {"required": { + "image": ("IMAGE",), + "upscale_method": (cls.upscale_methods, {"default": "lanczos"}), + "megapixels": ("FLOAT", { + "default": 1.0, "min": 0.01, "max": 16.0, "step": 0.01, + "tooltip": "Downscale toward this target; smaller images are only aligned to resolution_steps. Uses 1024 × 1024 pixels per MP, like core.", + }), + "resolution_steps": ("INT", { + "default": 1, "min": 1, "max": 256, "advanced": True, + "tooltip": "Align both dimensions to this multiple without enlarging. Rounding may put the result slightly above the megapixel target.", + }), + }} + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "downscale" + CATEGORY = "MXD/Upscaling" + + def downscale(self, image, upscale_method, megapixels, resolution_steps=1): + _, height, width, _ = image.shape + # No positive aligned size fits when a source dimension is below a step. + if min(height, width) < resolution_steps: + return (image,) + + scale = min(1.0, math.sqrt(megapixels * 1024 * 1024 / (width * height))) + + def aligned_size(size): + nearest = max(1, round(size * scale / resolution_steps)) + return min(nearest, size // resolution_steps) * resolution_steps + + new_width, new_height = aligned_size(width), aligned_size(height) + if (new_width, new_height) == (width, height): + return (image,) + samples = comfy.utils.common_upscale( + image.movedim(-1, 1), new_width, new_height, upscale_method, "disabled" + ) + return (samples.movedim(1, -1),) + ######################################################################################################################## # Image Scale To Total Pixels (SDXL Safe) class SDXLImageScaleToTotalPixelsSafe: @@ -334,6 +379,7 @@ class ResolutionSelectorMXD: ######################################################################################################################## NODE_CLASS_MAPPINGS = { + "DownscaleImageToTotalPixelsMXD": DownscaleImageToTotalPixelsMXD, "Image Scale To Total Pixels (SDXL Safe)": SDXLImageScaleToTotalPixelsSafe, "Flux Image Scale To Total Pixels (Flux Safe)": FluxImageScaleToTotalPixelsSafe, "FluxResolutionMatcher": FluxResolutionMatcher, @@ -342,6 +388,7 @@ NODE_CLASS_MAPPINGS = { } NODE_DISPLAY_NAME_MAPPINGS = { + "DownscaleImageToTotalPixelsMXD": "Downscale Image MXD", "Image Scale To Total Pixels (SDXL Safe)": "Scale SDXL Image MXD", "Flux Image Scale To Total Pixels (Flux Safe)": "Scale Flux Image MXD", "FluxResolutionMatcher": "Flux Resolution Matcher MXD", diff --git a/nodes/wan22/video_ops.py b/nodes/wan22/video_ops.py index c829112..9e51d61 100644 --- a/nodes/wan22/video_ops.py +++ b/nodes/wan22/video_ops.py @@ -23,6 +23,7 @@ import torch import folder_paths import comfy.model_management +from comfy.utils import ProgressBar from comfy.cli_args import args # Comfy API @@ -49,6 +50,29 @@ from .latent_io import _merge_prior_workflow_into_current VIDEO_EXTS = {".mp4", ".mov", ".mkv", ".webm", ".avi"} +class _VideoEncodingFrames: + """Report frames consumed by core's encoder without copying the image batch. + + Used only for saving: the returned VIDEO retains its original image tensor. + Core's VideoFromComponents encoder reads shape and iterates these frames. + """ + + def __init__(self, images, progress): + self.images = images + self.shape = images.shape + self.progress = progress + + def __len__(self): + return len(self.images) + + def __iter__(self): + for frame in self.images: + comfy.model_management.throw_exception_if_processing_interrupted() + yield frame + # Resuming the iterator means core has encoded/muxed this frame. + self.progress.update(1) + + def _frame_window(total, count, offset, mode): offset = max(1, min(offset, total)) count = max(1, min(count, total - offset + 1)) @@ -684,13 +708,26 @@ if HAVE_COMFY_API: saved_metadata = metadata file = f"{filename}_{counter:05}_.{VideoContainer.get_extension(format)}" - video.save_to( + # Reserve the last step for encoder flush, audio and container close. + # Keep the progress adapter out of the VIDEO output and downstream nodes. + progress = ProgressBar(len(images) + 1) + progress.update_absolute(0) + saving_video = VideoFromComponents( + VideoComponents( + images=_VideoEncodingFrames(images, progress), + audio=audio, + frame_rate=Fraction(fps), + ), + bit_depth=bit_depth, + ) + saving_video.save_to( os.path.join(full_output_folder, file), format=VideoContainer(format), codec=codec_name, metadata=saved_metadata, crf=encoding.get("crf"), ) + progress.update_absolute(progress.total) return io.NodeOutput( video, diff --git a/pyproject.toml b/pyproject.toml index 533ad17..ec6587e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "ComfyUI-MaxedOut" description = "Custom ComfyUI nodes used in Maxed Out workflows (SDXL, Flux, Wan 2.2, etc.)" -version = "3.3.0" +version = "3.3.1" license = {file = "LICENSE"} # classifiers = [ # # For OS-independent nodes (works on all operating systems) diff --git a/scripts/check_krea_release.py b/scripts/check_krea_release.py new file mode 100644 index 0000000..8cc080f --- /dev/null +++ b/scripts/check_krea_release.py @@ -0,0 +1,33 @@ +"""Release checks using synthetic images only; no model weights or personal media.""" +import ast +from pathlib import Path +import torch +import torch.nn.functional as F + +def functions(path,names): + tree=ast.parse(Path(path).read_text(encoding='utf-8')) + return compile(ast.Module(body=[n for n in tree.body if isinstance(n,ast.FunctionDef) and n.name in names],type_ignores=[]),path,'exec') +ns={'F':F,'torch':torch} +exec(functions('nodes/krea2_edit_core.py',{'fit_encode_image','fit_boost_mask'}),ns) +class VAE: + def encode(self,image): + self.image=image + return image.movedim(-1,1) +for h,w,th,tw in [(100,100,96,88),(120,240,96,96),(240,120,96,96),(96,96,96,96)]: + image=torch.zeros(1,h,w,3);image[:,h//4:3*h//4,w//4:3*w//4]=1 + mask=image[...,0].clone();vae=VAE() + ns['fit_encode_image'](image,vae,th//8,tw//8,{},(0,),'fit') + cropped=ns['fit_boost_mask'](mask,image,th,tw) + aligned=F.interpolate(cropped[:,None],size=vae.image.shape[1:3],mode='bicubic',antialias=True).clamp(0,1)[:,0] + torch.testing.assert_close(aligned,vae.image[...,0]) + assert ns['fit_boost_mask'](None,image,th,tw) is None +print('PASS: boost-mask crop matches encoded image for square, near-square, portrait and landscape references') +# The original published model class remains unchanged. +import subprocess +old=ast.parse(subprocess.check_output(['git','show','HEAD:nodes/krea.py']).decode());new=ast.parse(Path('nodes/krea.py').read_text()) +get=lambda tree,name:next(n for n in tree.body if isinstance(n,ast.ClassDef) and n.name==name) +assert ast.dump(get(old,'Krea2EditModelPatchMXD'))==ast.dump(get(new,'Krea2EditModelPatchMXD')) +for tree in [old,new]: + c=get(tree,'Krea2EditGroundedEncodeMXD');tree.schema=ast.dump(next(n for n in c.body if isinstance(n,ast.FunctionDef) and n.name=='INPUT_TYPES')) +assert old.schema==new.schema +print('PASS: published legacy model and legacy prompt input schema preserved') diff --git a/scripts/check_set_get.mjs b/scripts/check_set_get.mjs new file mode 100644 index 0000000..97c2d8b --- /dev/null +++ b/scripts/check_set_get.mjs @@ -0,0 +1,87 @@ +// Run with node scripts/check_set_get.mjs. No server or personal files needed. +import assert from 'node:assert/strict'; +import fs from 'node:fs'; +import vm from 'node:vm'; + +const source = fs.readFileSync(new URL('../web/nodes/mxd_set_get.js', import.meta.url), 'utf8') + .replace(/^import .*;\r?\n/gm, ''); +const run = context => vm.runInContext(`(async () => {${source}\n})()`, context); + +function environment(extensions = [], existing = {}) { + const registered = []; + class Node { + constructor() { this.properties = {}; this.inputs = []; this.outputs = []; this.widgets = []; } + addInput(name, type) { this.inputs.push({ name, type, link: null }); } + addOutput(name, type) { this.outputs.push({ name, type, links: [] }); } + addWidget(type, name, value, callback, options) { + const widget = { type, name, value, callback, options }; + this.widgets.push(widget); return widget; + } + serialize() { return { type: this.type, widgets_values: this.widgets.map(w => w.value) }; } + } + const LiteGraph = { + LGraphNode: Node, registered_node_types: { ...existing }, + registerNodeType(key, cls) { + assert.ok(!this.registered_node_types[key], `Duplicate ${key}`); + cls.prototype.type = key; this.registered_node_types[key] = cls; + } + }; + const context = vm.createContext({ + LiteGraph, window: {}, console, setTimeout, clearTimeout, + api: { getExtensions: async () => extensions }, + app: { + registerExtension: ext => registered.push(ext), + canvas: { setDirty() {} }, + ui: { settings: { getSettingValue() { return false; } } } + } + }); + return { context, registered, LiteGraph }; +} + +// Detection is based on the served file list, even before KJ runs its code. +for (const url of [ + '/extensions/comfyui-kjnodes/js/setgetnodes.js', + '/extensions/renamed%20KJ%20pack/js/setgetnodes.js?version=2' +]) { + const env = environment([url]); + await run(env.context); + assert.equal(env.registered.length, 0); + assert.equal(env.context.window.kjNodes, undefined); +} +const existing = environment([], { SetNode: class {}, GetNode: class {} }); +await run(existing.context); +assert.equal(existing.registered.length, 0); +const failed = environment(); +failed.context.api.getExtensions = async () => { throw new Error('simulated offline'); }; +failed.context.console = { warn() {} }; +await run(failed.context); +assert.equal(failed.registered.length, 0); + +// Two MXD installations racing must install only one provider. +const env = environment(['/extensions/ComfyUI-MaxedOut/nodes/mxd_set_get.js']); +await Promise.all([run(env.context), run(env.context)]); +assert.equal(new Set(env.registered.map(e => e.name)).size, env.registered.length); +assert.equal(env.registered.filter(e => e.name === 'KJNodes.SetNode').length, 1); +for (const ext of env.registered) ext.registerCustomNodes?.(); +assert.deepEqual(Object.keys(env.LiteGraph.registered_node_types).sort(), ['GetNode', 'SetNode']); +const { SetNode, GetNode } = env.LiteGraph.registered_node_types; +const setter = new SetNode(); +const getter = new GetNode(); +setter.widgets[0].value = getter.widgets[0].value = 'VAE'; +setter.inputs[0] = { name: 'VAE', type: 'VAE', link: 7 }; +const upstream = { id: 1, type: 'VAELoader' }; +const link = { id: 7, origin_id: 1, origin_slot: 0, target_id: 2, target_slot: 0, type: 'VAE' }; +const graph = { _nodes: [upstream, setter, getter], _links: { 7: link }, getNodeById: id => id === 1 ? upstream : undefined }; +setter.graph = getter.graph = graph; +getter.onRename(); +assert.equal(getter.findSetter(graph), setter); +assert.equal(getter.outputs[0].type, 'VAE'); +assert.equal(getter.getInputLink(0), link); +assert.equal(getter.serialize().widgets_values[0], 'VAE'); +// Renaming a setter must update the paired getter's saved name. +setter.properties.previousName = 'VAE'; +setter.widgets[0].value = 'Shared VAE'; +setter.update(); +assert.equal(getter.widgets[0].value, 'Shared VAE'); +assert.equal(getter.getInputLink(0), link); +console.log('OK: Set/Get provider coexistence, singleton, routing, serialization and rename checks'); diff --git a/scripts/dump_node_schema.py b/scripts/dump_node_schema.py index 6a31c64..11afc20 100644 --- a/scripts/dump_node_schema.py +++ b/scripts/dump_node_schema.py @@ -49,7 +49,7 @@ def _norm_combo(values, force_dynamic=False): and not any(_PATHY.search(v) for v in vals) ): return {"combo": vals} - return {"combo": "nonempty" if vals else "empty", "dynamic": True} + return {"combo": "dynamic", "dynamic": True} def _norm(obj, in_type_position=False): @@ -197,7 +197,7 @@ def dump(comfy_root, out_path): errors = [k for k, v in result["nodes"].items() if "__error__" in v.get("input_types", {})] if errors: print(f"[dump_node_schema] WARNING: INPUT_TYPES failed for: {errors}") - return 0 + return 1 if errors or not class_map else 0 def compare(a_path, b_path): diff --git a/scripts/node_schema_baseline.json b/scripts/node_schema_baseline.json index fe13b57..cfe3639 100644 --- a/scripts/node_schema_baseline.json +++ b/scripts/node_schema_baseline.json @@ -1,5 +1,5 @@ { - "node_count": 70, + "node_count": 80, "nodes": { "BboxDetectorCombinedBatchMXD": { "category": "MXD/Detector", @@ -175,7 +175,8 @@ "auto", "mp4", "mkv", - "webm" + "webm", + "mov" ], "tooltip": "The format to save the video as." }, @@ -231,6 +232,58 @@ "IMAGE" ] }, + "DownscaleImageToTotalPixelsMXD": { + "category": "MXD/Upscaling", + "class": "DownscaleImageToTotalPixelsMXD", + "display_name": "Downscale Image MXD", + "function": "downscale", + "input_types": { + "required": { + "image": { + "type": "IMAGE" + }, + "megapixels": { + "config": { + "default": 1.0, + "max": 16.0, + "min": 0.01, + "step": 0.01, + "tooltip": "Downscale toward this target; smaller images are only aligned to resolution_steps. Uses 1024 × 1024 pixels per MP, like core." + }, + "type": "FLOAT" + }, + "resolution_steps": { + "config": { + "advanced": true, + "default": 1, + "max": 256, + "min": 1, + "tooltip": "Align both dimensions to this multiple without enlarging. Rounding may put the result slightly above the megapixel target." + }, + "type": "INT" + }, + "upscale_method": { + "config": { + "default": "lanczos" + }, + "type": { + "combo": [ + "nearest-exact", + "bilinear", + "area", + "bicubic", + "lanczos" + ] + } + } + } + }, + "output_node": false, + "return_names": null, + "return_types": [ + "IMAGE" + ] + }, "Extract Workflow From Image MXD": { "category": "MXD/Image", "class": "ExtractWorkflowFromImageMXD", @@ -620,6 +673,75 @@ "IMAGE" ] }, + "GIF Stitch MXD": { + "category": "Maxed Out/GIF", + "class": "GifStitchMXD", + "display_name": "GIF Stitch MXD", + "function": "stitch", + "input_types": { + "optional": { + "gif2": { + "type": "MXD_GIF" + } + }, + "required": { + "direction": { + "config": { + "default": "right" + }, + "type": { + "combo": [ + "right", + "down", + "left", + "up" + ] + } + }, + "gif1": { + "type": "MXD_GIF" + }, + "match_image_size": { + "config": { + "default": true + }, + "type": "BOOLEAN" + }, + "spacing_color": { + "config": { + "default": "white" + }, + "type": { + "combo": [ + "white", + "black", + "red", + "green", + "blue" + ] + } + }, + "spacing_width": { + "config": { + "default": 0, + "max": 1024, + "min": 0, + "step": 2 + }, + "type": "INT" + } + } + }, + "output_node": false, + "return_names": [ + "gif", + "images" + ], + "return_types": [ + "MXD_GIF", + "IMAGE" + ] + }, "Get Latent Size": { "category": "MXD/Latent", "class": "GetLatentSizeMXD", @@ -820,10 +942,110 @@ "IMAGE" ] }, + "Krea2EditAdvancedMXD": { + "category": "MXD/Krea", + "class": "Krea2EditAdvancedMXD", + "display_name": "Krea2 2 Img Edit Adv MXD", + "function": "patch", + "input_types": { + "optional": { + "image_1_boost": { + "config": { + "default": 1.0, + "max": 1000.0, + "min": 0.0, + "round": 0.001, + "step": 0.01, + "tooltip": "Attention strength for Image 1. 1.0 = unchanged." + }, + "type": "FLOAT" + }, + "image_1_boost_mask": { + "config": { + "tooltip": "Limit the boost for Image 1 to white areas. This is not an edit or inpainting mask." + }, + "type": "MASK" + }, + "image_1_position": { + "config": { + "default": "Base", + "tooltip": "Base: follow this image's layout and framing. Ref: use its appearance with freer positioning. This does not change which image supplies source_latent." + }, + "type": { + "combo": [ + "Base", + "Ref" + ] + } + }, + "image_2_boost": { + "config": { + "default": 2.0, + "max": 1000.0, + "min": 0.0, + "round": 0.001, + "step": 0.01, + "tooltip": "Attention strength for Image 2. 1.0 = unchanged." + }, + "type": "FLOAT" + }, + "image_2_boost_mask": { + "config": { + "tooltip": "Limit the boost for Image 2 to white areas. This is not an edit or inpainting mask." + }, + "type": "MASK" + }, + "image_2_position": { + "config": { + "default": "Ref", + "tooltip": "Base: follow this image's layout and framing. Ref: use its appearance with freer positioning. This does not change which image supplies source_latent." + }, + "type": { + "combo": [ + "Base", + "Ref" + ] + } + } + }, + "required": { + "image_1": { + "config": { + "tooltip": "Primary reference image, usually the scene or image to edit." + }, + "type": "IMAGE" + }, + "image_2": { + "config": { + "tooltip": "Optional second reference, usually a subject to place into Image 1." + }, + "type": "IMAGE" + }, + "model": { + "type": "MODEL" + }, + "vae": { + "config": { + "tooltip": "Required. VAE used to encode the reference images." + }, + "type": "VAE" + } + } + }, + "output_node": false, + "return_names": [ + "model", + "source_latent" + ], + "return_types": [ + "MODEL", + "LATENT" + ] + }, "Krea2EditGroundedEncodeMXD": { "category": "MXD/Krea", "class": "Krea2EditGroundedEncodeMXD", - "display_name": "Krea2 Edit MXD", + "display_name": "Krea2 Edit Prompt MXD", "function": "encode", "input_types": { "optional": { @@ -942,6 +1164,250 @@ "LATENT" ] }, + "Krea2ImgEditAdvMXD": { + "category": "MXD/Krea", + "class": "Krea2ImgEditAdvMXD", + "display_name": "Krea2 Img Edit Adv MXD", + "function": "patch", + "input_types": { + "optional": { + "boost": { + "config": { + "default": 1.0, + "max": 1000.0, + "min": 0.0, + "round": 0.001, + "step": 0.01, + "tooltip": "Attention strength for the image. 1.0 = unchanged." + }, + "type": "FLOAT" + }, + "boost_mask": { + "config": { + "tooltip": "Limit the boost for the image to white areas. This is not an edit or inpainting mask." + }, + "type": "MASK" + }, + "position": { + "config": { + "default": "Base", + "tooltip": "Base: follow this image's layout and framing. Ref: use its appearance with freer positioning. This does not change which image supplies source_latent." + }, + "type": { + "combo": [ + "Base", + "Ref" + ] + } + } + }, + "required": { + "image": { + "config": { + "tooltip": "The image to edit; supplies source_latent." + }, + "type": "IMAGE" + }, + "model": { + "type": "MODEL" + }, + "vae": { + "config": { + "tooltip": "Required. VAE used to encode the reference images." + }, + "type": "VAE" + } + } + }, + "output_node": false, + "return_names": [ + "model", + "source_latent" + ], + "return_types": [ + "MODEL", + "LATENT" + ] + }, + "Krea2ImgEditMXD": { + "category": "MXD/Krea", + "class": "Krea2ImgEditMXD", + "display_name": "Krea2 Img Edit MXD", + "function": "patch", + "input_types": { + "optional": { + "boost": { + "config": { + "default": 1.0, + "max": 1000.0, + "min": 0.0, + "round": 0.001, + "step": 0.01, + "tooltip": "Attention strength for the image. 1.0 = unchanged." + }, + "type": "FLOAT" + } + }, + "required": { + "image": { + "config": { + "tooltip": "The image to edit; supplies source_latent." + }, + "type": "IMAGE" + }, + "model": { + "type": "MODEL" + }, + "vae": { + "config": { + "tooltip": "Required. VAE used to encode the reference images." + }, + "type": "VAE" + } + } + }, + "output_node": false, + "return_names": [ + "model", + "source_latent" + ], + "return_types": [ + "MODEL", + "LATENT" + ] + }, + "Krea2ImgPromptMXD": { + "category": "MXD/Krea", + "class": "Krea2ImgPromptMXD", + "display_name": "Krea2 Img Edit Prompt MXD", + "function": "encode", + "input_types": { + "required": { + "clip": { + "type": "CLIP" + }, + "image": { + "config": { + "tooltip": "Connect the same image as the single-image edit node." + }, + "type": "IMAGE" + }, + "prompt": { + "config": { + "default": "", + "multiline": true + }, + "type": "STRING" + } + } + }, + "output_node": false, + "return_names": null, + "return_types": [ + "CONDITIONING" + ] + }, + "Krea2TwoImgEditMXD": { + "category": "MXD/Krea", + "class": "Krea2TwoImgEditMXD", + "display_name": "Krea2 2 Img Edit MXD", + "function": "patch", + "input_types": { + "optional": { + "img_base_boost": { + "config": { + "default": 1.0, + "max": 1000.0, + "min": 0.0, + "round": 0.001, + "step": 0.01, + "tooltip": "Attention strength for the base image. 1.0 = unchanged." + }, + "type": "FLOAT" + }, + "img_ref_boost": { + "config": { + "default": 2.0, + "max": 1000.0, + "min": 0.0, + "round": 0.001, + "step": 0.01, + "tooltip": "Attention strength for the reference image. 1.0 = unchanged." + }, + "type": "FLOAT" + } + }, + "required": { + "img_base": { + "config": { + "tooltip": "Image to edit. Guides layout and framing; supplies source_latent." + }, + "type": "IMAGE" + }, + "img_ref": { + "config": { + "tooltip": "Appearance reference with freer placement than the base image." + }, + "type": "IMAGE" + }, + "model": { + "type": "MODEL" + }, + "vae": { + "config": { + "tooltip": "Required. VAE used to encode the reference images." + }, + "type": "VAE" + } + } + }, + "output_node": false, + "return_names": [ + "model", + "source_latent" + ], + "return_types": [ + "MODEL", + "LATENT" + ] + }, + "Krea2TwoImgPromptMXD": { + "category": "MXD/Krea", + "class": "Krea2TwoImgPromptMXD", + "display_name": "Krea2 2 Img Edit Prompt MXD", + "function": "encode", + "input_types": { + "required": { + "clip": { + "type": "CLIP" + }, + "img_base": { + "config": { + "tooltip": "The base image, or Image 1 when using the advanced editor." + }, + "type": "IMAGE" + }, + "img_ref": { + "config": { + "tooltip": "The reference image, or Image 2 when using the advanced editor." + }, + "type": "IMAGE" + }, + "prompt": { + "config": { + "default": "", + "multiline": true + }, + "type": "STRING" + } + } + }, + "output_node": false, + "return_names": null, + "return_types": [ + "CONDITIONING" + ] + }, "KreaLayerVarianceMXD": { "category": "MXD/conditioning", "class": "KreaLayerVarianceMXD", @@ -1346,6 +1812,36 @@ "STRING" ] }, + "Load GIF MXD": { + "category": "Maxed Out/GIF", + "class": "LoadGifMXD", + "display_name": "Load GIF MXD", + "function": "load", + "input_types": { + "required": { + "gif": { + "config": { + "animated_image_upload": true, + "default": "", + "image_upload": true + }, + "type": { + "combo": "dynamic", + "dynamic": true + } + } + } + }, + "output_node": false, + "return_names": [ + "gif", + "images" + ], + "return_types": [ + "MXD_GIF", + "IMAGE" + ] + }, "Load Image Batch MXD": { "category": "MXD/Image", "class": "LoadImageBatchMXD", @@ -1358,7 +1854,7 @@ "default": "" }, "type": { - "combo": "nonempty", + "combo": "dynamic", "dynamic": true } }, @@ -1401,7 +1897,7 @@ "default": "" }, "type": { - "combo": "nonempty", + "combo": "dynamic", "dynamic": true } }, @@ -1435,7 +1931,7 @@ "required": { "ckpt_name": { "type": { - "combo": "nonempty", + "combo": "dynamic", "dynamic": true } } @@ -1458,7 +1954,7 @@ "required": { "image": { "type": { - "combo": "nonempty", + "combo": "dynamic", "dynamic": true } }, @@ -1511,7 +2007,7 @@ "image_upload": true }, "type": { - "combo": "nonempty", + "combo": "dynamic", "dynamic": true } } @@ -1540,7 +2036,7 @@ "required": { "latent": { "type": { - "combo": "empty", + "combo": "dynamic", "dynamic": true } }, @@ -1599,7 +2095,7 @@ "required": { "latent": { "type": { - "combo": "empty", + "combo": "dynamic", "dynamic": true } }, @@ -1636,7 +2132,7 @@ "required": { "latent": { "type": { - "combo": "empty", + "combo": "dynamic", "dynamic": true } }, @@ -1695,7 +2191,7 @@ "required": { "subfolder": { "type": { - "combo": "nonempty", + "combo": "dynamic", "dynamic": true } } @@ -1738,7 +2234,7 @@ "required": { "subfolder": { "type": { - "combo": "nonempty", + "combo": "dynamic", "dynamic": true } } @@ -1761,7 +2257,7 @@ "required": { "subfolder": { "type": { - "combo": "nonempty", + "combo": "dynamic", "dynamic": true } } @@ -1853,7 +2349,7 @@ }, "video": { "type": { - "combo": "nonempty", + "combo": "dynamic", "dynamic": true } } @@ -2351,6 +2847,32 @@ "BOOLEAN" ] }, + "Save GIF MXD": { + "category": "Maxed Out/GIF", + "class": "SaveGifMXD", + "display_name": "Save GIF MXD", + "function": "save", + "input_types": { + "required": { + "filename_prefix": { + "config": { + "default": "MaxedOut/GIF" + }, + "type": "STRING" + }, + "gif": { + "type": "MXD_GIF" + } + } + }, + "output_node": true, + "return_names": [ + "gif" + ], + "return_types": [ + "MXD_GIF" + ] + }, "Save Image MXD": { "category": "MXD/Image", "class": "SaveImage_MXD", diff --git a/web/index.js b/web/index.js index 6e65877..75fc329 100644 --- a/web/index.js +++ b/web/index.js @@ -31,3 +31,4 @@ import './nodes/checkpoint_loader.js'; import './nodes/prompt_spellcheck.js'; import './nodes/video_preview.js'; import './nodes/live_preview_panel.js'; +import './nodes/mxd_set_get.js'; diff --git a/web/lib/mxd_utils_menu.js b/web/lib/mxd_utils_menu.js index 04a1be7..d8aceb4 100644 --- a/web/lib/mxd_utils_menu.js +++ b/web/lib/mxd_utils_menu.js @@ -183,6 +183,20 @@ function normalizeMenuCallbackValue(value) { return value; } +function positionChooserAfterLayout(menuEl, event) { + if (!menuEl?.isConnected || !Number.isFinite(event?.clientY)) return; + + // ContextMenu positions the full, flat file list before we collapse folders. + // Re-anchor the finished tree now, rather than waiting for the first keystroke + // to make ComfyUI's filter correct the stale (often far too high) position. + const rect = menuEl.getBoundingClientRect(); + const viewportHeight = document.documentElement.clientHeight || window.innerHeight; + const top = Math.max(10, Math.min(event.clientY - 10, viewportHeight - rect.height - 10)); + // Account for ContextMenu's zoom transform, including its transform origin. + const transformOffset = rect.top - (parseFloat(menuEl.style.top) || 0); + menuEl.style.top = `${top - transformOffset}px`; +} + /** Generic tree/search-enabled item chooser context menu, used for LoRAs, checkpoints, and * the smart UNET/CLIP model lists. */ export async function showModelChooser(event, callback, parentMenu, items, title = "Choose an Item") { @@ -203,6 +217,7 @@ export async function showModelChooser(event, callback, parentMenu, items, title const menuEl = getContextMenuRootElement(menu); decorateChooserAsTree(menuEl); attachSmartChooserFilter(menuEl); + positionChooserAfterLayout(menuEl, event); }); } diff --git a/web/nodes/mxd_set_get.js b/web/nodes/mxd_set_get.js new file mode 100644 index 0000000..c58a0a6 --- /dev/null +++ b/web/nodes/mxd_set_get.js @@ -0,0 +1,1747 @@ +// Adapted from kijai/ComfyUI-KJNodes web/js/setgetnodes.js (GPL-3.0). +// See THIRD_PARTY_LICENSES.md. Serialized node keys intentionally stay compatible. +import { app } from "../../../scripts/app.js"; +import { api } from "../../../scripts/api.js"; + +// Use the server's enabled extension list, not execution order: KJ's module may +// not have run yet. Skip ALL fallback side effects when another provider ships +// setgetnodes.js. A singleton also protects installations containing two MXD packs. +async function registerFallback() { + let extensions; + try { + extensions = await api.getExtensions(); + } catch (error) { + console.warn("[MXD] Cannot check Set/Get provider; leaving registration to other packs.", error); + return; + } + if (extensions.some(url => /(?:^|\/)setgetnodes\.js(?:[?#]|$)/i.test(decodeURIComponent(url)))) return; + if (globalThis.__mxdSetGetProvider || LiteGraph.registered_node_types.SetNode || LiteGraph.registered_node_types.GetNode) return; + globalThis.__mxdSetGetProvider = true; + installSetGet(); +} + +function installSetGet() { + + +// originally based on diffus3's SetGet: https://github.com/diffus3/ComfyUI-extensions + +// Nodes that allow you to hide connections for cleaner graphs + +let _typeColorMap; +function setColorAndBgColor(node, type) { + if (!_typeColorMap) { + _typeColorMap = { + "DEFAULT": LGraphCanvas.node_colors.gray, + "MODEL": LGraphCanvas.node_colors.blue, + "LATENT": LGraphCanvas.node_colors.purple, + "VAE": LGraphCanvas.node_colors.red, + "WANVAE": LGraphCanvas.node_colors.red, + "CONDITIONING": LGraphCanvas.node_colors.brown, + "IMAGE": LGraphCanvas.node_colors.pale_blue, + "CLIP": LGraphCanvas.node_colors.yellow, + "FLOAT": LGraphCanvas.node_colors.green, + "MASK": { color: "#1c5715", bgcolor: "#1f401b"}, + "INT": { color: "#1b4669", bgcolor: "#29699c"}, + "CONTROL_NET": { color: "#156653", bgcolor: "#1c453b"}, + "NOISE": { color: "#2e2e2e", bgcolor: "#242121"}, + "GUIDER": { color: "#3c7878", bgcolor: "#1c453b"}, + "SAMPLER": { color: "#614a4a", bgcolor: "#3b2c2c"}, + "SIGMAS": { color: "#485248", bgcolor: "#272e27"}, + }; + } + const colors = _typeColorMap[type] || LGraphCanvas.node_colors?.gray; + if (colors) { + node.color = colors.color; + node.bgcolor = colors.bgcolor; + } +} +function getDisablePrefix() { + return app.ui.settings.getSettingValue("KJNodes.disablePrefix") ?? false; +} +function prefixedTitle(prefix, name) { + return (getDisablePrefix() ? "" : prefix + "_") + name; +} +function autoColor(node, type) { + if (!app.ui.settings.getSettingValue("KJNodes.nodeAutoColor")) return; + if (type === '*') { node.color = null; node.bgcolor = null; } + else setColorAndBgColor(node, type); +} +function addNodeToSelectedOrCursor(nodeType, side) { + const selected = Object.values(app.canvas.selected_nodes || {}); + if (selected.length > 0) { + for (const source of selected) { + const graph = source.graph || app.graph; + const node = LiteGraph.createNode(nodeType); + if (!node) continue; + graph.add(node); + node.pos = [side === "right" ? source.pos[0] + source.size[0] + 30 : source.pos[0] - node.size[0] - 30, source.pos[1]]; + if (side === "right" && source.outputs?.length) source.connect(0, node, 0); + else if (side === "left" && source.inputs?.length) node.connect(0, source, 0); + graph.change(); + } + } else { + const canvas = app.canvas, graph = canvas.graph || app.graph; + const node = LiteGraph.createNode(nodeType); + if (!node) return; + node.pos = [canvas.graph_mouse[0], canvas.graph_mouse[1]]; + graph.add(node); + canvas.selectNode(node, false); + canvas.setDirty(true, true); + } +} +// Temporary map for paste rename coordination between Set and Get nodes. +// Key: old name, Value: new name. Cleared via setTimeout(0) after each paste cycle. +// This works because both onConfigure calls (Set + Get) fire synchronously within +// the same paste operation, before the timeout clears the entry. +const _pasteRenameMap = new Map(); + +let _forceShowAllLinks = false; + +function getShowLinksMode() { + if (_forceShowAllLinks) return "always"; + return app.ui.settings.getSettingValue("KJNodes.showSetGetLinks") ?? "never"; +} +const LGraphNode = LiteGraph.LGraphNode + +// Cross-graph traversal utilities with lexical scoping for subgraph support. +// Set nodes propagate downward: a Set in a parent graph is visible to all descendant subgraphs. +// Get nodes look upward: a Get searches its own graph first, then parent, then grandparent, etc. +// Duplicate names are allowed across unrelated (sibling) subgraphs. + +function findRootGraph(graph) { + if (!graph) return null; + return graph.rootGraph || graph; +} + +// Find which SubgraphNode in parentGraph wraps the given subgraph +function findSubgraphNodeFor(parentGraph, innerNode) { + if (!parentGraph?._nodes || !innerNode?.graph) return null; + for (const n of parentGraph._nodes) { + if (n.subgraph && n.subgraph === innerNode.graph) return n; + } + return null; +} + +// Walk from a subgraph up to root, returning [graph, parent, grandparent, ..., root] +function getGraphAncestors(graph) { + if (!graph) return []; + const root = findRootGraph(graph); + if (!root || graph === root) return [root]; + + const chain = [graph]; + const visited = new Set([graph]); + let current = graph; + // Walk up: find which SubgraphNode wraps current, then get its parent graph + while (current !== root) { + let found = false; + // Search root nodes + for (const n of root._nodes) { + if (n.subgraph === current) { + chain.push(root); + current = root; + found = true; + break; + } + } + if (found) break; + // Search sibling subgraphs (for nested subgraphs) + const subgraphs = root._subgraphs || root.subgraphs; + if (subgraphs) { + for (const sg of subgraphs.values()) { + if (sg === current || !sg._nodes) continue; + for (const n of sg._nodes) { + if (n.subgraph === current) { + if (visited.has(sg)) { found = false; break; } + visited.add(sg); + chain.push(sg); + current = sg; + found = true; + break; + } + } + if (found) break; + } + } + if (!found) { + // Can't find parent, add root as fallback + if (!chain.includes(root)) chain.push(root); + break; + } + } + return chain; +} + +// Get all descendant subgraphs of a graph (children, grandchildren, etc.) +function getGraphDescendants(graph, _visited) { + if (!graph?._nodes) return []; + const visited = _visited || new Set(); + if (visited.has(graph)) return []; + visited.add(graph); + const descendants = []; + for (const n of graph._nodes) { + if (n.subgraph && !visited.has(n.subgraph)) { + descendants.push(n.subgraph); + descendants.push(...getGraphDescendants(n.subgraph, visited)); + } + } + return descendants; +} + +// Collect nodes of a type from specific graphs +function collectNodesOfType(graphs, type) { + const results = []; + for (const g of graphs) { + if (!g?._nodes) continue; + for (const node of g._nodes) { + if (node.type === type) results.push({ node, graph: g }); + } + } + return results; +} + +// Find all nodes of type across ALL graphs (root + all subgraphs). Used for global operations. +function findAllNodesOfType(graph, type) { + const root = findRootGraph(graph); + if (!root) return []; + const allGraphs = [root]; + const subgraphs = root._subgraphs || root.subgraphs; + if (subgraphs) { + for (const sg of subgraphs.values()) allGraphs.push(sg); + } + return collectNodesOfType(allGraphs, type); +} + +// Scoped setter lookup: search current graph, then ancestors (look up). +function findSetterByName(graph, name) { + if (!name) return null; + for (const g of getGraphAncestors(graph)) { + if (!g?._nodes) continue; + for (const node of g._nodes) { + if (node.type === 'SetNode' && node.widgets[0].value === name) { + return { node, graph: g }; + } + } + } + return null; +} + +// Scoped getter lookup: search current graph + descendants (propagate down). +function findGettersByName(graph, name) { + if (!name) return []; + const graphs = [graph, ...getGraphDescendants(graph)]; + return collectNodesOfType(graphs, 'GetNode') + .filter(entry => entry.node.widgets[0].value === name); +} + +// Get all visible SetNode names for a GetNode's combo dropdown. +// Shows names from current graph + ancestors (what's in scope). +// Also builds a source map for getOptionLabel. This module-level variable is safe +// because the values getter and getOptionLabel run in the same synchronous render pass. +let _setNameSourceMap = new Map(); + +function getVisibleSetNames(graph, filterType) { + const sourceMap = new Map(); + const ancestors = getGraphAncestors(graph); + const entries = collectNodesOfType(ancestors, 'SetNode'); + for (const e of entries) { + const name = e.node.widgets[0].value; + if (!name) continue; + if (filterType && filterType !== '*') { + const setType = e.node.inputs[0]?.type; + if (setType && setType !== '*') { + const filterTypes = String(filterType).split(","); + if (!filterTypes.some(ft => ft === setType || setType.split(",").includes(ft))) continue; + } + } + if (!sourceMap.has(name)) { + sourceMap.set(name, e.graph === graph ? "local" : "parent"); + } + } + _setNameSourceMap = sourceMap; + return [...sourceMap.keys()].sort(); +} + +// Force every GetNode (root + subgraphs) to swap its widget.options reference so Vue re-reads values. +function refreshAllGetNodeCombos(graph) { + const root = findRootGraph(graph); + if (!root) return; + const allGraphs = [root]; + const subgraphs = root._subgraphs || root.subgraphs; + if (subgraphs) { + for (const sg of subgraphs.values()) allGraphs.push(sg); + } + for (const g of allGraphs) { + if (!g?._nodes) continue; + for (const node of g._nodes) { + if (node.type === 'GetNode') node._refreshComboOptions?.(); + } + } +} + +// Exposed globally for use in contextmenu.js +window.kjNodes = window.kjNodes || {}; +window.kjNodes.convertOutputsToSetGet = convertOutputsToSetGet; +window.kjNodes.snapshotSelectedNodes = snapshotSelectedNodes; +function convertOutputsToSetGet(graph, node) { + if (!graph || !node) return; + for (let slotIdx = 0; slotIdx < node.outputs.length; slotIdx++) { + const output = node.outputs[slotIdx]; + if (!output.links || output.links.length === 0) continue; + + const linkType = output.type || "*"; + const linkName = output.name || linkType; + + // Collect targets and link IDs to remove, skipping existing Set/Get nodes + const targets = []; + const linksToRemove = []; + for (const linkId of [...output.links]) { + const link = getLink(graph, linkId); + if (link) { + const targetNode = graph.getNodeById(link.target_id); + if (targetNode && targetNode.type !== 'SetNode' && targetNode.type !== 'GetNode') { + targets.push({ targetId: link.target_id, targetSlot: link.target_slot }); + linksToRemove.push(linkId); + } + } + } + if (targets.length === 0) continue; + + // Create Set node + const setNode = LiteGraph.createNode("SetNode"); + if (!setNode) continue; + setNode.pos = [ + node.pos[0] + node.size[0] + 30, + node.pos[1] + slotIdx * 60 + ]; + graph.add(setNode); + setNode.flags.collapsed = true; + + for (const linkId of linksToRemove) graph.removeLink(linkId); + + // Connect source → Set node input + node.connect(slotIdx, setNode, 0); + + // Set the name widget + setNode.widgets[0].value = linkName; + setNode.title = prefixedTitle("Set", linkName); + setNode.validateName(graph); + const finalName = setNode.widgets[0].value; + setNode.properties.previousName = finalName; + + // Create a Get node for each target + for (const target of targets) { + const targetNode = graph.getNodeById(target.targetId); + if (!targetNode) continue; + + const getNode = LiteGraph.createNode("GetNode"); + if (!getNode) continue; + getNode.pos = [ + targetNode.pos[0] - getNode.size[0] - 30, + targetNode.pos[1] + ]; + graph.add(getNode); + getNode.flags.collapsed = true; + + getNode.widgets[0].value = finalName; + getNode.onRename(); + + getNode.connect(0, targetNode, target.targetSlot); + } + } + app.canvas.setDirty(true, true); +} + +// Snapshot selection immediately (before right-click changes it) +function snapshotSelectedNodes(node, typeFilter) { + const selected = Object.values(app.canvas.selected_nodes || {}); + // Always include the right-clicked node, even if selection state is inconsistent + const nodeSet = new Set(selected); + nodeSet.add(node); + let nodes = [...nodeSet]; + if (typeFilter) nodes = nodes.filter(n => n.type === typeFilter); + return nodes; +} + +// Tracks nodes flagged with has_errors so they can be cleared on the next alert. +// Always cleared before repopulating, so it only holds nodes from the most recent alert. +const _errorNodes = new Set(); +function showAlert(message, nodes) { + const nodeList = nodes ? (Array.isArray(nodes) ? nodes : [nodes]) : []; + const nodeInfo = nodeList.map(n => + `${n.title || n.type} [${Math.round(n.pos[0])}, ${Math.round(n.pos[1])}]` + ).join(', '); + if (nodeList.length) { + // Clear previous error flags + for (const n of _errorNodes) n.has_errors = false; + _errorNodes.clear(); + setTimeout(() => { + for (const n of nodeList) { + n.has_errors = true; + _errorNodes.add(n); + } + app.canvas?.setDirty(true, true); + }, 100); + } + app.extensionManager.toast.add({ + severity: 'warn', + summary: "KJ Set/Get", + detail: nodeInfo ? `${message} — ${nodeInfo}` : message, + life: 5000, + }); +} +function convertAllSetGetToLinks(graph) { + if (!graph) return; + const rootGraph = findRootGraph(graph); + + // First pass: handle cross-graph pairs (must run before same-graph + // because converting cross-graph creates real links that replace + // the Set/Get, and same-graph pass would miss the removed nodes) + const allSetEntries = findAllNodesOfType(rootGraph, 'SetNode'); + for (const { node: setNode, graph: setGraph } of [...allSetEntries]) { + const name = setNode.widgets[0].value; + if (!name) continue; + const allGetEntries = findGettersByName(rootGraph, name); + const crossGraphGetters = allGetEntries.filter(e => e.graph !== setGraph); + if (crossGraphGetters.length > 0) { + convertCrossGraphSetGet(setNode, setGraph, crossGraphGetters); + } + } + + // Second pass: handle remaining same-graph pairs + for (const { node: setNode, graph: g } of [...findAllNodesOfType(rootGraph, 'SetNode')]) { + convertSetGetToLinks(g, setNode); + } +} + +// Compat shim: older frontends may lack graph.getLink(), and _links may be a Map or +// plain object depending on the litegraph version. The == null check catches both null +// and undefined intentionally. +function getLink(graph, linkId) { + if (linkId == null) return null; + if (graph.getLink) return graph.getLink(linkId); + return graph._links instanceof Map ? graph._links.get(linkId) : graph._links?.[linkId] ?? null; +} + +// Collect {targetId, targetSlot} from an output slot's links +function collectOutputConnections(graph, output) { + const connections = []; + if (output?.links) { + for (const linkId of [...output.links]) { + const link = getLink(graph, linkId); + if (link) connections.push({ targetId: link.target_id, targetSlot: link.target_slot }); + } + } + return connections; +} + +function convertCrossGraphSetGet(setNode, setGraph, crossGraphGetters) { + // Find what's connected to the SetNode's input + const setInput = setNode.inputs[0]; + if (setInput.link == null) return; + const sourceLink = getLink(setGraph, setInput.link); + if (!sourceLink) return; + const sourceNode = setGraph.getNodeById(sourceLink.origin_id); + if (!sourceNode) return; + const sourceSlot = sourceLink.origin_slot; + const linkType = setNode.inputs[0].type || '*'; + + for (const { node: getter, graph: getterGraph } of crossGraphGetters) { + // Collect getter's downstream connections before removing it + const connections = collectOutputConnections(getterGraph, getter.outputs[0]); + if (connections.length === 0) { + getterGraph.remove(getter); + continue; + } + + // Determine direction and create SubgraphInput or SubgraphOutput + const rootGraph = findRootGraph(setGraph); + const sgNodeForGetter = findSubgraphNodeFor(rootGraph, getter); + const sgNodeForSetter = findSubgraphNodeFor(rootGraph, setNode); + + if (sgNodeForGetter && setGraph === rootGraph) { + // Set in root, Get in subgraph → create SubgraphInput + const subgraph = sgNodeForGetter.subgraph; + const inputName = setNode.widgets[0].value || linkType; + const newInput = subgraph.addInput(inputName, linkType); + const inputIndex = subgraph.inputs.indexOf(newInput); + + // Connect source → SubgraphNode's new input in root graph + sourceNode.connect(sourceSlot, sgNodeForGetter, inputIndex); + + // Inside subgraph: connect SubgraphInput slot → getter's targets + for (const conn of connections) { + if (conn.targetId === subgraph.outputNode?.id) { + // GetNode fed a SubgraphOutput — bypass it and connect source directly in parent graph. + const sgNodeOutput = sgNodeForGetter.outputs[conn.targetSlot]; + if (sgNodeOutput?.links) { + for (const parentLinkId of [...sgNodeOutput.links]) { + const parentLink = getLink(rootGraph, parentLinkId); + if (parentLink) { + const parentTarget = rootGraph.getNodeById(parentLink.target_id); + if (parentTarget) { + sourceNode.connect(sourceSlot, parentTarget, parentLink.target_slot); + } + } + } + } + // Remove the now-unnecessary SubgraphOutput + const sgOutput = subgraph.outputs[conn.targetSlot]; + if (sgOutput) { + subgraph.removeOutput(sgOutput); + } + } else { + const targetNode = getterGraph.getNodeById(conn.targetId); + if (targetNode) { + newInput.connect(targetNode.inputs[conn.targetSlot], targetNode); + } + } + } + + // If no connections used the SubgraphInput, remove it + if (newInput.linkIds.length === 0) { + subgraph.removeInput(newInput); + } + + getterGraph.remove(getter); + + } else if (sgNodeForSetter && getterGraph === rootGraph) { + // Set in subgraph, Get in root → create SubgraphOutput + const subgraph = sgNodeForSetter.subgraph; + const outputName = setNode.widgets[0].value || linkType; + const newOutput = subgraph.addOutput(outputName, linkType); + const outputIndex = subgraph.outputs.indexOf(newOutput); + + // Inside subgraph: connect source output → SubgraphOutput slot + newOutput.connect(sourceNode.outputs[sourceSlot], sourceNode); + + // In root graph: connect SubgraphNode's new output → getter's targets + for (const conn of connections) { + const targetNode = getterGraph.getNodeById(conn.targetId); + if (targetNode) { + sgNodeForSetter.connect(outputIndex, targetNode, conn.targetSlot); + } + } + + getterGraph.remove(getter); + + } else { + // Both in different subgraphs (sibling) — would need both input and output + // Skip for now, too complex + console.warn(`KJNodes: Cannot convert cross-graph Set/Get "${setNode.widgets[0].value}" between sibling subgraphs`); + } + } + + // Remove SetNode if all its getters have been handled (check if any same-graph getters remain) + const remainingGetters = setGraph._nodes.filter( + n => n.type === 'GetNode' && n.widgets[0].value === setNode.widgets[0].value + ); + if (remainingGetters.length === 0) { + setGraph.remove(setNode); + } +} + +function convertSetGetToLinks(graph, setNode) { + if (!graph || !setNode) return; + + const name = setNode.widgets[0].value; + let sameGraphGetters = []; + if (name) { + const allGetters = findGettersByName(graph, name); + const crossGetters = []; + for (const entry of allGetters) { + if (entry.graph === graph) sameGraphGetters.push(entry.node); + else crossGetters.push(entry); + } + if (crossGetters.length > 0) { + convertCrossGraphSetGet(setNode, graph, crossGetters); + } + } + + // Find the source connected to the Set node's input + const setInput = setNode.inputs[0]; + if (setInput.link == null) return; + const sourceLink = getLink(graph, setInput.link); + if (!sourceLink) return; + const sourceNode = graph.getNodeById(sourceLink.origin_id); + if (!sourceNode) return; + const sourceSlot = sourceLink.origin_slot; + + // Collect all consumer connections from Get nodes and SetNode's own output passthrough + const connections = []; + for (const getter of sameGraphGetters) { + connections.push(...collectOutputConnections(graph, getter.outputs[0])); + } + connections.push(...collectOutputConnections(graph, setNode.outputs[0])); + + // Remove all Get nodes (this also removes their links) + for (const getter of sameGraphGetters) { + graph.remove(getter); + } + // Remove the Set node + graph.remove(setNode); + + // Create direct links from source to each consumer + for (const conn of connections) { + const targetNode = graph.getNodeById(conn.targetId); + if (targetNode) { + sourceNode.connect(sourceSlot, targetNode, conn.targetSlot); + } + } + + app.canvas?.setDirty(true, true); +} + +// region SetNode +app.registerExtension({ + name: "KJNodes.SetNode", + registerCustomNodes() { + class SetNode extends LGraphNode { + static title = "Set"; + static category = "MXD/Utilities"; + serialize_widgets = true; + drawConnection = false; + slotColor = "#FFF"; + canvas = app.canvas; + + constructor(title) { + super(title) + if (!this.properties) { + this.properties = { + "previousName": "" + }; + } + this.properties["Node name for S&R"] = "SetNode"; + this.properties["aux_id"] = "Maxed-Out-99/ComfyUI-MaxedOut"; + this.isVirtualNode = true; // This node is purely frontend and does not impact the resulting prompt so should not be serialized + + this.addWidget( + "text", + "Constant", + '', + () => { + if (!this.graph || app.configuringGraph) return; + this.validateName(this.graph); + if (this.widgets[0].value !== '') { + this.title = prefixedTitle("Set", this.widgets[0].value); + } + this.update(); + this.properties.previousName = this.widgets[0].value; + }, + {} + ) + + this.addInput("*", "*"); + this.addOutput("*", '*'); + } + + onConnectionsChange( + slotType, + slot, + isChangeConnect, + link_info + ) { + // During graph load, slots are restored by configure — skip side effects + if (app.configuringGraph) return; + //On Disconnect + if (slotType === LiteGraph.INPUT && !isChangeConnect) { + const outputConnected = this.outputs[0]?.links?.length > 0; + if (outputConnected) { + this.inputs[slot].type = this.outputs[0].type; + this.inputs[slot].name = this.outputs[0].name; + } else { + this.inputs[slot].type = '*'; + this.inputs[slot].name = '*'; + this.outputs[0].type = '*'; + this.outputs[0].name = '*'; + this.title = "Set"; + this.color = null; + this.bgcolor = null; + } + this.update(); + } + if (slotType === LiteGraph.OUTPUT && !isChangeConnect) { + if (this.outputs && this.outputs[slot]) { + // Keep type if input has a real connection + const inputConnected = this.inputs[0]?.link != null; + if (inputConnected) { + this.outputs[slot].type = this.inputs[0].type; + this.outputs[slot].name = this.inputs[0].name; + } else { + this.inputs[0].type = '*'; + this.inputs[0].name = '*'; + this.outputs[slot].type = '*'; + this.outputs[slot].name = '*'; + this.color = null; + this.bgcolor = null; + } + } + } + //On Connect + if (link_info && this.graph && slotType === LiteGraph.INPUT && isChangeConnect) { + const resolve = link_info.resolve(this.graph) + const resolvedSlot = resolve?.subgraphInput ?? resolve?.output; + const type = resolvedSlot?.type; + if (type) { + if (this.title === "Set"){ + this.title = prefixedTitle("Set", type); + } + if (this.widgets[0].value === '' || this.widgets[0].value === '*'){ + // Determine the initial widget value based on naming setting + const namingMode = app.ui.settings.getSettingValue("KJNodes.setGetNaming") ?? "empty"; + if (namingMode !== "empty") { + const link = getLink(this.graph, this.inputs[0]?.link); + const sourceNode = link ? this.graph.getNodeById(link.origin_id) : null; + const slotName = sourceNode?.outputs?.[link?.origin_slot]?.name || type; + switch (namingMode) { + case "slot name": this.widgets[0].value = slotName; break; + case "slot name (lowercase)": this.widgets[0].value = slotName.toLowerCase(); break; + case "slot name (UPPERCASE)": this.widgets[0].value = slotName.toUpperCase(); break; + } + } + } + + this.validateName(this.graph); + this.properties.previousName = this.widgets[0].value; + this.inputs[0].type = type; + this.inputs[0].name = type; + this.outputs[0].type = type; + this.outputs[0].name = type; + + autoColor(this, type); + } else { + showAlert(`node ${this.title} input undefined.`, this) + } + } + if (link_info && this.graph && slotType === LiteGraph.OUTPUT && isChangeConnect) { + const inputType = this.inputs[0]?.type; + if (inputType && inputType !== '*') { + this.outputs[0].type = inputType; + this.outputs[0].name = inputType; + } else { + const resolve = link_info.resolve(this.graph); + const type = resolve?.input?.type; + if (type && type !== '*') { + this.inputs[0].type = type; + this.inputs[0].name = type; + this.outputs[0].type = type; + this.outputs[0].name = type; + autoColor(this, type); + } + } + } + + //Update either way + this.update(); + } + + // Returns true if the name was changed + validateName(graph, sameGraphOnly) { + let widgetValue = this.widgets[0].value; + + if (widgetValue !== '') { + let tries = 0; + const existingValues = new Set(); + + // sameGraphOnly: only check the immediate graph (for paste/clone) + // otherwise check full scope (own graph + ancestors) + const scopeGraphs = sameGraphOnly ? [graph] : getGraphAncestors(graph); + const scopedSetNodes = collectNodesOfType(scopeGraphs, 'SetNode'); + scopedSetNodes.forEach(entry => { + if (entry.node !== this) { + existingValues.add(entry.node.widgets[0].value); + } + }); + + const originalValue = widgetValue; + // Only strip _N suffix during paste to avoid FOO_0_1_2 accumulation. + // For manual renames, keep the full name as base (user may intend FOO_3). + const baseName = this._justAdded ? widgetValue.replace(/_\d+$/, '') : widgetValue; + while (existingValues.has(widgetValue)) { + widgetValue = baseName + "_" + tries; + tries++; + } + + this.widgets[0].value = widgetValue; + this.title = prefixedTitle("Set", widgetValue); + return widgetValue !== originalValue; + } + return false; + } + + clone() { + const cloned = super.clone(); + cloned.inputs[0].name = '*'; + cloned.inputs[0].type = '*'; + cloned.properties.previousName = ''; + cloned.size = cloned.computeSize(); + return cloned; + } + + onAdded() { + this._justAdded = true; + // Vue mode doesn't always re-extract GetNode options on SetNode add; skip during workflow load. + if (LiteGraph.vueNodesMode && this.graph && !app.configuringGraph) { + refreshAllGetNodeCombos(this.graph); + } + } + + onRemoved() { + // Only Vue mode needs the refresh — legacy re-reads the values getter on every click. + if (!LiteGraph.vueNodesMode) return; + const g = this.graph; + if (!g) return; + // Defer: onRemoved fires before _nodes is spliced, so getters still see this SetNode. + setTimeout(() => refreshAllGetNodeCombos(g), 0); + } + + onConfigure() { + // Only run paste logic when actually pasting, not during workflow load + if (this._justAdded && this.graph && !app.configuringGraph) { + const oldName = this.widgets[0].value; + this.validateName(this.graph, true); + this._justAdded = false; + const newName = this.widgets[0].value; + if (newName !== oldName) { + _pasteRenameMap.set(oldName, newName); + // Clear the map after this paste cycle + setTimeout(() => _pasteRenameMap.delete(oldName), 0); + } + // Reset type and color on paste — nothing is connected yet + if (this.inputs[0]?.link == null) { + this.inputs[0].type = '*'; + this.inputs[0].name = '*'; + this.outputs[0].type = '*'; + this.outputs[0].name = '*'; + this.color = null; + this.bgcolor = null; + } + } + this._justAdded = false; + } + + update() { + if (!this.graph) return; + + const getters = this.findGetters(this.graph); + getters.forEach(getter => { + getter.setType(this.inputs[0].type); + }); + + if (this.widgets[0].value && this.properties.previousName) { + const gettersWithPreviousName = this.findGetters(this.graph, true); + gettersWithPreviousName.forEach(getter => { + getter.setName(this.widgets[0].value); + }); + } + app.canvas?.setDirty(true, true); + } + + findGetters(graph, checkForPreviousName) { + const name = checkForPreviousName ? this.properties.previousName : this.widgets[0].value; + if (!name || name === '') return []; + // Scoped: searches own graph + descendant subgraphs + return findGettersByName(graph, name).map(entry => entry.node); + } + + getExtraMenuOptions(_, options) { + options.unshift( + { + content: this.drawConnection ? "Hide connections" : "Show connections", + callback: () => { + this.drawConnection = !this.drawConnection; + const linkType = this.inputs[0].type; + this.slotColor = this.canvas.default_connection_color_byType[linkType]; + this.canvas.setDirty(true, true); + }, + }, + { + content: "Hide all connections", + callback: () => { + for (const n of this.graph._nodes) { + if (n.type === "GetNode" || n.type === "SetNode") n.drawConnection = false; + } + this.canvas.setDirty(true, true); + }, + + }, + ); + options.unshift({ + content: "Convert to links", + callback: () => { + const graph = this.graph; + for (const n of snapshotSelectedNodes(this, 'SetNode')) convertSetGetToLinks(graph, n); + }, + }); + options.unshift({ + content: "Add paired GetNode", + callback: () => { + const graph = this.graph; + const getNode = LiteGraph.createNode("GetNode"); + if (!getNode) return; + getNode.pos = [this.pos[0] + this.size[0] + 30, this.pos[1]]; + graph.add(getNode); + // Set the widget value to match — this drives type, color, and connection + const name = this.widgets[0].value; + if (getNode.widgets?.[0]) { + getNode.widgets[0].value = name; + } + getNode.onRename?.(); + app.canvas.selectNode(getNode, false); + app.canvas.setDirty(true, true); + }, + }); + // Dynamically add a submenu for all getters + const allGettersForMenu = this.findGetters(this.graph); + if (allGettersForMenu.length) { + let gettersSubmenu = allGettersForMenu.map(getter => { + const sameGraph = getter.graph === this.graph; + const sgNode = !sameGraph ? findSubgraphNodeFor(this.graph, getter) : null; + const label = sameGraph + ? `${getter.title} id: ${getter.id}` + : `${getter.title} (in subgraph${sgNode ? ': ' + (sgNode.title || sgNode.type) : ''})`; + return { + content: label, + callback: () => { + if (sameGraph) { + this.canvas.centerOnNode(getter); + this.canvas.selectNode(getter, false); + } else if (sgNode) { + this.canvas.openSubgraph(sgNode.subgraph, sgNode); + // Center on the getter after entering the subgraph + setTimeout(() => { + this.canvas.centerOnNode(getter); + this.canvas.selectNode(getter, false); + }, 0); + } + this.canvas.setDirty(true, true); + }, + }; + }); + + options.unshift({ + content: "Getters", + has_submenu: true, + submenu: { + title: "GetNodes", + options: gettersSubmenu, + } + }); + } + } + } + + LiteGraph.registerNodeType("SetNode", SetNode); + }, +}); + +// region GetNode +app.registerExtension({ + name: "KJNodes.GetNode", + registerCustomNodes() { + class GetNode extends LGraphNode { + static title = "Get"; + static category = "MXD/Utilities"; + serialize_widgets = true; + drawConnection = false; + currentSetter = null; + canvas = app.canvas; + + constructor(title) { + super(title) + if (!this.properties) { + this.properties = {}; + } + this.properties["Node name for S&R"] = "GetNode"; + this.properties["aux_id"] = "Maxed-Out-99/ComfyUI-MaxedOut"; + this.isVirtualNode = true; // This node is purely frontend and does not impact the resulting prompt so should not be serialized + const comboOptions = { + getOptionLabel: (value) => { + if (!value) return ""; + const source = _setNameSourceMap.get(value); + if (!source || source === "local") return value; + return `${value} (${source})`; + }, + }; + Object.defineProperty(comboOptions, 'values', { + get: () => { + if (!this.graph) return []; + let filterType = null; + if (app.ui.settings.getSettingValue("KJNodes.filterGetNodeOptions") !== false + && this.outputs[0]?.links?.length) { + const linkId = this.outputs[0].links[0]; + const link = getLink(this.graph, linkId); + if (link) { + const targetNode = this.graph.getNodeById(link.target_id); + filterType = targetNode?.inputs?.[link.target_slot]?.type || null; + } + } + return getVisibleSetNames(this.graph, filterType); + }, + enumerable: true, + configurable: true + }); + const constantWidget = this.addWidget("combo", "Constant", "", () => { if (!app.configuringGraph) this.onRename(); }, comboOptions); + // Legacy-only fix: pass pre-labeled values to ContextMenu so the core filter shows the search input. + // Sidesteps ComboWidget's empty-array + getOptionLabel bug. Vue mode handles it natively. + const origOnClick = constantWidget.onClick?.bind(constantWidget); + constantWidget.onClick = (params) => { + if (LiteGraph.vueNodesMode) return origOnClick?.(params); + const { e, canvas, node } = params; + const x = e.canvasX - node.pos[0]; + const width = constantWidget.width || node.size[0]; + if (x < 40) return constantWidget.decrementValue({ e, node, canvas }); + if (x > width - 40) return constantWidget.incrementValue({ e, node, canvas }); + const rawValues = comboOptions.values; + const labels = rawValues.map(v => comboOptions.getOptionLabel(v) || v); + const menu = new LiteGraph.ContextMenu(labels, { + scale: Math.max(1, canvas.ds.scale), + event: e, + className: 'dark', + callback: (selectedLabel) => { + const idx = labels.indexOf(selectedLabel); + if (idx >= 0) constantWidget.setValue(rawValues[idx], { e, node, canvas }); + } + }); + // Color each entry's left border by its SetNode's type. + if (!_typeColorMap) setColorAndBgColor({}, ''); // ensure lazy map is initialized + const entries = menu.root?.querySelectorAll('.litemenu-entry'); + rawValues.forEach((name, i) => { + if (!entries?.[i]) return; + const setter = findSetterByName(this.graph, name); + const type = setter?.node?.inputs?.[0]?.type; + const c = _typeColorMap?.[type]; + const borderColor = canvas.default_connection_color_byType?.[type] + || LGraphCanvas.link_type_colors?.[type] + || (c && (c.groupcolor || c.bgcolor || c.color)) + || '#888'; + entries[i].style.borderLeft = `4px solid ${borderColor}`; + entries[i].style.paddingLeft = '8px'; + }); + }; + // Fresh options object (live getter preserved) + remove/re-add to force Vue re-extraction. + this._refreshComboOptions = () => { + const w = this.widgets?.[0]; + if (!w) return; + const newOpts = { getOptionLabel: comboOptions.getOptionLabel }; + Object.defineProperty(newOpts, 'values', + Object.getOwnPropertyDescriptor(comboOptions, 'values')); + w.options = newOpts; + const idx = this.widgets.indexOf(w); + if (idx >= 0) { + this.widgets.splice(idx, 1); + this.widgets.splice(idx, 0, w); + } + }; + this.addOutput("*", '*'); + } + + onConnectionsChange() { + if (app.configuringGraph) return; + this.validateLinks(); + } + + setName(name) { + this.widgets[0].value = name; + this.onRename(); + this.serialize(); + } + + onRename() { + const setter = this.findSetter(this.graph); + if (setter) { + this.setType(setter.inputs[0].type); + this.title = prefixedTitle("Get", setter.widgets[0].value); + } else { + this.setType('*'); + const name = this.widgets[0].value; + this.title = name ? prefixedTitle("Get", name) : "Get"; + } + app.canvas?.setDirty(true, true); + } + + clone() { + const cloned = super.clone(); + cloned.size = cloned.computeSize(); + return cloned; + } + + validateLinks() { + if (this.outputs[0].type !== '*' && this.outputs[0].links && this.graph) { + this.outputs[0].links.filter(linkId => { + const link = getLink(this.graph, linkId); + if (!link || !link.type) return false; + if (link.type === '*') return false; + const targetNode = this.graph.getNodeById(link.target_id); + const targetType = targetNode?.inputs?.[link.target_slot]?.type; + if (targetType === '*') return false; + if (targetType) { + const targetTypes = String(targetType).split(","); + if (targetTypes.includes(this.outputs[0].type)) return false; + } + return !link.type.split(",").includes(this.outputs[0].type); + }).forEach(linkId => { + this.graph.removeLink(linkId); + }); + } + } + + setType(type) { + this.outputs[0].name = type; + this.outputs[0].type = type; + this.validateLinks(); + autoColor(this, type); + } + + findSetter(graph) { + const name = this.widgets[0].value; + // Scoped: searches own graph first, then ancestors + const result = findSetterByName(graph, name); + return result ? result.node : undefined; + } + + goToSetter() { + if (!this.currentSetter) return; + const setterGraph = this.currentSetter.graph; + if (setterGraph && setterGraph !== this.graph) { + this.canvas.setGraph(setterGraph); + setTimeout(() => { + this.canvas.centerOnNode(this.currentSetter); + this.canvas.selectNode(this.currentSetter, false); + this.canvas.setDirty(true, true); + }, 0); + } else { + this.canvas.centerOnNode(this.currentSetter); + this.canvas.selectNode(this.currentSetter, false); + } + } + + onAdded() { + this._justAdded = true; + } + + onConfigure() { + if (this._justAdded && !app.configuringGraph) { + const name = this.widgets[0].value; + if (name) { + // Check if our paired SetNode was renamed during this paste + const newName = _pasteRenameMap.get(name); + if (newName) { + this.widgets[0].value = newName; + } + // Restore type/color from setter after paste + setTimeout(() => this.onRename(), 0); + } + } + this._justAdded = false; + } + + getInputLink(slot) { + // Same-graph: standard resolveOutput path. + // Cross-graph: handled by resolveVirtualOutput. + const name = this.widgets[0].value; + if (!name || name === '') return null; + const setter = this.graph?._nodes?.find( + n => n.type === 'SetNode' && n.widgets[0].value === name + ); + if (setter) { + const slotInfo = setter.inputs[slot]; + if (!slotInfo || slotInfo.link == null) return null; + return getLink(this.graph, slotInfo.link); + } + // Scoped lookup: no setter found in own graph or ancestors + if (name && !findSetterByName(this.graph, name)) { + showAlert("No SetNode found for " + name + "(" + this.type + ")", this); + } + return null; + } + + resolveVirtualOutput(slot) { + const name = this.widgets[0].value; + // Scoped lookup: own graph first, then ancestors + const result = findSetterByName(this.graph, name); + if (!result) return undefined; + + // Same graph — let the standard getInputLink path handle it + if (result.graph === this.graph) return undefined; + + // Warn if multiple SetNodes with this name exist in scope + const scopeGraphs = getGraphAncestors(this.graph); + const scopedSetters = collectNodesOfType(scopeGraphs, 'SetNode') + .filter(e => e.node.widgets[0].value === name); + if (scopedSetters.length > 1) { + showAlert(`Multiple SetNodes named "${name}" found in scope. Rename duplicates or use "Convert to links" to resolve`, [this, ...scopedSetters.map(e => e.node)]); + return undefined; + } + + const { node: setter, graph: setterGraph } = result; + const slotInfo = setter.inputs[slot]; + if (!slotInfo || slotInfo.link == null) return undefined; + + const link = getLink(setterGraph, slotInfo.link); + if (!link) return undefined; + + const sourceNode = setterGraph.getNodeById(link.origin_id); + if (!sourceNode) return undefined; + + return { node: sourceNode, slot: link.origin_slot }; + } + getExtraMenuOptions(_, options) { + this.currentSetter = this.findSetter(this.graph) + if (!this.currentSetter) return + const sameGraph = this.currentSetter.graph === this.graph; + if (sameGraph) { + let menuEntry = this.drawConnection ? "Hide connections" : "Show connections"; + options.unshift( + { + content: "Convert to links", + callback: () => { + const graph = this.graph; + const setters = new Set(snapshotSelectedNodes(this, 'GetNode') + .map(n => n.findSetter?.(graph)).filter(Boolean)); + for (const s of setters) convertSetGetToLinks(graph, s); + }, + }, + { + content: "Go to setter", + callback: () => { + this.goToSetter(); + }, + }, + { + content: menuEntry, + callback: () => { + if (!this.currentSetter) return; + const linkType = this.currentSetter.inputs[0].type; + // Toggle on the SetNode so the canvas-level hook draws it + this.currentSetter.drawConnection = !this.currentSetter.drawConnection; + this.currentSetter.slotColor = this.canvas.default_connection_color_byType[linkType]; + this.drawConnection = this.currentSetter.drawConnection; + this.canvas.setDirty(true, true); + }, + }, + ); + } else { + // Cross-graph setter — navigate to its graph + const setterGraph = this.currentSetter.graph; + const isRoot = setterGraph === findRootGraph(this.graph); + options.unshift( + { + content: `Go to setter (in ${isRoot ? 'parent graph' : 'subgraph'})`, + callback: () => { + const canvas = this.canvas; + canvas.setGraph(setterGraph); + setTimeout(() => { + canvas.centerOnNode(this.currentSetter); + canvas.selectNode(this.currentSetter, false); + canvas.setDirty(true, true); + }, 0); + }, + }, + ); + } + } + } + + LiteGraph.registerNodeType("GetNode", GetNode); + }, +}); + +// region UI: commands, keybindings, canvas menu, draw hook, settings +app.registerExtension({ + name: "KJNodes.SetGetUI", + settings: [ + { + id: "KJNodes.setGetNaming", + name: "Default SetNode widget value", + category: ["KJNodes", "Set & Get", "Default SetNode widget value"], + tooltip: "Initial Constant value when a Set node is first connected to a slot", + type: "combo", + options: ["empty", "slot name", "slot name (lowercase)", "slot name (UPPERCASE)"], + defaultValue: "empty", + }, + { + id: "KJNodes.showSetGetLinks", + name: "Show links", + category: ["KJNodes", "Set & Get", "Show links"], + tooltip: "When to show virtual links between Set/Get pairs", + type: "combo", + options: ["never", "selected", "always"], + defaultValue: "never", + onChange: () => app.canvas?.setDirty(true, true), + }, + { + id: "KJNodes.nodeAutoColor", + name: "Auto-color nodes", + category: ["KJNodes", "Set & Get", "Auto-color nodes"], + tooltip: "Automatically color Set/Get nodes based on their connection type", + type: "boolean", + defaultValue: true, + }, + { + id: "KJNodes.disablePrefix", + name: "Disable Set_/Get_ prefix", + category: ["KJNodes", "Set & Get", "Disable Set_/Get_ prefix"], + tooltip: "Prevents automatically adding Set_ and Get_ prefixes to node titles", + defaultValue: false, + type: "boolean", + }, + { + id: "KJNodes.shiftMiddleClickSetGet", + name: "Shift+middle-click creates Set/Get", + category: ["KJNodes", "Set & Get", "Shift+middle-click creates Set/Get"], + tooltip: "Shift+middle-click on a slot creates a connected SetNode (output) or GetNode (input) instead of Reroute", + type: "boolean", + defaultValue: true, + }, + { + id: "KJNodes.middleClickSetGet", + name: "Middle-click creates Set/Get", + category: ["KJNodes", "Set & Get", "Middle-click creates Set/Get"], + tooltip: "Middle-click on a slot creates Set/Get instead of Reroute (overrides core behavior)", + type: "boolean", + defaultValue: false, + }, + ], + commands: [ + { + id: "KJNodes.AddSetNodeToSelected", + label: "Add Set node to selected / at cursor", + function: () => addNodeToSelectedOrCursor("SetNode", "right"), + }, + { + id: "KJNodes.AddGetNodeAtCursor", + label: "Add Get node to selected / at cursor", + function: () => addNodeToSelectedOrCursor("GetNode", "left"), + }, + { + id: "KJNodes.ToggleForceShowSetGetLinks", + label: "Toggle force-show all Set/Get connections", + function: () => { + _forceShowAllLinks = !_forceShowAllLinks; + app.canvas.setDirty(true, true); + }, + }, + ], + keybindings: [ + { + commandId: "KJNodes.AddSetNodeToSelected", + combo: { key: "s", ctrl: true, shift: true }, + targetElementId: "graph-canvas", + }, + { + commandId: "KJNodes.AddGetNodeAtCursor", + combo: { key: "g", ctrl: true, shift: true }, + targetElementId: "graph-canvas", + }, + { + commandId: "KJNodes.ToggleForceShowSetGetLinks", + combo: { key: "l", ctrl: true, shift: true }, + targetElementId: "graph-canvas", + }, + ], + getCanvasMenuItems() { + return [ + { + content: "KJNodes", + has_submenu: true, + submenu: { + options: [ + { + content: "Convert outputs on all selected nodes to Set/Get", + callback: () => { + const selected = Object.values(app.canvas.selected_nodes || {}); + if (selected.length === 0) return; + for (const n of selected) convertOutputsToSetGet(app.graph, n); + }, + }, + { + content: "Convert selected Set/Get to links", + callback: () => { + const graph = app.graph; + const selected = Object.values(app.canvas.selected_nodes || {}); + const setNodes = selected.filter(n => n.type === 'SetNode'); + const getNodes = selected.filter(n => n.type === 'GetNode'); + for (const n of setNodes) convertSetGetToLinks(graph, n); + const setters = new Set(getNodes.map(n => n.findSetter?.(graph)).filter(Boolean)); + for (const s of setters) convertSetGetToLinks(graph, s); + }, + }, + ], + }, + }, + ]; + }, + setup() { + // Double-click GetNode to jump to its SetNode (works in both legacy and Vue modes) + document.addEventListener("dblclick", (e) => { + // Don't interfere when a ContextMenu (e.g. combo dropdown) is open + if (document.querySelector(".litecontextmenu")) return; + const canvas = app.canvas; + if (!canvas) return; + const selected = Object.values(canvas.selected_nodes || {}); + if (selected.length !== 1) return; + const node = selected[0]; + if (node.type === "GetNode") { + node.currentSetter = node.findSetter(node.graph); + if (node.currentSetter) node.goToSetter(); + } + }); + // Suppress title rename on collapsed GetNodes so double-click navigates instead + document.addEventListener("litegraph:canvas", (e) => { + if (e.detail?.subType !== 'node-double-click') return; + const node = e.detail.node; + if (node?.type === "GetNode" && node.flags?.collapsed) { + e.stopImmediatePropagation(); + } + }); + + // Monkey-patch: no extension hook exists for filtering visible nodes. + // Keep Set/Get nodes with visible virtual connections in the visible set + // even when offscreen, so the link lines don't get clipped. + const originalComputeVisibleNodes = LGraphCanvas.prototype.computeVisibleNodes; + LGraphCanvas.prototype.computeVisibleNodes = function () { + const visible = originalComputeVisibleNodes.apply(this, arguments); + for (const node of this.graph._nodes) { + if ((node.type === "SetNode" || node.type === "GetNode") && node.drawConnection && !visible.includes(node)) { + visible.push(node); + } + } + return visible; + }; + + // Monkey-patch: no extension hook exists for custom canvas drawing. + // Draws virtual links between Set/Get pairs. + // Handles per-node "Show connections", "selected" mode, and "always" mode. + const origCanvasOnDrawBackground = app.canvas.onDrawBackground; + app.canvas.onDrawBackground = function(ctx, visibleArea) { + origCanvasOnDrawBackground?.call(this, ctx, visibleArea); + + const graph = this.graph || app.graph; + if (!graph?._nodes) return; + const mode = getShowLinksMode(); + + // Build set of selected Set/Get names for "selected" mode + let selectedNames = null; + if (mode === "selected") { + const sel = Object.values(this.selected_nodes || {}); + if (!sel.length) { + // No selection — still honor per-node drawConnection + if (!graph._nodes.some(n => n.type === 'SetNode' && n.drawConnection)) return; + } else { + selectedNames = new Set(); + for (const n of sel) { + if (n.type === 'SetNode' || n.type === 'GetNode') { + selectedNames.add(n.widgets?.[0]?.value ?? n.title); + } + } + if (!selectedNames.size && !graph._nodes.some(n => n.type === 'SetNode' && n.drawConnection)) return; + } + } else if (mode === "never") { + if (!graph._nodes.some(n => n.type === 'SetNode' && n.drawConnection)) return; + } + + // Collect all GetNodes once across current graph + descendants, grouped by name. + // This avoids re-discovering descendant graphs and re-scanning nodes per SetNode. + const gettersByName = new Map(); + const scopeGraphs = [graph, ...getGraphDescendants(graph)]; + for (const g of scopeGraphs) { + if (!g?._nodes) continue; + for (const node of g._nodes) { + if (node.type !== 'GetNode') continue; + const n = node.widgets?.[0]?.value; + if (!n) continue; + let list = gettersByName.get(n); + if (!list) { list = []; gettersByName.set(n, list); } + list.push(node); + } + } + + for (const setNode of graph._nodes) { + if (setNode.type !== 'SetNode') continue; + const name = setNode.widgets?.[0]?.value ?? setNode.title; + const showByMode = mode === "always" || (mode === "selected" && selectedNames?.has(name)); + if (!showByMode && !setNode.drawConnection) continue; + + const allGetters = gettersByName.get(name) || []; + const drawTargets = []; + const seenSubgraphs = new Set(); + for (const getter of allGetters) { + if (getter.graph === graph) { + drawTargets.push(getter); + } else { + const sgNode = findSubgraphNodeFor(graph, getter); + if (sgNode && !seenSubgraphs.has(sgNode)) { + seenSubgraphs.add(sgNode); + drawTargets.push(sgNode); + } + } + } + if (!drawTargets.length) continue; + + const linkType = setNode.inputs[0]?.type; + const slotColor = app.canvas.default_connection_color_byType?.[linkType] + || LGraphCanvas.link_type_colors?.[linkType] + || setNode.bgcolor + || (setNode.slotColor !== "#FFF" ? setNode.slotColor : null) || "#AAA"; + + const startPos = setNode.getConnectionPos(false, 0); + + for (const target of drawTargets) { + const endPos = target.getConnectionPos(true, 0); + const highlighted = setNode.is_selected || target.is_selected; + const color = highlighted ? "#FFF" : slotColor; + app.canvas.renderLink( + ctx, + startPos, + endPos, + null, + false, + null, + color, + LiteGraph.RIGHT, + LiteGraph.LEFT + ); + } + } + }; + + app.ui.settings.addSetting({ + id: "KJNodes.filterGetNodeOptions", + name: "Filter Get node options by type", + category: ["KJNodes", "Set & Get", "Filter Get node options by type"], + tooltip: "When a Get node is connected, only show Set nodes with compatible types in the dropdown", + type: "boolean", + defaultValue: true, + }); + app.ui.settings.addSetting({ + id: "KJNodes.convertAllSetGet", + name: "Convert ALL Set/Get to links", + category: ["KJNodes", "Set & Get", "Convert ALL Set/Get to links"], + tooltip: "Replaces all Set/Get node pairs with direct links, including across subgraph boundaries (irreversible)", + type: () => { + const btn = document.createElement("button"); + btn.textContent = "Convert ALL Set/Get to links"; + btn.style.cssText = "padding: 6px 12px; cursor: pointer;"; + btn.onclick = () => { + if (confirm("This will replace ALL Set/Get pairs with direct links. This is irreversible. Continue?")) { + convertAllSetGetToLinks(app.graph); + app.canvas?.setDirty(true, true); + app.extensionManager.toast.add({ + severity: 'info', + summary: "KJ Set/Get", + detail: "All Set/Get nodes converted to direct links", + life: 3000, + }); + } + }; + return btn; + }, + defaultValue: false, + }); + + // Middle-click Set/Get: let the core do slot hit detection, then intercept + // createDefaultNodeForSlot to create Set/Get instead of reroute. + let _fromMiddleClick = false; + let _lastMiddleClickShift = false; + const origProcessMiddle = LGraphCanvas.prototype._processMiddleButton; + LGraphCanvas.prototype._processMiddleButton = function(e, node) { + _fromMiddleClick = true; + _lastMiddleClickShift = !!e.shiftKey; + const wantSetGet = _lastMiddleClickShift + ? app.ui.settings.getSettingValue("KJNodes.shiftMiddleClickSetGet") ?? true + : app.ui.settings.getSettingValue("KJNodes.middleClickSetGet") ?? false; + const saved = LiteGraph.middle_click_slot_add_default_node; + if (wantSetGet) { + LiteGraph.middle_click_slot_add_default_node = true; + } else { + // Workaround: core slotDefaults.init() resets this flag to true + const coreSetting = app.ui.settings.getSettingValue("Comfy.Node.MiddleClickRerouteNode"); + if (coreSetting === false) LiteGraph.middle_click_slot_add_default_node = false; + } + const result = origProcessMiddle.call(this, e, node); + LiteGraph.middle_click_slot_add_default_node = saved; + return result; + }; + const origCreateDefault = LGraphCanvas.prototype.createDefaultNodeForSlot; + LGraphCanvas.prototype.createDefaultNodeForSlot = function(optPass) { + if (!_fromMiddleClick) return origCreateDefault.call(this, optPass); + _fromMiddleClick = false; + const wantSetGet = _lastMiddleClickShift + ? app.ui.settings.getSettingValue("KJNodes.shiftMiddleClickSetGet") ?? true + : app.ui.settings.getSettingValue("KJNodes.middleClickSetGet") ?? false; + if (!wantSetGet) return origCreateDefault.call(this, optPass); + + const isFrom = optPass.nodeFrom && optPass.slotFrom !== null; + const sourceNode = isFrom ? optPass.nodeFrom : optPass.nodeTo; + const slotIndex = isFrom ? optPass.slotFrom : optPass.slotTo; + if (!sourceNode || slotIndex == null) return origCreateDefault.call(this, optPass); + const graph = this.graph || app.graph; + if (!graph) return origCreateDefault.call(this, optPass); + + const newNode = LiteGraph.createNode(isFrom ? "SetNode" : "GetNode"); + if (!newNode) return false; + const bounding = sourceNode.getBounding(); + newNode.pos = [ + isFrom ? bounding[0] + bounding[2] + 30 : bounding[0] - 200, + optPass.position?.[1] ?? bounding[1] + ]; + graph.add(newNode); + if (isFrom) sourceNode.connect(slotIndex, newNode, 0); + else newNode.connect(0, sourceNode, slotIndex); + graph.change(); + this.setDirty(true, true); + return true; + }; + } +}); + +// Temporary backwards-compat patch: makes cross-subgraph Set/Get work on frontends +// that don't natively support resolveVirtualOutput. Remove when no longer needed. +app.registerExtension({ + name: "KJNodes.CrossGraphSetGet", + setup() { + let patched = false; + + const originalGraphToPrompt = app.graphToPrompt.bind(app); + app.graphToPrompt = async function(...args) { + if (!patched) { + try { + const subgraphNode = app.graph._nodes.find(n => typeof n.getInnerNodes === 'function'); + if (subgraphNode) { + const tempMap = new Map(); + const dtos = subgraphNode.getInnerNodes(tempMap, []); + if (dtos.length > 0) { + const proto = Object.getPrototypeOf(dtos[0]); + const DtoClass = proto.constructor; + const nativeSource = proto.resolveOutput.toString(); + const hasNativeSupport = nativeSource.includes('resolveVirtualOutput'); + console.log(`[KJNodes] Cross-graph Set/Get: frontend native support ${hasNativeSupport ? 'detected, skipping patch' : 'not found, applying patch'}`); + if (!hasNativeSupport) { + const origResolveOutput = proto.resolveOutput; + proto.resolveOutput = function(slot, type, visited) { + if (typeof this.node?.resolveVirtualOutput === 'function') { + const virtualSource = this.node.resolveVirtualOutput(slot); + if (virtualSource) { + const inputNodeDto = [...this.nodesByExecutionId.values()] + .find(dto => dto instanceof DtoClass && dto.node === virtualSource.node); + if (inputNodeDto) { + return inputNodeDto.resolveOutput(virtualSource.slot, type, visited); + } + throw new Error(`KJNodes: No DTO found for cross-graph source node [${virtualSource.node.id}]`); + } + } + return origResolveOutput.call(this, slot, type, visited); + }; + } + patched = true; + } + } + } catch (e) { + console.warn('[KJNodes] Failed to probe ExecutableNodeDTO for cross-graph patch:', e); + } + } + return originalGraphToPrompt(...args); + }; + } +}); + +// Monkey-patch: no hook exists to extend the link context menu. +// Appends a "Convert to Set/Get" entry by counting DOM menus before/after +// the original call to find the newly created one. Fragile but unavoidable. +app.registerExtension({ + name: "KJNodes.LinkToSetGet", + setup() { + const originalShowLinkMenu = LGraphCanvas.prototype.showLinkMenu; + LGraphCanvas.prototype.showLinkMenu = function(segment, e) { + const menusBefore = document.querySelectorAll(".litecontextmenu").length; + const result = originalShowLinkMenu.call(this, segment, e); + + const graph = this.graph; + if (!graph) return result; + + // Get the full link object to access target info + const link = getLink(graph, segment.id); + if (!link || link.origin_id == null || link.target_id == null) return result; + + // Find the newly created context menu + const menus = document.querySelectorAll(".litecontextmenu"); + if (menus.length <= menusBefore) return result; + const lastMenu = menus[menus.length - 1]; + if (!lastMenu) return result; + + const entries = lastMenu.querySelector(".litemenu-entry")?.parentElement; + if (!entries) return result; + + // Create separator + const separator = document.createElement("div"); + separator.className = "litemenu-entry separator"; + entries.appendChild(separator); + + // Create menu entry + const menuItem = document.createElement("div"); + menuItem.className = "litemenu-entry submenu"; + menuItem.textContent = "Convert to Set/Get"; + entries.appendChild(menuItem); + + const canvas = this; + menuItem.addEventListener("click", function() { + // Close menu + lastMenu.remove(); + + const originNode = graph.getNodeById(link.origin_id); + const targetNode = graph.getNodeById(link.target_id); + if (!originNode || !targetNode) return; + + const outputSlot = originNode.outputs[link.origin_slot]; + const linkType = outputSlot?.type || "*"; + const linkName = outputSlot?.name || linkType; + + // Create Set node near the source + const setNode = LiteGraph.createNode("SetNode"); + if (!setNode) return; + setNode.pos = [ + originNode.pos[0] + originNode.size[0] + 30, + originNode.pos[1] + ]; + graph.add(setNode); + + // Create Get node near the target + const getNode = LiteGraph.createNode("GetNode"); + if (!getNode) return; + graph.add(getNode); + getNode.pos = [ + targetNode.pos[0] - getNode.size[0] - 30, + targetNode.pos[1] + ]; + + // Remove the original link + graph.removeLink(link.id); + + // Connect source -> Set node input + originNode.connect(link.origin_slot, setNode, 0); + + // Set the name widget on the Set node + setNode.widgets[0].value = linkName; + setNode.title = prefixedTitle("Set", linkName); + setNode.validateName(graph); + setNode.properties.previousName = setNode.widgets[0].value; + + // Update Get node to match + const finalName = setNode.widgets[0].value; + getNode.widgets[0].value = finalName; + getNode.onRename(); + + // Connect Get node output -> target node input + getNode.connect(0, targetNode, link.target_slot); + + canvas.setDirty(true, true); + }); + + return result; + }; + } +}); + +} + +await registerFallback();