Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
a7c4b20049 | ||
|
|
8e72d7a1a0 |
@@ -2,6 +2,10 @@
|
||||
|
||||
Data is kept separate from node logic so new model families can be added
|
||||
without changing execution code.
|
||||
|
||||
min_res / max_res are the clamp range for *custom* sizes. Curated presets may go
|
||||
beyond max_res (e.g. 4 MP tiers); their own reliability/warning fields say how
|
||||
safe they are, and the node reports that warning in its info output.
|
||||
"""
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
"""Arithmetic expression evaluator shared by Math Eval (NH) and layer position expressions.
|
||||
|
||||
Accepted: numbers, named variables, + - * / // % **, unary + and -, and an optional
|
||||
whitelist of functions. Everything is computed in floating point, so no expression can
|
||||
build a huge integer and freeze the worker (``9**9**9`` simply overflows and is reported).
|
||||
"""
|
||||
|
||||
import ast
|
||||
import math
|
||||
import numbers
|
||||
import operator
|
||||
|
||||
# Python's ** returns a complex number for a negative base with a fractional exponent;
|
||||
# math.pow raises instead, which is what a node input needs.
|
||||
_BINARY = {
|
||||
ast.Add: operator.add,
|
||||
ast.Sub: operator.sub,
|
||||
ast.Mult: operator.mul,
|
||||
ast.Div: operator.truediv,
|
||||
ast.FloorDiv: operator.floordiv,
|
||||
ast.Mod: operator.mod,
|
||||
ast.Pow: math.pow,
|
||||
}
|
||||
_UNARY = {
|
||||
ast.UAdd: operator.pos,
|
||||
ast.USub: operator.neg,
|
||||
}
|
||||
_DIVISIONS = (ast.Div, ast.FloorDiv, ast.Mod)
|
||||
|
||||
MATH_FUNCTIONS = {
|
||||
"min": min,
|
||||
"max": max,
|
||||
"abs": abs,
|
||||
"round": lambda value, digits=0: round(value, int(digits)),
|
||||
"sqrt": math.sqrt,
|
||||
"floor": math.floor,
|
||||
"ceil": math.ceil,
|
||||
"clamp": lambda value, low, high: max(low, min(high, value)),
|
||||
}
|
||||
|
||||
|
||||
class ExpressionError(ValueError):
|
||||
"""Raised for any expression that cannot be evaluated to a finite number."""
|
||||
|
||||
|
||||
def evaluate(expression, variables=None, functions=None, *, zero_division="zero",
|
||||
case_insensitive=False, on_warning=None):
|
||||
"""Evaluate ``expression`` and return a finite float.
|
||||
|
||||
zero_division: "zero" makes x/0, x//0 and x%0 evaluate to 0 (reported through
|
||||
``on_warning``); "error" raises ExpressionError.
|
||||
case_insensitive: match variable names with str.casefold (keys must already be folded).
|
||||
"""
|
||||
variables = variables or {}
|
||||
functions = functions or {}
|
||||
text = str(expression).strip()
|
||||
if not text:
|
||||
raise ExpressionError("Expression is empty")
|
||||
|
||||
def warn(message):
|
||||
if on_warning is not None:
|
||||
on_warning(message)
|
||||
|
||||
def number(value, what):
|
||||
# numpy scalars and 0-d torch tensors arrive from upstream nodes on FLOAT/INT sockets.
|
||||
if getattr(value, "ndim", None) == 0 and callable(getattr(value, "item", None)):
|
||||
value = value.item()
|
||||
if not isinstance(value, numbers.Real): # bool is a Real (True + 1 == 2, as before)
|
||||
raise ExpressionError(f"{what} must be a number, got {type(value).__name__}")
|
||||
try:
|
||||
return float(value)
|
||||
except OverflowError as exc:
|
||||
raise ExpressionError(f"{what} is too large") from exc
|
||||
|
||||
def lookup(name):
|
||||
key = name.casefold() if case_insensitive else name
|
||||
if key in variables:
|
||||
return number(variables[key], f"Variable '{name}'")
|
||||
raise ExpressionError(f"Unknown variable '{name}'")
|
||||
|
||||
def visit(node):
|
||||
if isinstance(node, ast.Constant):
|
||||
return number(node.value, f"Constant {node.value!r}")
|
||||
if isinstance(node, ast.Name):
|
||||
return lookup(node.id)
|
||||
if isinstance(node, ast.BinOp):
|
||||
op = _BINARY.get(type(node.op))
|
||||
if op is None:
|
||||
raise ExpressionError(f"Operator '{type(node.op).__name__}' is not allowed")
|
||||
left, right = visit(node.left), visit(node.right)
|
||||
if isinstance(node.op, _DIVISIONS) and right == 0:
|
||||
if zero_division == "zero":
|
||||
warn("division by zero, using 0")
|
||||
return 0.0
|
||||
raise ExpressionError("Division by zero")
|
||||
return float(op(left, right))
|
||||
if isinstance(node, ast.UnaryOp):
|
||||
op = _UNARY.get(type(node.op))
|
||||
if op is None:
|
||||
raise ExpressionError(f"Operator '{type(node.op).__name__}' is not allowed")
|
||||
return op(visit(node.operand))
|
||||
if isinstance(node, ast.Call):
|
||||
name = node.func.id if isinstance(node.func, ast.Name) else None
|
||||
if name not in functions or node.keywords:
|
||||
raise ExpressionError(f"Function '{name}' is not allowed")
|
||||
return number(functions[name](*[visit(arg) for arg in node.args]), f"{name}()")
|
||||
raise ExpressionError(f"'{type(node).__name__}' is not allowed in an expression")
|
||||
|
||||
try:
|
||||
result = visit(ast.parse(text, mode="eval").body)
|
||||
except ExpressionError:
|
||||
raise
|
||||
except SyntaxError as exc:
|
||||
raise ExpressionError(f"Invalid expression: {text}") from exc
|
||||
except OverflowError as exc:
|
||||
raise ExpressionError(f"Result is too large: {text}") from exc
|
||||
except (ValueError, TypeError, ZeroDivisionError) as exc:
|
||||
# math domain errors, wrong argument counts, 0 ** -1 ...
|
||||
raise ExpressionError(f"Cannot evaluate {text}: {exc}") from exc
|
||||
except (RecursionError, MemoryError) as exc:
|
||||
# The parser raises MemoryError ("parser stack overflowed") for absurdly nested input.
|
||||
raise ExpressionError("Expression is nested too deeply") from exc
|
||||
|
||||
if not math.isfinite(result):
|
||||
raise ExpressionError(f"Result is not a finite number: {text}")
|
||||
return result
|
||||
+32
-4
@@ -12,10 +12,21 @@ def resize_image(image, target_size=768):
|
||||
return image.resize((new_w, new_h), Image.LANCZOS)
|
||||
|
||||
def tensor_to_pil(tensor_image):
|
||||
"""Convert tensor to PIL image"""
|
||||
"""Convert ONE ComfyUI image ([H,W,C] or [1,H,W,C]) to PIL.
|
||||
|
||||
Values are clipped to [0, 1] first, so overshoot from resampling cannot wrap around
|
||||
in uint8. 1/3/4 channels become L/RGB/RGBA. Batches must be converted per image.
|
||||
"""
|
||||
if tensor_image.dim() == 4:
|
||||
tensor_image = tensor_image.squeeze(0)
|
||||
np_image = (tensor_image.cpu().numpy() * 255).astype(np.uint8)
|
||||
if tensor_image.shape[0] != 1:
|
||||
raise ValueError(
|
||||
f"tensor_to_pil expects a single image but got a batch of {tensor_image.shape[0]}; "
|
||||
"convert each image of the batch separately"
|
||||
)
|
||||
tensor_image = tensor_image[0]
|
||||
np_image = np.clip(tensor_image.detach().cpu().float().numpy() * 255.0, 0, 255).astype(np.uint8)
|
||||
if np_image.ndim == 3 and np_image.shape[2] == 1:
|
||||
np_image = np_image[:, :, 0]
|
||||
return Image.fromarray(np_image)
|
||||
|
||||
def pil_to_tensor(pil_image):
|
||||
@@ -32,4 +43,21 @@ def pil_to_mask(pil_image):
|
||||
if len(np_image.shape) == 3: # Ensure single channel
|
||||
np_image = np_image[:,:,0]
|
||||
tensor = torch.from_numpy(np_image).unsqueeze(0)
|
||||
return tensor
|
||||
return tensor
|
||||
|
||||
def parse_hex_color(value, default=(0, 0, 0, 255)):
|
||||
"""Parse '#RGB', '#RRGGBB' or '#RRGGBBAA' (the '#' is optional) into an (r, g, b, a) tuple of 0-255 ints.
|
||||
|
||||
Without an alpha part the colour is opaque. Text that is not a valid colour returns ``default``.
|
||||
"""
|
||||
text = str(value or "").strip().lstrip("#")
|
||||
if len(text) == 3:
|
||||
text = "".join(ch * 2 for ch in text)
|
||||
if len(text) == 6:
|
||||
text += "ff"
|
||||
if len(text) != 8:
|
||||
return tuple(default)
|
||||
try:
|
||||
return tuple(int(text[i:i + 2], 16) for i in range(0, 8, 2))
|
||||
except ValueError:
|
||||
return tuple(default)
|
||||
|
||||
+83
-95
@@ -15,11 +15,27 @@ def remove_small(binary_mask, min_area=10000):
|
||||
small_white_regions = [region for region in regions if region.area < min_area]
|
||||
|
||||
for region in small_white_regions:
|
||||
min_row, min_col, max_row, max_col = region.bbox
|
||||
binary_mask[min_row:max_row, min_col:max_col] = 0
|
||||
# Only this region's pixels: its bounding box can overlap a large region.
|
||||
binary_mask[tuple(region.coords.T)] = 0
|
||||
return binary_mask
|
||||
|
||||
|
||||
def _merge_bound(keypoint_bound, cloth_bound, pick):
|
||||
"""Combine a keypoint bound (None when the whole keypoint group is missing) with the garment contour bound."""
|
||||
if keypoint_bound is None:
|
||||
return cloth_bound
|
||||
return pick(keypoint_bound, cloth_bound)
|
||||
|
||||
|
||||
def _fallback_box(x1, y1, x2, y2, shape):
|
||||
"""Whole-image edges only for bounds that neither keypoints nor the garment contour provided."""
|
||||
height, width = shape[:2]
|
||||
return (0 if x1 is None else x1, 0 if y1 is None else y1,
|
||||
width if x2 is None else x2, height if y2 is None else y2)
|
||||
|
||||
|
||||
# pose_data rows: body 0-17, foot 18-23, face 24-91, hands 92-133. A keypoint that was not detected is
|
||||
# negative (DWPose marks it -1 and the node then scales it by the width/height), so only values > 0 count.
|
||||
def get_img_agnostic_upper_rectangle(im_parse, pose_data, offset_top, offset_bottom, offset_left, offset_right):
|
||||
|
||||
faces = pose_data[24:92]
|
||||
@@ -69,48 +85,37 @@ def get_img_agnostic_upper_rectangle(im_parse, pose_data, offset_top, offset_bot
|
||||
except Exception:
|
||||
cloth_exist = False
|
||||
|
||||
x1 = [body[i, 0] for i in [2, 3, 4] if body[i, 0] != 0]+\
|
||||
[hands2[i, 0] for i in [5,9,13] if hands2[i, 0] != 0]+\
|
||||
[hands1[i, 0] for i in [5,9,13] if hands1[i, 0] != 0]
|
||||
if x1:
|
||||
x1 = np.min(x1)
|
||||
else:
|
||||
x1 = 0
|
||||
y1 = [body[i, 1] for i in [2, 5] if body[i, 1] != 0]+\
|
||||
[hands2[i, 1] for i in [5,9,13] if hands2[i, 1] != 0]+\
|
||||
[hands1[i, 1] for i in [5,9,13] if hands1[i, 1] != 0]
|
||||
if y1:
|
||||
y1 = np.min(y1)
|
||||
else:
|
||||
y1 = 0
|
||||
x2 = [body[i, 0] for i in [5, 6, 7] if body[i, 0] != 0]+\
|
||||
[hands1[i, 0] for i in [5,9,13] if hands1[i, 0] != 0]+\
|
||||
[hands2[i, 0] for i in [5,9,13] if hands2[i, 0] != 0]
|
||||
if x2:
|
||||
x2 = np.max(x2)
|
||||
else:
|
||||
x2 = parse_array.shape[1]
|
||||
y2 = [body[i, 1] for i in [8, 11] if body[i, 1] != 0]+\
|
||||
[hands2[i, 1] for i in [5,9,13] if hands2[i, 1] != 0]+\
|
||||
[hands1[i, 1] for i in [5,9,13] if hands1[i, 1] != 0]
|
||||
if y2:
|
||||
y2 = np.max(y2)
|
||||
else:
|
||||
y2 = parse_array.shape[0]
|
||||
x1 = [body[i, 0] for i in [2, 3, 4] if body[i, 0] > 0]+\
|
||||
[hands2[i, 0] for i in [5,9,13] if hands2[i, 0] > 0]+\
|
||||
[hands1[i, 0] for i in [5,9,13] if hands1[i, 0] > 0]
|
||||
x1 = np.min(x1) if x1 else None
|
||||
y1 = [body[i, 1] for i in [2, 5] if body[i, 1] > 0]+\
|
||||
[hands2[i, 1] for i in [5,9,13] if hands2[i, 1] > 0]+\
|
||||
[hands1[i, 1] for i in [5,9,13] if hands1[i, 1] > 0]
|
||||
y1 = np.min(y1) if y1 else None
|
||||
x2 = [body[i, 0] for i in [5, 6, 7] if body[i, 0] > 0]+\
|
||||
[hands1[i, 0] for i in [5,9,13] if hands1[i, 0] > 0]+\
|
||||
[hands2[i, 0] for i in [5,9,13] if hands2[i, 0] > 0]
|
||||
x2 = np.max(x2) if x2 else None
|
||||
y2 = [body[i, 1] for i in [8, 11] if body[i, 1] > 0]+\
|
||||
[hands2[i, 1] for i in [5,9,13] if hands2[i, 1] > 0]+\
|
||||
[hands1[i, 1] for i in [5,9,13] if hands1[i, 1] > 0]
|
||||
y2 = np.max(y2) if y2 else None
|
||||
|
||||
pad_y1 = 20
|
||||
pad_y2 = 10
|
||||
pad_x1 = pad_x2 = 25
|
||||
|
||||
if cloth_exist:
|
||||
x1 = min(x1, _x1)
|
||||
x2 = max(x2, _x2)
|
||||
y1 = min(y1, _y1)
|
||||
if y2>_y2:
|
||||
x1 = _merge_bound(x1, _x1, min)
|
||||
x2 = _merge_bound(x2, _x2, max)
|
||||
y1 = _merge_bound(y1, _y1, min)
|
||||
if y2 is not None and y2>_y2:
|
||||
pad_y2 = 0
|
||||
y2 = max(y2, _y2)
|
||||
y2 = _merge_bound(y2, _y2, max)
|
||||
x1, y1, x2, y2 = _fallback_box(x1, y1, x2, y2, parse_array.shape)
|
||||
|
||||
y_face = [faces[i, 1] for i in [5, 11] if faces[i, 1] != 0]
|
||||
y_face = [faces[i, 1] for i in [5, 11] if faces[i, 1] > 0]
|
||||
|
||||
if y_face:
|
||||
y_face = np.mean(y_face)
|
||||
@@ -180,33 +185,27 @@ def get_img_agnostic_lower_rectangle(im_parse, pose_data, offset_top, offset_bot
|
||||
except Exception:
|
||||
all_cloth_exist = False
|
||||
|
||||
x1 = [body[i, 0] for i in [8, 9, 10] if body[i, 0] != 0]+[foot2[i, 0] for i in [0, 1, 2] if foot2[i, 0] != 0]+\
|
||||
[body[i, 0] for i in [11, 12, 13] if body[i, 0] != 0]+[foot1[i, 0] for i in [0, 1, 2] if foot1[i, 0] != 0]
|
||||
if x1:
|
||||
x1 = np.min(x1)
|
||||
else:
|
||||
x1=0
|
||||
y1 = min(body[8, 1], body[11, 1])
|
||||
x2 = [body[i, 0] for i in [11, 12, 13] if body[i, 0] != 0]+[foot1[i, 0] for i in [0, 1, 2] if foot1[i, 0] != 0]+\
|
||||
[body[i, 0] for i in [8, 9, 10] if body[i, 0] != 0]+[foot2[i, 0] for i in [0, 1, 2] if foot2[i, 0] != 0]
|
||||
if x2:
|
||||
x2 = np.max(x2)
|
||||
else:
|
||||
x2 = parse_array.shape[1]
|
||||
y2 = [body[i, 1] for i in [10, 13] if body[i, 1] != 0]+[foot1[i, 1] for i in [2] if foot2[i, 1] != 0]+\
|
||||
[foot2[i, 1] for i in [2] if foot2[i, 1] != 0]
|
||||
if y2:
|
||||
y2 = np.max(y2)
|
||||
else:
|
||||
y2 = parse_array.shape[0]
|
||||
x1 = [body[i, 0] for i in [8, 9, 10] if body[i, 0] > 0]+[foot2[i, 0] for i in [0, 1, 2] if foot2[i, 0] > 0]+\
|
||||
[body[i, 0] for i in [11, 12, 13] if body[i, 0] > 0]+[foot1[i, 0] for i in [0, 1, 2] if foot1[i, 0] > 0]
|
||||
x1 = np.min(x1) if x1 else None
|
||||
# Keypoints that were not detected are negative (DWPose marks them -1): use only the detected hips.
|
||||
y1 = [body[i, 1] for i in [8, 11] if body[i, 1] > 0]
|
||||
y1 = np.min(y1) if y1 else None
|
||||
x2 = [body[i, 0] for i in [11, 12, 13] if body[i, 0] > 0]+[foot1[i, 0] for i in [0, 1, 2] if foot1[i, 0] > 0]+\
|
||||
[body[i, 0] for i in [8, 9, 10] if body[i, 0] > 0]+[foot2[i, 0] for i in [0, 1, 2] if foot2[i, 0] > 0]
|
||||
x2 = np.max(x2) if x2 else None
|
||||
y2 = [body[i, 1] for i in [10, 13] if body[i, 1] > 0]+[foot1[i, 1] for i in [2] if foot1[i, 1] > 0]+\
|
||||
[foot2[i, 1] for i in [2] if foot2[i, 1] > 0]
|
||||
y2 = np.max(y2) if y2 else None
|
||||
if cloth_exist:
|
||||
x1 = min(x1, _x1)
|
||||
x2 = max(x2, _x2)
|
||||
y1 = min(y1, _y1)
|
||||
y2 = max(y2, _y2)
|
||||
x1 = _merge_bound(x1, _x1, min)
|
||||
x2 = _merge_bound(x2, _x2, max)
|
||||
y1 = _merge_bound(y1, _y1, min)
|
||||
y2 = _merge_bound(y2, _y2, max)
|
||||
elif all_cloth_exist:
|
||||
x1 = min(x1, _x1)
|
||||
x2 = max(x2, _x2)
|
||||
x1 = _merge_bound(x1, _x1, min)
|
||||
x2 = _merge_bound(x2, _x2, max)
|
||||
x1, y1, x2, y2 = _fallback_box(x1, y1, x2, y2, parse_array.shape)
|
||||
|
||||
pad_y1 = 5 if cloth_exist else 50
|
||||
pad_y2 = 10
|
||||
@@ -266,48 +265,37 @@ def get_img_agnostic_dresses_rectangle(im_parse, pose_data, offset_top, offset_b
|
||||
except Exception:
|
||||
cloth_exist = False
|
||||
|
||||
x1 = [body[i, 0] for i in [2, 3, 4] if body[i, 0] != 0]+\
|
||||
[body[i, 0] for i in [8, 9, 10] if body[i, 0] != 0]+\
|
||||
[foot2[i, 0] for i in [0, 1, 2] if foot2[i, 0] != 0]+\
|
||||
[hands2[i, 0] for i in [5, 9, 13, 17] if hands2[i, 0] != 0]
|
||||
if x1:
|
||||
x1 = np.min(x1)
|
||||
else:
|
||||
x1 = 0
|
||||
y1 = [body[2, 1], body[5, 1]]+\
|
||||
[hands2[i, 1] for i in [5, 9, 13, 17] if hands2[i, 1] != 0]+\
|
||||
[hands1[i, 1] for i in [5, 9, 13, 17] if hands1[i, 1] != 0]
|
||||
if y1:
|
||||
y1 = np.min(y1)
|
||||
else:
|
||||
y1 = 0
|
||||
x2 = [body[i, 0] for i in [5, 6, 7] if body[i, 0] != 0]+\
|
||||
[body[i, 0] for i in [11, 12, 13] if body[i, 0] != 0]+\
|
||||
[foot1[i, 0] for i in [0, 1, 2] if foot1[i, 0] != 0]+\
|
||||
[hands1[i, 0] for i in [5, 9, 13, 17] if hands1[i, 0] != 0]
|
||||
if x2:
|
||||
x2 = np.max(x2)
|
||||
else:
|
||||
x2 = parse_array.shape[1]
|
||||
y2 = [body[i, 1] for i in [10, 13] if body[i, 1] != 0]+[foot1[i, 1] for i in [0, 1, 2] if foot1[i, 1] != 0]+\
|
||||
[foot2[i, 1] for i in [0, 1, 2] if foot2[i, 1] != 0]
|
||||
if y2:
|
||||
y2 = np.max(y2)
|
||||
else:
|
||||
y2 = parse_array.shape[0]
|
||||
x1 = [body[i, 0] for i in [2, 3, 4] if body[i, 0] > 0]+\
|
||||
[body[i, 0] for i in [8, 9, 10] if body[i, 0] > 0]+\
|
||||
[foot2[i, 0] for i in [0, 1, 2] if foot2[i, 0] > 0]+\
|
||||
[hands2[i, 0] for i in [5, 9, 13, 17] if hands2[i, 0] > 0]
|
||||
x1 = np.min(x1) if x1 else None
|
||||
y1 = [body[i, 1] for i in [2, 5] if body[i, 1] > 0]+\
|
||||
[hands2[i, 1] for i in [5, 9, 13, 17] if hands2[i, 1] > 0]+\
|
||||
[hands1[i, 1] for i in [5, 9, 13, 17] if hands1[i, 1] > 0]
|
||||
y1 = np.min(y1) if y1 else None
|
||||
x2 = [body[i, 0] for i in [5, 6, 7] if body[i, 0] > 0]+\
|
||||
[body[i, 0] for i in [11, 12, 13] if body[i, 0] > 0]+\
|
||||
[foot1[i, 0] for i in [0, 1, 2] if foot1[i, 0] > 0]+\
|
||||
[hands1[i, 0] for i in [5, 9, 13, 17] if hands1[i, 0] > 0]
|
||||
x2 = np.max(x2) if x2 else None
|
||||
y2 = [body[i, 1] for i in [10, 13] if body[i, 1] > 0]+[foot1[i, 1] for i in [0, 1, 2] if foot1[i, 1] > 0]+\
|
||||
[foot2[i, 1] for i in [0, 1, 2] if foot2[i, 1] > 0]
|
||||
y2 = np.max(y2) if y2 else None
|
||||
|
||||
if cloth_exist:
|
||||
x1 = min(x1, _x1)
|
||||
x2 = max(x2, _x2)
|
||||
y1 = min(y1, _y1)
|
||||
y2 = max(y2, _y2)
|
||||
x1 = _merge_bound(x1, _x1, min)
|
||||
x2 = _merge_bound(x2, _x2, max)
|
||||
y1 = _merge_bound(y1, _y1, min)
|
||||
y2 = _merge_bound(y2, _y2, max)
|
||||
x1, y1, x2, y2 = _fallback_box(x1, y1, x2, y2, parse_array.shape)
|
||||
|
||||
|
||||
pad_y1 = 20
|
||||
pad_y2 = 10
|
||||
pad_x1 = pad_x2 = 50
|
||||
|
||||
y_face = [faces[i, 1] for i in [5, 11] if faces[i, 1] != 0]
|
||||
y_face = [faces[i, 1] for i in [5, 11] if faces[i, 1] > 0]
|
||||
|
||||
if y_face:
|
||||
y_face = np.mean(y_face)
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
# Changelog
|
||||
|
||||
All notable changes to NH-Nodes. Versions follow the Comfy Registry releases.
|
||||
"Behaviour change" marks fixes that can change the output of an existing workflow.
|
||||
|
||||
## 1.5.0 — Correctness fixes
|
||||
|
||||
### Image
|
||||
- **Batches everywhere:** `VTON Ultimate Processor`, `NH Simple Face Paste` and `NH Agnostic Image Generator` now process every image of a batch (smaller inputs with batch 1 are reused for every item). Behaviour change: Face Paste returns the full destination batch instead of one image. VTON stops with a clear error naming the batch index when no person is detected (it used to return a meaningless mask or crash).
|
||||
- **Values are clipped before 8-bit conversion**, so overshoot from resampling no longer wraps bright pixels to black (Face Paste, Mask-Aware Resize, shared helpers).
|
||||
- `NH Mask-Aware Resize`: an empty mask now respects the mode (pad = fit and pad with `pad_color_hex`, crop = cover and centre) instead of stretching.
|
||||
- `NH Agnostic Image Generator`: masks of a different size are resized to the image, 2-D masks are accepted, soft masks keep their gradient, and the `noise` fill is deterministic (behaviour change).
|
||||
- `Layer Layout Composite` / `Layer Stack Composite`: RGBA layers are pasted with their own alpha instead of as opaque rectangles. RGBA backgrounds give RGBA output.
|
||||
- Padding colours accept `#RRGGBBAA`; without an alpha part the padding is opaque (it used to copy the blue value into alpha).
|
||||
- `Save Image Path (NH)`: `filename_prefix` and `new_folder_name` can no longer write outside the chosen folder (on Windows a prefix may not contain `:`, which blocks drive prefixes such as `D:name`); a prefix such as `renders/img` creates the sub-folder. An absolute `folder_path` is still allowed.
|
||||
- `NH Save Large Image`: WEBP (16383 px) and JPEG (65500 px) size limits are checked before anything is written; `--disable-metadata` is respected.
|
||||
- `Image Label (NH)`: much faster font fitting (cached fonts + binary search), same sizes as before.
|
||||
- `Image Grid Composite (NH)`: refuses a canvas that would not fit in the available memory with a clear message instead of crashing.
|
||||
- `NH Large Preview Cache Manager`: actions run on every queue (it used to be cached).
|
||||
|
||||
### Resolution
|
||||
- `NH Remove Resize Padding`: images resized with `cover_crop` are no longer stretched back (that distorted them). With `original_size` the picture is scaled back uniformly into the original frame and the cropped-away borders are filled from the edges; use `NH Restore Image Geometry` to composite onto the original instead.
|
||||
- Resize nodes keep float precision for Lanczos/box/hamming (no 8-bit quantisation).
|
||||
- `NH Smart Resolution Picker`: a stale preset or model name is reported as a warning in `info`.
|
||||
|
||||
### Logic, text and lists
|
||||
- `Math Eval (NH)`: uses a shared safe evaluator. Behaviour change: an invalid expression now shows an error that names the expression instead of silently returning 0. `x / 0`, `x // 0` and `x % 0` give 0 with a warning. numpy/torch scalar inputs work.
|
||||
- `Random Choice (NH)`: the weight list follows the input slots (`weights` item *i* belongs to `input_i`). Behaviour change when some slots are empty.
|
||||
- Seeds are used literally in `Random Choice`, `Prompt Scheduler`, `Text Split Lines` and `Text Random Line`, so seed 0 is reproducible (behaviour change).
|
||||
- `Prompt Template (NH)`: an empty variable no longer glues the neighbouring words together.
|
||||
- `Regex Extract (NH)`: optional groups that did not match give an empty string instead of "no match".
|
||||
- `Pack Universal (NH)`: values are stored by slot, so an empty `input_0` no longer shifts the other indices (behaviour change).
|
||||
- `Compare (NH)` operator `in`: plain membership / substring test (no more fuzzy numeric matching).
|
||||
- `Batch Merge (NH)`: RGB and RGBA batches can be merged; letterbox bars are transparent for both; very thin inputs no longer crash `pad`.
|
||||
- `NH Mask Aspect Ratio Match`: `pad`/`crop` placement fixed (no crash at the image edge, requested size honoured) and the whole batch is processed.
|
||||
- VTON garment masks: undetected keypoints and missing keypoint groups no longer stretch the mask to the image edges; small parsing islands are removed without cutting rectangles out of the garment.
|
||||
- `Universal Slider Builder (NH)` is marked deprecated (it never produced output); existing workflows still load.
|
||||
|
||||
## 1.4.3 — Tooling, licensing and dependencies
|
||||
- CI runs ruff and the test suite on every pull request; publishing to the Comfy Registry only happens when it passes.
|
||||
- `onnxruntime` is no longer installed automatically (the CPU wheel overwrote `onnxruntime-gpu`). `VTON Ultimate Processor` explains how to install it and falls back to the CPU with a warning.
|
||||
- License notices: the VTON code and parsing models are CC BY-NC-SA 4.0 and the face-parsing model is non-commercial — see `THIRD_PARTY_NOTICES.md`.
|
||||
|
||||
## 1.4.2
|
||||
- `NH Smart Resolution Picker`: FLUX.2 latents are 128 channels at 1/16 (images came out twice the size).
|
||||
- `Load Clip Index (NH)`: works with current ComfyUI (`CLIPType`) and lists every CLIP type.
|
||||
- `Counter (NH)`: counts on every queue.
|
||||
- `NH Mask-Aware Resize`: `crop` keeps the whole mask.
|
||||
@@ -72,13 +72,15 @@ class NH_BatchMerge:
|
||||
CATEGORY = "NH-Nodes/Batch"
|
||||
|
||||
def merge(self, batch_a, batch_b=None, batch_c=None, resize_mode="stretch"):
|
||||
# Target size tu batch_a
|
||||
target_h, target_w = batch_a.shape[1], batch_a.shape[2]
|
||||
# Target size va so kenh tu batch_a
|
||||
target_h, target_w, channels = batch_a.shape[1], batch_a.shape[2], batch_a.shape[3]
|
||||
batches = [batch_a]
|
||||
|
||||
for batch in [batch_b, batch_c]:
|
||||
if batch is None:
|
||||
continue
|
||||
# Channels first, so pad bars of an RGB batch merged into RGBA get the same (transparent) alpha.
|
||||
batch = self._match_channels(batch, channels)
|
||||
if batch.shape[1] != target_h or batch.shape[2] != target_w:
|
||||
batch = self._resize_batch(batch, target_h, target_w, resize_mode)
|
||||
batches.append(batch)
|
||||
@@ -86,6 +88,19 @@ class NH_BatchMerge:
|
||||
result = torch.cat(batches, dim=0)
|
||||
return (result, result.shape[0])
|
||||
|
||||
@staticmethod
|
||||
def _match_channels(batch, channels):
|
||||
"""RGB <-> RGBA de torch.cat chay duoc: bo kenh alpha thua, thieu thi them alpha dac (1.0)."""
|
||||
have = batch.shape[-1]
|
||||
if have == channels:
|
||||
return batch
|
||||
if have not in (3, 4) or channels not in (3, 4):
|
||||
raise ValueError(f"Batch Merge: cannot merge a {have}-channel batch with the {channels}-channel batch_a")
|
||||
if have > channels:
|
||||
return batch[..., :channels]
|
||||
opaque = torch.ones((*batch.shape[:-1], 1), dtype=batch.dtype, device=batch.device)
|
||||
return torch.cat([batch, opaque], dim=-1)
|
||||
|
||||
@staticmethod
|
||||
def _resize_batch(batch, th, tw, mode):
|
||||
"""Resize batch ve target size."""
|
||||
@@ -99,7 +114,7 @@ class NH_BatchMerge:
|
||||
# Scale de fill, roi crop
|
||||
_, _, sh, sw = b.shape
|
||||
scale = max(th / sh, tw / sw)
|
||||
new_h, new_w = round(sh * scale), round(sw * scale)
|
||||
new_h, new_w = max(1, round(sh * scale)), max(1, round(sw * scale))
|
||||
b = F.interpolate(b, size=(new_h, new_w), mode='bilinear', align_corners=False)
|
||||
# Center crop
|
||||
y1 = (new_h - th) // 2
|
||||
@@ -110,7 +125,7 @@ class NH_BatchMerge:
|
||||
# Scale de fit, roi pad
|
||||
_, _, sh, sw = b.shape
|
||||
scale = min(th / sh, tw / sw)
|
||||
new_h, new_w = round(sh * scale), round(sw * scale)
|
||||
new_h, new_w = max(1, round(sh * scale)), max(1, round(sw * scale))
|
||||
b = F.interpolate(b, size=(new_h, new_w), mode='bilinear', align_corners=False)
|
||||
# Center pad
|
||||
pad_top = (th - new_h) // 2
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
@@ -10,6 +11,10 @@ class NH_AgnosticImageGenerator:
|
||||
- agnostic_img: ảnh đã fill vùng mask (dùng làm input cho inpaint/VTON)
|
||||
- masked_img: ảnh gốc với vùng mask bị xóa trắng (transparent)
|
||||
- composite: ảnh gốc với viền mask overlay để preview
|
||||
|
||||
Mask là alpha mềm: giá trị trung gian được blend một phần, feathering làm mềm thêm từ chính các giá trị đó.
|
||||
Mask 2-D [H,W] coi như [1,H,W]; mask khác kích thước ảnh được resize về kích thước ảnh.
|
||||
Fill "noise" tất định: ảnh thứ i trong batch dùng seed = i.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
@@ -36,19 +41,30 @@ class NH_AgnosticImageGenerator:
|
||||
results_masked = []
|
||||
results_composite = []
|
||||
|
||||
# Chuẩn hoá mask về [B,H,W] đúng kích thước ảnh (batch 1 được dùng lại cho mọi ảnh ở vòng lặp dưới).
|
||||
if mask.dim() == 2:
|
||||
mask = mask.unsqueeze(0)
|
||||
if mask.shape[-2:] != image.shape[1:3]:
|
||||
mask = F.interpolate(
|
||||
mask.unsqueeze(1).float(), size=tuple(image.shape[1:3]), mode="bilinear", align_corners=False, antialias=True,
|
||||
).squeeze(1)
|
||||
|
||||
for i in range(image.shape[0]):
|
||||
img_np = image[i].cpu().numpy() # (H, W, C) float32 [0,1]
|
||||
m_np = mask[i].cpu().numpy() if i < mask.shape[0] else mask[0].cpu().numpy() # (H, W)
|
||||
|
||||
h, w, c = img_np.shape
|
||||
|
||||
# Tạo alpha mask [0,1]
|
||||
alpha = (m_np > 0.5).astype(np.float32)
|
||||
# Alpha giữ nguyên giá trị mềm của mask (không ngưỡng hoá), kẹp về [0,1]
|
||||
alpha = np.clip(m_np, 0.0, 1.0).astype(np.float32)
|
||||
# Viền preview cần hình dạng rõ nét: lấy ở mức 0.5 của mask gốc (như trước đây), không dùng alpha mềm
|
||||
outline_mask = (m_np > 0.5).astype(np.float32)
|
||||
|
||||
# Feathering: làm mềm biên mask
|
||||
if feathering > 0:
|
||||
kernel_size = feathering * 2 + 1
|
||||
alpha = cv2.GaussianBlur(alpha, (kernel_size, kernel_size), 0)
|
||||
outline_mask = cv2.GaussianBlur(outline_mask, (kernel_size, kernel_size), 0)
|
||||
|
||||
alpha_3ch = alpha[:, :, np.newaxis] # (H, W, 1)
|
||||
|
||||
@@ -57,7 +73,8 @@ class NH_AgnosticImageGenerator:
|
||||
fill = np.full_like(img_np, gray_value)
|
||||
|
||||
elif fill_mode == "noise":
|
||||
noise = np.random.rand(h, w, c).astype(np.float32)
|
||||
# Seed cố định theo index trong batch: cùng input luôn ra cùng noise
|
||||
noise = np.random.default_rng(i).random((h, w, c), dtype=np.float32)
|
||||
# Blend noise với gray trung tâm để bớt harsh
|
||||
fill = noise * noise_strength + gray_value * (1.0 - noise_strength)
|
||||
|
||||
@@ -75,7 +92,7 @@ class NH_AgnosticImageGenerator:
|
||||
masked = np.clip(masked, 0.0, 1.0)
|
||||
|
||||
# --- composite: overlay viền mask lên ảnh gốc ---
|
||||
composite = self._create_composite(img_np, alpha, agnostic)
|
||||
composite = self._create_composite(img_np, outline_mask, agnostic)
|
||||
|
||||
results_agnostic.append(agnostic)
|
||||
results_masked.append(masked)
|
||||
@@ -87,12 +104,12 @@ class NH_AgnosticImageGenerator:
|
||||
return (out_agnostic, out_masked, out_composite)
|
||||
|
||||
@staticmethod
|
||||
def _create_composite(img_np, alpha, agnostic):
|
||||
def _create_composite(img_np, outline_mask, agnostic):
|
||||
"""Tạo preview: nửa trái = ảnh gốc với viền mask, nửa phải = agnostic."""
|
||||
h, w = img_np.shape[:2]
|
||||
|
||||
# Tìm contour của mask để vẽ viền
|
||||
mask_uint8 = np.ascontiguousarray((alpha * 255).astype(np.uint8))
|
||||
mask_uint8 = np.ascontiguousarray((outline_mask * 255).astype(np.uint8))
|
||||
contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
# Nửa trái: ảnh gốc + viền xanh
|
||||
|
||||
+33
-10
@@ -11,6 +11,14 @@ def get_bounding_box(mask_np):
|
||||
xmin, xmax = np.where(cols)[0][[0, -1]]
|
||||
return xmin, ymin, xmax, ymax
|
||||
|
||||
def _to_uint8(image):
|
||||
# Clip trước khi ép uint8: giá trị lệch khỏi [0,1] (vd. do resample) sẽ bị wrap về gần 0 nếu không clip.
|
||||
return (np.clip(image.cpu().numpy(), 0.0, 1.0) * 255).astype(np.uint8)
|
||||
|
||||
def _as_mask_batch(mask):
|
||||
# MASK 2-D [H,W] không có chiều batch: coi là [1,H,W].
|
||||
return mask.unsqueeze(0) if mask.dim() == 2 else mask
|
||||
|
||||
class NH_SimpleFacePaste:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -29,20 +37,37 @@ class NH_SimpleFacePaste:
|
||||
CATEGORY = "NH-Nodes/Image"
|
||||
|
||||
def paste_face(self, dest_image, dest_mask, source_image, source_mask, feathering):
|
||||
src_img_np = (source_image[0].cpu().numpy() * 255).astype(np.uint8)
|
||||
src_mask_np = source_mask[0].cpu().numpy()
|
||||
|
||||
dest_img_np = (dest_image[0].cpu().numpy() * 255).astype(np.uint8)
|
||||
dest_mask_np = dest_mask[0].cpu().numpy()
|
||||
dest_mask = _as_mask_batch(dest_mask)
|
||||
source_mask = _as_mask_batch(source_mask)
|
||||
# Output theo batch của dest_image; input khác có batch 1 (hoặc ngắn hơn) được lặp lại theo vòng.
|
||||
output = dest_image.clone()
|
||||
for index in range(dest_image.shape[0]):
|
||||
pasted = self._paste_single(
|
||||
dest_image[index], dest_mask[index % dest_mask.shape[0]],
|
||||
source_image[index % source_image.shape[0]], source_mask[index % source_mask.shape[0]],
|
||||
feathering,
|
||||
)
|
||||
if pasted is not None: # mask rỗng: giữ nguyên ảnh đích của item này
|
||||
output[index] = pasted
|
||||
return (output,)
|
||||
|
||||
@staticmethod
|
||||
def _paste_single(dest_image, dest_mask, source_image, source_mask, feathering):
|
||||
"""Dán mặt của một ảnh nguồn [H,W,C] lên một ảnh đích [H,W,C]; trả None nếu một trong hai mask rỗng."""
|
||||
src_img_np = _to_uint8(source_image)
|
||||
src_mask_np = source_mask.cpu().numpy()
|
||||
|
||||
dest_img_np = _to_uint8(dest_image)
|
||||
dest_mask_np = dest_mask.cpu().numpy()
|
||||
|
||||
src_x1, src_y1, src_x2, src_y2 = get_bounding_box(src_mask_np > 0.5)
|
||||
dest_x1, dest_y1, dest_x2, dest_y2 = get_bounding_box(dest_mask_np > 0.5)
|
||||
|
||||
if src_x1 is None or dest_x1 is None:
|
||||
return (dest_image,)
|
||||
return None
|
||||
|
||||
src_face = src_img_np[src_y1:src_y2+1, src_x1:src_x2+1]
|
||||
src_face_mask = (src_mask_np[src_y1:src_y2+1, src_x1:src_x2+1] * 255).astype(np.uint8)
|
||||
src_face_mask = (np.clip(src_mask_np[src_y1:src_y2+1, src_x1:src_x2+1], 0.0, 1.0) * 255).astype(np.uint8)
|
||||
|
||||
dest_width = dest_x2 - dest_x1 + 1
|
||||
dest_height = dest_y2 - dest_y1 + 1
|
||||
@@ -61,10 +86,8 @@ class NH_SimpleFacePaste:
|
||||
|
||||
output_img_np = dest_img_np.copy()
|
||||
output_img_np[dest_y1:dest_y1+dest_height, dest_x1:dest_x1+dest_width] = blended_area.astype(np.uint8)
|
||||
|
||||
output_image = torch.from_numpy(output_img_np.astype(np.float32) / 255.0).unsqueeze(0)
|
||||
|
||||
return (output_image,)
|
||||
return torch.from_numpy(output_img_np.astype(np.float32) / 255.0)
|
||||
|
||||
# --- PHẦN BỊ THIẾU TRƯỚC ĐÂY ĐÃ ĐƯỢC THÊM VÀO ĐÂY ---
|
||||
# Ánh xạ để ComfyUI nhận diện class
|
||||
|
||||
@@ -2,9 +2,73 @@ import numpy as np
|
||||
import torch
|
||||
from PIL import Image, ImageColor
|
||||
|
||||
try:
|
||||
import psutil
|
||||
except ImportError: # ComfyUI installs psutil; without it the memory guard below is simply skipped
|
||||
psutil = None
|
||||
|
||||
|
||||
_RESAMPLING = getattr(Image, "Resampling", Image)
|
||||
|
||||
# Size guard for the output canvas. It is about RAM, not about Pillow's decompression-bomb limit
|
||||
# (Image.MAX_IMAGE_PIXELS): that limit is process-wide, other packs raise or disable it, and it only applies to
|
||||
# Image.open, never to an IMAGE that stays in memory.
|
||||
#
|
||||
# Peak RAM per canvas pixel, reached inside _pil_to_tensor (measured 28.1-28.2 B/px on Windows, Pillow 11, NumPy 2):
|
||||
# 4 B the PIL canvas (Pillow stores an RGB image as 4 bytes per pixel)
|
||||
# 12 B the float32 array from astype(np.float32)
|
||||
# 12 B the float32 result of "/ 255.0" (NumPy only reuses that temporary on Linux)
|
||||
# By then the fitted cells and the uint8 array from np.array() are freed, and torch.from_numpy() shares the float32
|
||||
# result, so the IMAGE tensor adds nothing (its 12 B/px are what stays alive after the node returns). The input
|
||||
# images were converted before the check, so psutil's "available" already leaves their memory out.
|
||||
_PEAK_BYTES_PER_PIXEL = 4 + 12 + 12
|
||||
# Share of the available RAM the node may plan to use; the rest is left to the OS and the other programs.
|
||||
_USABLE_MEMORY_FRACTION = 0.85
|
||||
# Image.new raises a bare MemoryError for an image wider than INT_MAX // 4 - 1 px, whatever the memory; the node's
|
||||
# own limits never get near it. The same ceiling is applied to the height to keep one rule.
|
||||
_PIL_MAX_SIDE = (2**31 - 1) // 4 - 1
|
||||
|
||||
|
||||
def _estimate_canvas_bytes(canvas_width, canvas_height):
|
||||
"""Peak RAM in bytes the node needs to turn a canvas of this size into an IMAGE tensor."""
|
||||
return _PEAK_BYTES_PER_PIXEL * canvas_width * canvas_height
|
||||
|
||||
|
||||
def _available_memory_bytes():
|
||||
"""Free physical RAM in bytes, or None when it cannot be read (the guard is advisory, so it then stays quiet)."""
|
||||
if psutil is None:
|
||||
return None
|
||||
try:
|
||||
return int(psutil.virtual_memory().available)
|
||||
except Exception: # whatever the platform backend raises must not stop a canvas that may well fit
|
||||
return None
|
||||
|
||||
|
||||
def _format_size(nbytes):
|
||||
return f"{nbytes / 2**30:.1f} GiB" if nbytes >= 2**30 else f"{nbytes / 2**20:.0f} MiB"
|
||||
|
||||
|
||||
def _check_canvas_size(canvas_width, canvas_height, columns, rows, cell_width, cell_height):
|
||||
"""Raise ValueError, before anything is allocated, for a canvas the node cannot build in the free memory."""
|
||||
layout = f"{canvas_width}x{canvas_height} px canvas ({columns}x{rows} cells of {cell_width}x{cell_height} px)"
|
||||
hint = "Lower columns/rows, the cell size or the gaps."
|
||||
|
||||
if max(canvas_width, canvas_height) > _PIL_MAX_SIDE:
|
||||
raise ValueError(
|
||||
f"Image Grid Composite: the {layout} is too large: Pillow cannot create an image wider or taller "
|
||||
f"than {_PIL_MAX_SIDE:,} px. {hint}"
|
||||
)
|
||||
|
||||
available = _available_memory_bytes()
|
||||
if available is None:
|
||||
return
|
||||
needed = _estimate_canvas_bytes(canvas_width, canvas_height)
|
||||
if needed > available * _USABLE_MEMORY_FRACTION:
|
||||
raise ValueError(
|
||||
f"Image Grid Composite: the {layout} is too large for the free memory: it needs about "
|
||||
f"{_format_size(needed)} of RAM and {_format_size(available)} is available. {hint}"
|
||||
)
|
||||
|
||||
|
||||
def _parse_hex_color(value, default):
|
||||
try:
|
||||
@@ -157,6 +221,7 @@ class NH_ImageGridComposite:
|
||||
|
||||
canvas_width = (columns * cell_width) + ((columns - 1) * gap_x)
|
||||
canvas_height = (rows * cell_height) + ((rows - 1) * gap_y)
|
||||
_check_canvas_size(canvas_width, canvas_height, columns, rows, cell_width, cell_height)
|
||||
fill_rgb = _parse_hex_color(background_color_hex, (0, 0, 0))
|
||||
canvas = Image.new("RGB", (canvas_width, canvas_height), fill_rgb)
|
||||
|
||||
@@ -173,8 +238,8 @@ class NH_ImageGridComposite:
|
||||
x = column * (cell_width + gap_x)
|
||||
y = row * (cell_height + gap_y)
|
||||
source = source_images[image_index]
|
||||
fitted = _fit_image_to_cell(source, cell_width, cell_height, resize_mode, fill_rgb)
|
||||
canvas.paste(fitted, (x, y))
|
||||
# No local name for the fitted cell: it must be freed here, not kept until the tensor conversion.
|
||||
canvas.paste(_fit_image_to_cell(source, cell_width, cell_height, resize_mode, fill_rgb), (x, y))
|
||||
info_lines.append(
|
||||
f"{image_index}: x={x}, y={y}, w={cell_width}, h={cell_height}, source={source.width}x{source.height}"
|
||||
)
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
import functools
|
||||
import os
|
||||
import re
|
||||
import unicodedata
|
||||
import uuid
|
||||
import ast
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
|
||||
@@ -15,8 +15,12 @@ import comfy.utils
|
||||
import folder_paths
|
||||
import node_helpers
|
||||
|
||||
from ...core.safe_eval import evaluate
|
||||
from ...core.utils import parse_hex_color, tensor_to_pil
|
||||
|
||||
|
||||
_RESAMPLING = getattr(Image, "Resampling", Image)
|
||||
_WINDOWS = os.name == "nt"
|
||||
_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".avif", ".bmp", ".tif", ".tiff"}
|
||||
_FONT_EXTENSIONS = {".ttf", ".otf", ".ttc"}
|
||||
_MAX_LAYER_STACK_INPUTS = 64
|
||||
@@ -64,10 +68,19 @@ _FONT_CHOICES = _font_choices()
|
||||
|
||||
|
||||
def _parse_hex_color(value, default):
|
||||
"""Parse a colour into an (r, g, b, a) tuple of 0-255 ints; ``default`` when the text is not a colour.
|
||||
|
||||
'#RGB', '#RRGGBB' and '#RRGGBBAA' go through core.utils.parse_hex_color (opaque unless AA is given). Other CSS
|
||||
colours Pillow understands ('red', 'rgb(1, 2, 3)') keep working, as they did before.
|
||||
"""
|
||||
text = str(value or "").strip()
|
||||
color = parse_hex_color(text, ())
|
||||
if color:
|
||||
return color
|
||||
try:
|
||||
return ImageColor.getrgb((value or "").strip())
|
||||
except Exception:
|
||||
return default
|
||||
return ImageColor.getcolor(text, "RGBA")
|
||||
except ValueError:
|
||||
return tuple(default) + (255,) * (4 - len(default))
|
||||
|
||||
|
||||
def _tensor_to_pil(image_tensor):
|
||||
@@ -75,24 +88,71 @@ def _tensor_to_pil(image_tensor):
|
||||
return Image.fromarray(array).convert("RGB")
|
||||
|
||||
|
||||
def _layer_to_pil(image_tensor):
|
||||
"""Like _tensor_to_pil, but a 4-channel image stays RGBA so its transparency is not lost."""
|
||||
pil_image = tensor_to_pil(image_tensor)
|
||||
return pil_image if pil_image.mode == "RGBA" else pil_image.convert("RGB")
|
||||
|
||||
|
||||
def _pil_to_tensor(image):
|
||||
array = np.array(image).astype(np.float32) / 255.0
|
||||
dtype = comfy.model_management.intermediate_dtype()
|
||||
return torch.from_numpy(array).unsqueeze(0).to(dtype=dtype)
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=64)
|
||||
def _cached_truetype(path, size):
|
||||
"""Open a TrueType font once per (path, size), as the Label node asks for many sizes. None if it cannot be opened."""
|
||||
try:
|
||||
return ImageFont.truetype(path, size)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _default_font():
|
||||
return ImageFont.load_default()
|
||||
|
||||
|
||||
def _load_font(font_name, font_size):
|
||||
if font_name and font_name != "default":
|
||||
font_path = Path(os.environ.get("WINDIR", "C:\\Windows")) / "Fonts" / font_name
|
||||
if font_path.exists():
|
||||
try:
|
||||
return ImageFont.truetype(str(font_path), font_size)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
return ImageFont.truetype("arial.ttf", font_size)
|
||||
except Exception:
|
||||
return ImageFont.load_default()
|
||||
font = _cached_truetype(str(font_path), font_size)
|
||||
if font is not None:
|
||||
return font
|
||||
font = _cached_truetype("arial.ttf", font_size)
|
||||
return font if font is not None else _default_font()
|
||||
|
||||
|
||||
_MIN_FIT_FONT_SIZE = 8
|
||||
|
||||
|
||||
def _fit_font_size(font_size, fits):
|
||||
"""Largest size in [8, font_size] for which ``fits(size)`` is true; 8 when none does.
|
||||
|
||||
A label gets wider as the font grows, so a binary search lands on the same size as shrinking one point at a
|
||||
time, after about log2(font_size) measurements instead of up to font_size. Hinting can make a text a pixel or
|
||||
two narrower at a few isolated sizes; then the size returned still fits but may sit a point or two below what
|
||||
the one-point-at-a-time loop would have found.
|
||||
"""
|
||||
if font_size <= _MIN_FIT_FONT_SIZE or fits(font_size):
|
||||
return font_size
|
||||
low, high = _MIN_FIT_FONT_SIZE, font_size # `low` fits (or is the floor), `high` does not
|
||||
while high - low > 1:
|
||||
middle = (low + high) // 2
|
||||
if fits(middle):
|
||||
low = middle
|
||||
else:
|
||||
high = middle
|
||||
return low
|
||||
|
||||
|
||||
def _label_text_width(draw, text, font_name, font_size, stroke_width):
|
||||
left, _top, right, _bottom = draw.multiline_textbbox(
|
||||
(0, 0), text, font=_load_font(font_name, font_size), stroke_width=stroke_width
|
||||
)
|
||||
return right - left
|
||||
|
||||
|
||||
def _resize_image(image, size, method):
|
||||
@@ -134,11 +194,16 @@ def _common_upscale_method(resize_method):
|
||||
|
||||
|
||||
def _pad_samples(samples, target_width, target_height, fill_rgb):
|
||||
"""Centre ``samples`` ([B, C, H, W]) on a canvas of the padding colour.
|
||||
|
||||
``fill_rgb`` is (r, g, b) or (r, g, b, a) in 0-255. The fourth channel of RGBA images takes the alpha value, which
|
||||
is opaque when the colour has none.
|
||||
"""
|
||||
batch, channels, resized_h, resized_w = samples.shape
|
||||
canvas = torch.empty((batch, channels, target_height, target_width), dtype=samples.dtype, device=samples.device)
|
||||
fill_values = [channel / 255.0 for channel in fill_rgb]
|
||||
fill_values = [channel / 255.0 for channel in fill_rgb[:3]] + [fill_rgb[3] / 255.0 if len(fill_rgb) > 3 else 1.0]
|
||||
for channel in range(channels):
|
||||
canvas[:, channel, :, :] = fill_values[min(channel, len(fill_values) - 1)]
|
||||
canvas[:, channel, :, :] = fill_values[min(channel, 3)]
|
||||
|
||||
offset_x = max(0, (target_width - resized_w) // 2)
|
||||
offset_y = max(0, (target_height - resized_h) // 2)
|
||||
@@ -163,6 +228,18 @@ def _resize_contain(image, target_size):
|
||||
return _resize_image(image, (max(1, round(width * scale)), max(1, round(height * scale))), "lanczos")
|
||||
|
||||
|
||||
def _paste_layer(canvas, layer, position):
|
||||
"""Paste ``layer`` onto ``canvas`` through its own alpha, so transparent areas leave the canvas untouched."""
|
||||
if layer.mode != "RGBA":
|
||||
canvas.paste(layer, position) # opaque, whatever the canvas mode
|
||||
elif canvas.mode == "RGBA":
|
||||
# Both have alpha: a straight paste would blend the alpha channels as if they were colours.
|
||||
box = (position[0], position[1], position[0] + layer.width, position[1] + layer.height)
|
||||
canvas.paste(Image.alpha_composite(canvas.crop(box), layer), position)
|
||||
else:
|
||||
canvas.paste(layer, position, layer)
|
||||
|
||||
|
||||
def _fit_layer_image(image, target_size, resize_mode):
|
||||
target_width, target_height = target_size
|
||||
if target_width <= 0 or target_height <= 0:
|
||||
@@ -289,42 +366,12 @@ def _layout_slots(layout_mode, layout_preset, canvas_width, canvas_height, colum
|
||||
|
||||
|
||||
def _eval_position_expr(expression, variables):
|
||||
operators = {
|
||||
ast.Add: lambda left, right: left + right,
|
||||
ast.Sub: lambda left, right: left - right,
|
||||
ast.Mult: lambda left, right: left * right,
|
||||
ast.Div: lambda left, right: left / right,
|
||||
ast.FloorDiv: lambda left, right: left // right,
|
||||
ast.Mod: lambda left, right: left % right,
|
||||
ast.Pow: lambda left, right: left ** right,
|
||||
}
|
||||
unary_operators = {
|
||||
ast.UAdd: lambda value: value,
|
||||
ast.USub: lambda value: -value,
|
||||
}
|
||||
|
||||
def visit(node):
|
||||
if isinstance(node, ast.Expression):
|
||||
return visit(node.body)
|
||||
if isinstance(node, ast.Constant) and isinstance(node.value, (int, float)):
|
||||
return node.value
|
||||
if isinstance(node, ast.Num):
|
||||
return node.n
|
||||
if isinstance(node, ast.Name):
|
||||
key = node.id.casefold()
|
||||
if key not in variables:
|
||||
raise ValueError(f"Unknown position variable: {node.id}")
|
||||
return variables[key]
|
||||
if isinstance(node, ast.BinOp) and type(node.op) in operators:
|
||||
return operators[type(node.op)](visit(node.left), visit(node.right))
|
||||
if isinstance(node, ast.UnaryOp) and type(node.op) in unary_operators:
|
||||
return unary_operators[type(node.op)](visit(node.operand))
|
||||
raise ValueError(f"Unsupported position expression: {expression}")
|
||||
|
||||
# Variable keys are already casefolded. core.safe_eval works in floating point, so 9**9**9 overflows into an
|
||||
# error instead of building a huge integer and freezing the worker.
|
||||
try:
|
||||
return visit(ast.parse(str(expression), mode="eval"))
|
||||
except SyntaxError as exc:
|
||||
raise ValueError(f"Invalid position expression: {expression}") from exc
|
||||
return evaluate(expression, variables, zero_division="error", case_insensitive=True)
|
||||
except ValueError as exc: # core.safe_eval.ExpressionError is a ValueError
|
||||
raise ValueError(f"Invalid position expression '{expression}': {exc}") from exc
|
||||
|
||||
|
||||
def _split_position_line(line):
|
||||
@@ -377,16 +424,48 @@ def _save_dir_from_text(folder_path):
|
||||
return os.path.join(folder_paths.get_output_directory(), folder_path)
|
||||
|
||||
|
||||
def _contained_path(base_dir, relative, what):
|
||||
"""Absolute path of ``relative`` below ``base_dir``; ValueError (naming ``what``) if it would leave ``base_dir``.
|
||||
|
||||
Both slash styles count as separators on every OS, so a workflow written on Windows is judged the same elsewhere.
|
||||
"""
|
||||
base_dir = os.path.abspath(base_dir)
|
||||
path = os.path.abspath(os.path.join(base_dir, str(relative).replace("\\", "/")))
|
||||
try:
|
||||
inside = os.path.commonpath([base_dir, path]) == base_dir
|
||||
except ValueError: # another drive or a UNC share
|
||||
inside = False
|
||||
if not inside:
|
||||
raise ValueError(f"{what} must stay inside the target folder: {relative!r}")
|
||||
return path
|
||||
|
||||
|
||||
def _save_child_dir(save_dir, folder_name):
|
||||
folder_name = str(folder_name or "").strip()
|
||||
if not folder_name:
|
||||
return save_dir
|
||||
return _contained_path(save_dir, folder_name, "new_folder_name")
|
||||
|
||||
base_dir = os.path.abspath(save_dir)
|
||||
child_dir = os.path.abspath(os.path.join(base_dir, folder_name))
|
||||
if os.path.commonpath([base_dir, child_dir]) != base_dir:
|
||||
raise ValueError("new_folder_name must stay inside folder_path")
|
||||
return child_dir
|
||||
|
||||
def _split_filename_prefix(save_dir, prefix):
|
||||
"""('sub/dir/', 'name', directory) for a filename_prefix such as 'sub/dir/name'.
|
||||
|
||||
Subfolders are allowed, but the folder part must stay inside ``save_dir`` (no '..', no absolute path) and the
|
||||
file name must be a plain name. Without a separator the name is not looked at by the folder check, and on
|
||||
Windows a colon turns it into a path: os.path.join(save_dir, 'D:x0001.png') is a file on drive D:, 'C:x' is
|
||||
relative to the current folder of drive C: (join drops the 'C:', so the reported name is wrong and the
|
||||
numbering never finds the file again) and 'name:stream' writes an NTFS alternate data stream. A Windows file
|
||||
name cannot contain a colon anyway, so such a prefix is refused whole, before anything is created.
|
||||
"""
|
||||
if _WINDOWS and ":" in prefix:
|
||||
raise ValueError(
|
||||
"filename_prefix cannot contain ':' on Windows (it selects a drive or an NTFS stream, "
|
||||
f"not a file name): {prefix!r}"
|
||||
)
|
||||
folder_part, name = re.match(r"(.*[\\/])?([^\\/]*)\Z", prefix, re.DOTALL).groups()
|
||||
if not folder_part:
|
||||
return "", name, save_dir
|
||||
return folder_part, name, _contained_path(save_dir, folder_part, "filename_prefix")
|
||||
|
||||
|
||||
def _resolve_load_dir(folder_path, image="", upload_probe=""):
|
||||
@@ -738,16 +817,19 @@ class NH_SaveImagePath:
|
||||
save_dir = _save_dir_from_text(folder_path)
|
||||
if create_new_folder:
|
||||
save_dir = _save_child_dir(save_dir, new_folder_name)
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
# A prefix may name subfolders ("renders/img"); they live below save_dir and are where files are counted.
|
||||
prefix = (filename_prefix or "NH").strip() or "NH"
|
||||
numbered_pattern = re.compile(rf"^{re.escape(prefix)}(\d+)\.(?:png|jpe?g|webp|avif)$", re.IGNORECASE)
|
||||
legacy_numbered_pattern = re.compile(rf"^(?:{re.escape(prefix)})?(\d+)\.(?:png|jpe?g|webp|avif)$", re.IGNORECASE)
|
||||
bare_pattern = re.compile(rf"^{re.escape(prefix)}\.(?:png|jpe?g|webp|avif)$", re.IGNORECASE)
|
||||
prefix_folder, prefix_name, target_dir = _split_filename_prefix(save_dir, prefix)
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
|
||||
numbered_pattern = re.compile(rf"^{re.escape(prefix_name)}(\d+)\.(?:png|jpe?g|webp|avif)$", re.IGNORECASE)
|
||||
legacy_numbered_pattern = re.compile(rf"^(?:{re.escape(prefix_name)})?(\d+)\.(?:png|jpe?g|webp|avif)$", re.IGNORECASE)
|
||||
bare_pattern = re.compile(rf"^{re.escape(prefix_name)}\.(?:png|jpe?g|webp|avif)$", re.IGNORECASE)
|
||||
|
||||
used_numbers = set()
|
||||
bare_exists = False
|
||||
for name in os.listdir(save_dir):
|
||||
for name in os.listdir(target_dir):
|
||||
if bare_pattern.match(name):
|
||||
bare_exists = True
|
||||
continue
|
||||
@@ -765,22 +847,22 @@ class NH_SaveImagePath:
|
||||
if number_suffix:
|
||||
while next_number in used_numbers:
|
||||
next_number += 1
|
||||
filename = f"{prefix}{next_number:04d}.{file_format}"
|
||||
filename = f"{prefix_name}{next_number:04d}.{file_format}"
|
||||
used_numbers.add(next_number)
|
||||
next_number += 1
|
||||
elif not bare_exists and not used_numbers:
|
||||
filename = f"{prefix}.{file_format}"
|
||||
filename = f"{prefix_name}.{file_format}"
|
||||
bare_exists = True
|
||||
else:
|
||||
while next_number in used_numbers:
|
||||
next_number += 1
|
||||
filename = f"{prefix}{next_number}.{file_format}"
|
||||
filename = f"{prefix_name}{next_number}.{file_format}"
|
||||
used_numbers.add(next_number)
|
||||
next_number += 1
|
||||
|
||||
pil_image = _tensor_to_pil(image)
|
||||
_save_pil_image(pil_image, os.path.join(save_dir, filename), file_format, int(dpi))
|
||||
saved_files.append(filename)
|
||||
_save_pil_image(pil_image, os.path.join(target_dir, filename), file_format, int(dpi))
|
||||
saved_files.append(prefix_folder + filename)
|
||||
|
||||
if preview_image:
|
||||
preview_filename = f"NH_preview_{uuid.uuid4().hex}_{batch_index}.png"
|
||||
@@ -1048,16 +1130,18 @@ class NH_LayerLayoutComposite:
|
||||
if right <= left or bottom <= top:
|
||||
return
|
||||
crop_box = (left - int(x), top - int(y), right - int(x), bottom - int(y))
|
||||
canvas.paste(layer.crop(crop_box), (left, top))
|
||||
_paste_layer(canvas, layer.crop(crop_box), (left, top))
|
||||
|
||||
def _create_canvas(self, canvas_width, canvas_height, background_color, background_fit, background_image=None):
|
||||
fill_rgb = _parse_hex_color(background_color, (255, 255, 255))
|
||||
canvas = Image.new("RGB", (canvas_width, canvas_height), fill_rgb)
|
||||
fill = _parse_hex_color(background_color, (255, 255, 255))
|
||||
# An opaque colour gives the RGB canvas (and RGB output) this node always had; a colour with transparency
|
||||
# ('#RRGGBBAA' with AA below FF) gives an RGBA canvas so that the transparency survives.
|
||||
canvas = Image.new("RGBA" if fill[3] < 255 else "RGB", (canvas_width, canvas_height), fill)
|
||||
|
||||
if background_image is None:
|
||||
return canvas
|
||||
|
||||
background = _tensor_to_pil(background_image[0])
|
||||
background = _layer_to_pil(background_image[0])
|
||||
fitted, offset_x, offset_y = _fit_layer_image(background, (canvas_width, canvas_height), background_fit)
|
||||
self._paste_clipped(canvas, fitted, offset_x, offset_y)
|
||||
return canvas
|
||||
@@ -1105,7 +1189,7 @@ class NH_LayerLayoutComposite:
|
||||
|
||||
for image_index in range(used_count):
|
||||
x, y, width, height = slots[image_index]
|
||||
layer = _tensor_to_pil(images[image_index])
|
||||
layer = _layer_to_pil(images[image_index])
|
||||
fitted, offset_x, offset_y = _fit_layer_image(layer, (int(width), int(height)), resize_mode)
|
||||
paste_x = int(x) + offset_x
|
||||
paste_y = int(y) + offset_y
|
||||
@@ -1157,11 +1241,12 @@ class NH_LayerStackComposite:
|
||||
return False
|
||||
|
||||
crop_box = (left - x, top - y, right - x, bottom - y)
|
||||
canvas.paste(layer.crop(crop_box), (left, top))
|
||||
_paste_layer(canvas, layer.crop(crop_box), (left, top))
|
||||
return True
|
||||
|
||||
def composite(self, background_image, layer_count, positions, **kwargs):
|
||||
background = _tensor_to_pil(background_image[0])
|
||||
# The output has the channels of the background: RGB stays RGB, an RGBA background gives RGBA.
|
||||
background = _layer_to_pil(background_image[0])
|
||||
canvas = background.copy()
|
||||
requested_count = max(0, min(int(layer_count), _MAX_LAYER_STACK_INPUTS))
|
||||
layer_images = {}
|
||||
@@ -1171,7 +1256,7 @@ class NH_LayerStackComposite:
|
||||
if layer is None:
|
||||
layer_sizes.append((0, 0))
|
||||
continue
|
||||
layer_image = _tensor_to_pil(layer[0])
|
||||
layer_image = _layer_to_pil(layer[0])
|
||||
layer_images[index] = layer_image
|
||||
layer_sizes.append((layer_image.width, layer_image.height))
|
||||
|
||||
@@ -1416,17 +1501,17 @@ class NH_ImageLabel:
|
||||
band_top = height
|
||||
|
||||
draw = ImageDraw.Draw(canvas)
|
||||
font = _load_font(font_name, font_size)
|
||||
bbox = draw.multiline_textbbox((0, 0), label_text, font=font, stroke_width=outline_width)
|
||||
current_size = font_size
|
||||
|
||||
if fit_text:
|
||||
max_width = max(1, width - (padding_x * 2))
|
||||
while (bbox[2] - bbox[0]) > max_width and current_size > 8:
|
||||
current_size -= 1
|
||||
font = _load_font(font_name, current_size)
|
||||
bbox = draw.multiline_textbbox((0, 0), label_text, font=font, stroke_width=outline_width)
|
||||
current_size = _fit_font_size(
|
||||
font_size,
|
||||
lambda size: _label_text_width(draw, label_text, font_name, size, outline_width) <= max_width,
|
||||
)
|
||||
|
||||
font = _load_font(font_name, current_size)
|
||||
bbox = draw.multiline_textbbox((0, 0), label_text, font=font, stroke_width=outline_width)
|
||||
text_width = bbox[2] - bbox[0]
|
||||
text_height = bbox[3] - bbox[1]
|
||||
if align == "left":
|
||||
|
||||
@@ -5,9 +5,23 @@ import os
|
||||
|
||||
import folder_paths
|
||||
import numpy as np
|
||||
from comfy.cli_args import args
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
# Cạnh dài nhất (px) mà từng định dạng ghi được; PNG/TIFF không bị giới hạn trong thực tế.
|
||||
# JPEG: spec cho tới 65535 nhưng libjpeg (Pillow) từ chối > 65500 (JPEG_MAX_DIMENSION).
|
||||
_MAX_SIDE_BY_EXTENSION = {"webp": 16383, "jpg": 65500}
|
||||
|
||||
|
||||
def _check_size_limit(extension, width, height):
|
||||
limit = _MAX_SIDE_BY_EXTENSION.get(extension)
|
||||
if limit is not None and max(width, height) > limit:
|
||||
raise ValueError(
|
||||
f"NH Save Large Image: {width}x{height} px is too large for {extension.upper()} "
|
||||
f"(max {limit} px per side). Use png or tiff for images this large."
|
||||
)
|
||||
|
||||
|
||||
def _extension_for_save_format(file_format):
|
||||
file_format = str(file_format or "png").lower()
|
||||
@@ -85,6 +99,10 @@ class NH_SaveLargeImage:
|
||||
|
||||
def save(self, image, filename_prefix="NH_Large", file_format="png", quality=95, embed_metadata=True, prompt=None, extra_pnginfo=None):
|
||||
extension = _extension_for_save_format(file_format)
|
||||
# Kiểm tra giới hạn trước khi ghi bất kỳ file nào (mọi ảnh trong batch cùng kích thước), kể cả tạo thư mục.
|
||||
_check_size_limit(extension, int(image.shape[2]), int(image.shape[1]))
|
||||
# --disable-metadata của ComfyUI thắng ô embed_metadata: không nhúng prompt/workflow.
|
||||
embed_metadata = bool(embed_metadata) and not args.disable_metadata
|
||||
first = image[0]
|
||||
full_output_folder, filename, counter, _subfolder, _filename_prefix = folder_paths.get_save_image_path(
|
||||
filename_prefix,
|
||||
@@ -99,7 +117,7 @@ class NH_SaveLargeImage:
|
||||
file_name = f"{filename_with_batch_num}_{counter:05}_.{extension}"
|
||||
file_path = os.path.join(full_output_folder, file_name)
|
||||
pil_image = _tensor_to_pil(image_tensor)
|
||||
_save_image(pil_image, file_path, extension, quality, bool(embed_metadata), prompt, extra_pnginfo)
|
||||
_save_image(pil_image, file_path, extension, quality, embed_metadata, prompt, extra_pnginfo)
|
||||
saved_paths.append(file_path)
|
||||
counter += 1
|
||||
|
||||
|
||||
@@ -24,6 +24,12 @@ class NH_LargePreviewCacheManager:
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "NH-Nodes/Image"
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, **kwargs):
|
||||
# Info/cleanup đọc và xoá file trên đĩa, input không đổi vẫn phải chạy lại mỗi lần queue;
|
||||
# NaN không bằng chính nó nên ComfyUI không bao giờ dùng lại output cache.
|
||||
return float("nan")
|
||||
|
||||
def run(self, action="info", max_age_hours=24.0, max_total_mb=2048.0, cache_id=""):
|
||||
action = str(action or "info")
|
||||
before = cache_info()
|
||||
|
||||
@@ -8,6 +8,7 @@ class NH_MaskResizeImage:
|
||||
|
||||
- Crop: scale để mask vừa khít target, crop phần ảnh thừa (pad nếu ảnh nhỏ hơn target).
|
||||
- Pad: scale vừa đủ để mask nằm trong target, pad phần thừa bằng màu tùy chọn.
|
||||
- Mask rỗng: coi như mask phủ cả ảnh (crop: lấp đầy target rồi cắt giữa; pad: cả ảnh nằm gọn trong target).
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
@@ -44,7 +45,8 @@ class NH_MaskResizeImage:
|
||||
results_img.append(out_img)
|
||||
results_mask.append(out_mask)
|
||||
|
||||
out_images = torch.from_numpy(np.stack(results_img, axis=0)).float()
|
||||
# Lanczos có thể vọt ra ngoài [0,1] (ringing) quanh cạnh tương phản cao -> kẹp lại
|
||||
out_images = torch.from_numpy(np.clip(np.stack(results_img, axis=0), 0.0, 1.0)).float()
|
||||
out_masks = torch.from_numpy(np.stack(results_mask, axis=0)).float()
|
||||
return (out_images, out_masks)
|
||||
|
||||
@@ -144,10 +146,23 @@ class NH_MaskResizeImage:
|
||||
return canvas_img, canvas_mask
|
||||
|
||||
def _simple_resize(self, img_np, mask_np, tw, th, mode, pad_color):
|
||||
"""Fallback khi mask rỗng."""
|
||||
img_resized = cv2.resize(img_np, (tw, th), interpolation=cv2.INTER_LANCZOS4)
|
||||
mask_resized = cv2.resize(mask_np, (tw, th), interpolation=cv2.INTER_NEAREST)
|
||||
return img_resized, mask_resized
|
||||
"""Fallback khi mask rỗng: coi như mask phủ cả ảnh, không méo ảnh.
|
||||
|
||||
crop = scale lớn nhất (lấp đầy target) rồi cắt giữa; pad = scale nhỏ nhất (cả ảnh nằm trong target), phần thừa pad_color.
|
||||
"""
|
||||
sh, sw = img_np.shape[:2]
|
||||
scale = max(tw / sw, th / sh) if mode == "crop" else min(tw / sw, th / sh)
|
||||
new_sw = max(1, round(sw * scale))
|
||||
new_sh = max(1, round(sh * scale))
|
||||
img_scaled = cv2.resize(img_np, (new_sw, new_sh), interpolation=cv2.INTER_LANCZOS4)
|
||||
mask_scaled = cv2.resize(mask_np, (new_sw, new_sh), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
# Tâm ảnh sau khi scale
|
||||
cx = round(sw / 2.0 * scale)
|
||||
cy = round(sh / 2.0 * scale)
|
||||
if mode == "crop":
|
||||
return self._crop_around_center(img_scaled, mask_scaled, tw, th, cx, cy, pad_color)
|
||||
return self._pad_around_center(img_scaled, mask_scaled, tw, th, cx, cy, pad_color)
|
||||
|
||||
@staticmethod
|
||||
def _hex_to_rgb(hex_str):
|
||||
|
||||
+14
-21
@@ -1,7 +1,5 @@
|
||||
"""Phase 1: Logic Core — compare, branching, and gating nodes."""
|
||||
|
||||
import re
|
||||
|
||||
from comfy_execution.graph_utils import ExecutionBlocker
|
||||
|
||||
|
||||
@@ -28,6 +26,11 @@ class NH_Compare:
|
||||
RETURN_NAMES = ("result", "a_passthrough", "b_passthrough")
|
||||
FUNCTION = "compare"
|
||||
CATEGORY = "NH-Nodes/Logic"
|
||||
DESCRIPTION = (
|
||||
'Compare a with b. "in" is plain membership: when b is a list, tuple, set or dict, a must be one of '
|
||||
"its items or keys; otherwise str(a) must be a substring of str(b), case-sensitive "
|
||||
'(13 is in "a13b" but not in "1024x768 3").'
|
||||
)
|
||||
|
||||
def compare(self, a, b, op, type_cast="auto"):
|
||||
if a is None or b is None:
|
||||
@@ -73,29 +76,19 @@ class NH_Compare:
|
||||
|
||||
@staticmethod
|
||||
def _contains(a, b):
|
||||
"""`a in b` as plain membership / substring, nothing fuzzier:
|
||||
|
||||
- b is a dict: a is one of its keys.
|
||||
- b is a list, tuple or set: a equals one of its items.
|
||||
- anything else: str(a) is a substring of str(b), case-sensitive, so 13 in "a13b" is True
|
||||
and 13 in "1024x768 3" is False.
|
||||
Values also match by their text (7 matches "7"), because widgets hand over strings.
|
||||
"""
|
||||
if isinstance(b, dict):
|
||||
return a in b or str(a) in [str(x) for x in b.keys()]
|
||||
if isinstance(b, (list, tuple, set)):
|
||||
return a in b or str(a) in [str(x) for x in b]
|
||||
needle = str(a)
|
||||
haystack = str(b)
|
||||
if needle in haystack:
|
||||
return True
|
||||
|
||||
if not NH_Compare._supports_subsequence(a, needle):
|
||||
return False
|
||||
|
||||
idx = 0
|
||||
for char in haystack:
|
||||
if idx < len(needle) and char == needle[idx]:
|
||||
idx += 1
|
||||
return idx == len(needle)
|
||||
|
||||
@staticmethod
|
||||
def _supports_subsequence(a, needle):
|
||||
if isinstance(a, (int, float)):
|
||||
return True
|
||||
return bool(re.fullmatch(r"[+-]?\d+(?:\.\d+)?", needle.strip()))
|
||||
return str(a) in str(b)
|
||||
|
||||
|
||||
class NH_LogicGate:
|
||||
|
||||
+25
-82
@@ -1,74 +1,15 @@
|
||||
"""Phase 2: Math & Random — NH_MathEval, NH_RandomChoice"""
|
||||
|
||||
import ast
|
||||
import operator
|
||||
import math
|
||||
import random
|
||||
|
||||
from ...core import safe_eval
|
||||
|
||||
# --- Safe eval engine ---
|
||||
_SAFE_OPS = {
|
||||
ast.Add: operator.add,
|
||||
ast.Sub: operator.sub,
|
||||
ast.Mult: operator.mul,
|
||||
ast.Div: operator.truediv,
|
||||
ast.FloorDiv: operator.floordiv,
|
||||
ast.Mod: operator.mod,
|
||||
ast.Pow: operator.pow,
|
||||
ast.USub: operator.neg,
|
||||
}
|
||||
|
||||
_SAFE_FUNCS = {
|
||||
"min": min, "max": max, "abs": abs,
|
||||
"round": round,
|
||||
"sqrt": math.sqrt, "floor": math.floor, "ceil": math.ceil,
|
||||
"clamp": lambda v, lo, hi: max(lo, min(hi, v)),
|
||||
}
|
||||
|
||||
|
||||
def _safe_eval(node, variables):
|
||||
"""De quy eval AST node, chi cho phep operators va functions an toan."""
|
||||
if isinstance(node, ast.Constant):
|
||||
return node.value
|
||||
elif isinstance(node, ast.Name):
|
||||
if node.id in variables:
|
||||
return variables[node.id]
|
||||
if node.id in _SAFE_FUNCS:
|
||||
return _SAFE_FUNCS[node.id]
|
||||
raise ValueError(f"Unknown variable: {node.id}")
|
||||
elif isinstance(node, ast.BinOp):
|
||||
left = _safe_eval(node.left, variables)
|
||||
right = _safe_eval(node.right, variables)
|
||||
op = _SAFE_OPS.get(type(node.op))
|
||||
if op is None:
|
||||
raise ValueError(f"Unsupported operator: {type(node.op).__name__}")
|
||||
if isinstance(node.op, ast.Div) and right == 0:
|
||||
print("[NH-Nodes] Warning: Division by zero, returning 0")
|
||||
return 0.0
|
||||
if isinstance(node.op, ast.Pow):
|
||||
# Gioi han power de tranh treo
|
||||
if abs(right) > 100:
|
||||
print("[NH-Nodes] Warning: Power too large, clamped to 100")
|
||||
right = 100 if right > 0 else -100
|
||||
return op(left, right)
|
||||
elif isinstance(node, ast.UnaryOp):
|
||||
operand = _safe_eval(node.operand, variables)
|
||||
op = _SAFE_OPS.get(type(node.op))
|
||||
if op is None:
|
||||
raise ValueError(f"Unsupported unary op: {type(node.op).__name__}")
|
||||
return op(operand)
|
||||
elif isinstance(node, ast.Call):
|
||||
func_name = node.func.id if isinstance(node.func, ast.Name) else None
|
||||
if func_name not in _SAFE_FUNCS:
|
||||
raise ValueError(f"Unknown function: {func_name}")
|
||||
args = [_safe_eval(arg, variables) for arg in node.args]
|
||||
return _SAFE_FUNCS[func_name](*args)
|
||||
else:
|
||||
raise ValueError(f"Unsupported expression type: {type(node).__name__}")
|
||||
# Random Choice co 6 input: input_0 .. input_5
|
||||
_CHOICE_SLOTS = 6
|
||||
|
||||
|
||||
class NH_MathEval:
|
||||
"""Tinh toan bieu thuc toan hoc an toan bang AST parsing."""
|
||||
"""Tinh toan bieu thuc toan hoc an toan bang AST parsing (core/safe_eval.py)."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -94,15 +35,18 @@ class NH_MathEval:
|
||||
CATEGORY = "NH-Nodes/Logic"
|
||||
|
||||
def evaluate(self, expression, a=0.0, b=0.0, c=0.0, d=0.0, round_to=-1):
|
||||
variables = {"a": a, "b": b, "c": c, "d": d}
|
||||
# Bieu thuc sai hoac ket qua khong huu han -> ValueError (ComfyUI hien loi, khong tra 0 im lang).
|
||||
# x / 0, x // 0, x % 0 -> 0 kem canh bao.
|
||||
try:
|
||||
tree = ast.parse(expression, mode='eval')
|
||||
result = _safe_eval(tree.body, variables)
|
||||
except Exception as e:
|
||||
print(f"[NH-Nodes] MathEval error: {e}")
|
||||
result = 0.0
|
||||
|
||||
result = float(result)
|
||||
result = safe_eval.evaluate(
|
||||
expression,
|
||||
{"a": a, "b": b, "c": c, "d": d},
|
||||
safe_eval.MATH_FUNCTIONS,
|
||||
zero_division="zero",
|
||||
on_warning=lambda message: print(f"[NH-Nodes] MathEval warning: {message}"),
|
||||
)
|
||||
except safe_eval.ExpressionError as exc:
|
||||
raise ValueError(f"Math Eval: {exc} (expression: {expression!r})") from exc
|
||||
if round_to >= 0:
|
||||
result = round(result, round_to)
|
||||
|
||||
@@ -133,7 +77,7 @@ class NH_RandomChoice:
|
||||
|
||||
def choose(self, seed, weights, **kwargs):
|
||||
connected = []
|
||||
for i in range(6):
|
||||
for i in range(_CHOICE_SLOTS):
|
||||
key = f"input_{i}"
|
||||
if key in kwargs and kwargs[key] is not None:
|
||||
connected.append((i, kwargs[key]))
|
||||
@@ -141,23 +85,22 @@ class NH_RandomChoice:
|
||||
if not connected:
|
||||
return (None, -1, "")
|
||||
|
||||
# Parse weights
|
||||
w = []
|
||||
for s in weights.split(","):
|
||||
# Weight thu i thuoc input_i (khong phai input thu i dang noi), slot thieu weight = 1
|
||||
slot_weights = []
|
||||
for s in weights.split(",")[:_CHOICE_SLOTS]:
|
||||
try:
|
||||
w.append(max(0.0, float(s.strip())))
|
||||
slot_weights.append(max(0.0, float(s.strip())))
|
||||
except (ValueError, TypeError):
|
||||
w.append(1.0)
|
||||
while len(w) < len(connected):
|
||||
w.append(1.0)
|
||||
w = w[:len(connected)]
|
||||
slot_weights.append(1.0)
|
||||
slot_weights += [1.0] * (_CHOICE_SLOTS - len(slot_weights))
|
||||
w = [slot_weights[i] for i, _ in connected]
|
||||
|
||||
# Normalize
|
||||
total = sum(w) or 1.0
|
||||
probs = [x / total for x in w]
|
||||
|
||||
# Seed
|
||||
rng = random.Random(seed if seed > 0 else None)
|
||||
# Seed dung dung gia tri (0 cung tat dinh): ComfyUI cache output theo input
|
||||
rng = random.Random(seed)
|
||||
|
||||
# Weighted pick
|
||||
r = rng.random()
|
||||
|
||||
@@ -1,25 +1,82 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
def get_bounding_box_and_ratio(mask_tensor):
|
||||
"""Tính toán bounding box và tỷ lệ (width / height) từ một mask tensor."""
|
||||
m_np = mask_tensor.cpu().numpy()
|
||||
rows, cols = np.where(m_np > 0.5)
|
||||
|
||||
|
||||
if len(rows) == 0:
|
||||
return None, 1.0 # Trả về None và tỷ lệ 1:1 nếu mask rỗng
|
||||
|
||||
ymin, ymax = np.min(rows), np.max(rows)
|
||||
xmin, xmax = np.min(cols), np.max(cols)
|
||||
|
||||
|
||||
bbox = (xmin, ymin, xmax, ymax)
|
||||
width = xmax - xmin + 1
|
||||
height = ymax - ymin + 1
|
||||
|
||||
|
||||
aspect_ratio = width / height if height > 0 else 1.0
|
||||
return bbox, aspect_ratio
|
||||
|
||||
|
||||
def _fit_window(start, length, limit):
|
||||
"""Cửa sổ [start, start + length) trượt vào trong [0, limit); lớn hơn limit thì cắt ở mép cuối."""
|
||||
start = max(0, min(start, limit - length))
|
||||
return start, min(limit, start + length)
|
||||
|
||||
|
||||
def _adjust_mask(mask_np, bbox, target_size, mode):
|
||||
"""Chỉnh hộp (bbox) của một mask [H, W] về đúng tỷ lệ target_size = (rộng, cao).
|
||||
Kích thước hộp mới tính bằng số nguyên nên không bị lệch 1px do số thực. Ngoài hộp luôn là 0."""
|
||||
height, width = mask_np.shape
|
||||
xmin, ymin, xmax, ymax = (int(v) for v in bbox)
|
||||
cur_w, cur_h = xmax - xmin + 1, ymax - ymin + 1
|
||||
target_w, target_h = target_size
|
||||
target_is_wider = target_w * cur_h > cur_w * target_h
|
||||
output = np.zeros_like(mask_np)
|
||||
|
||||
if mode in ("stretch", "pad"):
|
||||
# Nới một chiều: target rộng hơn thì nới ngang, ngược lại nới dọc
|
||||
new_w, new_h = cur_w, cur_h
|
||||
if target_is_wider:
|
||||
new_w = max(1, cur_h * target_w // target_h)
|
||||
else:
|
||||
new_h = max(1, cur_w * target_h // target_w)
|
||||
|
||||
# Hộp mới nới đối xứng quanh hộp cũ. Chạm mép ảnh thì trượt vào trong để vẫn đủ kích thước.
|
||||
x0, x1 = _fit_window(xmin - (new_w - cur_w) // 2, new_w, width)
|
||||
y0, y1 = _fit_window(ymin - (new_h - cur_h) // 2, new_h, height)
|
||||
|
||||
if mode == "stretch":
|
||||
# Kéo dãn nội dung (nearest) cho vừa hộp mới. Hộp lớn hơn cả ảnh thì chỉ hiện phần giữa.
|
||||
off_x = (new_w - (x1 - x0)) // 2
|
||||
off_y = (new_h - (y1 - y0)) // 2
|
||||
cols = (np.arange(x1 - x0) + off_x) * cur_w // new_w
|
||||
rows = (np.arange(y1 - y0) + off_y) * cur_h // new_h
|
||||
output[y0:y1, x0:x1] = mask_np[ymin:ymax + 1, xmin:xmax + 1][np.ix_(rows, cols)]
|
||||
else:
|
||||
# Pad: phần đệm thuộc vùng mask (1.0), nội dung gốc giữ nguyên vị trí cũ
|
||||
output[y0:y1, x0:x1] = 1.0
|
||||
output[ymin:ymax + 1, xmin:xmax + 1] = mask_np[ymin:ymax + 1, xmin:xmax + 1]
|
||||
|
||||
elif mode == "crop":
|
||||
# Cắt bớt một chiều quanh tâm; phần giữ lại không bị dịch chỗ
|
||||
new_w, new_h = cur_w, cur_h
|
||||
if target_is_wider:
|
||||
new_h = max(1, cur_w * target_h // target_w)
|
||||
else:
|
||||
new_w = max(1, cur_h * target_w // target_h)
|
||||
x0 = xmin + (cur_w - new_w) // 2
|
||||
y0 = ymin + (cur_h - new_h) // 2
|
||||
output[y0:y0 + new_h, x0:x0 + new_w] = mask_np[y0:y0 + new_h, x0:x0 + new_w]
|
||||
|
||||
else:
|
||||
raise ValueError(f"NH Mask Aspect Ratio Match: unknown mode '{mode}'")
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class NH_MaskAspectRatioMatch:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -35,88 +92,37 @@ class NH_MaskAspectRatioMatch:
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "match_aspect_ratio"
|
||||
CATEGORY = "NH-Nodes/Mask"
|
||||
DESCRIPTION = (
|
||||
"Resizes the bounding box of mask_to_adjust to the aspect ratio of target_ratio_mask, around the same centre. "
|
||||
"stretch: scale the content to the new box. "
|
||||
"pad: keep the content and add mask pixels (1.0) around it until the box has the ratio. "
|
||||
"crop: keep the middle part of the content. "
|
||||
"A box that would leave the image slides back inside it. Every mask of the batch is processed; "
|
||||
"a shorter target batch repeats."
|
||||
)
|
||||
|
||||
def match_aspect_ratio(self, target_ratio_mask, mask_to_adjust, mode):
|
||||
mask_a = target_ratio_mask[0]
|
||||
mask_b = mask_to_adjust[0]
|
||||
|
||||
_ , ratio_a = get_bounding_box_and_ratio(mask_a)
|
||||
bbox_b, ratio_b = get_bounding_box_and_ratio(mask_b)
|
||||
|
||||
if bbox_b is None:
|
||||
# Chấp nhận mask [H, W] hoặc [B, H, W]; xử lý từng mask trong batch
|
||||
targets = target_ratio_mask.reshape(-1, *target_ratio_mask.shape[-2:])
|
||||
masks = mask_to_adjust.reshape(-1, *mask_to_adjust.shape[-2:])
|
||||
if targets.shape[0] == 0 or masks.shape[0] == 0:
|
||||
return (mask_to_adjust,)
|
||||
|
||||
xmin, ymin, xmax, ymax = bbox_b
|
||||
current_width = xmax - xmin + 1
|
||||
current_height = ymax - ymin + 1
|
||||
center_x = (xmin + xmax) // 2
|
||||
center_y = (ymin + ymax) // 2
|
||||
adjusted = []
|
||||
for index in range(masks.shape[0]):
|
||||
# Batch của mask A ngắn hơn thì lặp lại vòng tròn
|
||||
bbox_a, _ = get_bounding_box_and_ratio(targets[index % targets.shape[0]])
|
||||
bbox_b, _ = get_bounding_box_and_ratio(masks[index])
|
||||
if bbox_b is None:
|
||||
adjusted.append(masks[index]) # Mask B rỗng: giữ nguyên
|
||||
continue
|
||||
|
||||
mask_b_np = mask_b.cpu().numpy()
|
||||
output_mask_np = np.zeros_like(mask_b_np)
|
||||
# Mask A rỗng thì coi như tỷ lệ 1:1
|
||||
target_size = (1, 1) if bbox_a is None else (bbox_a[2] - bbox_a[0] + 1, bbox_a[3] - bbox_a[1] + 1)
|
||||
output_np = _adjust_mask(masks[index].cpu().numpy(), bbox_b, target_size, mode)
|
||||
adjusted.append(torch.from_numpy(output_np).to(masks.device))
|
||||
|
||||
# Cắt nội dung thực sự của Mask B
|
||||
cropped_content_b = mask_b_np[ymin:ymax+1, xmin:xmax+1]
|
||||
|
||||
# ==========================================================
|
||||
# LOGIC CHO TỪNG CHẾ ĐỘ
|
||||
# ==========================================================
|
||||
|
||||
if mode == "stretch":
|
||||
new_height, new_width = current_height, current_width
|
||||
if ratio_a > ratio_b: # Target rộng hơn
|
||||
new_width = int(current_height * ratio_a)
|
||||
else: # Target cao hơn
|
||||
new_height = int(current_width / ratio_a)
|
||||
|
||||
# Kéo dãn nội dung đã cắt
|
||||
stretched_content = cv2.resize(cropped_content_b, (new_width, new_height), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
# Dán vào canvas mới tại vị trí trung tâm
|
||||
paste_x = center_x - new_width // 2
|
||||
paste_y = center_y - new_height // 2
|
||||
|
||||
# Giới hạn tọa độ để không bị tràn
|
||||
p_x_start, p_y_start = max(0, paste_x), max(0, paste_y)
|
||||
p_x_end, p_y_end = min(mask_b_np.shape[1], paste_x + new_width), min(mask_b_np.shape[0], paste_y + new_height)
|
||||
|
||||
s_x_start, s_y_start = p_x_start - paste_x, p_y_start - paste_y
|
||||
s_x_end, s_y_end = s_x_start + (p_x_end - p_x_start), s_y_start + (p_y_end - p_y_start)
|
||||
|
||||
output_mask_np[p_y_start:p_y_end, p_x_start:p_x_end] = stretched_content[s_y_start:s_y_end, s_x_start:s_x_end]
|
||||
|
||||
elif mode == "pad":
|
||||
new_height, new_width = current_height, current_width
|
||||
if ratio_a > ratio_b: # Target rộng hơn -> Thêm đệm ngang
|
||||
new_width = int(current_height * ratio_a)
|
||||
else: # Target cao hơn -> Thêm đệm dọc
|
||||
new_height = int(current_width / ratio_a)
|
||||
|
||||
# Dán nội dung gốc vào trung tâm của hộp mới lớn hơn
|
||||
paste_x = center_x - current_width // 2
|
||||
paste_y = center_y - current_height // 2
|
||||
output_mask_np[paste_y:paste_y + current_height, paste_x:paste_x + current_width] = cropped_content_b
|
||||
|
||||
elif mode == "crop":
|
||||
new_height, new_width = current_height, current_width
|
||||
if ratio_a > ratio_b: # Target rộng hơn -> Cắt bớt chiều dọc
|
||||
new_height = int(current_width / ratio_a)
|
||||
else: # Target cao hơn -> Cắt bớt chiều ngang
|
||||
new_width = int(current_height * ratio_a)
|
||||
|
||||
# Cắt bớt nội dung
|
||||
crop_y = (current_height - new_height) // 2
|
||||
crop_x = (current_width - new_width) // 2
|
||||
cropped_content = cropped_content_b[crop_y:crop_y + new_height, crop_x:crop_x + new_width]
|
||||
|
||||
# Dán nội dung đã cắt vào vị trí trung tâm
|
||||
paste_x = center_x - new_width // 2
|
||||
paste_y = center_y - new_height // 2
|
||||
output_mask_np[paste_y:paste_y + new_height, paste_x:paste_x + new_width] = cropped_content
|
||||
|
||||
# Chuyển lại sang định dạng tensor
|
||||
output_mask = torch.from_numpy(output_mask_np).unsqueeze(0)
|
||||
return (output_mask,)
|
||||
return (torch.stack(adjusted),)
|
||||
|
||||
|
||||
# Ánh xạ để ComfyUI nhận diện class
|
||||
@@ -127,4 +133,4 @@ NODE_CLASS_MAPPINGS = {
|
||||
# Tên sẽ hiển thị trong menu của ComfyUI
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NH_MaskAspectRatioMatch": "NH Mask Aspect Ratio Match"
|
||||
}
|
||||
}
|
||||
|
||||
@@ -19,6 +19,7 @@ from ...core.resolution_data import (
|
||||
PRESET_LABELS_BY_MODEL,
|
||||
PRESET_LOOKUP_BY_MODEL,
|
||||
)
|
||||
from ...core.utils import parse_hex_color
|
||||
|
||||
|
||||
TARGET_RESOLUTION_LEVELS = [f"{mp} MP" for mp in range(1, 13)]
|
||||
@@ -46,6 +47,9 @@ _PIL_RESAMPLE_METHODS = {
|
||||
"box": _RESAMPLING.BOX,
|
||||
"hamming": _RESAMPLING.HAMMING,
|
||||
}
|
||||
# Lanczos has negative lobes, so it overshoots beside hard edges. Box and Hamming are plain weighted averages (their
|
||||
# weights are all positive), so they stay inside [0, 1] and need no clipping between the width and the height pass.
|
||||
_RINGING_METHODS = {"lanczos"}
|
||||
|
||||
|
||||
def _clamp(value, lower, upper):
|
||||
@@ -68,18 +72,6 @@ def _ratio_label(width, height):
|
||||
return f"{ratio_w}:{ratio_h}"
|
||||
|
||||
|
||||
def _parse_hex_color(value, default=(255, 255, 255)):
|
||||
text = str(value or "").strip().lstrip("#")
|
||||
try:
|
||||
if len(text) == 3:
|
||||
text = "".join(char * 2 for char in text)
|
||||
if len(text) != 6:
|
||||
raise ValueError
|
||||
return tuple(int(text[index:index + 2], 16) for index in (0, 2, 4))
|
||||
except Exception:
|
||||
return default
|
||||
|
||||
|
||||
def _needs_warning(entry, extra_warnings):
|
||||
if extra_warnings:
|
||||
return True
|
||||
@@ -89,14 +81,21 @@ def _needs_warning(entry, extra_warnings):
|
||||
return any(marker in text for marker in ("community", "extrapolation", "rủi ro", "risk"))
|
||||
|
||||
|
||||
def _build_info(model_spec, tier, ratio, width, height, entry=None, extra_warnings=None):
|
||||
def _build_info(model_spec, tier, ratio, width, height, entry=None, extra_warnings=None, warnings=None):
|
||||
"""Info text of the resolution nodes. ``warnings`` become 'Warning:' lines.
|
||||
|
||||
Presets carry their own reliability warning; max_res only bounds custom sizes (see resolution_data).
|
||||
"""
|
||||
extra_warnings = extra_warnings or []
|
||||
warnings = list(warnings or [])
|
||||
mp = (int(width) * int(height)) / 1_000_000
|
||||
prefix = "⚠️ " if _needs_warning(entry, extra_warnings) else ""
|
||||
prefix = "⚠️ " if warnings or _needs_warning(entry, extra_warnings) else ""
|
||||
info = f"{prefix}{model_spec['name']} | {tier} | {ratio} | {int(width)}×{int(height)} | {mp:.2f}MP"
|
||||
warning = entry.get("warning") if entry else ""
|
||||
if warning:
|
||||
info += f"\nWarning: {warning}"
|
||||
for message in warnings:
|
||||
info += f"\nWarning: {message}"
|
||||
note = model_spec.get("note")
|
||||
if note:
|
||||
info += f"\nNote: {note}"
|
||||
@@ -335,25 +334,57 @@ def _build_resize_data(source_width, source_height, target_width, target_height,
|
||||
return data
|
||||
|
||||
|
||||
def _resize_planes_pil(planes, size, resample, premultiplied):
|
||||
"""Resize [C, H, W] float planes with one Pillow call and clip the result to [0, 1].
|
||||
|
||||
Pillow's 8-bit path clips after each of its two passes (width, then height); on float planes it clips nothing.
|
||||
Lanczos overshoots beside hard edges, so the caller resizes one direction per call to clip in between: otherwise
|
||||
the second pass amplifies the overshoot of the first (15-33/255 off on text and line art). Premultiplied colour
|
||||
cannot be brighter than its alpha, so it is clipped to [0, alpha].
|
||||
"""
|
||||
resized = np.empty((planes.shape[0], size[1], size[0]), dtype=np.float32)
|
||||
for index in (3, 0, 1, 2) if premultiplied else range(planes.shape[0]): # alpha first, it bounds the colour
|
||||
upper = resized[3] if premultiplied and index != 3 else 1.0
|
||||
np.clip(np.asarray(Image.fromarray(planes[index]).resize(size, resample)), 0.0, upper, out=resized[index])
|
||||
return resized
|
||||
|
||||
|
||||
def _resize_bchw_pil(samples, height, width, resize_method):
|
||||
"""Resize [B, C, H, W] with Pillow's Lanczos/box/Hamming filters, in float.
|
||||
|
||||
Every channel goes through a 32-bit float ("F") image, so the result keeps its precision instead of being rounded
|
||||
to 8 bits. Four-channel images are resized with premultiplied alpha, as Pillow does for RGBA, so colour does not
|
||||
bleed in from transparent pixels. Lanczos resizes the width and the height in two passes, as Pillow does inside one
|
||||
resize, but clips each pass to [0, 1] (colour to [0, alpha] when premultiplied), like Pillow's 8-bit path does.
|
||||
"""
|
||||
original_dtype = samples.dtype
|
||||
original_device = samples.device
|
||||
resample = _PIL_RESAMPLE_METHODS[resize_method]
|
||||
size = (int(width), int(height))
|
||||
resized_items = []
|
||||
|
||||
for sample in samples:
|
||||
array = sample.movedim(0, -1).detach().cpu().numpy()
|
||||
array = np.clip(array * 255.0, 0, 255).astype(np.uint8)
|
||||
if array.shape[-1] == 1:
|
||||
pil_image = Image.fromarray(array[:, :, 0], mode="L")
|
||||
else:
|
||||
pil_image = Image.fromarray(array)
|
||||
|
||||
pil_image = pil_image.resize((int(width), int(height)), resample)
|
||||
resized_array = np.array(pil_image).astype(np.float32) / 255.0
|
||||
if resized_array.ndim == 2:
|
||||
resized_array = resized_array[:, :, None]
|
||||
resized_items.append(torch.from_numpy(resized_array).movedim(-1, 0))
|
||||
# [C, H, W] float32 in 0-1 (NaN and out-of-range values are cleaned up first, as the uint8 round trip did).
|
||||
# Done in torch, which is several times faster than numpy on large images.
|
||||
cleaned = torch.nan_to_num(sample.detach().cpu().float(), nan=0.0, posinf=1.0, neginf=0.0).clamp_(0.0, 1.0)
|
||||
if (cleaned.shape[2], cleaned.shape[1]) == size:
|
||||
# Nothing to resample; skipping the premultiply round trip keeps the colour under transparent pixels.
|
||||
resized_items.append(cleaned)
|
||||
continue
|
||||
planes = cleaned.contiguous().numpy()
|
||||
premultiply = planes.shape[0] == 4
|
||||
if premultiply:
|
||||
planes = np.concatenate((planes[:3] * planes[3:], planes[3:]))
|
||||
# Lanczos: the width first, then the height (Pillow's own order), clipped after each. The others: one call.
|
||||
passes = ((size[0], planes.shape[1]), size) if resize_method in _RINGING_METHODS else (size,)
|
||||
for pass_size in passes:
|
||||
if pass_size != (planes.shape[2], planes.shape[1]):
|
||||
planes = _resize_planes_pil(planes, pass_size, resample, premultiply)
|
||||
if premultiply:
|
||||
alpha = planes[3:]
|
||||
colour = np.divide(planes[:3], alpha, out=np.zeros_like(planes[:3]), where=alpha > 1e-6)
|
||||
planes = np.concatenate((colour, alpha))
|
||||
resized_items.append(torch.from_numpy(planes))
|
||||
|
||||
return torch.stack(resized_items, dim=0).to(device=original_device, dtype=original_dtype)
|
||||
|
||||
@@ -401,14 +432,15 @@ def _resize_image_batch_with_data(images, target_width, target_height, resize_mo
|
||||
cropped = resized[:, :, top:top + int(target_height), left:left + int(target_width)]
|
||||
return cropped.movedim(1, -1), resize_data
|
||||
|
||||
fill_rgb = _parse_hex_color(pad_color_hex, (255, 255, 255))
|
||||
fill_rgba = parse_hex_color(pad_color_hex, (255, 255, 255, 255))
|
||||
canvas = torch.empty(
|
||||
(batch, channels, int(target_height), int(target_width)),
|
||||
dtype=images.dtype,
|
||||
device=images.device,
|
||||
)
|
||||
for channel in range(channels):
|
||||
canvas[:, channel, :, :] = fill_rgb[min(channel, 2)] / 255.0
|
||||
# r, g, b, then alpha for the fourth channel of RGBA images (opaque unless the colour is #RRGGBBAA).
|
||||
canvas[:, channel, :, :] = fill_rgba[min(channel, 3)] / 255.0
|
||||
|
||||
left = int(resize_data["content_left"])
|
||||
top = int(resize_data["content_top"])
|
||||
@@ -450,6 +482,76 @@ def _has_padding(resize_data):
|
||||
return any(int(resize_data.get(key, 0)) > 0 for key in ("pad_left", "pad_top", "pad_right", "pad_bottom"))
|
||||
|
||||
|
||||
def _has_crop(resize_data):
|
||||
return any(int(resize_data.get(key, 0)) > 0 for key in ("crop_left", "crop_top", "crop_right", "crop_bottom"))
|
||||
|
||||
|
||||
def _cropped_footprint(resize_data):
|
||||
"""Where the picture of a ``cover_crop`` result sits in the source frame.
|
||||
|
||||
Returns ``(source_width, source_height, left, top, right, bottom)`` in source pixels, or None when ``resize_data``
|
||||
does not say. The resize scaled the source uniformly by resized/source and the crop kept ``target`` pixels from
|
||||
offset ``crop_left``/``crop_top`` of that, so the kept window maps back by source/resized. A source pixel belongs
|
||||
to it when its centre is inside the window, the rule NH Restore Image Geometry uses for its mask.
|
||||
"""
|
||||
try:
|
||||
source_width, source_height, resized_width, resized_height, target_width, target_height, crop_left, crop_top = (
|
||||
int(resize_data[key])
|
||||
for key in (
|
||||
"source_width", "source_height", "resized_width", "resized_height",
|
||||
"target_width", "target_height", "crop_left", "crop_top",
|
||||
)
|
||||
)
|
||||
except (KeyError, TypeError, ValueError):
|
||||
return None
|
||||
if min(source_width, source_height, resized_width, resized_height, target_width, target_height) <= 0:
|
||||
return None
|
||||
|
||||
def span(crop, target, resized, source):
|
||||
start = _clamp(math.ceil(crop * source / resized - 0.5), 0, source - 1)
|
||||
end = _clamp(math.ceil((crop + target) * source / resized - 0.5), start + 1, source)
|
||||
return start, end
|
||||
|
||||
left, right = span(crop_left, target_width, resized_width, source_width)
|
||||
top, bottom = span(crop_top, target_height, resized_height, source_height)
|
||||
return source_width, source_height, left, top, right, bottom
|
||||
|
||||
|
||||
def _scale_into_source_frame(image, footprint, resize_method):
|
||||
"""Scale [B, H, W, C] into ``footprint`` and fill the rest of the source frame with the nearest edge pixels."""
|
||||
source_width, source_height, left, top, right, bottom = footprint
|
||||
scaled = _resize_image_batch(image, right - left, bottom - top, "stretch", resize_method, "#000000")
|
||||
padding = (left, source_width - right, top, source_height - bottom)
|
||||
return F.pad(scaled.movedim(-1, 1), padding, mode="replicate").movedim(1, -1)
|
||||
|
||||
|
||||
def _cropped_resize_info(resize_data, image_width, image_height, footprint=None):
|
||||
"""Explain what happened to a ``cover_crop`` result; ``footprint`` is None when it was returned unchanged."""
|
||||
source_width = int(resize_data.get("source_width", image_width))
|
||||
source_height = int(resize_data.get("source_height", image_height))
|
||||
target_width = int(resize_data.get("target_width", image_width))
|
||||
target_height = int(resize_data.get("target_height", image_height))
|
||||
crop_text = " ".join(f"{side[0].upper()}{int(resize_data.get(f'crop_{side}', 0))}" for side in ("left", "top", "right", "bottom"))
|
||||
details = (
|
||||
f"source {source_width}×{source_height}; resize target {target_width}×{target_height}; "
|
||||
f"input {image_width}×{image_height}; crop {crop_text}"
|
||||
)
|
||||
if footprint is None:
|
||||
return (
|
||||
"cover_crop resize detected: the image was cropped, so there is no padding to remove and the cropped "
|
||||
"borders cannot be recovered; returned unchanged. output_size=original_size scales it back into the "
|
||||
"source frame (it needs complete resize_data), and NH Restore Image Geometry composites it onto the "
|
||||
f"original image; {details}"
|
||||
)
|
||||
_, _, left, top, right, bottom = footprint
|
||||
return (
|
||||
"cover_crop resize detected: the cropped borders cannot be recovered. The image was scaled back uniformly "
|
||||
f"into the source frame (x {left}-{right}, y {top}-{bottom}); the borders that were cut off repeat its nearest "
|
||||
"edge pixels. NH Restore Image Geometry composites onto the original image instead, which keeps its real "
|
||||
f"borders; {details}; output {source_width}×{source_height}"
|
||||
)
|
||||
|
||||
|
||||
class NH_SmartResolutionPicker:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -486,13 +588,21 @@ class NH_SmartResolutionPicker:
|
||||
swap_orientation=False,
|
||||
batch_size=1,
|
||||
):
|
||||
# VALIDATE_INPUTS accepts any value, so a saved workflow can carry a model or preset label that no longer
|
||||
# exists. The default is used for it, and the info says so.
|
||||
warnings = []
|
||||
if model_family not in MODEL_ID_BY_LABEL:
|
||||
warnings.append(f"unknown model_family '{model_family}', using {DEFAULT_MODEL_LABEL}")
|
||||
model_family = DEFAULT_MODEL_LABEL
|
||||
|
||||
preset_lookup = PRESET_LOOKUP_BY_MODEL.get(model_family, {})
|
||||
model_presets = PRESET_LABELS_BY_MODEL.get(model_family, [])
|
||||
default_preset = DEFAULT_PRESET if model_family == DEFAULT_MODEL_LABEL else model_presets[0]
|
||||
entry = preset_lookup.get(preset) or preset_lookup[default_preset]
|
||||
entry = preset_lookup.get(preset)
|
||||
if entry is None:
|
||||
entry = preset_lookup[default_preset]
|
||||
if not use_custom: # a custom size does not use the preset's size
|
||||
warnings.append(f"preset '{preset}' not found for {model_family}, using '{default_preset}'")
|
||||
model_spec = MODEL_SPECS[entry["model_id"]]
|
||||
extra_warnings = []
|
||||
|
||||
@@ -530,7 +640,7 @@ class NH_SmartResolutionPicker:
|
||||
"downscale_ratio_spacial": downscale,
|
||||
}
|
||||
aspect_ratio = float(width) / max(float(height), 1.0)
|
||||
info = _build_info(model_spec, tier, ratio, width, height, entry_for_warning, extra_warnings)
|
||||
info = _build_info(model_spec, tier, ratio, width, height, entry_for_warning, extra_warnings, warnings)
|
||||
|
||||
return (int(width), int(height), latent, info, aspect_ratio)
|
||||
|
||||
@@ -754,6 +864,17 @@ class NH_RemoveResizePadding:
|
||||
f"No resize_data metadata found; passed through {image_width}×{image_height}.",
|
||||
)
|
||||
|
||||
if resize_data.get("resize_mode") == "cover_crop" and _has_crop(resize_data):
|
||||
# A cover_crop resize has no padding, and the part it cut off is gone. Stretching what is left to the
|
||||
# source size would distort it (about 2.6% for 1500x1000 -> 1216x832), so for the original size the image
|
||||
# is scaled back by the same uniform factor, put where it came from, and the lost borders repeat its edges.
|
||||
footprint = _cropped_footprint(resize_data) if output_size == "original_size" else None
|
||||
if footprint is None:
|
||||
return (image, image_width, image_height, _cropped_resize_info(resize_data, image_width, image_height))
|
||||
restored = _scale_into_source_frame(image, footprint, resize_method)
|
||||
info = _cropped_resize_info(resize_data, image_width, image_height, footprint)
|
||||
return (restored, footprint[0], footprint[1], info)
|
||||
|
||||
left, top, right, bottom = _scaled_content_bounds(resize_data, image_width, image_height)
|
||||
cropped = image[:, top:bottom, left:right, :]
|
||||
|
||||
|
||||
@@ -13,12 +13,15 @@ def _as_prompt_text(value):
|
||||
|
||||
|
||||
def _remove_placeholder_preserve_lines(text, name):
|
||||
# Chi bo placeholder + khoang trang truoc no (dau dong: ca khoang trang sau no).
|
||||
# Dau phay de _cleanup_prompt_line gop: "a, {x}, b" -> "a,, b" -> "a, b" (khong dinh thanh "ab").
|
||||
placeholder = r"\{" + re.escape(name) + r"\}"
|
||||
horizontal_ws = r"[^\S\r\n]*"
|
||||
return re.sub(
|
||||
horizontal_ws + r",?" + horizontal_ws + placeholder + horizontal_ws + r",?" + horizontal_ws,
|
||||
r"^" + horizontal_ws + placeholder + horizontal_ws + r"|" + horizontal_ws + placeholder,
|
||||
"",
|
||||
text,
|
||||
flags=re.MULTILINE,
|
||||
)
|
||||
|
||||
|
||||
@@ -147,7 +150,8 @@ class NH_PromptScheduler:
|
||||
idx = cycle_len - pos
|
||||
|
||||
elif mode == "random":
|
||||
rng = random.Random(seed + current_step if seed > 0 else None)
|
||||
# Seed dung dung gia tri (0 cung tat dinh): ComfyUI cache output theo input
|
||||
rng = random.Random(seed + current_step)
|
||||
idx = rng.randint(0, n - 1)
|
||||
|
||||
else:
|
||||
|
||||
@@ -191,7 +191,8 @@ class NH_RegexExtract:
|
||||
if mode == "match":
|
||||
m = re.search(pattern, text)
|
||||
if m:
|
||||
groups = "\n".join(m.groups()) if m.groups() else m.group(0)
|
||||
# Group tuy chon khong khop la None -> chuoi rong
|
||||
groups = "\n".join(g or "" for g in m.groups()) if m.groups() else m.group(0)
|
||||
return (m.group(0), True, groups, 1)
|
||||
return ("", False, "", 0)
|
||||
|
||||
|
||||
@@ -70,8 +70,8 @@ class NH_TextSplitLines:
|
||||
if count == 0:
|
||||
return ([""], "", 0, "(empty)", True, True)
|
||||
|
||||
# Reorder by mode
|
||||
rng = random.Random(seed if seed > 0 else None)
|
||||
# Reorder by mode. The seed is used as given (0 too): ComfyUI caches the output per input.
|
||||
rng = random.Random(seed)
|
||||
|
||||
if mode == "shuffle":
|
||||
lines = lines[:]
|
||||
@@ -132,7 +132,8 @@ class NH_TextRandomLine:
|
||||
if total == 0:
|
||||
return ("", -1, 0)
|
||||
|
||||
rng = random.Random(seed if seed > 0 else None)
|
||||
# The seed is used as given (0 too): ComfyUI caches the output per input.
|
||||
rng = random.Random(seed)
|
||||
idx = rng.randint(0, total - 1)
|
||||
|
||||
return (lines[idx], idx, total)
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
# Kiểu dữ liệu tùy chỉnh cho đường ống đa năng
|
||||
NH_UNIVERSAL_PIPE = ("NH_UNIVERSAL_PIPE",)
|
||||
|
||||
# Số ô input của Pack (input_0 .. input_9) cũng là độ dài của package
|
||||
_PIPE_SLOTS = 10
|
||||
|
||||
class NH_PackUniversal:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -19,9 +22,9 @@ class NH_PackUniversal:
|
||||
CATEGORY = "NH-Nodes/Utils/Pipe"
|
||||
|
||||
def pack(self, **kwargs):
|
||||
# Đóng gói tất cả các giá trị đầu vào thành một list
|
||||
# kwargs sẽ là một dictionary chứa tất cả các input đã được kết nối
|
||||
pipe = list(kwargs.values())
|
||||
# Đóng gói theo số ô: package[i] là input_i, ô chưa nối là None.
|
||||
# (list(kwargs.values()) chỉ giữ ô đã nối nên index bị lệch khi bỏ trống input_0.)
|
||||
pipe = [kwargs.get(f"input_{i}") for i in range(_PIPE_SLOTS)]
|
||||
return (pipe,)
|
||||
|
||||
class NH_UnpackUniversal:
|
||||
@@ -29,14 +32,14 @@ class NH_UnpackUniversal:
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"index": ("INT", {"default": 0, "min": 0, "max": 9}),
|
||||
"index": ("INT", {"default": 0, "min": 0, "max": _PIPE_SLOTS - 1}),
|
||||
},
|
||||
"optional": {
|
||||
"package": NH_UNIVERSAL_PIPE,
|
||||
"pipe": NH_UNIVERSAL_PIPE,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
# Đầu ra là kiểu dữ liệu bất kỳ
|
||||
RETURN_TYPES = ("*",)
|
||||
RETURN_NAMES = ("item",)
|
||||
@@ -48,8 +51,8 @@ class NH_UnpackUniversal:
|
||||
if pipe_data is None:
|
||||
return (None,)
|
||||
|
||||
# Kiểm tra xem index có hợp lệ không
|
||||
if index < len(pipe_data):
|
||||
# index = số ô input của Pack; ngoài phạm vi (kể cả âm) thì trả None
|
||||
if 0 <= index < len(pipe_data):
|
||||
return (pipe_data[index],)
|
||||
else:
|
||||
# Trả về None nếu index nằm ngoài phạm vi
|
||||
|
||||
@@ -22,6 +22,9 @@ class NH_UniversalSliderBuilder:
|
||||
FUNCTION = "build_sliders"
|
||||
CATEGORY = "NH-Nodes/Utils"
|
||||
OUTPUT_NODE = True # Đánh dấu đây là node có đầu ra động
|
||||
# Node này không có input/output thật nên không dùng được. DEPRECATED ẩn nó khỏi search,
|
||||
# còn workflow cũ đã lưu vẫn load được. Dùng Multi-Slider (FLOAT/INT) thay thế.
|
||||
DEPRECATED = True
|
||||
|
||||
def build_sliders(self, config, **kwargs):
|
||||
cls = self.__class__
|
||||
|
||||
@@ -65,15 +65,38 @@ class NH_VTonUltimateProcessor:
|
||||
}
|
||||
return LOADED_VTON_MODELS
|
||||
|
||||
def process(self, human_image, category, mask_feathering, cover_shoes, refine_hands, refine_hair, device,
|
||||
def process(self, human_image, category, mask_feathering, cover_shoes, refine_hands, refine_hair, device,
|
||||
offset_top=0, offset_bottom=0, offset_left=0, offset_right=0):
|
||||
# Import trước load_model: thiếu scikit-image thì báo lỗi ngay, không phải tải model xong mới lỗi.
|
||||
from ...core.utils_mask import get_mask_location
|
||||
|
||||
if human_image.shape[0] == 0:
|
||||
raise ValueError("VTON Ultimate Processor: human_image is an empty batch")
|
||||
vton_model = self.load_model(device)
|
||||
human_img_pil = tensor_to_pil(human_image)
|
||||
# Mỗi ảnh trong batch chạy độc lập với cùng setting, rồi nối các output theo chiều batch.
|
||||
results = [
|
||||
self._process_image(
|
||||
vton_model, get_mask_location, human_image[index], index, category, mask_feathering, cover_shoes,
|
||||
refine_hands, refine_hair, offset_top, offset_bottom, offset_left, offset_right,
|
||||
)
|
||||
for index in range(human_image.shape[0])
|
||||
]
|
||||
return tuple(torch.cat(outputs, dim=0) for outputs in zip(*results))
|
||||
|
||||
def _process_image(self, vton_model, get_mask_location, human_image, index, category, mask_feathering, cover_shoes,
|
||||
refine_hands, refine_hair, offset_top, offset_bottom, offset_left, offset_right):
|
||||
"""Xử lý đúng một ảnh [H,W,C]; `index` là vị trí trong batch (chỉ dùng cho thông báo lỗi)."""
|
||||
# DWPose và parsing cần RGB 3 kênh: ảnh xám làm hỏng bước đảo BGR, RGBA đưa 4 kênh vào DWPose (alpha bị bỏ).
|
||||
human_img_pil = tensor_to_pil(human_image).convert("RGB")
|
||||
human_img_resized = resize_image(human_img_pil)
|
||||
|
||||
pose_img_np, _, _, candidate = vton_model['dwprocessor'](np.array(human_img_resized)[:,:,::-1])
|
||||
# DWPose không trả mảng rỗng khi không thấy người: nó vẫn chạy pose trên cả ảnh và đánh dấu mọi keypoint là -1.
|
||||
if len(candidate) == 0 or not np.any(candidate[0] != -1):
|
||||
raise ValueError(
|
||||
f"VTON Ultimate Processor: DWPose found no person in the image at batch index {index}. "
|
||||
"Use a photo where the person is clearly visible."
|
||||
)
|
||||
pose_image_pil = Image.fromarray(pose_img_np[:,:,::-1])
|
||||
parsing_map_pil, parsing_map_colored_pil = vton_model['parsing_model'](human_img_resized)
|
||||
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "nh-nodes"
|
||||
description = "ComfyUI custom nodes for image workflows, tiling, masks, smart resize, logic routing, prompt and text tools, batch automation, garment/portrait segmentation, and VTON preprocessing."
|
||||
version = "1.4.3"
|
||||
version = "1.5.0"
|
||||
license = { file = "LICENSE" }
|
||||
dependencies = [
|
||||
"numpy",
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
"""Shared helpers: safe arithmetic evaluation, tensor/PIL conversion, hex colours."""
|
||||
|
||||
import importlib
|
||||
import math
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
safe_eval = importlib.import_module("NH-Nodes.core.safe_eval")
|
||||
utils = importlib.import_module("NH-Nodes.core.utils")
|
||||
|
||||
evaluate = safe_eval.evaluate
|
||||
ExpressionError = safe_eval.ExpressionError
|
||||
|
||||
|
||||
# --- safe_eval ---
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"expression,expected",
|
||||
[
|
||||
("a + b * 2", 7.0),
|
||||
("(a + b) / 2", 2.0),
|
||||
("a ** 2", 1.0),
|
||||
("-a + +b", 2.0),
|
||||
("b // 2", 1.0),
|
||||
("b % 2", 1.0),
|
||||
("clamp(10, 0, 5)", 5.0),
|
||||
("round(3.14159, 2)", 3.14),
|
||||
("sqrt(16) + floor(1.7) + ceil(1.2) + abs(-1) + min(a, b) + max(a, b)", 4 + 1 + 2 + 1 + 1 + 3),
|
||||
],
|
||||
)
|
||||
def test_math_expressions(expression, expected):
|
||||
assert evaluate(expression, {"a": 1, "b": 3}, safe_eval.MATH_FUNCTIONS) == pytest.approx(expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expression", ["9**9**9", "10**400", "(10**100)**100**100", "1e308 * 10"])
|
||||
def test_runaway_numbers_fail_fast_instead_of_freezing(expression):
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate(expression)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"expression",
|
||||
[
|
||||
"'a' * 10**10", # string constants are not numbers
|
||||
"__import__('os')", # calls outside the whitelist
|
||||
"a.real", # attribute access
|
||||
"[1, 2][0]", # containers / subscripts
|
||||
"lambda: 1",
|
||||
"(-8) ** 0.5", # would be a complex number
|
||||
"sqrt(-1)",
|
||||
"1e308 * 10 - 1e308 * 10", # inf - inf = nan
|
||||
"unknown + 1",
|
||||
"",
|
||||
"1 +",
|
||||
],
|
||||
)
|
||||
def test_rejected_expressions(expression):
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate(expression, {"a": 1}, safe_eval.MATH_FUNCTIONS)
|
||||
|
||||
|
||||
def test_division_by_zero_modes():
|
||||
warnings = []
|
||||
assert evaluate("a / 0 + 1", {"a": 5}, on_warning=warnings.append) == 1.0
|
||||
assert evaluate("a // 0", {"a": 5}) == 0.0
|
||||
assert evaluate("a % 0", {"a": 5}) == 0.0
|
||||
assert warnings == ["division by zero, using 0"]
|
||||
with pytest.raises(ExpressionError, match="Division by zero"):
|
||||
evaluate("a / 0", {"a": 5}, zero_division="error")
|
||||
|
||||
|
||||
def test_case_insensitive_variables():
|
||||
assert evaluate("BG_W / 2 - w1 / 2", {"bg_w": 100, "w1": 20}, case_insensitive=True) == 40.0
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate("BG_W", {"bg_w": 100})
|
||||
|
||||
|
||||
def test_non_numeric_variables_are_rejected():
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate("a", {"a": "5"})
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate("a", {"a": None})
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate("a", {"a": torch.ones(2)}) # not a scalar
|
||||
|
||||
|
||||
def test_bools_numpy_and_torch_scalars_count_as_numbers():
|
||||
import numpy as np
|
||||
|
||||
assert evaluate("True + 1") == 2.0
|
||||
assert evaluate("a + b", {"a": np.float32(2), "b": np.int64(3)}) == 5.0
|
||||
assert evaluate("a * 2", {"a": torch.tensor(2.5)}) == 5.0
|
||||
assert evaluate("a + 1", {"a": np.bool_(True)}) == 2.0
|
||||
|
||||
|
||||
def test_absurd_nesting_is_an_expression_error_not_a_crash():
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate("-" * 100000 + "1")
|
||||
with pytest.raises(ExpressionError):
|
||||
evaluate("(" * 5000 + "1" + ")" * 5000)
|
||||
|
||||
|
||||
def test_result_is_always_a_finite_float():
|
||||
result = evaluate("2 ** 10")
|
||||
assert isinstance(result, float) and math.isfinite(result) and result == 1024.0
|
||||
|
||||
|
||||
# --- tensor_to_pil ---
|
||||
|
||||
def test_tensor_to_pil_clips_instead_of_wrapping():
|
||||
image = torch.tensor([[[1.2, -0.1, 0.5]]]) # 1x1 RGB with out-of-range values
|
||||
pixel = utils.tensor_to_pil(image).getpixel((0, 0))
|
||||
assert pixel == (255, 0, 127)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("channels,mode", [(1, "L"), (3, "RGB"), (4, "RGBA")])
|
||||
def test_tensor_to_pil_channel_modes(channels, mode):
|
||||
image = torch.rand((1, 8, 6, channels))
|
||||
pil = utils.tensor_to_pil(image)
|
||||
assert pil.mode == mode and pil.size == (6, 8)
|
||||
|
||||
|
||||
def test_tensor_to_pil_rejects_batches_with_clear_message():
|
||||
with pytest.raises(ValueError, match="batch of 2"):
|
||||
utils.tensor_to_pil(torch.rand((2, 4, 4, 3)))
|
||||
|
||||
|
||||
# --- parse_hex_color ---
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text,expected",
|
||||
[
|
||||
("#FF8000", (255, 128, 0, 255)),
|
||||
("ff8000", (255, 128, 0, 255)),
|
||||
("#F80", (255, 136, 0, 255)),
|
||||
("#11223344", (17, 34, 51, 68)),
|
||||
(" #000000 ", (0, 0, 0, 255)),
|
||||
],
|
||||
)
|
||||
def test_parse_hex_color(text, expected):
|
||||
assert utils.parse_hex_color(text) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("text", ["", "#12", "#GGGGGG", "red", None])
|
||||
def test_parse_hex_color_invalid_returns_default(text):
|
||||
assert utils.parse_hex_color(text, default=(1, 2, 3, 4)) == (1, 2, 3, 4)
|
||||
@@ -0,0 +1,279 @@
|
||||
"""Regression tests (2026-10-18 review): the Image Grid Composite size guard and the
|
||||
file-name part of Save Image Path."""
|
||||
|
||||
import importlib
|
||||
import os
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import psutil
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
grid = importlib.import_module("NH-Nodes.nodes.image.image_grid_nodes")
|
||||
tools = importlib.import_module("NH-Nodes.nodes.image.image_tools_nodes")
|
||||
|
||||
GIB = 2**30
|
||||
|
||||
|
||||
# --- Image Grid Composite: the size guard follows the free memory, not Pillow's decompression-bomb limit ---
|
||||
|
||||
@pytest.fixture
|
||||
def free_memory(monkeypatch):
|
||||
"""Make psutil report the given available RAM. The tests only exercise the decision: nothing big is allocated."""
|
||||
|
||||
def report(nbytes):
|
||||
monkeypatch.setattr(psutil, "virtual_memory", lambda: types.SimpleNamespace(available=int(nbytes)))
|
||||
|
||||
return report
|
||||
|
||||
|
||||
def _check(width, height, columns=1, rows=1):
|
||||
grid._check_canvas_size(width, height, columns, rows, width // columns, height // rows)
|
||||
|
||||
|
||||
def _composite(images, **overrides):
|
||||
params = dict(
|
||||
columns=2, rows=2, cell_size_mode="manual", cell_width=8, cell_height=8,
|
||||
resize_mode="stretch", gap_x=0, gap_y=0, background_color_hex="#000000",
|
||||
)
|
||||
params.update(overrides)
|
||||
return grid.NH_ImageGridComposite().composite(images, **params)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pillow_limit", [89_478_485, 1000, None]) # Pillow's default, a low value, "disabled"
|
||||
@pytest.mark.parametrize(
|
||||
"columns,rows,cell",
|
||||
[(40, 1, 1024), (4, 4, 4096), (1, 1, 16384)], # a 40960x1024 filmstrip; 16384x16384 as 4x4 cells and as one cell
|
||||
)
|
||||
def test_grid_guard_accepts_big_canvases_whatever_pillow_limit_other_packs_set(
|
||||
monkeypatch, free_memory, pillow_limit, columns, rows, cell
|
||||
):
|
||||
# Image.MAX_IMAGE_PIXELS is process-wide and only guards Image.open; an in-memory IMAGE never goes through it.
|
||||
monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", pillow_limit)
|
||||
free_memory(64 * GIB)
|
||||
|
||||
grid._check_canvas_size(columns * cell, rows * cell, columns, rows, cell, cell)
|
||||
|
||||
|
||||
def test_grid_estimate_is_the_peak_of_the_canvas_to_tensor_conversion():
|
||||
# PIL keeps an RGB canvas as 4 bytes per pixel; _pil_to_tensor then holds two float32 RGB arrays (astype, / 255.0).
|
||||
pil_rgb, float32_rgb = 4, 3 * 4
|
||||
peak_per_pixel = pil_rgb + 2 * float32_rgb
|
||||
|
||||
assert grid._estimate_canvas_bytes(1, 1) == peak_per_pixel
|
||||
assert grid._estimate_canvas_bytes(40960, 1024) == 40960 * 1024 * peak_per_pixel
|
||||
assert grid._estimate_canvas_bytes(16384, 16384) == 7 * GIB
|
||||
|
||||
|
||||
def test_grid_estimate_covers_what_the_node_keeps_alive():
|
||||
canvas = Image.new("RGB", (64, 48))
|
||||
tensor = grid._pil_to_tensor(canvas)
|
||||
|
||||
held = tensor.numel() * tensor.element_size() + 64 * 48 * 4 # the IMAGE tensor and the PIL canvas
|
||||
assert grid._estimate_canvas_bytes(64, 48) >= held
|
||||
|
||||
|
||||
def test_grid_guard_needs_headroom_on_top_of_the_estimate(free_memory):
|
||||
width = height = 8192
|
||||
need = grid._estimate_canvas_bytes(width, height)
|
||||
|
||||
free_memory(4 * need)
|
||||
_check(width, height)
|
||||
|
||||
for free in (need * 1.05, need, need // 2, 0): # barely enough is not enough: some memory stays free
|
||||
free_memory(free)
|
||||
with pytest.raises(ValueError, match=r"(?i)too large"):
|
||||
_check(width, height)
|
||||
|
||||
|
||||
def test_grid_guard_switches_exactly_at_the_usable_share_of_the_free_memory(free_memory):
|
||||
need = grid._estimate_canvas_bytes(4096, 4096)
|
||||
threshold = need / grid._USABLE_MEMORY_FRACTION
|
||||
|
||||
free_memory(threshold + 2)
|
||||
_check(4096, 4096)
|
||||
|
||||
free_memory(threshold - 2)
|
||||
with pytest.raises(ValueError, match=r"(?i)too large"):
|
||||
_check(4096, 4096)
|
||||
|
||||
|
||||
def test_grid_guard_error_names_the_grid_and_the_memory(free_memory):
|
||||
free_memory(2 * GIB)
|
||||
|
||||
with pytest.raises(ValueError) as error:
|
||||
grid._check_canvas_size(16384, 16384, 4, 4, 4096, 4096)
|
||||
|
||||
message = str(error.value)
|
||||
assert "16384x16384" in message and "4x4 cells of 4096x4096" in message
|
||||
assert "7.0 GiB" in message and "2.0 GiB" in message # what it needs / what is free
|
||||
|
||||
|
||||
def test_grid_node_checks_the_memory_before_it_allocates_the_canvas(monkeypatch, free_memory):
|
||||
free_memory(1024) # 1 KiB free: even this 16x16 canvas (2x2 cells of 8x8) cannot be planned
|
||||
created_sizes = []
|
||||
real_new = Image.new
|
||||
|
||||
def spy(mode, size, *args, **kwargs):
|
||||
created_sizes.append(size)
|
||||
return real_new(mode, size, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(Image, "new", spy)
|
||||
|
||||
with pytest.raises(ValueError, match=r"(?i)too large"):
|
||||
_composite([torch.rand((1, 8, 8, 3))])
|
||||
|
||||
assert (16, 16) not in created_sizes # Pillow itself makes 8x8 images while converting the input
|
||||
|
||||
|
||||
def test_grid_node_ignores_the_pillow_pixel_limit(monkeypatch):
|
||||
monkeypatch.setattr(Image, "MAX_IMAGE_PIXELS", 1000) # the first guard refused every canvas over 2 x 1000 px
|
||||
|
||||
image = _composite([torch.rand((1, 8, 8, 3))], cell_width=32, cell_height=32)[0]
|
||||
|
||||
assert image.shape == (1, 64, 64, 3) # 4096 px
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"columns,rows,cell_width,cell_height,shape",
|
||||
[
|
||||
(3, 1, 16384, 1, (1, 1, 49152, 3)), # 49152 px wide
|
||||
(1, 3, 1, 16384, (1, 49152, 1, 3)), # 49152 px tall
|
||||
],
|
||||
)
|
||||
def test_grid_node_builds_strips_longer_than_the_old_32768_px_cap(columns, rows, cell_width, cell_height, shape):
|
||||
image = _composite(
|
||||
[torch.rand((1, 2, 2, 3))], columns=columns, rows=rows, cell_width=cell_width, cell_height=cell_height
|
||||
)[0]
|
||||
|
||||
assert image.shape == shape
|
||||
|
||||
|
||||
def test_grid_guard_does_nothing_without_psutil(monkeypatch):
|
||||
monkeypatch.setattr(grid, "psutil", None)
|
||||
|
||||
_check(16384, 16384) # no information about the memory: no guard, as before the guard existed
|
||||
|
||||
|
||||
def test_grid_guard_does_nothing_when_psutil_cannot_read_the_memory(monkeypatch):
|
||||
def broken():
|
||||
raise OSError("memory status is not available")
|
||||
|
||||
monkeypatch.setattr(psutil, "virtual_memory", broken)
|
||||
|
||||
_check(16384, 16384)
|
||||
|
||||
|
||||
def test_grid_guard_refuses_a_side_pillow_cannot_create(free_memory):
|
||||
free_memory(1024 * GIB)
|
||||
ceiling = grid._PIL_MAX_SIDE # Image.new raises a bare MemoryError for a width above INT_MAX // 4 - 1
|
||||
|
||||
_check(ceiling, 1)
|
||||
for width, height in ((ceiling + 1, 1), (1, ceiling + 1)):
|
||||
with pytest.raises(ValueError, match=r"(?i)too large.*Pillow"):
|
||||
_check(width, height)
|
||||
|
||||
|
||||
# --- Save Image Path: the file-name part of filename_prefix is checked too ---
|
||||
|
||||
@pytest.fixture
|
||||
def written(monkeypatch):
|
||||
"""Record the files Save Image Path would write instead of writing them, so that a prefix that escapes the
|
||||
target folder cannot leave a stray file on the real disk."""
|
||||
paths = []
|
||||
monkeypatch.setattr(tools, "_save_pil_image", lambda image, path, file_format, dpi: paths.append(path))
|
||||
return paths
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def as_windows(monkeypatch):
|
||||
"""Judge prefixes by the Windows rules even where the tests run on another OS (a no-op on Windows itself)."""
|
||||
monkeypatch.setattr(tools, "_WINDOWS", True, raising=False)
|
||||
|
||||
|
||||
def _save_image_path(tmp_path, prefix, **kwargs):
|
||||
settings = dict(folder_path=str(tmp_path / "out"), create_new_folder=False, new_folder_name="")
|
||||
settings.update(kwargs)
|
||||
return tools.NH_SaveImagePath().save_images(torch.rand((1, 4, 4, 3)), filename_prefix=prefix, **settings)
|
||||
|
||||
|
||||
def _png_files_outside_out(tmp_path):
|
||||
return [p for p in tmp_path.rglob("*.png") if p.relative_to(tmp_path).parts[0] != "out"]
|
||||
|
||||
|
||||
# "D:x" and "x:" are paths on another drive, "C:x" is relative to the current folder of drive C: (os.path.join keeps
|
||||
# it below the target but drops the "C:", so the reported name is wrong and every run overwrites x0001.png), and
|
||||
# "name:stream" writes an NTFS alternate data stream. A colon cannot be part of a Windows file name at all.
|
||||
COLON_PREFIXES = [
|
||||
"D:x", "C:x", "c:x", "a:b", "x:", ":", "12:30", "name:stream",
|
||||
"sub/D:x", "sub\\C:x", "C:sub/name", "D:\\x", "D:/x",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("as_windows")
|
||||
@pytest.mark.parametrize("number_suffix", [True, False])
|
||||
@pytest.mark.parametrize("prefix", COLON_PREFIXES)
|
||||
def test_windows_prefix_with_a_colon_is_refused_before_anything_is_written(tmp_path, written, prefix, number_suffix):
|
||||
with pytest.raises(ValueError, match="filename_prefix"):
|
||||
_save_image_path(tmp_path, prefix, number_suffix=number_suffix)
|
||||
|
||||
assert written == []
|
||||
assert not (tmp_path / "out").exists() # not even the target folder is created
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"prefix",
|
||||
["/x", "\\x", "//server/share/x", "\\\\server\\share\\x", "\\\\?\\C:\\x", "../x", "..\\x", "sub/../../x"],
|
||||
)
|
||||
def test_rooted_unc_and_parent_prefixes_are_refused_on_every_os(tmp_path, written, prefix):
|
||||
with pytest.raises(ValueError, match="filename_prefix"):
|
||||
_save_image_path(tmp_path, prefix)
|
||||
|
||||
assert written == []
|
||||
assert not (tmp_path / "out").exists()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"prefix,expected",
|
||||
[
|
||||
("name", "name0001.png"),
|
||||
("sub/name", "sub/name0001.png"),
|
||||
("sub\\name", "sub\\name0001.png"),
|
||||
("a/b/name", "a/b/name0001.png"),
|
||||
("my..name", "my..name0001.png"),
|
||||
("sub/", "sub/0001.png"),
|
||||
],
|
||||
)
|
||||
def test_plain_and_subfolder_prefixes_keep_working(tmp_path, prefix, expected):
|
||||
filenames, path, count = _save_image_path(tmp_path, prefix)
|
||||
|
||||
assert (filenames, path, count) == (expected, str(tmp_path / "out"), 1)
|
||||
assert (tmp_path / "out" / Path(expected.replace("\\", "/"))).is_file()
|
||||
assert _png_files_outside_out(tmp_path) == []
|
||||
|
||||
|
||||
def test_plain_prefix_is_numbered_across_runs(tmp_path):
|
||||
first = _save_image_path(tmp_path, "sub/name")[0]
|
||||
second = _save_image_path(tmp_path, "sub\\name")[0]
|
||||
|
||||
assert (first, second) == ("sub/name0001.png", "sub\\name0002.png")
|
||||
assert (tmp_path / "out" / "sub" / "name0002.png").is_file()
|
||||
|
||||
|
||||
@pytest.mark.skipif(os.name == "nt", reason="a colon is not a valid character in a Windows file name")
|
||||
def test_colons_stay_allowed_outside_windows(tmp_path):
|
||||
filenames, _, _ = _save_image_path(tmp_path, "render 12:30")
|
||||
|
||||
assert filenames == "render 12:300001.png"
|
||||
assert (tmp_path / "out" / "render 12:300001.png").is_file()
|
||||
|
||||
|
||||
def test_same_size_rgba_resize_keeps_colour_under_transparent_pixels():
|
||||
picker = importlib.import_module("NH-Nodes.nodes.resolution.smart_resolution_picker")
|
||||
image = torch.zeros((1, 4, 8, 8))
|
||||
image[:, :3] = torch.tensor([0.2, 0.6, 0.9]).view(1, 3, 1, 1) # colour kept under alpha 0
|
||||
for method in ("lanczos", "box", "hamming"):
|
||||
resized = picker._resize_bchw(image, 8, 8, method)
|
||||
assert torch.allclose(resized, image), method
|
||||
@@ -0,0 +1,462 @@
|
||||
"""Regression tests: image core nodes (VTON batches, face paste, agnostic image, mask-aware
|
||||
resize, large image save, cache manager, grid composite)."""
|
||||
|
||||
import importlib
|
||||
import importlib.util
|
||||
import os
|
||||
|
||||
import comfy.cli_args
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
vton = importlib.import_module("NH-Nodes.nodes.vton.vton_preprocessor_nodes")
|
||||
face_paste = importlib.import_module("NH-Nodes.nodes.image.face_paste")
|
||||
agnostic = importlib.import_module("NH-Nodes.nodes.image.agnostic_image")
|
||||
mask_resize = importlib.import_module("NH-Nodes.nodes.image.mask_resize_image")
|
||||
large_save = importlib.import_module("NH-Nodes.nodes.image.large_image_save")
|
||||
cache_manager = importlib.import_module("NH-Nodes.nodes.image.large_preview_cache_manager")
|
||||
grid = importlib.import_module("NH-Nodes.nodes.image.image_grid_nodes")
|
||||
|
||||
RED, GREEN, BLUE = (1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0)
|
||||
YELLOW, WHITE, CYAN = (1.0, 1.0, 0.0), (1.0, 1.0, 1.0), (0.0, 1.0, 1.0)
|
||||
|
||||
|
||||
def _solid(color, h=32, w=32):
|
||||
"""One image [H,W,C] filled with a colour."""
|
||||
return torch.tensor(color, dtype=torch.float32).expand(h, w, len(color)).clone()
|
||||
|
||||
|
||||
def _boxes(boxes, h=32, w=32):
|
||||
"""MASK batch [B,H,W]; each box is (x1, y1, x2, y2) with an exclusive end."""
|
||||
masks = torch.zeros((len(boxes), h, w))
|
||||
for index, (x1, y1, x2, y2) in enumerate(boxes):
|
||||
masks[index, y1:y2, x1:x2] = 1.0
|
||||
return masks
|
||||
|
||||
|
||||
def _bbox(condition):
|
||||
"""(top, bottom, left, right) of the True pixels of a 2-D bool tensor."""
|
||||
rows, cols = torch.nonzero(condition, as_tuple=True)
|
||||
return int(rows.min()), int(rows.max()), int(cols.min()), int(cols.max())
|
||||
|
||||
|
||||
# --- VTON Ultimate Processor: batches, RGB conversion, no-person error ---
|
||||
|
||||
_MODEL_FILES = (
|
||||
"dwpose/yolox_l.onnx", "dwpose/dw-ll_ucoco_384.onnx",
|
||||
"humanparsing/parsing_atr.onnx", "humanparsing/parsing_lip.onnx",
|
||||
)
|
||||
|
||||
|
||||
def _vton_ready():
|
||||
if importlib.util.find_spec("onnxruntime") is None or importlib.util.find_spec("skimage") is None:
|
||||
return False
|
||||
return all(os.path.isfile(os.path.join(vton.MODELS_DIR, name)) for name in _MODEL_FILES)
|
||||
|
||||
|
||||
class _FakePose:
|
||||
"""DWposeDetector stand-in. Like the real one it never returns an empty array for a person-less
|
||||
image: it marks every keypoint -1. A completely black image counts as "no person"."""
|
||||
|
||||
def __init__(self):
|
||||
self.seen = []
|
||||
self.empty_when_no_person = False
|
||||
|
||||
def __call__(self, bgr):
|
||||
self.seen.append(bgr)
|
||||
height, width = bgr.shape[:2]
|
||||
pose_image = np.zeros((height, width, 3), np.uint8)
|
||||
if bgr.max() == 0:
|
||||
candidate = np.zeros((0, 134, 2), np.float32) if self.empty_when_no_person else np.full((1, 134, 2), -1.0, np.float32)
|
||||
else:
|
||||
candidate = np.full((1, 134, 2), 0.5, np.float32)
|
||||
return pose_image, None, None, candidate
|
||||
|
||||
|
||||
class _FakeParsing:
|
||||
"""Parsing stand-in: the left part of the image becomes upper clothes (label 5), wider when brighter."""
|
||||
|
||||
def __call__(self, image):
|
||||
pixels = np.array(image)
|
||||
labels = np.zeros(pixels.shape[:2], np.uint8)
|
||||
labels[:, : max(1, int(pixels.shape[1] * pixels.mean() / 255.0))] = 5
|
||||
return Image.fromarray(labels), Image.fromarray(labels)
|
||||
|
||||
|
||||
def _fake_get_mask_location(category, parsing_map, keypoints, width, height, *offsets):
|
||||
return Image.fromarray(((np.array(parsing_map) == 5) * 255).astype(np.uint8)), None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_vton(monkeypatch):
|
||||
pytest.importorskip("skimage") # core.utils_mask imports it at module level
|
||||
utils_mask = importlib.import_module("NH-Nodes.core.utils_mask")
|
||||
pose = _FakePose()
|
||||
monkeypatch.setattr(utils_mask, "get_mask_location", _fake_get_mask_location)
|
||||
monkeypatch.setattr(vton, "LOADED_VTON_MODELS", {"dwprocessor": pose, "parsing_model": _FakeParsing(), "device": "cpu"})
|
||||
return pose
|
||||
|
||||
|
||||
def _process(images, category="Upper-body"):
|
||||
return vton.NH_VTonUltimateProcessor().process(images, category, 4, False, True, True, "cpu")
|
||||
|
||||
|
||||
def _flat_batch(levels, channels=3, h=48, w=64):
|
||||
return torch.stack([torch.full((h, w, channels), level) for level in levels])
|
||||
|
||||
|
||||
def test_vton_processes_every_batch_image_independently(fake_vton):
|
||||
batch = _flat_batch([0.2, 0.5, 0.8])
|
||||
outputs = _process(batch)
|
||||
|
||||
assert len(outputs) == 8
|
||||
assert all(output.shape[0] == 3 for output in outputs)
|
||||
for index in range(3):
|
||||
single = _process(batch[index:index + 1])
|
||||
for got, want in zip(outputs, single):
|
||||
assert torch.equal(got[index:index + 1], want)
|
||||
# The items really differ, so the comparison above is not trivially true.
|
||||
assert not torch.equal(outputs[0][0], outputs[0][2])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("channels", [1, 4])
|
||||
def test_vton_converts_grayscale_and_rgba_input_to_rgb(fake_vton, channels):
|
||||
outputs = _process(_flat_batch([0.5], channels=channels))
|
||||
|
||||
assert fake_vton.seen[0].ndim == 3 and fake_vton.seen[0].shape[2] == 3 # what DWPose receives
|
||||
assert outputs[4].shape == (1, 48, 64, 3) # masked_image
|
||||
|
||||
|
||||
@pytest.mark.parametrize("empty_candidate", [False, True])
|
||||
def test_vton_reports_batch_index_when_no_person_is_found(fake_vton, empty_candidate):
|
||||
fake_vton.empty_when_no_person = empty_candidate
|
||||
batch = _flat_batch([0.5, 0.0, 0.5]) # the middle image is black: nothing to detect
|
||||
|
||||
with pytest.raises(ValueError, match="batch index 1"):
|
||||
_process(batch)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _vton_ready(), reason="VTON models or onnxruntime not available")
|
||||
def test_vton_batch_of_two_matches_single_results_on_cpu(monkeypatch):
|
||||
from skimage import data
|
||||
|
||||
monkeypatch.setattr(vton, "LOADED_VTON_MODELS", None)
|
||||
astronaut = torch.from_numpy(data.astronaut().astype(np.float32) / 255.0)
|
||||
batch = torch.stack([astronaut, astronaut.flip(1)]) # the second image is mirrored left-right
|
||||
node = vton.NH_VTonUltimateProcessor()
|
||||
|
||||
outputs = node.process(batch, "Upper-body", 4, False, True, True, "cpu")
|
||||
|
||||
assert len(outputs) == 8
|
||||
assert all(outputs[i].shape == (2, 512, 512) for i in (0, 1, 5, 6, 7))
|
||||
assert all(outputs[i].shape == (2, 512, 512, 3) for i in (2, 3, 4))
|
||||
for index in range(2):
|
||||
single = node.process(batch[index:index + 1], "Upper-body", 4, False, True, True, "cpu")
|
||||
for got, want in zip(outputs, single):
|
||||
torch.testing.assert_close(got[index:index + 1], want)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _vton_ready(), reason="VTON models or onnxruntime not available")
|
||||
def test_vton_real_models_report_batch_index_for_person_less_image(monkeypatch):
|
||||
from skimage import data
|
||||
|
||||
monkeypatch.setattr(vton, "LOADED_VTON_MODELS", None)
|
||||
astronaut = torch.from_numpy(data.astronaut().astype(np.float32) / 255.0)
|
||||
blank = torch.full_like(astronaut, 0.5)
|
||||
|
||||
with pytest.raises(ValueError, match="batch index 1"):
|
||||
vton.NH_VTonUltimateProcessor().process(torch.stack([astronaut, blank]), "Upper-body", 4, False, True, True, "cpu")
|
||||
|
||||
|
||||
# --- NH Simple Face Paste ---
|
||||
|
||||
def _paste(dest, dest_mask, source, source_mask, feathering=0):
|
||||
return face_paste.NH_SimpleFacePaste().paste_face(dest, dest_mask, source, source_mask, feathering)[0]
|
||||
|
||||
|
||||
def test_face_paste_processes_every_item_and_broadcasts_single_inputs():
|
||||
dest = torch.stack([_solid(RED), _solid(GREEN)])
|
||||
|
||||
out = _paste(dest, _boxes([(8, 8, 24, 24)]), torch.stack([_solid(BLUE)]), _boxes([(0, 0, 16, 16)]))
|
||||
|
||||
assert out.shape == (2, 32, 32, 3)
|
||||
for index, color in enumerate((RED, GREEN)):
|
||||
assert torch.equal(out[index, 16, 16], torch.tensor(BLUE)) # pasted
|
||||
assert torch.equal(out[index, 2, 2], torch.tensor(color)) # destination untouched
|
||||
|
||||
|
||||
def test_face_paste_pairs_batch_items_by_index():
|
||||
dest = torch.stack([_solid(RED), _solid(GREEN)])
|
||||
source = torch.stack([_solid(BLUE), _solid(YELLOW)])
|
||||
|
||||
out = _paste(dest, _boxes([(0, 0, 8, 8), (24, 24, 32, 32)]), source, _boxes([(0, 0, 16, 16)]))
|
||||
|
||||
assert torch.equal(out[0, 4, 4], torch.tensor(BLUE)) and torch.equal(out[0, 28, 28], torch.tensor(RED))
|
||||
assert torch.equal(out[1, 28, 28], torch.tensor(YELLOW)) and torch.equal(out[1, 4, 4], torch.tensor(GREEN))
|
||||
|
||||
|
||||
def test_face_paste_clips_overshoot_instead_of_wrapping():
|
||||
# Resampled images can leave [0, 1]; uint8 conversion must saturate, not wrap around to black.
|
||||
dest = torch.stack([_solid((1.2, 1.2, 1.2))])
|
||||
source = torch.stack([_solid((-0.2, 1.2, 1.0))])
|
||||
|
||||
out = _paste(dest, _boxes([(8, 8, 24, 24)]), source, _boxes([(0, 0, 16, 16)]))
|
||||
|
||||
assert torch.equal(out[0, 16, 16], torch.tensor([0.0, 1.0, 1.0]))
|
||||
assert torch.equal(out[0, 2, 2], torch.tensor(WHITE))
|
||||
|
||||
|
||||
def test_face_paste_leaves_items_without_a_mask_untouched():
|
||||
dest = torch.rand((2, 32, 32, 3))
|
||||
dest_mask = _boxes([(8, 8, 24, 24), (0, 0, 0, 0)]) # the second mask is empty
|
||||
|
||||
out = _paste(dest, dest_mask, torch.stack([_solid(BLUE)]), _boxes([(0, 0, 16, 16)]))
|
||||
|
||||
assert out.shape == dest.shape
|
||||
assert torch.equal(out[0, 16, 16], torch.tensor(BLUE))
|
||||
assert torch.equal(out[1], dest[1]) # exact, not re-quantised through uint8
|
||||
|
||||
|
||||
def test_face_paste_accepts_2d_masks():
|
||||
out = _paste(
|
||||
torch.stack([_solid(RED)]), _boxes([(8, 8, 24, 24)])[0],
|
||||
torch.stack([_solid(BLUE)]), _boxes([(0, 0, 16, 16)])[0],
|
||||
)
|
||||
assert torch.equal(out[0, 16, 16], torch.tensor(BLUE))
|
||||
|
||||
|
||||
# --- NH Agnostic Image Generator ---
|
||||
|
||||
def _generate(image, mask, fill_mode="gray", blur_radius=65, noise_strength=1.0, gray_value=0.5, feathering=0):
|
||||
return agnostic.NH_AgnosticImageGenerator().generate(
|
||||
image, mask, fill_mode, blur_radius, noise_strength, gray_value, feathering,
|
||||
)
|
||||
|
||||
|
||||
def test_agnostic_resizes_a_mask_of_a_different_size():
|
||||
image = torch.rand((1, 16, 16, 3))
|
||||
mask = torch.zeros((1, 8, 8))
|
||||
mask[:, 2:6, 2:6] = 1.0
|
||||
|
||||
agnostic_img, masked_img, composite = _generate(image, mask)
|
||||
|
||||
assert agnostic_img.shape == masked_img.shape == composite.shape == (1, 16, 16, 3)
|
||||
assert torch.allclose(agnostic_img[0, 8, 8], torch.full((3,), 0.5)) # inside the (rescaled) mask
|
||||
assert torch.equal(agnostic_img[0, 0, 0], image[0, 0, 0]) # outside
|
||||
|
||||
|
||||
def test_agnostic_treats_a_2d_mask_as_a_batch_of_one():
|
||||
image = torch.rand((1, 16, 16, 3))
|
||||
mask = torch.zeros((16, 16))
|
||||
mask[4:12, 4:12] = 1.0
|
||||
|
||||
agnostic_img = _generate(image, mask)[0]
|
||||
|
||||
assert agnostic_img.shape == (1, 16, 16, 3)
|
||||
assert torch.allclose(agnostic_img[0, 8, 8], torch.full((3,), 0.5))
|
||||
assert torch.equal(agnostic_img[0, 0, 0], image[0, 0, 0])
|
||||
|
||||
|
||||
def test_agnostic_single_mask_broadcasts_over_the_image_batch():
|
||||
image = torch.rand((3, 16, 16, 3))
|
||||
mask = torch.zeros((1, 16, 16))
|
||||
mask[:, 4:12, 4:12] = 1.0
|
||||
|
||||
agnostic_img = _generate(image, mask)[0]
|
||||
|
||||
assert agnostic_img.shape == (3, 16, 16, 3)
|
||||
assert torch.allclose(agnostic_img[:, 8, 8], torch.full((3, 3), 0.5))
|
||||
|
||||
|
||||
def test_agnostic_binary_mask_result_is_unchanged():
|
||||
image = torch.rand((1, 16, 16, 3))
|
||||
mask = torch.zeros((1, 16, 16))
|
||||
mask[:, 4:12, 4:12] = 1.0
|
||||
|
||||
agnostic_img, masked_img, _ = _generate(image, mask)
|
||||
|
||||
expected = image.clone()
|
||||
expected[:, 4:12, 4:12] = 0.5
|
||||
assert torch.equal(agnostic_img, expected)
|
||||
expected[:, 4:12, 4:12] = 1.0
|
||||
assert torch.equal(masked_img, expected)
|
||||
|
||||
|
||||
def test_agnostic_keeps_soft_mask_values():
|
||||
image = torch.rand((1, 16, 16, 3))
|
||||
mask = torch.full((1, 16, 16), 0.25) # a soft mask must blend, not be thresholded away
|
||||
|
||||
agnostic_img = _generate(image, mask)[0]
|
||||
|
||||
assert torch.allclose(agnostic_img, image * 0.75 + 0.5 * 0.25, atol=1e-6)
|
||||
|
||||
|
||||
def test_agnostic_feathering_starts_from_the_soft_mask():
|
||||
image = torch.rand((1, 16, 16, 3))
|
||||
mask = torch.full((1, 16, 16), 0.5)
|
||||
|
||||
agnostic_img = _generate(image, mask, feathering=3)[0]
|
||||
|
||||
assert torch.allclose(agnostic_img, image * 0.5 + 0.5 * 0.5, atol=1e-5)
|
||||
|
||||
|
||||
def test_agnostic_preview_outline_follows_the_half_level_of_a_soft_mask():
|
||||
image = torch.zeros((1, 64, 64, 3))
|
||||
box = np.zeros((64, 64), np.float32)
|
||||
box[20:44, 20:44] = 1.0
|
||||
soft = torch.from_numpy(cv2.GaussianBlur(box, (31, 31), 5)).unsqueeze(0)
|
||||
|
||||
composite = _generate(image, soft)[2]
|
||||
|
||||
assert float(composite[0, 32, 8].sum()) == 0.0 # faint fringe (alpha ~0.01) is not outlined
|
||||
assert float(composite[0, 32, 20, 1]) > 0.5 # the box edge (alpha 0.5) is
|
||||
|
||||
|
||||
def test_agnostic_noise_is_deterministic_per_batch_index():
|
||||
image = torch.rand((2, 8, 8, 3))
|
||||
mask = torch.ones((2, 8, 8))
|
||||
|
||||
first = _generate(image, mask, "noise")[0]
|
||||
second = _generate(image, mask, "noise")[0]
|
||||
|
||||
assert torch.equal(first, second)
|
||||
assert not torch.equal(first[0], first[1])
|
||||
for index in range(2):
|
||||
expected = np.random.default_rng(index).random((8, 8, 3), dtype=np.float32)
|
||||
assert torch.allclose(first[index], torch.from_numpy(expected))
|
||||
|
||||
|
||||
# --- NH Mask-Aware Resize ---
|
||||
|
||||
def _marker_image(h=40, w=80):
|
||||
"""Blue background, red/green end stripes (12.5% wide each) and a white 8x8 marker in the middle."""
|
||||
image = torch.zeros((h, w, 3))
|
||||
image[..., 2] = 1.0
|
||||
image[:, :10] = torch.tensor(RED)
|
||||
image[:, 70:] = torch.tensor(GREEN)
|
||||
image[16:24, 36:44] = 1.0
|
||||
return image.unsqueeze(0)
|
||||
|
||||
|
||||
def _is_white(image):
|
||||
return (image[..., 0] > 0.5) & (image[..., 1] > 0.5) & (image[..., 2] > 0.5)
|
||||
|
||||
|
||||
def test_mask_resize_empty_mask_pad_fits_the_whole_image():
|
||||
image, mask = _marker_image(), torch.zeros((1, 40, 80))
|
||||
|
||||
out_img, out_mask = mask_resize.NH_MaskResizeImage().resize_by_mask(image, mask, 64, 64, "pad", "#00FFFF")
|
||||
|
||||
assert out_img.shape == (1, 64, 64, 3) and out_mask.shape == (1, 64, 64)
|
||||
top, bottom, left, right = _bbox((out_img[0] != torch.tensor(CYAN)).any(-1))
|
||||
assert (bottom - top + 1, right - left + 1) == (32, 64) # 80x40 contained in 64x64, padded top/bottom
|
||||
m_top, m_bottom, m_left, m_right = _bbox(_is_white(out_img[0]))
|
||||
assert abs((m_bottom - m_top) - (m_right - m_left)) <= 1 # the square marker stays square
|
||||
|
||||
|
||||
def test_mask_resize_empty_mask_crop_covers_the_target_with_a_centred_crop():
|
||||
image, mask = _marker_image(), torch.zeros((1, 40, 80))
|
||||
|
||||
out_img, out_mask = mask_resize.NH_MaskResizeImage().resize_by_mask(image, mask, 64, 64, "crop", "#00FFFF")
|
||||
|
||||
assert out_img.shape == (1, 64, 64, 3) and out_mask.shape == (1, 64, 64)
|
||||
red = (out_img[0, ..., 0] > 0.5) & (out_img[0, ..., 1] < 0.5) & (out_img[0, ..., 2] < 0.5)
|
||||
green = (out_img[0, ..., 1] > 0.5) & (out_img[0, ..., 0] < 0.5) & (out_img[0, ..., 2] < 0.5)
|
||||
assert not red.any() and not green.any() # the end stripes are cropped away, not squeezed in
|
||||
assert not (out_img[0] == torch.tensor(CYAN)).all(-1).any() # no padding: the image covers the target
|
||||
m_top, m_bottom, m_left, m_right = _bbox(_is_white(out_img[0]))
|
||||
assert abs((m_bottom - m_top) - (m_right - m_left)) <= 1
|
||||
|
||||
|
||||
def test_mask_resize_output_is_clipped_to_the_unit_range():
|
||||
image = torch.zeros((1, 64, 64, 3))
|
||||
image[:, :, 32:] = 1.0 # a hard edge makes Lanczos ring below 0 and above 1
|
||||
mask = torch.ones((1, 64, 64))
|
||||
|
||||
out_img, _ = mask_resize.NH_MaskResizeImage().resize_by_mask(image, mask, 100, 100, "pad", "#000000")
|
||||
|
||||
assert float(out_img.min()) >= 0.0 and float(out_img.max()) <= 1.0
|
||||
|
||||
|
||||
# --- NH Save Large Image ---
|
||||
|
||||
def _save(tmp_path, image, file_format, **kwargs):
|
||||
node = large_save.NH_SaveLargeImage()
|
||||
node.output_dir = str(tmp_path)
|
||||
return node.save(image, "nh_test", file_format, **kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"file_format,width",
|
||||
[("webp", 16384), ("jpg", 65501)],
|
||||
)
|
||||
def test_save_large_rejects_oversized_images_before_writing(tmp_path, file_format, width):
|
||||
batch = torch.zeros((2, 1, width, 3)) # 1 px tall keeps the test light
|
||||
|
||||
with pytest.raises(ValueError, match=r"(?is)too large.*png.*tiff"):
|
||||
_save(tmp_path, batch, file_format)
|
||||
|
||||
assert list(tmp_path.rglob("*")) == [] # not even the output folder was touched
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"file_format,width",
|
||||
[("webp", 16383), ("jpg", 65500), ("png", 70000), ("tiff", 70000)],
|
||||
)
|
||||
def test_save_large_writes_images_at_the_format_limit(tmp_path, file_format, width):
|
||||
result = _save(tmp_path, torch.zeros((1, 1, width, 3)), file_format)
|
||||
|
||||
saved_count = result["result"][2]
|
||||
assert saved_count == 1
|
||||
assert os.path.getsize(result["result"][1]) > 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("disabled", [False, True])
|
||||
def test_save_large_png_metadata_follows_disable_metadata(tmp_path, monkeypatch, disabled):
|
||||
monkeypatch.setattr(comfy.cli_args.args, "disable_metadata", disabled)
|
||||
|
||||
result = _save(
|
||||
tmp_path, torch.zeros((1, 4, 4, 3)), "png",
|
||||
prompt={"1": {"class_type": "Example"}}, extra_pnginfo={"workflow": {"nodes": []}},
|
||||
)
|
||||
|
||||
with Image.open(result["result"][1]) as saved:
|
||||
text = dict(saved.text)
|
||||
if disabled:
|
||||
assert "prompt" not in text and "workflow" not in text
|
||||
else:
|
||||
assert "prompt" in text and "workflow" in text
|
||||
|
||||
|
||||
# --- NH Large Preview Cache Manager ---
|
||||
|
||||
def test_cache_manager_runs_on_every_queue():
|
||||
changed = cache_manager.NH_LargePreviewCacheManager.IS_CHANGED(
|
||||
action="info", max_age_hours=24.0, max_total_mb=2048.0, cache_id="",
|
||||
)
|
||||
assert changed != changed # NaN never equals itself -> ComfyUI never serves the cached output
|
||||
|
||||
|
||||
# --- Image Grid Composite ---
|
||||
|
||||
def _grid(images, **overrides):
|
||||
params = dict(
|
||||
columns=2, rows=2, cell_size_mode="manual", cell_width=8, cell_height=8,
|
||||
resize_mode="stretch", gap_x=0, gap_y=0, background_color_hex="#000000",
|
||||
)
|
||||
params.update(overrides)
|
||||
return grid.NH_ImageGridComposite().composite(images, **params)
|
||||
|
||||
|
||||
def test_grid_builds_a_normal_canvas():
|
||||
image, slot_count, used_count, _info = _grid([torch.rand((1, 8, 8, 3))])
|
||||
|
||||
assert image.shape == (1, 16, 16, 3)
|
||||
assert (slot_count, used_count) == (4, 1)
|
||||
|
||||
|
||||
# The canvas size guard (free memory, not Pillow's decompression-bomb limit) is tested in
|
||||
# test_grid_save_path_resize_guards.py.
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,435 @@
|
||||
"""Regression tests: logic, text, batch, utils and mask nodes.
|
||||
|
||||
Each fix has tests that fail on the code before release 1.5.0; a few cases only pin behaviour that must
|
||||
keep working (variables, rounding, plain membership, filled template variables). Sections follow the
|
||||
1.5.0 entries of docs/CHANGELOG.md.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
math_nodes = importlib.import_module("NH-Nodes.nodes.logic.math_nodes")
|
||||
logic_nodes = importlib.import_module("NH-Nodes.nodes.logic.logic_nodes")
|
||||
prompt_nodes = importlib.import_module("NH-Nodes.nodes.text.prompt_nodes")
|
||||
text_nodes = importlib.import_module("NH-Nodes.nodes.text.text_nodes")
|
||||
split_lines = importlib.import_module("NH-Nodes.nodes.text.text_split_lines")
|
||||
batch_nodes = importlib.import_module("NH-Nodes.nodes.batch.batch_nodes")
|
||||
pipe_nodes = importlib.import_module("NH-Nodes.nodes.utils.universal_pipe_nodes")
|
||||
utils_nodes = importlib.import_module("NH-Nodes.nodes.utils.utils_nodes")
|
||||
aspect = importlib.import_module("NH-Nodes.nodes.mask.mask_aspect_match")
|
||||
utils_mask = importlib.import_module("NH-Nodes.core.utils_mask")
|
||||
|
||||
|
||||
# --- 1. Math Eval: bad expressions raise, x / 0 only zeroes its own term ---
|
||||
|
||||
def _evaluate(expression, **kwargs):
|
||||
return math_nodes.NH_MathEval().evaluate(expression, **kwargs)
|
||||
|
||||
|
||||
def test_math_eval_keeps_variables_functions_and_rounding():
|
||||
assert _evaluate("a + b * 2", a=1.0, b=3.0) == (7.0, 7, "7.0")
|
||||
assert _evaluate("clamp(a * 10, b, c) + max(d, 1)", a=0.5, b=1.0, c=3.0, d=0.0) == (4.0, 4, "4.0")
|
||||
assert _evaluate("a / 3", a=10.0, round_to=2) == (3.33, 3, "3.33")
|
||||
assert _evaluate("a", a=2.7, round_to=0) == (3.0, 3, "3.0") # result_int is int() of the rounded value
|
||||
assert _evaluate("a", a=2.7) == (2.7, 2, "2.7") # round_to = -1: no rounding
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"expression",
|
||||
["", "1 +", "unknown + 1", "__import__('os')", "'a' * 3", "(-8) ** 0.5", "1e308 * 10", "9 ** 9 ** 9", "sqrt(-1)"],
|
||||
)
|
||||
def test_math_eval_invalid_expression_raises_value_error(expression):
|
||||
with pytest.raises(ValueError):
|
||||
_evaluate(expression)
|
||||
|
||||
|
||||
def test_math_eval_division_and_modulo_by_zero_only_zero_their_own_term(capsys):
|
||||
node = math_nodes.NH_MathEval()
|
||||
assert node.evaluate("a + b / 0", a=1.0, b=5.0)[0] == 1.0
|
||||
assert node.evaluate("a + b // 0", a=1.0, b=5.0)[0] == 1.0 # used to zero the whole result
|
||||
assert node.evaluate("a + b % 0", a=1.0, b=5.0)[0] == 1.0
|
||||
assert capsys.readouterr().out.count("[NH-Nodes] MathEval warning: division by zero") == 3
|
||||
|
||||
|
||||
# --- 2. Random Choice: weight N belongs to input_N, the seed is used as given ---
|
||||
|
||||
def test_random_choice_weight_belongs_to_its_input_slot():
|
||||
node = math_nodes.NH_RandomChoice()
|
||||
# input_0 is not connected: input_1 keeps the weight of slot 1 and input_2 the weight of slot 2.
|
||||
assert node.choose(1, "1,1,10", input_1="B", input_2="C")[2] == "9.1%, 90.9%"
|
||||
picks = [node.choose(seed, "1,1,10", input_1="B", input_2="C")[1] for seed in range(1, 401)]
|
||||
assert picks.count(2) > 300 # about 91 % of 400; matching weights to connected inputs gave 50 %
|
||||
# A slot without a weight counts as 1: input_3 is the 4th weight, not the 2nd.
|
||||
assert node.choose(1, "1,2", input_0="A", input_3="D")[2] == "50.0%, 50.0%"
|
||||
|
||||
|
||||
def test_random_choice_zero_weight_input_is_never_picked():
|
||||
node = math_nodes.NH_RandomChoice()
|
||||
picks = {node.choose(seed, "0,1,1", input_0="A", input_1="B", input_2="C")[1] for seed in range(200)}
|
||||
assert picks == {1, 2}
|
||||
|
||||
|
||||
def test_random_choice_seed_zero_is_deterministic():
|
||||
node = math_nodes.NH_RandomChoice()
|
||||
inputs = {f"input_{index}": chr(ord("A") + index) for index in range(6)}
|
||||
weights = "1,1,1,1,1,1"
|
||||
assert len({node.choose(0, weights, **inputs)[1] for _ in range(30)}) == 1
|
||||
assert node.choose(5, weights, **inputs) == node.choose(5, weights, **inputs)
|
||||
assert {node.choose(seed, weights, **inputs)[1] for seed in range(60)} == set(range(6))
|
||||
|
||||
|
||||
# --- 3. Prompt Template: an empty variable never glues its neighbours together ---
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"template,names,values,expected,missing",
|
||||
[
|
||||
("red dress, {style}, 8k photo", "style", {}, "red dress, 8k photo", "style"),
|
||||
("{style}, red dress, 8k photo", "style", {}, "red dress, 8k photo", "style"), # at the start
|
||||
("red dress, 8k photo, {style}", "style", {}, "red dress, 8k photo", "style"), # at the end
|
||||
("red dress, {a}, {b}, 8k photo", "a,b", {}, "red dress, 8k photo", "a, b"), # consecutive
|
||||
("{a}, {b}, red dress", "a,b", {}, "red dress", "a, b"),
|
||||
("red dress, {a}, {b}", "a,b", {}, "red dress", "a, b"),
|
||||
("{a}, {b}", "a,b", {}, "", "a, b"),
|
||||
("{a}", "a", {}, "", "a"),
|
||||
("a {color} dress", "color", {}, "a dress", "color"), # inside a sentence
|
||||
("a {color} {garment}, {style}, 8k photo", "color,garment,style", {"var2": "dress"}, "a dress, 8k photo", "color, style"),
|
||||
("a red {garment}, {style}, 8k photo", "garment,style", {}, "a red, 8k photo", "garment, style"),
|
||||
("line one, {x}\nline two, {y}, tail\n{x}, last", "x,y", {}, "line one\nline two, tail\nlast", "x, y"),
|
||||
("a, {x}\r\nb", "x", {}, "a\r\nb", "x"),
|
||||
("a, {x}, b", "x", {"var1": None}, "a, b", "x"), # an unconnected input counts as empty
|
||||
],
|
||||
)
|
||||
def test_prompt_template_empty_variable_keeps_neighbours_apart(template, names, values, expected, missing):
|
||||
assert prompt_nodes.NH_PromptTemplate().render(template, names, **values) == (expected, missing)
|
||||
|
||||
|
||||
def test_prompt_template_filled_variables_are_substituted():
|
||||
node = prompt_nodes.NH_PromptTemplate()
|
||||
result = node.render("a {color} {garment}, {style}, 8k photo", "color,garment,style", var1="red", var2="dress", var3="studio")
|
||||
assert result == ("a red dress, studio, 8k photo", "")
|
||||
|
||||
|
||||
# --- 4. Seeds are used as given, so seed 0 is reproducible too ---
|
||||
|
||||
@pytest.mark.parametrize("seed", [0, 3])
|
||||
def test_prompt_scheduler_random_mode_uses_the_seed_literally(seed):
|
||||
node = prompt_nodes.NH_PromptScheduler()
|
||||
prompts = "\n".join("abcdef")
|
||||
for step in range(8):
|
||||
index = node.schedule(prompts, step, 10, "random", seed=seed)[2]
|
||||
assert index == random.Random(seed + step).randint(0, 5)
|
||||
assert index == node.schedule(prompts, step, 10, "random", seed=seed)[2]
|
||||
|
||||
|
||||
def test_text_split_lines_seed_zero_is_deterministic():
|
||||
node = split_lines.NH_TextSplitLines()
|
||||
text = "\n".join("abcdef")
|
||||
expected = list("abcdef")
|
||||
random.Random(0).shuffle(expected)
|
||||
assert node.split_lines(text, 0, "shuffle", 0)[0] == expected
|
||||
assert len({tuple(node.split_lines(text, 0, "shuffle", 0)[0]) for _ in range(20)}) == 1
|
||||
# index -1 picks a random line with the same seeded generator
|
||||
picks = {node.split_lines(text, -1, "sequential", 0)[1] for _ in range(30)}
|
||||
assert picks == {text.split("\n")[random.Random(0).randint(0, 5)]}
|
||||
|
||||
|
||||
def test_text_random_line_seed_zero_is_deterministic():
|
||||
node = split_lines.NH_TextRandomLine()
|
||||
text = "\n".join("abcdef")
|
||||
expected = random.Random(0).randint(0, 5)
|
||||
assert [node.pick(text, 0)[1] for _ in range(30)] == [expected] * 30
|
||||
|
||||
|
||||
# --- 5. Regex Extract: a group that did not take part is an empty string ---
|
||||
|
||||
def test_regex_extract_unmatched_optional_group_is_empty():
|
||||
node = text_nodes.NH_RegexExtract()
|
||||
assert node.extract("b", r"(a)|(b)", "match") == ("b", True, "\nb", 1)
|
||||
assert node.extract("x-y", r"(x)(-)?(z)?", "match") == ("x-", True, "x\n-\n", 1)
|
||||
assert node.extract("ab", r"(a)(c)?", "match") == ("a", True, "a\n", 1)
|
||||
assert node.extract("ab", r"a", "match") == ("a", True, "a", 1) # no group: the whole match
|
||||
|
||||
|
||||
# --- 6. Pack Universal: the package is indexed by slot ---
|
||||
|
||||
def test_pack_universal_keeps_slot_positions():
|
||||
pack = pipe_nodes.NH_PackUniversal().pack(input_1="ONE", input_3="THREE")[0]
|
||||
unpack = pipe_nodes.NH_UnpackUniversal()
|
||||
assert unpack.unpack(0, package=pack) == (None,) # input_0 was left empty
|
||||
assert unpack.unpack(1, package=pack) == ("ONE",)
|
||||
assert unpack.unpack(2, package=pack) == (None,)
|
||||
assert unpack.unpack(3, package=pack) == ("THREE",)
|
||||
assert unpack.unpack(9, package=pack) == (None,)
|
||||
assert unpack.unpack(1, pipe=pack) == ("ONE",) # the legacy `pipe` input reads the same package
|
||||
|
||||
|
||||
def test_pack_universal_ignores_argument_order_and_keeps_falsy_values():
|
||||
pack = pipe_nodes.NH_PackUniversal().pack(input_2="C", input_0=0, input_1="")[0]
|
||||
unpack = pipe_nodes.NH_UnpackUniversal()
|
||||
assert [unpack.unpack(index, package=pack)[0] for index in range(4)] == [0, "", "C", None]
|
||||
assert pipe_nodes.NH_UnpackUniversal().unpack(0, package=pipe_nodes.NH_PackUniversal().pack()[0]) == (None,)
|
||||
|
||||
|
||||
# --- 7. Compare "in": plain membership / substring ---
|
||||
|
||||
@pytest.mark.parametrize("needle,haystack", [(13, "1024x768 3"), ("13", "1024x768 3"), (15, "1,05"), (25, "2 and 5")])
|
||||
def test_compare_in_is_not_a_subsequence_match(needle, haystack):
|
||||
assert logic_nodes.NH_Compare().compare(needle, haystack, "in")[0] is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"needle,haystack",
|
||||
[
|
||||
("cat", "concatenate"),
|
||||
(13, "size 1024x13"),
|
||||
(7, [1, 7]),
|
||||
("7", [1, 7]),
|
||||
(7, ["1", "7"]),
|
||||
("k", {"k": 1}),
|
||||
(1, {"1": "x"}),
|
||||
("", "anything"),
|
||||
],
|
||||
)
|
||||
def test_compare_in_finds_substrings_and_members(needle, haystack):
|
||||
assert logic_nodes.NH_Compare().compare(needle, haystack, "in")[0] is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("needle,haystack", [("dog", "concatenate"), ("CAT", "concatenate"), (4, [1, 7]), ("z", {"k": 1})])
|
||||
def test_compare_in_rejects_non_members(needle, haystack):
|
||||
assert logic_nodes.NH_Compare().compare(needle, haystack, "in")[0] is False
|
||||
|
||||
|
||||
# --- 8. Batch Merge: channel counts follow batch_a ---
|
||||
|
||||
def test_batch_merge_aligns_channels_to_batch_a():
|
||||
node = batch_nodes.NH_BatchMerge()
|
||||
rgb = torch.full((2, 4, 4, 3), 0.25)
|
||||
rgba = torch.full((1, 4, 4, 4), 0.75)
|
||||
rgba[..., 3] = 0.5
|
||||
|
||||
merged, count = node.merge(rgb, rgba) # RGB first: the alpha channel is dropped
|
||||
assert merged.shape == (3, 4, 4, 3) and count == 3
|
||||
assert torch.equal(merged[:2], rgb) and torch.equal(merged[2], rgba[0, ..., :3])
|
||||
|
||||
merged, count = node.merge(rgba, rgb) # RGBA first: RGB images get a fully opaque alpha
|
||||
assert merged.shape == (3, 4, 4, 4) and count == 3
|
||||
assert torch.equal(merged[0], rgba[0])
|
||||
assert torch.equal(merged[1:, ..., :3], rgb)
|
||||
assert torch.all(merged[1:, ..., 3] == 1.0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("resize_mode", ["stretch", "crop", "pad"])
|
||||
def test_batch_merge_handles_size_and_channel_mismatch_together(resize_mode):
|
||||
node = batch_nodes.NH_BatchMerge()
|
||||
rgba = torch.rand((1, 4, 4, 4))
|
||||
rgb_big = torch.rand((2, 8, 6, 3))
|
||||
rgb_same = torch.rand((1, 4, 4, 3))
|
||||
merged, count = node.merge(rgba, rgb_big, rgb_same, resize_mode=resize_mode)
|
||||
assert merged.shape == (4, 4, 4, 4) and count == 4
|
||||
if resize_mode == "pad":
|
||||
# 8x6 fits as 4x3: the picture is opaque, the 1 px letterbox column is transparent (like RGBA input)
|
||||
assert torch.all(merged[1:3, :, :3, 3] == 1.0) and torch.all(merged[1:3, :, 3, 3] == 0.0)
|
||||
assert torch.all(merged[3, ..., 3] == 1.0)
|
||||
else:
|
||||
assert torch.all(merged[1:, ..., 3] == 1.0)
|
||||
|
||||
|
||||
def test_batch_merge_rejects_unsupported_channel_counts_clearly():
|
||||
node = batch_nodes.NH_BatchMerge()
|
||||
with pytest.raises(ValueError, match="channel"):
|
||||
node.merge(torch.rand((1, 4, 4, 3)), torch.rand((1, 4, 4, 2)))
|
||||
|
||||
|
||||
# --- 9. NH Mask Aspect Ratio Match ---
|
||||
|
||||
def _canvas(*boxes, size=100):
|
||||
"""MASK batch [B, size, size]; each box is (x1, y1, x2, y2) with an exclusive end."""
|
||||
masks = torch.zeros((len(boxes), size, size))
|
||||
for index, (x1, y1, x2, y2) in enumerate(boxes):
|
||||
masks[index, y1:y2, x1:x2] = 1.0
|
||||
return masks
|
||||
|
||||
|
||||
def _bounds(mask):
|
||||
"""(x1, y1, x2, y2), inclusive, of the pixels above 0.5."""
|
||||
rows, cols = torch.nonzero(mask > 0.5, as_tuple=True)
|
||||
return int(cols.min()), int(rows.min()), int(cols.max()), int(rows.max())
|
||||
|
||||
|
||||
TALL = _canvas((20, 10, 50, 70)) # 30 x 60 -> ratio 0.5
|
||||
WIDE = _canvas((10, 20, 70, 50)) # 60 x 30 -> ratio 2.0
|
||||
|
||||
|
||||
def _match(target, box, mode):
|
||||
return aspect.NH_MaskAspectRatioMatch().match_aspect_ratio(target, box, mode)[0]
|
||||
|
||||
|
||||
def test_aspect_match_pad_grows_the_box_to_the_requested_size():
|
||||
box = _canvas((30, 30, 50, 50)) # 20 x 20 in the middle
|
||||
out = _match(TALL, box, "pad")
|
||||
assert out.shape == box.shape
|
||||
assert _bounds(out[0]) == (30, 20, 49, 59) # 20 x 40, same centre, nothing shifted
|
||||
assert float(out.sum()) == 20 * 40
|
||||
|
||||
|
||||
def test_aspect_match_pad_keeps_the_original_pixels_where_they_were():
|
||||
box = _canvas((30, 30, 50, 50))
|
||||
box[0, 38:42, 38:42] = 0.0 # a hole must stay exactly there, the padding is added around it
|
||||
out = _match(TALL, box, "pad")[0]
|
||||
assert float(out[38:42, 38:42].sum()) == 0.0
|
||||
assert torch.equal(out[30:50, 30:50], box[0, 30:50, 30:50])
|
||||
assert float(out[20:30, 30:50].min()) == 1.0 and float(out[50:60, 30:50].min()) == 1.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode,expected", [("stretch", (0, 40, 39, 59)), ("pad", (0, 40, 39, 59)), ("crop", (0, 45, 19, 54))])
|
||||
def test_aspect_match_box_touching_the_left_edge(mode, expected):
|
||||
box = _canvas((0, 40, 20, 60)) # x1 = 0 used to crash pad and crop (negative paste offset)
|
||||
out = _match(WIDE, box, mode)
|
||||
assert out.shape == box.shape
|
||||
# The widened box slides inwards instead of being cut by the edge, so it keeps the requested ratio.
|
||||
assert _bounds(out[0]) == expected
|
||||
|
||||
|
||||
def test_aspect_match_stretch_is_centred_and_scaled_not_shifted():
|
||||
box = _canvas((30, 30, 50, 50))
|
||||
box[0, 38:42, 38:42] = 0.0
|
||||
out = _match(TALL, box, "stretch")[0]
|
||||
assert _bounds(out) == (30, 20, 49, 59) # the unchanged width must not move by a pixel
|
||||
assert float(out.sum()) == 20 * 40 - 4 * 8 # the 4 x 4 hole is stretched to 4 x 8
|
||||
|
||||
|
||||
def test_aspect_match_stretch_larger_than_the_image_shows_the_middle_of_the_content():
|
||||
box = _canvas((40, 40, 60, 60))
|
||||
box[0, 45:55, 40:44] = 0.0 # a hole on the far left of the box
|
||||
extreme = _canvas((0, 0, 1000, 10), size=1100) # ratio 100: the stretched box would be 2000 px wide
|
||||
out = _match(extreme, box, "stretch")[0]
|
||||
assert out.shape == box.shape[1:]
|
||||
# Only the middle of the stretched content fits into the 100 px wide image, so the hole is not in view.
|
||||
assert _bounds(out) == (0, 40, 99, 59)
|
||||
assert float(out.sum()) == 100 * 20
|
||||
|
||||
|
||||
def test_aspect_match_crop_trims_around_the_centre_without_shifting():
|
||||
out = _match(TALL, _canvas((30, 30, 50, 50)), "crop")
|
||||
assert _bounds(out[0]) == (35, 30, 44, 49) # 10 x 20 taken from the middle of 20 x 20
|
||||
assert float(out.sum()) == 10 * 20
|
||||
|
||||
|
||||
def test_aspect_match_processes_the_whole_batch():
|
||||
boxes = _canvas((30, 30, 50, 50), (10, 10, 40, 40), (0, 0, 0, 0)) # the last one is empty
|
||||
out = _match(TALL, boxes, "pad")
|
||||
assert out.shape == boxes.shape
|
||||
assert _bounds(out[0]) == (30, 20, 49, 59)
|
||||
assert _bounds(out[1]) == (10, 0, 39, 59) # 30 x 60; grown upwards by sliding back inside the canvas
|
||||
assert float(out[2].sum()) == 0.0 # an empty mask is returned as it is
|
||||
|
||||
|
||||
def test_aspect_match_pairs_target_and_box_by_batch_index():
|
||||
targets = torch.cat([TALL, WIDE])
|
||||
boxes = _canvas((30, 30, 50, 50), (30, 30, 50, 50))
|
||||
out = _match(targets, boxes, "pad")
|
||||
assert _bounds(out[0]) == (30, 20, 49, 59) # tall target -> 20 x 40
|
||||
assert _bounds(out[1]) == (20, 30, 59, 49) # wide target -> 40 x 20
|
||||
assert _match(TALL, boxes, "pad").shape == boxes.shape # a single target serves the whole batch
|
||||
|
||||
|
||||
def test_aspect_match_accepts_masks_without_batch_dimension():
|
||||
out = _match(TALL[0], _canvas((30, 30, 50, 50))[0], "pad")
|
||||
assert out.shape == (1, 100, 100)
|
||||
assert _bounds(out[0]) == (30, 20, 49, 59)
|
||||
|
||||
|
||||
# --- 10. core/utils_mask.py (FitDiT garment masks) ---
|
||||
|
||||
def test_remove_small_clears_only_the_pixels_of_small_regions():
|
||||
mask = np.zeros((60, 60), np.uint8)
|
||||
mask[11:50, 11:50] = 255 # large region, 39 x 39
|
||||
mask[8, 8:20] = 255 # thin L whose bounding box (rows/cols 8..19) overlaps the corner of the square
|
||||
mask[8:20, 8] = 255
|
||||
cleaned = utils_mask.remove_small(mask, min_area=100)
|
||||
assert np.array_equal(cleaned[11:50, 11:50], mask[11:50, 11:50]) # nothing cut out of the large region
|
||||
assert int((cleaned > 0).sum()) == 39 * 39
|
||||
assert not cleaned[8, 8:20].any() and not cleaned[8:20, 8].any()
|
||||
|
||||
|
||||
POSE_W, POSE_H = 200, 300
|
||||
|
||||
|
||||
def _pose(points):
|
||||
"""pose_data [134, 2] like the VTON node builds it: DWPose marks an undetected keypoint -1 and the node
|
||||
scales x by the width and y by the height, so a missing keypoint is (-W, -H).
|
||||
Rows: body 0-17, foot1 18-20 (right), foot2 21-23 (left), face 24-91, hands1 92-112, hands2 113-133."""
|
||||
pose = np.empty((134, 2), np.float32)
|
||||
pose[:] = (-POSE_W, -POSE_H)
|
||||
for row, xy in points.items():
|
||||
pose[row] = xy
|
||||
return pose
|
||||
|
||||
|
||||
def _garment_mask(category, pose):
|
||||
parse = Image.new("L", (POSE_W, POSE_H), 0) # no clothes, no head: only the keypoints decide
|
||||
mask, _ = utils_mask.get_mask_location(category, parse, pose, POSE_W, POSE_H, 0, 0, 0, 0)
|
||||
return mask
|
||||
|
||||
|
||||
def _rect(mask):
|
||||
"""(x1, y1, x2, y2), inclusive, of the white pixels."""
|
||||
rows, cols = np.nonzero(np.array(mask))
|
||||
return int(cols.min()), int(rows.min()), int(cols.max()), int(rows.max())
|
||||
|
||||
|
||||
def test_upper_mask_ignores_keypoints_that_were_not_detected():
|
||||
# Left shoulder, left hip and every hand/face point are missing (negative).
|
||||
pose = _pose({2: (80, 100), 3: (70, 150), 4: (65, 190), 6: (130, 150), 7: (135, 190), 8: (85, 200)})
|
||||
# x: 65 - 25 .. 135 + 25, y: 100 - 20 .. 200 + 10. A missing point used to pull the box to the image top/left.
|
||||
assert _rect(_garment_mask("Upper-body", pose)) == (40, 80, 160, 210)
|
||||
|
||||
|
||||
def test_upper_mask_uses_only_detected_face_points():
|
||||
pose = _pose({2: (80, 100), 5: (120, 100), 8: (85, 200), 11: (115, 200), 29: (100, 60)}) # face point 5 only
|
||||
assert _rect(_garment_mask("Upper-body", pose)) == (55, 60, 145, 210)
|
||||
|
||||
|
||||
def test_lower_mask_starts_at_the_detected_hip():
|
||||
# Left hip, left knee and the left foot are missing; the right hip is at y = 150.
|
||||
pose = _pose({8: (85, 150), 9: (84, 220), 10: (88, 240), 13: (115, 245), 20: (92, 262)})
|
||||
# top = 150 - 50 (no garment found), bottom = foot 262 + 10. min(right hip, missing left hip) put it at 0.
|
||||
assert _rect(_garment_mask("Lower-body", pose)) == (54, 100, 145, 272)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("left_foot", [(0, 0), (-POSE_W, -POSE_H)])
|
||||
def test_lower_mask_bottom_uses_the_right_foot_when_the_left_foot_is_missing(left_foot):
|
||||
pose = _pose({8: (85, 150), 11: (115, 150), 20: (92, 262), 23: left_foot}) # ankles not detected
|
||||
assert _rect(_garment_mask("Lower-body", pose)) == (55, 100, 145, 272)
|
||||
|
||||
|
||||
def test_dresses_mask_ignores_keypoints_that_were_not_detected():
|
||||
pose = _pose({
|
||||
2: (80, 100), 3: (70, 150), 6: (130, 150), # right shoulder side; left shoulder and both wrists missing
|
||||
8: (85, 170), 9: (84, 220), 10: (88, 250), 11: (115, 170), 12: (116, 220), 13: (118, 252),
|
||||
18: (95, 262), 21: (120, 264),
|
||||
})
|
||||
# x: 70 - 50 .. 130 + 50, y: 100 - 20 .. 264 + 10; the missing left shoulder used to put the top at 0.
|
||||
assert _rect(_garment_mask("Dresses", pose)) == (20, 80, 180, 274)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("category", ["Upper-body", "Lower-body", "Dresses"])
|
||||
def test_mask_without_any_detected_keypoint_falls_back_to_the_whole_image(category):
|
||||
mask = _garment_mask(category, _pose({}))
|
||||
assert np.array(mask).all()
|
||||
|
||||
|
||||
# --- 11. Universal Slider Builder cannot work, so it is hidden from the search ---
|
||||
|
||||
def test_universal_slider_builder_is_deprecated_but_still_loads_old_workflows():
|
||||
node_class = utils_nodes.NH_UniversalSliderBuilder
|
||||
assert node_class.DEPRECATED is True
|
||||
assert utils_nodes.NODE_CLASS_MAPPINGS["NH_UniversalSliderBuilder"] is node_class
|
||||
assert utils_nodes.NODE_DISPLAY_NAME_MAPPINGS["NH_UniversalSliderBuilder"] == "Universal Slider Builder (NH)"
|
||||
assert list(node_class.INPUT_TYPES()["required"]) == ["config"]
|
||||
assert not getattr(utils_nodes.NH_BooleanSwitch, "DEPRECATED", False)
|
||||
@@ -0,0 +1,91 @@
|
||||
"""Regression tests for Math Eval scalars, Batch Merge padding and VTON garment boxes."""
|
||||
|
||||
import importlib
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
math_nodes = importlib.import_module("NH-Nodes.nodes.logic.math_nodes")
|
||||
batch_nodes = importlib.import_module("NH-Nodes.nodes.batch.batch_nodes")
|
||||
utils_mask = importlib.import_module("NH-Nodes.core.utils_mask")
|
||||
|
||||
|
||||
# --- Math Eval accepts the scalar types upstream nodes actually send ---
|
||||
|
||||
@pytest.mark.parametrize("value", [np.float32(2), np.int64(2), torch.tensor(2.0), True])
|
||||
def test_math_eval_accepts_numpy_torch_and_bool_scalars(value):
|
||||
result_float, result_int, result_string = math_nodes.NH_MathEval().evaluate("a + b", a=value, b=1.0)
|
||||
assert result_float == float(value) + 1.0 and result_int == int(result_float) and result_string == str(result_float)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("expression", ["foo * 2", "__import__('os')", "a < b", "''"])
|
||||
def test_math_eval_errors_always_name_the_expression(expression):
|
||||
with pytest.raises(ValueError) as error:
|
||||
math_nodes.NH_MathEval().evaluate(expression, a=1.0, b=2.0)
|
||||
assert repr(expression) in str(error.value)
|
||||
|
||||
|
||||
# --- Batch Merge ---
|
||||
|
||||
def test_batch_merge_pad_never_rounds_a_side_to_zero():
|
||||
batch_a = torch.zeros((1, 64, 64, 3))
|
||||
thin = torch.ones((1, 1, 1000, 3)) # scale 0.064 -> 0.064 px tall before the fix
|
||||
result, count = batch_nodes.NH_BatchMerge().merge(batch_a, thin, resize_mode="pad")
|
||||
assert result.shape == (2, 64, 64, 3) and count == 2
|
||||
|
||||
|
||||
def test_batch_merge_pad_bars_get_the_same_alpha_for_rgb_and_rgba_inputs():
|
||||
batch_a = torch.ones((1, 8, 8, 4))
|
||||
rgb = torch.ones((1, 4, 8, 3)) # letterboxed: 2 px bars on top and bottom
|
||||
rgba = torch.ones((1, 4, 8, 4))
|
||||
merged, _ = batch_nodes.NH_BatchMerge().merge(batch_a, rgb, rgba, resize_mode="pad")
|
||||
bars_rgb, bars_rgba = merged[1, :2, :, 3], merged[2, :2, :, 3]
|
||||
assert torch.equal(bars_rgb, bars_rgba) and float(bars_rgb.max()) == 0.0
|
||||
assert float(merged[1, 2:6, :, 3].min()) == 1.0 # the picture itself is opaque
|
||||
|
||||
|
||||
# --- VTON garment boxes when a whole keypoint group is missing ---
|
||||
|
||||
W, H = 200, 300
|
||||
|
||||
|
||||
def _pose(points):
|
||||
pose = np.empty((134, 2), np.float32)
|
||||
pose[:] = (-W, -H) # DWPose marks undetected keypoints -1, the node scales them by W/H
|
||||
for row, xy in points.items():
|
||||
pose[row] = xy
|
||||
return pose
|
||||
|
||||
|
||||
def _rect(mask):
|
||||
rows, cols = np.nonzero(np.array(mask))
|
||||
return int(cols.min()), int(rows.min()), int(cols.max()), int(rows.max())
|
||||
|
||||
|
||||
def test_lower_mask_uses_the_trousers_when_both_hips_are_missing():
|
||||
parse = np.zeros((H, W), np.uint8)
|
||||
parse[150:280, 70:130] = 6 # trousers
|
||||
pose = _pose({9: (84, 220), 10: (88, 250), 12: (116, 220), 13: (118, 252)}) # knees + ankles only
|
||||
mask, _ = utils_mask.get_mask_location("Lower-body", Image.fromarray(parse), pose, W, H, 0, 0, 0, 0)
|
||||
x1, y1, x2, y2 = _rect(mask)
|
||||
assert 140 <= y1 <= 150 # trousers top minus the 5 px pad; the missing hips used to put it at 0
|
||||
assert x1 < 70 and x2 >= 129 and y2 >= 279
|
||||
|
||||
|
||||
def test_upper_mask_with_only_face_points_stays_around_the_shirt():
|
||||
parse = np.zeros((H, W), np.uint8)
|
||||
parse[90:200, 60:140] = 4 # upper clothes
|
||||
pose = _pose({29: (100, 60)}) # one face point, no body or hand keypoints
|
||||
mask, _ = utils_mask.get_mask_location("Upper-body", Image.fromarray(parse), pose, W, H, 0, 0, 0, 0)
|
||||
x1, y1, x2, y2 = _rect(mask)
|
||||
assert (x1, x2) == (35, 165) # shirt contour 60..140 (exclusive end) plus 25 px padding, not the whole width
|
||||
assert y1 == 60 and y2 < H - 1 # from the face line down to the shirt hem, not the whole height
|
||||
|
||||
|
||||
@pytest.mark.parametrize("category", ["Upper-body", "Lower-body", "Dresses"])
|
||||
def test_no_keypoints_and_no_garment_still_cover_the_whole_image(category):
|
||||
parse = Image.fromarray(np.zeros((H, W), np.uint8))
|
||||
mask, _ = utils_mask.get_mask_location(category, parse, _pose({}), W, H, 0, 0, 0, 0)
|
||||
assert np.array(mask).all()
|
||||
@@ -39,6 +39,46 @@ function addMissingLayerInputs(node, count) {
|
||||
}
|
||||
}
|
||||
|
||||
// graph.links is a Map in current frontends (a Proxy that still allows links[id]) and a plain object keyed by
|
||||
// link id in older ones.
|
||||
function findLink(graph, linkId) {
|
||||
if (linkId == null || !graph) {
|
||||
return null;
|
||||
}
|
||||
if (typeof graph.getLink === "function") {
|
||||
return graph.getLink(linkId) ?? null;
|
||||
}
|
||||
const links = graph.links;
|
||||
if (links instanceof Map) {
|
||||
return links.get(linkId) ?? null;
|
||||
}
|
||||
return links?.[linkId] ?? null;
|
||||
}
|
||||
|
||||
// The link attached to the input at `index`. Current frontends derive it from the index (input.link is only a
|
||||
// deprecated view of that), older ones keep the link id on the input.
|
||||
function getInputLink(node, index) {
|
||||
if (!node.graph) {
|
||||
return null; // removed from the graph before the timer fired: nothing can be linked to it
|
||||
}
|
||||
if (typeof node.getInputLink === "function") {
|
||||
return node.getInputLink(index) ?? null;
|
||||
}
|
||||
return findLink(node.graph, node.inputs[index]?.link);
|
||||
}
|
||||
|
||||
// A link remembers the index of the input it ends on (target_slot). Move every link onto its input's new index.
|
||||
function retargetLinks(moves, slotCount) {
|
||||
// Park the links on free slots first: current frontends refuse to move a link onto a slot that another link
|
||||
// still holds, which is what happens when two linked inputs change places.
|
||||
moves.forEach(({ link }, offset) => {
|
||||
link.target_slot = slotCount + offset;
|
||||
});
|
||||
moves.forEach(({ link, newIndex }) => {
|
||||
link.target_slot = newIndex;
|
||||
});
|
||||
}
|
||||
|
||||
function sortLayerInputs(node) {
|
||||
const inputs = node.inputs || [];
|
||||
const background = inputs.filter((input) => input.name === "background_image");
|
||||
@@ -48,7 +88,25 @@ function sortLayerInputs(node) {
|
||||
const others = inputs.filter(
|
||||
(input) => input.name !== "background_image" && !/^layer_\d+$/.test(input.name)
|
||||
);
|
||||
node.inputs = [...background, ...layers, ...others];
|
||||
const sorted = [...background, ...layers, ...others];
|
||||
if (sorted.every((input, index) => input === inputs[index])) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Read the links before reordering: current frontends find an input's link through its index.
|
||||
const moves = [];
|
||||
sorted.forEach((input, newIndex) => {
|
||||
const oldIndex = inputs.indexOf(input);
|
||||
const link = oldIndex === newIndex ? null : getInputLink(node, oldIndex);
|
||||
if (link) {
|
||||
moves.push({ link, newIndex });
|
||||
}
|
||||
});
|
||||
|
||||
node.inputs = sorted;
|
||||
retargetLinks(moves, sorted.length);
|
||||
node.setDirtyCanvas?.(true, true);
|
||||
(node.graph ?? app.graph)?.setDirtyCanvas?.(true, true);
|
||||
}
|
||||
|
||||
function updateLayerInputs(node) {
|
||||
|
||||
Reference in New Issue
Block a user