From 36feff88638b9c9a17dcf59bbaabdc7ac52e5f9f Mon Sep 17 00:00:00 2001 From: kijai Date: Mon, 20 Jul 2026 16:51:20 +0300 Subject: [PATCH] Smaller fixes and adjustments --- comfy/clip_vision.py | 2 +- comfy/ldm/trellis2/vae.py | 29 ------------------- comfy/model_detection.py | 2 +- comfy_extras/nodes_mesh_postprocess.py | 12 ++++---- comfy_extras/nodes_save_3d.py | 11 ++++++- comfy_extras/nodes_trellis2.py | 40 ++++++-------------------- 6 files changed, 26 insertions(+), 70 deletions(-) diff --git a/comfy/clip_vision.py b/comfy/clip_vision.py index 4d5ce024f..a84a350e6 100644 --- a/comfy/clip_vision.py +++ b/comfy/clip_vision.py @@ -143,7 +143,7 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False): json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino2_large.json") elif 'layer.0.mlp.gate_proj.weight' in sd and 'layer.31.norm1.weight' in sd: # Dinov3 ViT-H/16+ (SwiGLU gated MLP, 32 layers) json_config = comfy.image_encoders.dino3.DINOV3_VITH_CONFIG - elif 'layer.9.attention.o_proj.bias' in sd: # dinov3 large (24 layers); generic o_proj.bias key, so must come after the ViT-H check + elif 'layer.23.attention.o_proj.bias' in sd: # dinov3 large (24 layers) json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino3_large.json") else: return None diff --git a/comfy/ldm/trellis2/vae.py b/comfy/ldm/trellis2/vae.py index d43776ce2..2d446f0ea 100644 --- a/comfy/ldm/trellis2/vae.py +++ b/comfy/ldm/trellis2/vae.py @@ -665,35 +665,6 @@ class SparseTensor(VarLenTensor): spatial_cache=dict(self._spatial_cache), ) - def __merge_sparse_cache(self, other: 'SparseTensor') -> dict: - new_cache = {} - for k in set(list(self._spatial_cache.keys()) + list(other._spatial_cache.keys())): - if k in self._spatial_cache: - new_cache[k] = self._spatial_cache[k] - if k in other._spatial_cache: - if k not in new_cache: - new_cache[k] = other._spatial_cache[k] - else: - new_cache[k].update(other._spatial_cache[k]) - return new_cache - - def __elemwise__(self, other: Union[torch.Tensor, VarLenTensor], op: callable) -> 'SparseTensor': - if isinstance(other, torch.Tensor): - # Try per-batch [B, C] -> per-voxel [N, C] broadcast. RuntimeError - # fires for incompatible shapes; fall through and let op() handle. - try: - other = torch.broadcast_to(other, self.shape) - other = other[self.batch_boardcast_map] - except RuntimeError: - pass - if isinstance(other, VarLenTensor): - other = other.feats - new_feats = op(self.feats, other) - new_tensor = self.replace(new_feats) - if isinstance(other, SparseTensor): - new_tensor._spatial_cache = self.__merge_sparse_cache(other) - return new_tensor - def __getitem__(self, idx): if isinstance(idx, int): idx = [idx] diff --git a/comfy/model_detection.py b/comfy/model_detection.py index 6020a1bf5..7dda87880 100644 --- a/comfy/model_detection.py +++ b/comfy/model_detection.py @@ -120,7 +120,7 @@ def detect_unet_config(state_dict, key_prefix, metadata=None): has_tex = tex_key in state_dict_keys unet_config = { "image_model": "trellis2", - "resolution": 32 if (metadata is not None and "is_512" in metadata) else 64, + "resolution": 32 if (metadata or {}).get("is_512") else 64, "init_txt_model": has_tex, "txt_only": has_tex and not has_shape, } diff --git a/comfy_extras/nodes_mesh_postprocess.py b/comfy_extras/nodes_mesh_postprocess.py index 550dbd070..a2ff49af8 100644 --- a/comfy_extras/nodes_mesh_postprocess.py +++ b/comfy_extras/nodes_mesh_postprocess.py @@ -34,13 +34,13 @@ def paint_mesh_with_voxels(mesh, voxel_coords, voxel_colors, resolution): verts = mesh.vertices.to(device).squeeze(0) voxel_colors = voxel_colors.to(device) - voxel_pos_np = voxel_pos.numpy() - verts_np = verts.numpy() + voxel_pos_np = voxel_pos.cpu().numpy() + verts_np = verts.cpu().numpy() tree = cKDTree(voxel_pos_np) _, nearest_idx_np = tree.query(verts_np, k=1, workers=-1) - nearest_idx = torch.from_numpy(nearest_idx_np).long() + nearest_idx = torch.from_numpy(nearest_idx_np).long().to(voxel_colors.device) v_colors = voxel_colors[nearest_idx] # Voxel field may carry full PBR; vertex colors use only base_color RGB. if v_colors.shape[-1] > 3: @@ -2315,7 +2315,7 @@ class RemeshMesh(IO.ComfyNode): ), inputs=[ IO.Mesh.Input("mesh"), - IO.Int.Input("resolution", default=512, min=32, max=1024, + IO.Int.Input("resolution", default=512, min=32, max=2048, tooltip="Voxel grid resolution (output density). 256 ~ 100k faces, 512 ~ 1M. " "For an exact face count, follow with Decimate Mesh."), IO.DynamicCombo.Input("sign_mode", options=sign_mode_options, display_name="sign_mode", @@ -2926,14 +2926,14 @@ class FillHoles(IO.ComfyNode): ) @classmethod - def execute(cls, mesh, max_perimeter, weld_epsilon_rel, max_verts, fill_chains): + def execute(cls, mesh, max_perimeter, weld_epsilon_rel, max_vertices, fill_chains): def _fn(v, f, c): if max_perimeter > 0: v, f, c = fill_holes_v2_fn( v, f, max_perimeter=max_perimeter, colors=c, weld_epsilon_rel=weld_epsilon_rel, fill_chains=fill_chains, - max_verts=max_verts, + max_verts=max_vertices, ) return v, f, c return _process_mesh_batch(mesh, _fn) diff --git a/comfy_extras/nodes_save_3d.py b/comfy_extras/nodes_save_3d.py index 5b19c8838..b4ef4dfe2 100644 --- a/comfy_extras/nodes_save_3d.py +++ b/comfy_extras/nodes_save_3d.py @@ -719,10 +719,19 @@ class RotateMesh(IO.ComfyNode): out.vertices = [rotate(v) for v in mesh.vertices] else: out.vertices = rotate(mesh.vertices) - # Normals are directions; rotate them too (R is orthogonal) so they stay valid. + # Normals are directions, rotate them too (R is orthogonal). nrm = mesh.normals if nrm is not None: out.normals = [rotate(n) for n in nrm] if isinstance(nrm, list) else rotate(nrm) + # Tangents (xyz + handedness w) are directions too: rotate xyz, keep w + tng = mesh.tangents + if tng is not None: + def rotate_tangent(t: torch.Tensor) -> torch.Tensor: + rt = t.clone() + rt[..., :3] = rotate(t[..., :3]) + return rt + out.tangents = ([rotate_tangent(t) for t in tng] if isinstance(tng, list) + else rotate_tangent(tng)) return IO.NodeOutput(out) diff --git a/comfy_extras/nodes_trellis2.py b/comfy_extras/nodes_trellis2.py index 704f8efbb..0b3aae4a9 100644 --- a/comfy_extras/nodes_trellis2.py +++ b/comfy_extras/nodes_trellis2.py @@ -283,9 +283,7 @@ class Trellis2UpsampleStage(IO.ComfyNode): """Cascade-upsamples a 512-resolution shape latent into high-resolution sparse coords and sets up the second shape-stage sampling pass at the target resolution, attaching per-stage metadata to the conditioning for - the model to consume via extra_conds. target_resolution is reduced in - 128-step decrements until the unique upsampled coord count fits under - max_tokens (floor 1024).""" + the model to consume via extra_conds.""" @classmethod def define_schema(cls): return IO.Schema( @@ -297,12 +295,8 @@ class Trellis2UpsampleStage(IO.ComfyNode): IO.Conditioning.Input("negative"), IO.Latent.Input("shape_latent", tooltip="The 512-resolution shape latent output from the first shape-stage KSampler."), IO.Vae.Input("vae"), - IO.Combo.Input("target_resolution", options=["1024", "1536"], default="1024", tooltip="Controls output detail level for upsampling."), - IO.Int.Input("max_tokens", default=49152, min=1024, max=100000, - tooltip=( - "Maximum number of output elements (coordinates) allowed after upsampling. " - "Used to limit memory usage and control mesh density." - )), + IO.Int.Input("target_resolution", default=1024, min=1024, max=2048, step=128, + tooltip="Voxel resolution of the upsampled shape. Higher = more detail, more VRAM."), ], outputs=[ IO.Conditioning.Output(display_name="positive"), @@ -326,14 +320,13 @@ class Trellis2UpsampleStage(IO.ComfyNode): return quant.unique(dim=0) @classmethod - def execute(cls, positive, negative, shape_latent, vae, target_resolution, max_tokens): + def execute(cls, positive, negative, shape_latent, vae, target_resolution): device = comfy.model_management.get_torch_device() vae.prepare_decode(shape_latent["samples"].shape) coord_counts = shape_latent.get("coord_counts") shape_vae = vae.first_stage_model lr_resolution = 512 - target_resolution = int(target_resolution) proj_pack = _proj_pack_from_conditioning(positive) pixal3d_mode = proj_pack is not None @@ -356,28 +349,11 @@ class Trellis2UpsampleStage(IO.ComfyNode): slat_i = shape_norm(feats_i.to(device), coords_i) sample_hr_coords.append(shape_vae.upsample_shape(slat_i.to(vae.vae_dtype), upsample_times=4)) - # Resolution search — cache the final iteration's quantized unique tensors - # so we don't recompute .unique() per sample after picking hr_resolution. hr_resolution = target_resolution - quant_unique_list = [] - while True: - quant_unique_list = [] - exceeds_limit = False - for hr_coords_i in sample_hr_coords: - qu = cls._quantize_unique(hr_coords_i, lr_resolution, hr_resolution, pixal3d_mode) - quant_unique_list.append(qu) - if qu.shape[0] >= max_tokens: - exceeds_limit = True - break - if not exceeds_limit: - break - if hr_resolution <= 1024: - for k in range(len(quant_unique_list), len(sample_hr_coords)): - quant_unique_list.append( - cls._quantize_unique(sample_hr_coords[k], lr_resolution, hr_resolution, pixal3d_mode) - ) - break - hr_resolution -= 128 + quant_unique_list = [ + cls._quantize_unique(hr_coords_i, lr_resolution, hr_resolution, pixal3d_mode) + for hr_coords_i in sample_hr_coords + ] # Rewrite batch column to match per-sample offset and concat. per_sample_counts = []