Maxed-Out-99
f02f7a8c9a
Introduces the ZImageTurboEmptyLatentImage node for SD3-compatible latent generation with curated safe resolutions. Refactors SaveLatentMXD and SaveLatent_I2V_MXD to remove preview image logic and streamline latent/conditioning saving. Updates WAN22_I2V_Image_Scaler_MXD to add a 'Safe Auto' mode for robust video-extend workflows, improves bucket selection, and enhances error handling for unsupported resolutions. Also updates display names and increments the package version to 1.3.0.
1082 lines
40 KiB
Python
1082 lines
40 KiB
Python
from __future__ import annotations
|
||
import torch, math, comfy, os, folder_paths, node_helpers, comfy.model_management, comfy.utils, json, hashlib, re
|
||
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
|
||
import numpy as np
|
||
from PIL import Image, ImageOps, ImageSequence
|
||
|
||
########################################################################################################################
|
||
# Flux Empty Latent Image (SD3-compatible)
|
||
class FluxEmptyLatentImage:
|
||
DESCRIPTION = """
|
||
- Provides a wide selection of resolutions for Flux for easy selection.
|
||
|
||
- Meant to save time from manually entering the same
|
||
resolutions in the "Empty Latent Image" node over and over.
|
||
"""
|
||
TITLE = "Flux Empty Latent Image"
|
||
CATEGORY = "MXD/Latent"
|
||
|
||
RESOLUTIONS = {
|
||
"— High Resolutions —": None,
|
||
"Square (1:1) 1408x1408": (1408, 1408),
|
||
"Standard (4:3) 1664x1216": (1664, 1216),
|
||
"Landscape (3:2) 1728x1152": (1728, 1152),
|
||
"Widescreen (16:9) 1920x1088": (1920, 1088),
|
||
"Ultrawide (21:9) 2176x960": (2176, 960),
|
||
|
||
"— Standard Resolutions —": None,
|
||
"Square (1:1) 1024x1024": (1024, 1024),
|
||
"Standard (4:3) 1152x896": (1152, 896),
|
||
"Landscape (3:2) 1216x832": (1216, 832),
|
||
"Widescreen (16:9) 1344x768": (1344, 768),
|
||
"Ultrawide (21:9) 1536x640": (1536, 640),
|
||
|
||
"— Low Resolutions —": None,
|
||
"Square (1:1) 320x320": (320, 320),
|
||
"Standard (4:3) 448x320": (448, 320),
|
||
"Landscape (3:2) 384x256": (384, 256),
|
||
"Widescreen (16:9) 448x256": (448, 256),
|
||
"Ultrawide (21:9) 576x256": (576, 256),
|
||
}
|
||
|
||
def __init__(self):
|
||
self.device = comfy.model_management.intermediate_device()
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls) -> dict:
|
||
return {
|
||
"required": {
|
||
"resolution": (
|
||
list(cls.RESOLUTIONS.keys()),
|
||
{"default": "Square (1:1) 1024x1024"}
|
||
),
|
||
"vertical": ("BOOLEAN", {"default": False}),
|
||
"batch_size": (
|
||
"INT",
|
||
{
|
||
"default": 1,
|
||
"min": 1,
|
||
"max": 4096,
|
||
"tooltip": "The number of latent images in the batch."
|
||
}
|
||
)
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("LATENT",)
|
||
OUTPUT_TOOLTIPS = ("The empty latent image batch.",)
|
||
FUNCTION = "generate"
|
||
|
||
def generate(self, resolution, vertical, batch_size=1) -> tuple:
|
||
size = self.RESOLUTIONS.get(resolution)
|
||
if size is None:
|
||
raise ValueError(f"'{resolution}' is a header or invalid option.")
|
||
|
||
width, height = size
|
||
if vertical:
|
||
width, height = height, width
|
||
|
||
latent = torch.zeros([batch_size, 16, height // 8, width // 8], device=self.device)
|
||
return ({"samples": latent},)
|
||
|
||
########################################################################################################################
|
||
# Flux Resolution Selector (for feeding into FluxEmptyLatentImage)
|
||
class FluxResolutionSelector:
|
||
DESCRIPTION = """
|
||
- Provides a dropdown selection of Flux resolutions.
|
||
- Output connects directly to the 'resolution' input of FluxEmptyLatentImage.
|
||
- Useful for separating resolution selection from latent generation.
|
||
"""
|
||
TITLE = "Flux Resolution Selector"
|
||
CATEGORY = "MXD/Latent"
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls) -> dict:
|
||
return {
|
||
"required": {
|
||
"resolution": (
|
||
list(FluxEmptyLatentImage.RESOLUTIONS.keys()), # Include ALL keys including headers
|
||
{"default": "Square (1:1) 1024x1024"}
|
||
),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = (list(FluxEmptyLatentImage.RESOLUTIONS.keys()),)
|
||
RETURN_NAMES = ("resolution",)
|
||
OUTPUT_TOOLTIPS = ("The selected resolution string for FluxEmptyLatentImage.",)
|
||
FUNCTION = "select_resolution"
|
||
|
||
def select_resolution(self, resolution) -> tuple:
|
||
return (resolution,)
|
||
|
||
########################################################################################################################
|
||
# Sdxl Empty Latent Image
|
||
class SdxlEmptyLatentImage:
|
||
|
||
DESCRIPTION = """
|
||
- Provides all compatible SDXL resolutions easy selection.
|
||
|
||
- Meant to save time from manually entering the same
|
||
resolutions in the "Empty Latent Image" node over and over.
|
||
"""
|
||
TITLE = "Sdxl Empty Latent Image (With Resolutions)"
|
||
CATEGORY = "MXD/Latent"
|
||
|
||
# SDXL predefined resolutions (width, height)
|
||
RESOLUTIONS = {
|
||
"Square (1:1) 1024x1024": (1024, 1024),
|
||
"Standard (4:3) 1152x896": (1152, 896),
|
||
"Landscape (3:2) 1216x832": (1216, 832),
|
||
"Widescreen (16:9) 1344x768": (1344, 768),
|
||
"Ultra-Wide (21:9) 1536x640": (1536, 640),
|
||
}
|
||
|
||
def __init__(self):
|
||
# Retrieve the intermediate device (usually the GPU) from ComfyUI's model management.
|
||
self.device = comfy.model_management.intermediate_device()
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls) -> dict:
|
||
return {
|
||
"required": {
|
||
# Dropdown selection for one of the predefined SDXL resolutions.
|
||
"resolution": (list(cls.RESOLUTIONS.keys()),),
|
||
# Toggle for vertical mode (swaps width and height).
|
||
"vertical": ("BOOLEAN", {"default": False}),
|
||
# Number of latent images to create in the batch.
|
||
"batch_size": (
|
||
"INT",
|
||
{
|
||
"default": 1,
|
||
"min": 1,
|
||
"max": 4096,
|
||
"tooltip": "The number of latent images in the batch."
|
||
}
|
||
)
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("LATENT",)
|
||
OUTPUT_TOOLTIPS = ("The empty latent image batch.",)
|
||
FUNCTION = "generate"
|
||
|
||
def generate(self, resolution, vertical, batch_size=1) -> tuple:
|
||
# Get the selected resolution tuple (width, height)
|
||
width, height = self.RESOLUTIONS[resolution]
|
||
# If vertical mode is enabled, swap width and height.
|
||
if vertical:
|
||
width, height = height, width
|
||
|
||
# Create an empty latent tensor.
|
||
# Typically, the latent space has 4 channels and each spatial dimension is 1/8th of the image.
|
||
latent = torch.zeros([batch_size, 4, height // 8, width // 8], device=self.device)
|
||
return ({"samples": latent},)
|
||
|
||
########################################################################################################################
|
||
# Z-Image Turbo Empty Latent Image (SD3-compatible)
|
||
class ZImageTurboEmptyLatentImage:
|
||
DESCRIPTION = """
|
||
- Provides a curated set of SAFE and OPTIMAL resolutions for Z-Image Turbo.
|
||
|
||
- All presets stay within Z-Image Turbo's native comfort zone
|
||
to avoid composition breakage and latent corruption.
|
||
|
||
- Designed to save time and remove guesswork.
|
||
"""
|
||
TITLE = "Z-Image Turbo Empty Latent Image"
|
||
CATEGORY = "MXD/Latent"
|
||
|
||
RESOLUTIONS = {
|
||
"— Recommended (Best Overall) —": None,
|
||
"Square (1:1) 1024x1024": (1024, 1024),
|
||
"Landscape (16:9) 1920x1088": (1920, 1088),
|
||
"Portrait (2:3) 1024x1536": (1024, 1536),
|
||
|
||
"— High Detail —": None,
|
||
"Square (1:1) 1280x1280": (1280, 1280),
|
||
"Square (1:1) 1536x1536": (1536, 1536),
|
||
"Landscape (16:9) 2048x1152": (2048, 1152),
|
||
"Portrait (9:16) 1152x2048": (1152, 2048),
|
||
|
||
"— Cinematic / Wide —": None,
|
||
"Landscape (21:9) 2016x864": (2016, 864),
|
||
"Landscape (21:9) 1680x720": (1680, 720),
|
||
|
||
"— Fast / Draft —": None,
|
||
"Square (1:1) 768x768": (768, 768),
|
||
"Landscape (16:9) 1280x720": (1280, 720),
|
||
"Portrait (3:4) 832x1216": (832, 1216),
|
||
}
|
||
|
||
def __init__(self):
|
||
self.device = comfy.model_management.intermediate_device()
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls) -> dict:
|
||
return {
|
||
"required": {
|
||
"resolution": (
|
||
list(cls.RESOLUTIONS.keys()),
|
||
{"default": "Square (1:1) 1024x1024"}
|
||
),
|
||
"vertical": (
|
||
"BOOLEAN",
|
||
{
|
||
"default": False,
|
||
"tooltip": "Swap width and height."
|
||
}
|
||
),
|
||
"batch_size": (
|
||
"INT",
|
||
{
|
||
"default": 1,
|
||
"min": 1,
|
||
"max": 4096,
|
||
"tooltip": "Number of latent images in the batch."
|
||
}
|
||
)
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("LATENT",)
|
||
OUTPUT_TOOLTIPS = ("The empty Z-Image Turbo latent batch.",)
|
||
FUNCTION = "generate"
|
||
|
||
def generate(self, resolution, vertical, batch_size=1) -> tuple:
|
||
size = self.RESOLUTIONS.get(resolution)
|
||
if size is None:
|
||
raise ValueError(f"'{resolution}' is a header or invalid option.")
|
||
|
||
width, height = size
|
||
if vertical:
|
||
width, height = height, width
|
||
|
||
latent = torch.zeros(
|
||
[batch_size, 16, height // 8, width // 8],
|
||
device=self.device
|
||
)
|
||
|
||
return ({"samples": latent},)
|
||
|
||
########################################################################################################################
|
||
# SDXL Resolution Selector
|
||
class SdxlResolutionSelector:
|
||
DESCRIPTION = """
|
||
- Provides a dropdown selection of SDXL resolutions.
|
||
- Output connects directly to SDXL Empty Latent Image resolution input.
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls) -> dict:
|
||
return {
|
||
"required": {
|
||
"resolution": (
|
||
list(SdxlEmptyLatentImage.RESOLUTIONS.keys()),
|
||
{"default": "Square (1:1) 1024x1024"}
|
||
),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = (list(SdxlEmptyLatentImage.RESOLUTIONS.keys()),)
|
||
RETURN_NAMES = ("resolution",)
|
||
FUNCTION = "select_resolution"
|
||
CATEGORY = "MXD/Latent"
|
||
|
||
def select_resolution(self, resolution) -> tuple:
|
||
return (resolution,)
|
||
|
||
########################################################################################################################
|
||
# Image Scale To Total Pixels (SDXL Safe)
|
||
class SDXLImageScaleToTotalPixelsSafe:
|
||
DESCRIPTION = """
|
||
- Scales to target megapixel count, preserving aspect ratio.
|
||
|
||
- If image matches SDXL resolutions (e.g. those used in
|
||
"SDXL Empty Latent Image" node), scaling is skipped.
|
||
|
||
- Meant for SDXL workflows (e.g. image-to-image, inpainting)
|
||
to auto-scale random images but not images already made with SDXL.
|
||
"""
|
||
upscale_methods = ["bilinear", "bicubic", "lanczos", "nearest-exact", "area"]
|
||
|
||
# SDXL-safe resolutions (width, height) – store one orientation only,
|
||
# the code will check both (w, h) and (h, w)
|
||
SDXL_SAFE_RESOLUTIONS = [
|
||
(1024, 1024),
|
||
(1152, 896),
|
||
(1216, 832),
|
||
(1344, 768),
|
||
(1536, 640),
|
||
]
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"upscale_method": (cls.upscale_methods, {"default": "bilinear"}),
|
||
"total_megapixels": (
|
||
"FLOAT",
|
||
{
|
||
"default": 1.0,
|
||
"min": 0.01,
|
||
"max": 128.0,
|
||
"step": 0.01,
|
||
"tooltip": "Set the total megapixels (e.g., 1.0 = 1 MP)",
|
||
},
|
||
),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
FUNCTION = "upscale"
|
||
CATEGORY = "MXD/Upscaling"
|
||
|
||
def upscale(self, image, upscale_method, total_megapixels):
|
||
if upscale_method in ["nearest-exact", "area"]:
|
||
raise Exception(
|
||
f"❌ '{upscale_method}' gives poor results.\n\n"
|
||
f"👉 Go to the Scale SDXL Image MXD node and switch to another like 'lanczos'.\n\n"
|
||
f"Node may be hidden behind KSampler."
|
||
)
|
||
|
||
b, h, w, c = image.shape
|
||
|
||
# Skip scaling if the image already matches an SDXL-safe resolution
|
||
if (w, h) in self.SDXL_SAFE_RESOLUTIONS or (h, w) in self.SDXL_SAFE_RESOLUTIONS:
|
||
return (image,)
|
||
|
||
# ComfyUI-native megapixel math
|
||
samples = image.movedim(-1, 1)
|
||
orig_h, orig_w = samples.shape[2], samples.shape[3]
|
||
|
||
target_pixels = int(round(total_megapixels * 1024 * 1024))
|
||
scale_by = math.sqrt(target_pixels / (orig_w * orig_h))
|
||
|
||
new_w = max(1, round(orig_w * scale_by))
|
||
new_h = max(1, round(orig_h * scale_by))
|
||
|
||
scaled = comfy.utils.common_upscale(samples, new_w, new_h, upscale_method, "disabled")
|
||
scaled = scaled.movedim(1, -1)
|
||
return (scaled,)
|
||
|
||
########################################################################################################################
|
||
# Flux Image Scale To Total Pixels (Flux Safe)
|
||
class FluxImageScaleToTotalPixelsSafe:
|
||
DESCRIPTION = """
|
||
- Scales to target megapixel count, preserving aspect ratio.
|
||
|
||
- If image matches Flux-safe resolutions (e.g. those used in
|
||
the Flux Empty Latent Image node), scaling is skipped.
|
||
|
||
- Meant for image-to-image or inpainting workflows to auto-scale
|
||
arbitrary images, but skip images already matching Flux resolutions.
|
||
"""
|
||
upscale_methods = ["bilinear", "bicubic", "lanczos", "nearest-exact", "area"]
|
||
|
||
# Flux-safe resolutions (width, height) – stored in one orientation only
|
||
FLUX_SAFE_RESOLUTIONS = [
|
||
(1408, 1408),
|
||
(1728, 1152),
|
||
(1664, 1216),
|
||
(1920, 1088),
|
||
(2176, 960),
|
||
(1024, 1024),
|
||
(1216, 832),
|
||
(1152, 896),
|
||
(1344, 768),
|
||
(1536, 640),
|
||
(320, 320),
|
||
(384, 256),
|
||
(448, 320),
|
||
(448, 256),
|
||
(576, 256),
|
||
]
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE",),
|
||
"upscale_method": (cls.upscale_methods, {"default": "bilinear"}),
|
||
"total_megapixels": (
|
||
"FLOAT",
|
||
{
|
||
"default": 1.0,
|
||
"min": 0.01,
|
||
"max": 128.0,
|
||
"step": 0.01,
|
||
"tooltip": "Set the total megapixels (e.g., 1.0 = 1 MP)",
|
||
},
|
||
),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
FUNCTION = "upscale"
|
||
CATEGORY = "MXD/Upscaling"
|
||
|
||
def upscale(self, image, upscale_method, total_megapixels):
|
||
if upscale_method in ["nearest-exact", "area"]:
|
||
raise Exception(
|
||
f"❌ '{upscale_method}' gives poor results.\n\n"
|
||
f"👉 Go to the Scale Flux Image MXD node and switch to another like 'lanczos'.\n\n"
|
||
f"Node may be hidden behind KSampler."
|
||
)
|
||
|
||
b, h, w, c = image.shape
|
||
|
||
# Skip scaling if image matches any Flux-safe resolution
|
||
if (w, h) in self.FLUX_SAFE_RESOLUTIONS or (h, w) in self.FLUX_SAFE_RESOLUTIONS:
|
||
return (image,)
|
||
|
||
samples = image.movedim(-1, 1)
|
||
orig_h, orig_w = samples.shape[2], samples.shape[3]
|
||
|
||
target_pixels = int(round(total_megapixels * 1024 * 1024))
|
||
scale_by = math.sqrt(target_pixels / (orig_w * orig_h))
|
||
|
||
new_w = max(1, round(orig_w * scale_by))
|
||
new_h = max(1, round(orig_h * scale_by))
|
||
|
||
scaled = comfy.utils.common_upscale(samples, new_w, new_h, upscale_method, "disabled")
|
||
scaled = scaled.movedim(1, -1)
|
||
return (scaled,)
|
||
|
||
########################################################################################################################
|
||
# Prompt with Guidance (Flux)
|
||
class PromptWithGuidance(ComfyNodeABC):
|
||
DESCRIPTION = """
|
||
- Combines clip text encode with flux guidance to lower node count.
|
||
|
||
- Also removes the need to convert them into node group within ComfyUI.
|
||
"""
|
||
@classmethod
|
||
def INPUT_TYPES(cls) -> InputTypeDict:
|
||
return {
|
||
"required": {
|
||
"text": (IO.STRING, {"multiline": True, "dynamicPrompts": True}),
|
||
"clip": (IO.CLIP, {"tooltip": "The CLIP model used for encoding the text."}),
|
||
"guidance": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 100.0, "step": 0.1})
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = (IO.CONDITIONING,)
|
||
FUNCTION = "encode_and_guide"
|
||
CATEGORY = "MXD/conditioning"
|
||
|
||
def encode_and_guide(self, text, clip, guidance):
|
||
if clip is None:
|
||
raise RuntimeError("CLIP model is None. Your checkpoint may not contain a text encoder.")
|
||
|
||
tokens = clip.tokenize(text)
|
||
conditioning = clip.encode_from_tokens_scheduled(tokens)
|
||
conditioning = node_helpers.conditioning_set_values(conditioning, {"guidance": guidance})
|
||
return (conditioning,)
|
||
|
||
########################################################################################################################
|
||
class FluxResolutionMatcher:
|
||
DESCRIPTION = """
|
||
- Automatically matches the Flux Empty Latent Image node's resolution and orientation
|
||
to the input image.
|
||
|
||
- Only uses 'Standard Resolutions' internally, but remains fully compatible with
|
||
the Flux Empty Latent Image node output types.
|
||
"""
|
||
CATEGORY = "MXD/Latent"
|
||
FUNCTION = "match_resolution"
|
||
RETURN_NAMES = ("resolution", "vertical")
|
||
|
||
# Full set kept for compatibility (enum list must match FluxEmptyLatentImage)
|
||
RESOLUTIONS = {
|
||
"— High Resolutions —": None,
|
||
"Square (1:1) 1408x1408": (1408, 1408),
|
||
"Standard (4:3) 1664x1216": (1664, 1216),
|
||
"Landscape (3:2) 1728x1152": (1728, 1152),
|
||
"Widescreen (16:9) 1920x1088": (1920, 1088),
|
||
"Ultrawide (21:9) 2176x960": (2176, 960),
|
||
|
||
"— Standard Resolutions —": None,
|
||
"Square (1:1) 1024x1024": (1024, 1024),
|
||
"Standard (4:3) 1152x896": (1152, 896),
|
||
"Landscape (3:2) 1216x832": (1216, 832),
|
||
"Widescreen (16:9) 1344x768": (1344, 768),
|
||
"Ultrawide (21:9) 1536x640": (1536, 640),
|
||
|
||
"— Low Resolutions —": None,
|
||
"Square (1:1) 320x320": (320, 320),
|
||
"Standard (4:3) 448x320": (448, 320),
|
||
"Landscape (3:2) 384x256": (384, 256),
|
||
"Widescreen (16:9) 448x256": (448, 256),
|
||
"Ultrawide (21:9) 576x256": (576, 256),
|
||
}
|
||
|
||
# Keep same enum type so it connects to FluxEmptyLatentImage
|
||
RETURN_TYPES = (list(RESOLUTIONS.keys()), "BOOLEAN")
|
||
|
||
# Precompute aspect ratio groups (only for standard resolutions)
|
||
ASPECT_RATIO_GROUPS = {}
|
||
for res_str, dims in RESOLUTIONS.items():
|
||
if dims is None:
|
||
continue
|
||
# ✅ Skip high and low groups for logic
|
||
if "High" in res_str or "Low" in res_str:
|
||
continue
|
||
group_name = " ".join(res_str.split(' ')[:-1])
|
||
if group_name not in ASPECT_RATIO_GROUPS:
|
||
w, h = dims
|
||
ratio = w / h
|
||
ASPECT_RATIO_GROUPS[group_name] = {'ratio': ratio, 'resolutions': []}
|
||
ASPECT_RATIO_GROUPS[group_name]['resolutions'].append(res_str)
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {"required": {"image": ("IMAGE",)}}
|
||
|
||
def match_resolution(self, image: torch.Tensor):
|
||
if image.dim() < 4 or image.shape[1] < 1 or image.shape[2] < 1:
|
||
print("Warning: Invalid image tensor received. Falling back to default resolution.")
|
||
return ("Square (1:1) 1024x1024", False)
|
||
|
||
_batch, height, width, _channels = image.shape
|
||
is_vertical = height > width
|
||
img_aspect_ratio = (height / width) if is_vertical else (width / height)
|
||
img_area = height * width
|
||
|
||
best_ar_group_name = min(
|
||
self.ASPECT_RATIO_GROUPS.keys(),
|
||
key=lambda name: abs(img_aspect_ratio - self.ASPECT_RATIO_GROUPS[name]['ratio'])
|
||
)
|
||
|
||
candidate_res_strings = self.ASPECT_RATIO_GROUPS[best_ar_group_name]['resolutions']
|
||
|
||
best_res_string = min(
|
||
candidate_res_strings,
|
||
key=lambda res_str: abs(img_area - (self.RESOLUTIONS[res_str][0] * self.RESOLUTIONS[res_str][1]))
|
||
)
|
||
|
||
return (best_res_string, is_vertical)
|
||
########################################################################################################################
|
||
|
||
class LatentHalfMasks:
|
||
DESCRIPTION = """
|
||
- Splits a latent tensor into left and right half masks.
|
||
|
||
- Designed for dual-character setups (e.g. two Apply PuLID nodes).
|
||
|
||
- Simplifies workflows by removing manual math/masking steps.
|
||
"""
|
||
TITLE = "Latent Half Masks"
|
||
CATEGORY = "MXD/Latent"
|
||
|
||
RETURN_TYPES = ("MASK", "MASK")
|
||
RETURN_NAMES = ("mask_left", "mask_right")
|
||
OUTPUT_TOOLTIPS = (
|
||
"Mask covering the left half of the latent.",
|
||
"Mask covering the right half of the latent.",
|
||
)
|
||
FUNCTION = "make_masks"
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"latent": ("LATENT",),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("MASK", "MASK")
|
||
RETURN_NAMES = ("mask_left", "mask_right")
|
||
FUNCTION = "make_masks"
|
||
CATEGORY = "MXD/latent"
|
||
|
||
def make_masks(self, latent):
|
||
# Infer width/height from latent (assumes 8x scale)
|
||
samples = latent.get("samples", None)
|
||
if samples is None or not isinstance(samples, torch.Tensor):
|
||
raise ValueError("LatentHalfMasks: invalid latent or missing 'samples' tensor.")
|
||
h_lat, w_lat = samples.shape[-2], samples.shape[-1]
|
||
w, h = int(w_lat * 8), int(h_lat * 8)
|
||
|
||
# Always vertical, center split, no feather, no swap
|
||
split_px = w // 2
|
||
left = torch.zeros((h, w), dtype=torch.float32)
|
||
right = torch.zeros((h, w), dtype=torch.float32)
|
||
left[:, :split_px] = 1.0
|
||
right[:, split_px:] = 1.0
|
||
|
||
return left, right
|
||
|
||
########################################################################################################################
|
||
|
||
# --- Helper function to find the bounding box of a mask ---
|
||
def get_bounding_box(mask_tensor):
|
||
"""
|
||
Finds the bounding box of a non-zero region in a mask tensor.
|
||
The mask is expected to be a 2D tensor (H, W).
|
||
Returns a tuple (x_min, y_min, x_max, y_max) or None if the mask is empty.
|
||
"""
|
||
# Get non-zero coordinates from the mask
|
||
non_zero_coords = torch.nonzero(mask_tensor, as_tuple=False)
|
||
|
||
# If the mask is empty, there is no bounding box
|
||
if non_zero_coords.numel() == 0:
|
||
return None
|
||
|
||
# Find the min and max coordinates for y (dim 0) and x (dim 1)
|
||
min_y = non_zero_coords[:, 0].min().item()
|
||
max_y = non_zero_coords[:, 0].max().item()
|
||
min_x = non_zero_coords[:, 1].min().item()
|
||
max_x = non_zero_coords[:, 1].max().item()
|
||
|
||
# The bounding box for PIL needs (left, upper, right, lower).
|
||
# We add +1 to the max values because the upper bound is exclusive.
|
||
return (min_x, min_y, max_x + 1, max_y + 1)
|
||
|
||
# --- Tensor to PIL and PIL to Tensor conversion helpers ---
|
||
def tensor_to_pil(tensor):
|
||
"""Converts a torch tensor (B, H, W, C) to a list of PIL Images."""
|
||
if tensor is None:
|
||
return []
|
||
|
||
# Handle different tensor dimensions
|
||
if tensor.dim() == 4: # Batch of images
|
||
images = []
|
||
for i in range(tensor.shape[0]):
|
||
img_np = 255. * tensor[i].cpu().numpy()
|
||
images.append(Image.fromarray(np.clip(img_np, 0, 255).astype(np.uint8)))
|
||
return images
|
||
elif tensor.dim() == 3: # Single image
|
||
img_np = 255. * tensor.cpu().numpy()
|
||
return [Image.fromarray(np.clip(img_np, 0, 255).astype(np.uint8))]
|
||
else:
|
||
raise ValueError(f"Unsupported tensor dimension: {tensor.dim()}")
|
||
|
||
def pil_to_tensor(pil_images):
|
||
"""Converts a list of PIL Images back to a torch tensor (B, H, W, C)."""
|
||
if not isinstance(pil_images, list):
|
||
pil_images = [pil_images]
|
||
|
||
tensors = []
|
||
for img in pil_images:
|
||
# Convert to RGB, then to a numpy array, normalize, and create a tensor
|
||
img_np = np.array(img.convert("RGB")).astype(np.float32) / 255.0
|
||
tensors.append(torch.from_numpy(img_np).unsqueeze(0))
|
||
|
||
# Stack all tensors into a single batch tensor
|
||
return torch.cat(tensors, dim=0)
|
||
|
||
# --------------------------------------------------------------------
|
||
# ✨ The Main Node Class ✨
|
||
# --------------------------------------------------------------------
|
||
class PlaceImageByMask:
|
||
Description = """
|
||
- Overlay an image onto a base image.
|
||
|
||
- The position and scale of the overlay are determined by a mask's bounding box.
|
||
|
||
- Useful for placing images in specific regions of a base image.
|
||
"""
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"base_image": ("IMAGE",),
|
||
"mask": ("MASK",),
|
||
"overlay_image": ("IMAGE",),
|
||
},
|
||
"optional": {
|
||
"maintain_aspect_ratio": ("BOOLEAN", {"default": True}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
FUNCTION = "place_image"
|
||
CATEGORY = "MXD/Image"
|
||
|
||
def place_image(self, base_image, overlay_image, mask, maintain_aspect_ratio=True):
|
||
# Convert input tensors to lists of PIL Images
|
||
base_pils = tensor_to_pil(base_image)
|
||
overlay_pils = tensor_to_pil(overlay_image)
|
||
|
||
processed_images = []
|
||
|
||
# Process each image in the batch
|
||
for i, base_pil in enumerate(base_pils):
|
||
# Work with an RGBA version of the base image for clean pasting
|
||
composited_image = base_pil.convert("RGBA")
|
||
|
||
# Select the corresponding overlay and mask for the current base image
|
||
# Clamping the index prevents errors if batch sizes are mismatched
|
||
overlay_pil = overlay_pils[min(i, len(overlay_pils) - 1)].convert("RGBA")
|
||
current_mask = mask[min(i, mask.shape[0] - 1)]
|
||
|
||
# Find the bounding box from the mask
|
||
bbox = get_bounding_box(current_mask)
|
||
|
||
# If no mask is found, just use the original base image and skip to the next
|
||
if not bbox:
|
||
raise ValueError("The base image must be masked where you want the overlay to appear.")
|
||
|
||
x_min, y_min, x_max, y_max = bbox
|
||
box_width = x_max - x_min
|
||
box_height = y_max - y_min
|
||
|
||
# If the bounding box has no area, skip to the next image
|
||
if box_width <= 0 or box_height <= 0:
|
||
processed_images.append(base_pil)
|
||
continue
|
||
|
||
# --- Resize the overlay image using the specified method ---
|
||
if maintain_aspect_ratio:
|
||
# Resize to fit *within* the box, preserving aspect ratio (like a thumbnail)
|
||
resized_overlay = overlay_pil.copy()
|
||
resized_overlay.thumbnail((box_width, box_height), Image.Resampling.LANCZOS)
|
||
|
||
# Calculate position to center the resized overlay within the bounding box
|
||
paste_x = x_min + (box_width - resized_overlay.width) // 2
|
||
paste_y = y_min + (box_height - resized_overlay.height) // 2
|
||
paste_pos = (paste_x, paste_y)
|
||
else:
|
||
# As originally requested: stretch to fill the bounding box exactly
|
||
resized_overlay = overlay_pil.resize((box_width, box_height), resample=Image.Resampling.LANCZOS)
|
||
paste_pos = (x_min, y_min)
|
||
|
||
# --- Paste the resized overlay onto the base image ---
|
||
# The alpha channel of the overlay itself is used as the mask for pasting.
|
||
# This ensures transparent areas of the overlay are handled correctly.
|
||
composited_image.paste(resized_overlay, paste_pos, resized_overlay)
|
||
|
||
processed_images.append(composited_image)
|
||
|
||
# Convert the list of processed PIL images back to a single batch tensor for output
|
||
output_tensor = pil_to_tensor(processed_images)
|
||
return (output_tensor,)
|
||
|
||
######################################################################################################################################
|
||
|
||
class CropImageByMask:
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"image": ("IMAGE", ),
|
||
},
|
||
"optional": {
|
||
"mask": ("MASK", ),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", )
|
||
RETURN_NAMES = ("image", )
|
||
FUNCTION = "crop"
|
||
CATEGORY = "MXD/image"
|
||
|
||
def crop(self, image, mask=None):
|
||
# If no mask is provided or the mask is completely empty, return the original image
|
||
if mask is None or not torch.any(mask > 0):
|
||
return (image, )
|
||
|
||
B, H, W, C = image.shape
|
||
mask = mask.round()
|
||
|
||
# Find bounding box for each batch
|
||
crops = []
|
||
|
||
for b in range(B):
|
||
current_mask = mask[min(b, mask.shape[0]-1)]
|
||
|
||
# Check if the mask for this specific image is empty.
|
||
if not torch.any(current_mask > 0):
|
||
# If a specific mask in a batch is empty, we can't crop.
|
||
# To prevent errors with torch.cat later due to different sizes,
|
||
# we'll skip cropping for the whole batch and return the original.
|
||
# This ensures the output is always a valid tensor.
|
||
print("Warning: An empty mask was found in a batch. Returning original images.")
|
||
return (image, )
|
||
|
||
# Get coordinates of non-zero elements
|
||
rows = torch.any(current_mask > 0, dim=1)
|
||
cols = torch.any(current_mask > 0, dim=0)
|
||
|
||
# Find boundaries
|
||
y_min, y_max = torch.where(rows)[0][[0, -1]]
|
||
x_min, x_max = torch.where(cols)[0][[0, -1]]
|
||
|
||
# Crop image
|
||
crop = image[b:b+1, y_min:y_max+1, x_min:x_max+1, :]
|
||
crops.append(crop)
|
||
|
||
# Note: This will raise an error if the crops have different sizes.
|
||
# The original code had this limitation.
|
||
cropped_images = torch.cat(crops, dim=0)
|
||
|
||
return (cropped_images, )
|
||
|
||
########################################################################################################################
|
||
# ---------- Helpers (copied from latent loader style) ----------
|
||
def _safe_json_loads(s):
|
||
if s is None:
|
||
return None
|
||
if isinstance(s, bytes):
|
||
try:
|
||
s = s.decode("utf-8", "ignore")
|
||
except Exception:
|
||
return None
|
||
if not isinstance(s, str):
|
||
return None
|
||
try:
|
||
return json.loads(s)
|
||
except Exception:
|
||
try:
|
||
return json.loads(json.loads(s))
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _extract_params_from_prompt_json(prompt_json: dict):
|
||
"""
|
||
Returns (positive, negative) from saved Comfy prompt graph.
|
||
"""
|
||
pos = ""
|
||
neg = ""
|
||
if not isinstance(prompt_json, dict):
|
||
return pos, neg
|
||
|
||
# unwrap if saved as {"prompt": {...}}
|
||
graph = prompt_json.get("prompt", prompt_json)
|
||
if not isinstance(graph, dict):
|
||
return pos, neg
|
||
|
||
# try to find KSampler/KSamplerAdvanced node
|
||
ks = None
|
||
for _, v in graph.items():
|
||
if "KSampler" in v.get("class_type", ""):
|
||
ks = v
|
||
break
|
||
if not ks:
|
||
return pos, neg
|
||
|
||
kin = ks.get("inputs", {})
|
||
|
||
def _as_node_id(x):
|
||
return str(x[0]) if isinstance(x, (list, tuple)) and x else None
|
||
|
||
def _text_from_clip(node_id):
|
||
n = graph.get(str(node_id), {})
|
||
if n.get("class_type") == "CLIPTextEncode":
|
||
return str(n.get("inputs", {}).get("text", "")).strip()
|
||
return ""
|
||
|
||
pos = _text_from_clip(_as_node_id(kin.get("positive")))
|
||
neg = _text_from_clip(_as_node_id(kin.get("negative")))
|
||
|
||
return pos, neg
|
||
|
||
def _strip_counter(name: str) -> str:
|
||
# Only strip the trailing pattern we generate when saving: "_<5digits>_"
|
||
# Preserve numeric-only base names like "96".
|
||
stem, _ = os.path.splitext(name)
|
||
m = re.match(r"^(.*?)(?:_\d{5}_)$", stem)
|
||
return m.group(1) if m else stem
|
||
|
||
# ---------- Node ----------
|
||
class LoadImageBatchMXD:
|
||
DESCRIPTION = """
|
||
- Loads all images in a selected folder under outputs/.
|
||
- Automatically generates masks from alpha channel if present.
|
||
- Extracts positive/negative prompts from image metadata (if available).
|
||
"""
|
||
TITLE = "Load Image Batch (Outputs + Prompts)"
|
||
CATEGORY = "MXD/Image"
|
||
|
||
RETURN_TYPES = ("IMAGE", "MASK", "STRING", "STRING")
|
||
RETURN_NAMES = ("IMAGE", "MASK", "positive", "negative")
|
||
OUTPUT_IS_LIST = (True, True, True, True)
|
||
FUNCTION = "load_batch"
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
outputs_root = folder_paths.get_output_directory()
|
||
subdirs = [""] + sorted(
|
||
[d for d in os.listdir(outputs_root) if os.path.isdir(os.path.join(outputs_root, d))]
|
||
)
|
||
return {"required": {"folder": (tuple(subdirs), {"default": ""})}}
|
||
|
||
def _extract_prompts(self, image: Image.Image):
|
||
pos, neg = "", ""
|
||
try:
|
||
raw = image.info.get("prompt")
|
||
if raw:
|
||
prompt_json = _safe_json_loads(raw)
|
||
if prompt_json:
|
||
pos, neg = _extract_params_from_prompt_json(prompt_json)
|
||
else:
|
||
pos = raw
|
||
except Exception as e:
|
||
print(f"[LoadImageBatchMXD] Prompt parse failed: {e}")
|
||
return pos, neg
|
||
|
||
def load_batch(self, folder: str):
|
||
outputs_root = folder_paths.get_output_directory()
|
||
folder_path = os.path.join(outputs_root, folder) if folder else outputs_root
|
||
|
||
if not os.path.isdir(folder_path):
|
||
raise FileNotFoundError(f"No such folder: {folder_path}")
|
||
|
||
valid_exts = (".png", ".jpg", ".jpeg", ".webp")
|
||
files = [os.path.join(folder_path, f) for f in sorted(os.listdir(folder_path))
|
||
if f.lower().endswith(valid_exts)]
|
||
|
||
images, masks, positives, negatives, prefixes = [], [], [], [], []
|
||
|
||
for path in files:
|
||
i = Image.open(path)
|
||
i = ImageOps.exif_transpose(i)
|
||
|
||
pos, neg = self._extract_prompts(i)
|
||
positives.append(pos)
|
||
negatives.append(neg)
|
||
|
||
rgb = i.convert("RGB")
|
||
arr = np.array(rgb).astype(np.float32) / 255.0
|
||
img_t = torch.from_numpy(arr)[None, ...]
|
||
|
||
if 'A' in i.getbands():
|
||
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
|
||
mask_t = 1.0 - torch.from_numpy(mask).unsqueeze(0)
|
||
else:
|
||
h, w = arr.shape[:2]
|
||
mask_t = torch.zeros((1, h, w), dtype=torch.float32)
|
||
|
||
images.append(img_t)
|
||
masks.append(mask_t)
|
||
|
||
return (images, masks, positives, negatives)
|
||
|
||
class LoadImageWithPromptsMXD:
|
||
CATEGORY = "image"
|
||
|
||
RETURN_TYPES = ("IMAGE", "MASK", "STRING", "STRING")
|
||
RETURN_NAMES = ("IMAGE", "MASK", "positive", "negative")
|
||
FUNCTION = "load_image"
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
input_dir = folder_paths.get_input_directory()
|
||
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
|
||
files = folder_paths.filter_files_content_types(files, ["image"])
|
||
return {"required": {"image": (sorted(files), {"image_upload": True})}}
|
||
|
||
def _extract_prompts(self, img: Image.Image):
|
||
pos, neg = "", ""
|
||
raw = img.info.get("prompt")
|
||
if raw:
|
||
prompt_json = _safe_json_loads(raw)
|
||
if prompt_json:
|
||
pos, neg = _extract_params_from_prompt_json(prompt_json)
|
||
else:
|
||
pos = raw
|
||
return pos, neg
|
||
|
||
def load_image(self, image):
|
||
image_path = folder_paths.get_annotated_filepath(image)
|
||
img = node_helpers.pillow(Image.open, image_path)
|
||
|
||
output_images, output_masks = [], []
|
||
pos, neg = "", ""
|
||
w, h = None, None
|
||
|
||
excluded_formats = ['MPO']
|
||
|
||
for i in ImageSequence.Iterator(img):
|
||
i = node_helpers.pillow(ImageOps.exif_transpose, i)
|
||
|
||
if i.mode == 'I':
|
||
i = i.point(lambda i: i * (1 / 255))
|
||
frame = i.convert("RGB")
|
||
|
||
if len(output_images) == 0:
|
||
w, h = frame.size
|
||
# extract prompts only once (from first frame)
|
||
pos, neg = self._extract_prompts(i)
|
||
|
||
if frame.size != (w, h):
|
||
continue
|
||
|
||
arr = np.array(frame).astype(np.float32) / 255.0
|
||
tensor_img = torch.from_numpy(arr)[None, ...]
|
||
|
||
if 'A' in i.getbands():
|
||
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
|
||
mask = 1. - torch.from_numpy(mask)
|
||
elif i.mode == 'P' and 'transparency' in i.info:
|
||
mask = np.array(i.convert('RGBA').getchannel('A')).astype(np.float32) / 255.0
|
||
mask = 1. - torch.from_numpy(mask)
|
||
else:
|
||
mask = torch.zeros((1, 64, 64), dtype=torch.float32, device="cpu")
|
||
|
||
output_images.append(tensor_img)
|
||
output_masks.append(mask.unsqueeze(0))
|
||
|
||
if len(output_images) > 1 and img.format not in excluded_formats:
|
||
output_image = torch.cat(output_images, dim=0)
|
||
output_mask = torch.cat(output_masks, dim=0)
|
||
else:
|
||
output_image = output_images[0]
|
||
output_mask = output_masks[0]
|
||
|
||
return (output_image, output_mask, pos, neg)
|
||
|
||
@classmethod
|
||
def IS_CHANGED(s, image):
|
||
image_path = folder_paths.get_annotated_filepath(image)
|
||
m = hashlib.sha256()
|
||
with open(image_path, 'rb') as f:
|
||
m.update(f.read())
|
||
return m.digest().hex()
|
||
|
||
@classmethod
|
||
def VALIDATE_INPUTS(s, image):
|
||
if not folder_paths.exists_annotated_filepath(image):
|
||
return f"Invalid image file: {image}"
|
||
return True
|
||
|
||
########################################################################################################################
|
||
|
||
|
||
# NODE MAPPING
|
||
NODE_CLASS_MAPPINGS = {
|
||
"Flux Empty Latent Image": FluxEmptyLatentImage,
|
||
"Sdxl Empty Latent Image": SdxlEmptyLatentImage,
|
||
"Flux Resolution Selector": FluxResolutionSelector,
|
||
"Sdxl Resolution Selector": SdxlResolutionSelector,
|
||
"Image Scale To Total Pixels (SDXL Safe)": SDXLImageScaleToTotalPixelsSafe,
|
||
"Flux Image Scale To Total Pixels (Flux Safe)": FluxImageScaleToTotalPixelsSafe,
|
||
"Prompt With Guidance (Flux)": PromptWithGuidance,
|
||
"FluxResolutionMatcher": FluxResolutionMatcher,
|
||
"LatentHalfMasks": LatentHalfMasks,
|
||
"Place Image By Mask": PlaceImageByMask,
|
||
"Crop Image By Mask": CropImageByMask,
|
||
"Load Image Batch MXD": LoadImageBatchMXD,
|
||
"LoadImageWithPromptsMXD": LoadImageWithPromptsMXD,
|
||
"ZImageTurboEmptyLatentImage": ZImageTurboEmptyLatentImage,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"Flux Empty Latent Image": "Flux Empty Latent Image MXD",
|
||
"Sdxl Empty Latent Image": "SDXL Empty Latent Image MXD",
|
||
"Flux Resolution Selector": "Flux Resolution Selector MXD",
|
||
"Sdxl Resolution Selector": "SDXL Resolution Selector MXD",
|
||
"Image Scale To Total Pixels (SDXL Safe)": "Scale SDXL Image MXD",
|
||
"Flux Image Scale To Total Pixels (Flux Safe)": "Scale Flux Image MXD",
|
||
"Prompt With Guidance (Flux)": "Prompt with Flux Guidance MXD",
|
||
"FluxResolutionMatcher": "Flux Resolution Matcher MXD",
|
||
"LatentHalfMasks": "Latent to L/R Masks MXD",
|
||
"Place Image By Mask": "Place Image by Mask MXD",
|
||
"Crop Image By Mask": "Crop Image by Mask MXD",
|
||
"Load Image Batch MXD": "Load Image Batch MXD",
|
||
"LoadImageWithPromptsMXD": "Load Image MXD",
|
||
"ZImageTurboEmptyLatentImage": "ZImageTurbo Empty Latent Image MXD",
|
||
}
|