Smaller fixes and adjustments

This commit is contained in:
kijai 2026-07-20 16:51:20 +03:00
parent b7c6a61515
commit 36feff8863
6 changed files with 26 additions and 70 deletions

View File

@ -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

View File

@ -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]

View File

@ -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,
}

View File

@ -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)

View File

@ -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)

View File

@ -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 = []