From 1f23bcb27bef8a3f71ff10779cc75fd6b0f9b9af Mon Sep 17 00:00:00 2001 From: kijai Date: Fri, 17 Jul 2026 14:22:34 +0300 Subject: [PATCH] Resolve review comments and suggestions --- comfy_extras/nodes_images.py | 145 +++++++++++++++++++++++ comfy_extras/nodes_mesh_postprocess.py | 66 +++++------ comfy_extras/nodes_trellis2.py | 154 +------------------------ 3 files changed, 183 insertions(+), 182 deletions(-) diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py index fe1937ba5..6b2e5c8fb 100644 --- a/comfy_extras/nodes_images.py +++ b/comfy_extras/nodes_images.py @@ -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, diff --git a/comfy_extras/nodes_mesh_postprocess.py b/comfy_extras/nodes_mesh_postprocess.py index 1923f4c7b..4d08eed32 100644 --- a/comfy_extras/nodes_mesh_postprocess.py +++ b/comfy_extras/nodes_mesh_postprocess.py @@ -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)."), ], diff --git a/comfy_extras/nodes_trellis2.py b/comfy_extras/nodes_trellis2.py index bfffe3356..704f8efbb 100644 --- a/comfy_extras/nodes_trellis2.py +++ b/comfy_extras/nodes_trellis2.py @@ -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,