Smaller fixes and adjustments
This commit is contained in:
parent
b7c6a61515
commit
36feff8863
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
|
|
|
|||
Loading…
Reference in New Issue