wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-07-18 20:38:15 +08:00
Code Issues Actions 27 Packages Projects Releases Wiki Activity

Resolve review comments and suggestions

Browse Source
This commit is contained in:
kijai 2026-07-17 14:22:34 +03:00
parent 4cf6a777f6
commit 1f23bcb27b
3 changed files with 183 additions and 182 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

145
comfy_extras/nodes_images.py
View File

@ -10,6 +10,7 @@ import math
import numpy as np
import struct
import torch
import logging
import zlib
import comfy.utils
@ -89,6 +90,149 @@ class ImageCropV2(IO.ComfyNode):
return IO.NodeOutput(img, ui=UI.PreviewImage(img))
def _crop_image_with_mask(item_image, item_mask, max_image_size=1024, pad_factor=1.1,
mask_offset=0, mask_threshold=0.05, bg_rgb=(0.0, 0.0, 0.0),
aspect_ratio=1.0):
img = item_image.permute(2, 0, 1).unsqueeze(0).cpu().float().clamp(0, 1)
mask = item_mask.unsqueeze(0).unsqueeze(0).cpu().float().clamp(0, 1)
# Detect and correct an inverted mask, only when border and center have opposite polarity.
m2d = mask[0, 0]
h, w = m2d.shape
border = torch.cat([m2d[0, :], m2d[-1, :], m2d[:, 0], m2d[:, -1]])
center = m2d[h // 4:h - h // 4, w // 4:w - w // 4]
if float(border.mean()) > 0.5 and float(center.mean()) < 0.5:
mask = 1.0 - mask
if mask_offset > 0:
r = mask_offset
mask = torch.nn.functional.max_pool2d(mask, kernel_size=2 * r + 1, stride=1, padding=r)
elif mask_offset < 0:
r = -mask_offset
mask = 1.0 - torch.nn.functional.max_pool2d(1.0 - mask, kernel_size=2 * r + 1, stride=1, padding=r)
if mask_threshold > 0.0:
mask = torch.where(mask < mask_threshold, torch.zeros_like(mask), mask)
H, W = img.shape[-2:]
if max(H, W) > max_image_size:
scale = max_image_size / max(H, W)
new_w, new_h = int(W * scale), int(H * scale)
img = comfy.utils.common_upscale(img, new_w, new_h, "lanczos", "disabled")
mask = comfy.utils.common_upscale(mask, new_w, new_h, "lanczos", "disabled")
# common_upscale's lanczos path drops the singleton channel dim for masks (utils.py:1062).
if mask.ndim == 3:
mask = mask.unsqueeze(1)
H, W = new_h, new_w
scene_size = (W, H)
alpha_u8 = (mask[0, 0].clamp(0, 1) * 255.0).to(torch.uint8)
fg_pixels = (alpha_u8 > 204).nonzero()
if fg_pixels.numel() == 0:
# Try the inverted mask — auto-invert above may have been too conservative.
inv_fg = ((255 - alpha_u8) > 204).nonzero()
if inv_fg.numel() > 0:
logging.info("Trellis2 preprocess: mask bbox empty, using inverted mask.")
mask = 1.0 - mask
fg_pixels = inv_fg
if fg_pixels.numel() > 0:
y_min, x_min = fg_pixels.min(dim=0).values.tolist()
y_max, x_max = fg_pixels.max(dim=0).values.tolist()
center_y, center_x = (y_min + y_max) / 2.0, (x_min + x_max) / 2.0
bw = x_max - x_min
bh = y_max - y_min
# Grow the bbox so its aspect matches `aspect_ratio` (width/height),
# anchored on the max side. Then apply pad_factor.
if bw / max(bh, 1) >= aspect_ratio:
crop_w = int(bw * pad_factor)
crop_h = int(bw / aspect_ratio * pad_factor)
else:
crop_h = int(bh * pad_factor)
crop_w = int(bh * aspect_ratio * pad_factor)
half_w, half_h = crop_w // 2, crop_h // 2
crop_x1 = int(center_x - half_w)
crop_y1 = int(center_y - half_h)
crop_x2 = crop_x1 + 2 * half_w
crop_y2 = crop_y1 + 2 * half_h
else:
logging.warning("Mask for the image is empty; a clean foreground mask is required for best quality.")
crop_x1, crop_y1, crop_x2, crop_y2 = 0, 0, W, H
crop_bbox = (crop_x1, crop_y1, crop_x2, crop_y2)
# Zero-pad out-of-bounds slice (PIL.crop semantics).
pad_l = max(0, -crop_x1)
pad_t = max(0, -crop_y1)
pad_r = max(0, crop_x2 - W)
pad_b = max(0, crop_y2 - H)
if pad_l or pad_t or pad_r or pad_b:
img = torch.nn.functional.pad(img, (pad_l, pad_r, pad_t, pad_b), value=0.0)
mask = torch.nn.functional.pad(mask, (pad_l, pad_r, pad_t, pad_b), value=0.0)
crop_x1 += pad_l
crop_x2 += pad_l
crop_y1 += pad_t
crop_y2 += pad_t
cropped_img = img [..., crop_y1:crop_y2, crop_x1:crop_x2]
cropped_mask = mask[..., crop_y1:crop_y2, crop_x1:crop_x2]
bg = torch.tensor(bg_rgb, dtype=cropped_img.dtype, device=cropped_img.device).view(1, 3, 1, 1)
composite = (cropped_img * cropped_mask + bg * (1.0 - cropped_mask)).clamp(0, 1)
return composite, crop_bbox, scene_size
class ImageCropToMask(IO.ComfyNode):
"""Crop an image to its mask's bounding box (centered square, with pad_factor
margin), then composite `img * mask` and resize to a square. Handles OOB crops
with zero-padding. Useful for 3D pipelines that expect a centered, background-free
subject at a fixed input resolution (Trellis2, Pixal3D, Hunyuan3D, TripoSR, etc.)."""
@classmethod
def define_schema(cls):
return IO.Schema(
node_id="ImageCropToMask",
display_name="Crop Image to Mask",
category="image/transform",
search_aliases=["crop to mask", "mask crop", "crop mask", "mask crop resize", "crop mask resize", "trellis2", "pixal3d"],
inputs=[
IO.Image.Input("images"),
IO.Mask.Input("masks"),
IO.Int.Input("width", default=1024, min=64, max=4096, step=8, tooltip="Output width in pixels."),
IO.Int.Input("height", default=1024, min=64, max=4096, step=8, tooltip="Output height in pixels."),
IO.Float.Input("pad_factor", default=1.0, min=1.0, max=2.0, step=0.01, tooltip="Extra margin around the mask bounding box as a multiplier."),
IO.Int.Input("grow_mask", default=0, min=-32, max=32, step=1, tooltip="Grow or shrink the mask by this many pixels before cropping."),
IO.Color.Input("background", default="#000000", tooltip="Background color behind the masked subject."),
],
outputs=[IO.Image.Output(display_name="images")],
)
@classmethod
def execute(cls, images, masks, width, height, pad_factor, grow_mask, background) -> IO.NodeOutput:
h = background.lstrip("#")
bg_rgb = (int(h[0:2], 16) / 255.0, int(h[2:4], 16) / 255.0, int(h[4:6], 16) / 255.0) if len(h) == 6 else (0.0, 0.0, 0.0)
images = images[..., :3]
batch_size = images.shape[0]
if masks.shape[0] == 1 and batch_size > 1:
masks = masks.expand(batch_size, -1, -1)
elif masks.shape[0] != batch_size:
raise ValueError(f"Mask batch {masks.shape[0]} does not match image batch {batch_size}")
if masks.shape[-2:] != images.shape[1:3]:
masks = comfy.utils.common_upscale(masks.unsqueeze(1).float(), images.shape[2], images.shape[1], "bilinear", "disabled").squeeze(1)
out_images = []
for b in range(batch_size):
composite, _, _ = _crop_image_with_mask(
images[b], masks[b], max_image_size=max(width, height), pad_factor=pad_factor,
mask_offset=grow_mask, bg_rgb=bg_rgb, aspect_ratio=width / height,
)
composite = comfy.utils.common_upscale(composite, width, height, "lanczos", "disabled")
out_images.append(composite.movedim(-3, -1))
result = torch.cat(out_images, dim=0).to(
device=comfy.model_management.intermediate_device(),
dtype=comfy.model_management.intermediate_dtype(),
)
return IO.NodeOutput(result)
class BoundingBox(IO.ComfyNode):
@classmethod
def define_schema(cls):
@ -1233,6 +1377,7 @@ class ImagesExtension(ComfyExtension):
return [
ImageCrop,
ImageCropV2,
ImageCropToMask,
BoundingBox,
RepeatImageBatch,
ImageFromBatch,

66
comfy_extras/nodes_mesh_postprocess.py
View File

@ -1329,7 +1329,7 @@ class BakeTextureFromVoxel(IO.ComfyNode):
return IO.Schema(
node_id="BakeTextureFromVoxel",
display_name="Bake Texture From Voxel",
category="3d/mesh/texturing",
category="3d/texturing",
description=(
"Bakes PBR textures onto the mesh's existing UV layout (trilinear-sample the "
"sparse voxel volume). Does NOT unwrap — connect a UV unwrap node upstream. Outputs "
@ -1339,7 +1339,7 @@ class BakeTextureFromVoxel(IO.ComfyNode):
inputs=[
IO.Mesh.Input("mesh"),
IO.Voxel.Input("voxel_colors"),
IO.Int.Input("texture_size", default=2048, min=64, max=8192,
IO.Int.Input("texture_size", display_name="resolution", default=2048, min=64, max=8192,
tooltip="Square UV atlas resolution."),
IO.Mesh.Input("reference_mesh", optional=True,
tooltip=(
@ -1432,7 +1432,7 @@ class MeshTextureToImage(IO.ComfyNode):
return IO.Schema(
node_id="MeshTextureToImage",
display_name="Mesh Texture to Image",
category="3d/mesh/texturing",
category="3d/texturing",
description=(
"Extracts a mesh's baked textures as individual IMAGEs: base_color, metallic, "
"roughness, occlusion and normal_map. Channels with nothing baked come back "
@ -1490,7 +1490,7 @@ class ApplyTextureToMesh(IO.ComfyNode):
return IO.Schema(
node_id="ApplyTextureToMesh",
display_name="Apply Texture to Mesh",
category="3d/mesh/texturing",
category="3d/texturing",
description=(
"Attaches baked texture IMAGEs to a mesh's UV layout for SaveGLB. Feed the SAME mesh you baked"
),
@ -1583,12 +1583,12 @@ class BakeNormalMapFromMesh(IO.ComfyNode):
return IO.Schema(
node_id="BakeNormalMapFromMesh",
display_name="Bake Normal Map from Mesh",
category="3d/mesh/texturing",
category="3d/texturing",
description=(
"Bakes a tangent-space normal map (glTF/OpenGL +Y) from a high-poly mesh into a "
"low-poly's UV layout, capturing detail lost to decimation. Feed the UV-unwrapped "
"low_poly and the same-frame high_poly it was decimated from. Outputs an IMAGE for "
"ApplyTextureToMesh's normal_map input."
"low-poly's UV layout, capturing detail lost during decimation. Feed the UV-unwrapped "
"low_poly and the high_poly it was decimated from. Outputs an image for "
"Apply Texture To Mesh's normal_map input."
),
inputs=[
IO.Mesh.Input("low_poly"),
@ -1596,7 +1596,7 @@ class BakeNormalMapFromMesh(IO.ComfyNode):
IO.Int.Input("resolution", default=1024, min=64, max=8192, step=64,
display_name="resolution"),
IO.Float.Input("cage_distance", default=0.05, min=0.001, max=0.5, step=0.001,
tooltip="Surface search band, as a fraction of the bbox diagonal. "
tooltip="Surface search band, as a fraction of the bounding box diagonal. "
"Raise for wrong/missing patches under heavy decimation; "
"lower if it grabs across gaps."),
IO.Boolean.Input("ignore_backfaces", default=True,
@ -1664,12 +1664,12 @@ class BakeAmbientOcclusion(IO.ComfyNode):
return IO.Schema(
node_id="BakeAmbientOcclusion",
display_name="Bake Ambient Occlusion",
category="3d/mesh/texturing",
category="3d/texturing",
description=(
"Bakes an ambient-occlusion map from a high-poly mesh into a low-poly's UV "
"layout (white = open, dark = crevices). Feed the UV-unwrapped low_poly and the "
"high_poly it was decimated from. Outputs a grayscale IMAGE for "
"ApplyTextureToMesh's occlusion input (packed into the ORM map / occlusionTexture)."
"high_poly it was decimated from. Outputs a grayscale image for "
"Apply Texture To Mesh's occlusion input (packed into the ORM map / occlusionTexture)."
),
inputs=[
IO.Mesh.Input("low_poly"),
@ -1678,12 +1678,12 @@ class BakeAmbientOcclusion(IO.ComfyNode):
IO.Int.Input("samples", default=64, min=4, max=1024, step=4,
tooltip="Rays per texel. More = smoother, slower. Raise if grainy."),
IO.Float.Input("max_distance", default=0.5, min=0.01, max=2.0, step=0.01,
tooltip="Ray length, as a fraction of the bbox diagonal. "
tooltip="Ray length, as a fraction of the bounding box diagonal. "
"Smaller = tighter, more local occlusion."),
IO.Float.Input("strength", default=1.0, min=0.0, max=2.0, step=0.05,
tooltip="Scales the occlusion. >1 darkens, <1 lightens."),
IO.Float.Input("bias", default=0.01, min=0.0001, max=0.2, step=0.0005,
tooltip="Ray origin lift off the surface, as a fraction of the bbox "
tooltip="Ray origin lift off the surface, as a fraction of the bounding box "
"diagonal. Raise if even surfaces show dark blotches/holes."),
],
outputs=[IO.Image.Output(display_name="occlusion")],
@ -2202,7 +2202,7 @@ class DecimateMesh(IO.ComfyNode):
IO.Float.Input("feature_edge_quadric_weight", default=0.0, min=0.0, max=1000.0, step=1.0,
tooltip="Extra quadric weight on dihedral feature edges (creases). 0 = off."),
IO.Float.Input("feature_edge_min_dihedral_deg", default=30.0, min=0.0, max=180.0, step=1.0,
tooltip="Min dihedral angle (deg) to count an edge as a feature edge."),
tooltip="Minimum dihedral angle (degree) to count an edge as a feature edge."),
IO.Boolean.Input("clamp_v_to_edge", default=True,
tooltip="Project the QEM-optimal position onto the collapsed edge segment."),
]),
@ -2214,7 +2214,7 @@ class DecimateMesh(IO.ComfyNode):
description=(
"Simplifies a mesh to a target face count using QEM, on the active compute "
"device. 'midpoint' is the cumesh-faithful preset (best quality, preserves thin "
"features / hair); 'qem' places verts at the QEM optimum with line/feature-edge "
"features / hair); 'qem' places vertices at the QEM optimum with line/feature-edge "
"controls. Output stays welded."
),
inputs=[
@ -2285,7 +2285,7 @@ class RemeshMesh(IO.ComfyNode):
sign_mode_options = [
IO.DynamicCombo.Option(key="udf", inputs=[
IO.Boolean.Input("qef", default=False, advanced=True,
tooltip="QEF dual-vertex placement for sharper edges."),
tooltip="QEF (Quadratic Error Function) dual-vertex placement for sharper edges."),
IO.Boolean.Input("drop_inverted_components", default=False, advanced=True,
tooltip="Drop inward-normal (negative-volume) closed components — the UDF inner shell."),
IO.Boolean.Input("drop_enclosed_components", default=False, advanced=True,
@ -2293,7 +2293,7 @@ class RemeshMesh(IO.ComfyNode):
]),
IO.DynamicCombo.Option(key="sdf", inputs=[
IO.Boolean.Input("qef", default=True,
tooltip="QEF dual-vertex placement (recovers sharp features) vs edge-crossing centroid."),
tooltip="QEF (Quadratic Error Function) dual-vertex placement (recovers sharp features) vs edge-crossing centroid."),
IO.Boolean.Input("manifold", default=False,
tooltip="Manifold Dual Contouring: 1-4 dual verts/voxel for multi-sheet cases. Slower."),
]),
@ -2305,25 +2305,25 @@ class RemeshMesh(IO.ComfyNode):
description=(
"Re-extracts a uniformly tessellated mesh via a narrow-band distance field + Dual "
"Contouring, on the active compute device. Normalizes messy / non-manifold / "
"self-intersecting topology; run before DecimateMesh to hit an exact face count. "
"self-intersecting topology; run before Decimate Mesh to hit an exact face count. "
"Output stays welded."
),
inputs=[
IO.Mesh.Input("mesh"),
IO.Int.Input("resolution", default=512, min=32, max=1024,
tooltip="Voxel grid resolution (output density). 256 ~ 100k faces, 512 ~ 1M. "
"For an exact face count, follow with DecimateMesh."),
"For an exact face count, follow with Decimate Mesh."),
IO.DynamicCombo.Input("sign_mode", options=sign_mode_options, display_name="sign_mode",
tooltip="udf: robust to messy/non-manifold input. sdf: clean single "
"surface with QEF sharp-feature recovery, but needs consistent winding."),
"surface with QEF (Quadratic Error Function) sharp-feature recovery, but needs consistent winding."),
IO.Float.Input("band", default=1.0, min=0.5, max=4.0, step=0.1, advanced=True,
tooltip="Narrow-band width in voxel units. In UDF mode also offsets the surface."),
IO.Float.Input("project_back", default=0.0, min=0.0, max=1.0, step=0.05, advanced=True,
tooltip="Lerp verts toward the original surface (0 = pure DC, 1 = snapped)."),
tooltip="Linearly interpolate vertices toward the original surface (0 = pure DC, 1 = snapped)."),
IO.Boolean.Input("fix_poles", default=False, advanced=True,
tooltip="Collapse valence-3 vertex pairs (DC T-junction artifact)."),
IO.Int.Input("smooth_iters", default=0, min=0, max=20,
tooltip="Taubin smoothing iters (0 = off). 2-3 cleans DC stairstepping; higher rounds off QEF edges."),
tooltip="Taubin smoothing iterations (0 = off). 2-3 cleans DC staircase-like artifacts; higher over-smooth QEF edges."),
IO.Float.Input("drop_small_components", default=0.01, min=0.0, max=0.5, step=0.005, advanced=True,
tooltip="Drop components below this fraction of the largest's face count. 0 disables."),
IO.Int.Input("precluster_max_verts", default=8_000_000, min=0, max=50_000_000, advanced=True,
@ -2631,11 +2631,11 @@ class UnwrapMesh(IO.ComfyNode):
return IO.Schema(
node_id="UnwrapMesh",
display_name="Unwrap Mesh UVs",
category="3d/mesh/texturing",
category="3d/texturing",
description=(
"Generates a UV atlas (pure-torch, no xatlas): segments the surface into charts, "
"parameterizes each, packs into a [0,1] atlas. Seam verts duplicated. Run after "
"DecimateMesh/RemeshMesh, before texture baking."
"Generates a UV atlas: segments the surface into charts, parameterizes each, packs into a [0,1] atlas. "
"Vertices on chart seams are duplicated (one copy per chart, same position, own UV), so the output mesh "
"has more vertices than the input. Run after Remesh/Decimate, before texture baking."
),
inputs=[
IO.Mesh.Input("mesh"),
@ -2865,7 +2865,7 @@ class RenderUVAtlas(IO.ComfyNode):
return IO.Schema(
node_id="RenderUVAtlas",
display_name="Render UV Atlas",
category="3d/mesh/texturing",
category="3d/texturing",
description=("Renders a mesh's UV layout as an image."),
inputs=[
IO.Mesh.Input("mesh"),
@ -2909,10 +2909,10 @@ class FillHoles(IO.ComfyNode):
IO.Float.Input("max_perimeter", default=0.03, min=0.0, step=0.0001,
tooltip="Max hole perimeter to fill. 0 disables."),
IO.Float.Input("weld_epsilon_rel", default=1e-5, min=0.0, step=1e-6,
tooltip="Pre-weld tolerance (fraction of bbox diagonal); boundary detection "
"needs welded verts. 0 skips."),
IO.Int.Input("max_verts", default=16, min=3, max=1024,
tooltip="Cap boundary verts per cycle; centroid-fan only works for small "
tooltip="Pre-weld tolerance (fraction of bounding box diagonal); boundary detection "
"needs welded vertices. 0 skips."),
IO.Int.Input("max_vertices", default=16, min=3, max=1024,
tooltip="Cap boundary vertices per cycle; centroid-fan only works for small "
"near-planar holes. Keep ≤16."),
IO.Boolean.Input("fill_chains", default=False,
tooltip="Also fill open chains (not just cycles). Noisy; OFF matches cumesh."),
@ -2984,7 +2984,7 @@ class MeshSmoothNormals(IO.ComfyNode):
inputs=[
IO.Mesh.Input("mesh"),
IO.Float.Input("crease_angle", default=180.0, min=0.0, max=180.0, step=1.0,
tooltip="Edges whose dihedral angle exceeds this (degrees) stay "
tooltip="Edges whose dihedral angle exceeds this value (degrees) stay "
"hard (vertices are split). 180 = fully smooth; lower "
"preserves sharp edges (e.g. ~30-60 for hard-surface)."),
],

154
comfy_extras/nodes_trellis2.py
View File

@ -7,7 +7,6 @@ from comfy_extras.nodes_mesh_postprocess import pack_variable_mesh_batch
import comfy.latent_formats
import comfy.model_management
import comfy.utils
import logging
import math
import torch
@ -102,7 +101,7 @@ class VaeDecodeShapeTrellis(IO.ComfyNode):
def define_schema(cls):
return IO.Schema(
node_id="VaeDecodeShapeTrellis",
category="latent/3d",
category="model/latent/trellis",
inputs=[
IO.Latent.Input("samples"),
IO.Vae.Input("vae"),
@ -174,7 +173,7 @@ class VaeDecodeTextureTrellis(IO.ComfyNode):
def define_schema(cls):
return IO.Schema(
node_id="VaeDecodeTextureTrellis",
category="latent/3d",
category="model/latent/trellis",
inputs=[
IO.Latent.Input("samples"),
IO.Vae.Input("vae"),
@ -248,7 +247,7 @@ class VaeDecodeStructureTrellis2(IO.ComfyNode):
def define_schema(cls):
return IO.Schema(
node_id="VaeDecodeStructureTrellis2",
category="latent/3d",
category="model/latent/trellis",
inputs=[
IO.Latent.Input("samples"),
IO.Vae.Input("vae"),
@ -618,7 +617,7 @@ class EmptyTrellis2LatentStructure(IO.ComfyNode):
def define_schema(cls):
return IO.Schema(
node_id="EmptyTrellis2LatentStructure",
category="latent/3d",
category="model/latent/trellis",
inputs=[
IO.Int.Input("batch_size", default=1, min=1, max=4096, tooltip="The number of latent images in the batch."),
],
@ -648,95 +647,6 @@ def _dinov3_patches_to_2d(tokens, image_size, patch_size=16):
return patches.transpose(1, 2).reshape(tokens.shape[0], -1, h_p, w_p).contiguous()
def _crop_image_with_mask(item_image, item_mask, max_image_size=1024, pad_factor=1.1,
mask_offset=0, mask_threshold=0.05, bg_rgb=(0.0, 0.0, 0.0),
aspect_ratio=1.0):
img = item_image.permute(2, 0, 1).unsqueeze(0).cpu().float().clamp(0, 1)
mask = item_mask.unsqueeze(0).unsqueeze(0).cpu().float().clamp(0, 1)
# Detect and correct an inverted mask, only when border and center have opposite polarity.
m2d = mask[0, 0]
h, w = m2d.shape
border = torch.cat([m2d[0, :], m2d[-1, :], m2d[:, 0], m2d[:, -1]])
center = m2d[h // 4:h - h // 4, w // 4:w - w // 4]
if float(border.mean()) > 0.5 and float(center.mean()) < 0.5:
mask = 1.0 - mask
if mask_offset > 0:
r = mask_offset
mask = torch.nn.functional.max_pool2d(mask, kernel_size=2 * r + 1, stride=1, padding=r)
elif mask_offset < 0:
r = -mask_offset
mask = 1.0 - torch.nn.functional.max_pool2d(1.0 - mask, kernel_size=2 * r + 1, stride=1, padding=r)
if mask_threshold > 0.0:
mask = torch.where(mask < mask_threshold, torch.zeros_like(mask), mask)
H, W = img.shape[-2:]
if max(H, W) > max_image_size:
scale = max_image_size / max(H, W)
new_w, new_h = int(W * scale), int(H * scale)
img = comfy.utils.common_upscale(img, new_w, new_h, "lanczos", "disabled")
mask = comfy.utils.common_upscale(mask, new_w, new_h, "lanczos", "disabled")
# common_upscale's lanczos path drops the singleton channel dim for masks (utils.py:1062).
if mask.ndim == 3:
mask = mask.unsqueeze(1)
H, W = new_h, new_w
scene_size = (W, H)
alpha_u8 = (mask[0, 0].clamp(0, 1) * 255.0).to(torch.uint8)
fg_pixels = (alpha_u8 > 204).nonzero()
if fg_pixels.numel() == 0:
# Try the inverted mask — auto-invert above may have been too conservative.
inv_fg = ((255 - alpha_u8) > 204).nonzero()
if inv_fg.numel() > 0:
logging.info("Trellis2 preprocess: mask bbox empty, using inverted mask.")
mask = 1.0 - mask
fg_pixels = inv_fg
if fg_pixels.numel() > 0:
y_min, x_min = fg_pixels.min(dim=0).values.tolist()
y_max, x_max = fg_pixels.max(dim=0).values.tolist()
center_y, center_x = (y_min + y_max) / 2.0, (x_min + x_max) / 2.0
bw = x_max - x_min
bh = y_max - y_min
# Grow the bbox so its aspect matches `aspect_ratio` (width/height),
# anchored on the max side. Then apply pad_factor.
if bw / max(bh, 1) >= aspect_ratio:
crop_w = int(bw * pad_factor)
crop_h = int(bw / aspect_ratio * pad_factor)
else:
crop_h = int(bh * pad_factor)
crop_w = int(bh * aspect_ratio * pad_factor)
half_w, half_h = crop_w // 2, crop_h // 2
crop_x1 = int(center_x - half_w)
crop_y1 = int(center_y - half_h)
crop_x2 = crop_x1 + 2 * half_w
crop_y2 = crop_y1 + 2 * half_h
else:
logging.warning("Mask for the image is empty; a clean foreground mask is required for best quality.")
crop_x1, crop_y1, crop_x2, crop_y2 = 0, 0, W, H
crop_bbox = (crop_x1, crop_y1, crop_x2, crop_y2)
# Zero-pad out-of-bounds slice (PIL.crop semantics).
pad_l = max(0, -crop_x1)
pad_t = max(0, -crop_y1)
pad_r = max(0, crop_x2 - W)
pad_b = max(0, crop_y2 - H)
if pad_l or pad_t or pad_r or pad_b:
img = torch.nn.functional.pad(img, (pad_l, pad_r, pad_t, pad_b), value=0.0)
mask = torch.nn.functional.pad(mask, (pad_l, pad_r, pad_t, pad_b), value=0.0)
crop_x1 += pad_l
crop_x2 += pad_l
crop_y1 += pad_t
crop_y2 += pad_t
cropped_img = img [..., crop_y1:crop_y2, crop_x1:crop_x2]
cropped_mask = mask[..., crop_y1:crop_y2, crop_x1:crop_x2]
bg = torch.tensor(bg_rgb, dtype=cropped_img.dtype, device=cropped_img.device).view(1, 3, 1, 1)
composite = (cropped_img * cropped_mask + bg * (1.0 - cropped_mask)).clamp(0, 1)
return composite, crop_bbox, scene_size
def _dino_encode_batch(clip_vision_model, image, out_device, *, want_patches=False):
"""Encode an already-preprocessed image through DINOv3 at 512 and 1024.
@ -774,59 +684,6 @@ def _dino_encode_batch(clip_vision_model, image, out_device, *, want_patches=Fal
out["composites"] = composite_list
return out
class ImageCropToMask(IO.ComfyNode):
"""Crop an image to its mask's bounding box (centered square, with pad_factor
margin), then composite `img * mask` and resize to a square. Handles OOB crops
with zero-padding. Useful for 3D pipelines that expect a centered, background-free
subject at a fixed input resolution (Trellis2, Pixal3D, Hunyuan3D, TripoSR, etc.)."""
@classmethod
def define_schema(cls):
return IO.Schema(
node_id="ImageCropToMask",
display_name="Image Crop to Mask",
category="image/transform",
search_aliases=["crop to mask", "mask crop", "crop mask", "mask crop resize", "crop mask resize", "trellis2", "pixal3d"],
inputs=[
IO.Image.Input("image"),
IO.Mask.Input("mask"),
IO.Int.Input("width", default=1024, min=64, max=4096, step=8, tooltip="Output width in pixels."),
IO.Int.Input("height", default=1024, min=64, max=4096, step=8, tooltip="Output height in pixels."),
IO.Float.Input("pad_factor", default=1.0, min=1.0, max=2.0, step=0.01, tooltip="Extra margin around the mask bbox as a multiplier."),
IO.Int.Input("mask_offset", default=0, min=-32, max=32, step=1, tooltip="Grow or shrink the mask by this many pixels before cropping."),
IO.Color.Input("background", default="#000000", tooltip="Fill color behind the masked subject."),
],
outputs=[IO.Image.Output(display_name="image")],
)
@classmethod
def execute(cls, image, mask, width, height, pad_factor, mask_offset, background) -> IO.NodeOutput:
h = background.lstrip("#")
bg_rgb = (int(h[0:2], 16) / 255.0, int(h[2:4], 16) / 255.0, int(h[4:6], 16) / 255.0) if len(h) == 6 else (0.0, 0.0, 0.0)
image = image[..., :3]
batch_size = image.shape[0]
if mask.shape[0] == 1 and batch_size > 1:
mask = mask.expand(batch_size, -1, -1)
elif mask.shape[0] != batch_size:
raise ValueError(f"Mask batch {mask.shape[0]} does not match image batch {batch_size}")
out_images = []
for b in range(batch_size):
composite, _, _ = _crop_image_with_mask(
image[b], mask[b], max_image_size=max(width, height), pad_factor=pad_factor,
mask_offset=mask_offset, bg_rgb=bg_rgb, aspect_ratio=width / height,
)
composite = comfy.utils.common_upscale(composite, width, height, "lanczos", "disabled")
out_images.append(composite.movedim(-3, -1))
result = torch.cat(out_images, dim=0).to(
device=comfy.model_management.intermediate_device(),
dtype=comfy.model_management.intermediate_dtype(),
)
return IO.NodeOutput(result)
class Pixal3DConditioning(IO.ComfyNode):
@classmethod
@ -839,7 +696,7 @@ class Pixal3DConditioning(IO.ComfyNode):
IO.Image.Input("image", tooltip="Preprocessed image from ImageCropToMask (pad_factor=1.1 for Pixal3D)."),
IO.Float.Input(
"camera_angle_x", display_name="fov",
default=49.13, min=1.0, max=170.0, step=0.01, advanced=True,
default=49.13, min=1.0, max=170.0, step=0.01,
tooltip="Horizontal FOV in degrees. Wire a MoGeGeometryToFOV "
"(axis='horizontal', unit='degrees') for a per-image FoV (matches upstream default).",
),
@ -930,7 +787,6 @@ class Trellis2Extension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
return [
ImageCropToMask,
Trellis2Conditioning,
Pixal3DConditioning,
Trellis2ShapeStage,