mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-04-27 19:02:31 +08:00
Fix Trellis2 batched shape and texture semantics
This commit is contained in:
John Pollock
parent
880d7823e8
commit
c81ddf2349
@ -786,6 +786,7 @@ class Trellis2(nn.Module):
|
||||
# 32 -> 512px path, 64 -> 1024px path.
|
||||
uses_1024_conditioning = self.img2shape.resolution == 64
|
||||
coords = transformer_options.get("coords", None)
|
||||
coord_counts = transformer_options.get("coord_counts")
|
||||
mode = transformer_options.get("generation_mode", "structure_generation")
|
||||
is_512_run = False
|
||||
timestep = timestep.to(self.dtype)
|
||||
@ -811,40 +812,205 @@ class Trellis2(nn.Module):
|
||||
cond = context
|
||||
shape_rule = sigmas < self.guidance_interval[0] or sigmas > self.guidance_interval[1]
|
||||
txt_rule = sigmas < self.guidance_interval_txt[0] or sigmas > self.guidance_interval_txt[1]
|
||||
dense_out = None
|
||||
|
||||
if not_struct_mode:
|
||||
orig_bsz = x.shape[0]
|
||||
rule = txt_rule if mode == "texture_generation" else shape_rule
|
||||
|
||||
if rule and orig_bsz > 1:
|
||||
x_eval = x[1].unsqueeze(0)
|
||||
t_eval = timestep[1].unsqueeze(0) if timestep.shape[0] > 1 else timestep
|
||||
logical_batch = coord_counts.shape[0] if coord_counts is not None else 1
|
||||
if rule and orig_bsz > logical_batch:
|
||||
half = orig_bsz // 2
|
||||
x_eval = x[half:]
|
||||
t_eval = timestep[half:] if timestep.shape[0] > 1 else timestep
|
||||
c_eval = cond
|
||||
else:
|
||||
x_eval = x
|
||||
t_eval = timestep
|
||||
c_eval = context
|
||||
|
||||
x_eval_norms = [float(v) for v in x_eval.square().sum(dim=(1, 2)).detach().cpu().tolist()]
|
||||
c_eval_norms = [float(v) for v in c_eval.square().sum(dim=(1, 2)).detach().cpu().tolist()]
|
||||
print(
|
||||
"TRELLIS2_NOT_STRUCT_INPUT_TRACE",
|
||||
{
|
||||
"mode": mode,
|
||||
"orig_bsz": int(orig_bsz),
|
||||
"logical_batch": int(logical_batch),
|
||||
"rule": bool(rule),
|
||||
"coord_counts": coord_counts.tolist() if coord_counts is not None else None,
|
||||
"x_eval_norms": x_eval_norms,
|
||||
"c_eval_norms": c_eval_norms,
|
||||
},
|
||||
)
|
||||
|
||||
B, N, C = x_eval.shape
|
||||
|
||||
if mode in ["shape_generation", "texture_generation"]:
|
||||
feats_flat = x_eval.reshape(-1, C)
|
||||
if coord_counts is not None:
|
||||
logical_batch = coord_counts.shape[0]
|
||||
if B % logical_batch != 0:
|
||||
raise ValueError(
|
||||
f"Trellis2 coord_counts batch {logical_batch} doesn't divide latent batch {B}"
|
||||
)
|
||||
repeat_factor = B // logical_batch
|
||||
sparse_outs = []
|
||||
active_coord_counts = []
|
||||
if mode == "shape_generation" and repeat_factor > 1:
|
||||
grouped_outs = []
|
||||
grouped_counts = []
|
||||
for i in range(logical_batch):
|
||||
count = int(coord_counts[i].item())
|
||||
coords_i = coords[coords[:, 0] == i].clone()
|
||||
if coords_i.shape[0] != count:
|
||||
raise ValueError(
|
||||
f"Trellis2 coords rows for batch {i} expected {count}, got {coords_i.shape[0]}"
|
||||
)
|
||||
|
||||
# inflate coords [N, 4] -> [B*N, 4]
|
||||
coords_list = []
|
||||
for i in range(B):
|
||||
c = coords.clone()
|
||||
c[:, 0] = i
|
||||
coords_list.append(c)
|
||||
feat_batches = []
|
||||
coord_batches = []
|
||||
index_batch = []
|
||||
for rep in range(repeat_factor):
|
||||
out_index = rep * logical_batch + i
|
||||
feat_batches.append(x_eval[out_index, :count])
|
||||
coords_rep = coords_i.clone()
|
||||
coords_rep[:, 0] = rep
|
||||
coord_batches.append(coords_rep)
|
||||
index_batch.append(out_index)
|
||||
|
||||
batched_coords = torch.cat(coords_list, dim=0)
|
||||
print(
|
||||
"TRELLIS2_GROUPED_INPUT_TRACE",
|
||||
{
|
||||
"mode": mode,
|
||||
"sample_index": int(i),
|
||||
"coord_count": int(count),
|
||||
"feat_norms": [float(v.square().sum().detach().cpu().item()) for v in feat_batches],
|
||||
},
|
||||
)
|
||||
|
||||
x_st_i = SparseTensor(
|
||||
feats=torch.cat(feat_batches, dim=0),
|
||||
coords=torch.cat(coord_batches, dim=0).to(torch.int32),
|
||||
)
|
||||
index_tensor = torch.tensor(index_batch, device=x_eval.device, dtype=torch.long)
|
||||
if t_eval.shape[0] > 1:
|
||||
t_i = t_eval.index_select(0, index_tensor)
|
||||
else:
|
||||
t_i = t_eval
|
||||
if c_eval.shape[0] > 1:
|
||||
c_i = c_eval.index_select(0, index_tensor)
|
||||
else:
|
||||
c_i = c_eval
|
||||
|
||||
if is_512_run:
|
||||
sparse_out = self.img2shape_512(x_st_i, t_i, c_i)
|
||||
else:
|
||||
sparse_out = self.img2shape(x_st_i, t_i, c_i)
|
||||
|
||||
feats_group, coords_group = sparse_out.to_tensor_list()
|
||||
if len(feats_group) != repeat_factor:
|
||||
raise ValueError(
|
||||
f"Trellis2 expected {repeat_factor} sparse output groups for batch {i}, got {len(feats_group)}"
|
||||
)
|
||||
for rep, (feats_rep, coords_rep) in enumerate(zip(feats_group, coords_group)):
|
||||
if feats_rep.shape[0] != count:
|
||||
raise ValueError(
|
||||
f"Trellis2 sparse output rows for batch {i} rep {rep} expected {count}, got {feats_rep.shape[0]}"
|
||||
)
|
||||
if coords_rep.shape[0] != count:
|
||||
raise ValueError(
|
||||
f"Trellis2 sparse output coords for batch {i} rep {rep} expected {count}, got {coords_rep.shape[0]}"
|
||||
)
|
||||
grouped_outs.append(feats_group)
|
||||
grouped_counts.append(count)
|
||||
|
||||
for rep in range(repeat_factor):
|
||||
for i in range(logical_batch):
|
||||
sparse_outs.append(grouped_outs[i][rep])
|
||||
active_coord_counts.append(grouped_counts[i])
|
||||
else:
|
||||
for rep in range(repeat_factor):
|
||||
for i in range(logical_batch):
|
||||
out_index = rep * logical_batch + i
|
||||
count = int(coord_counts[i].item())
|
||||
coords_i = coords[coords[:, 0] == i].clone()
|
||||
if coords_i.shape[0] != count:
|
||||
raise ValueError(
|
||||
f"Trellis2 coords rows for batch {i} expected {count}, got {coords_i.shape[0]}"
|
||||
)
|
||||
coords_i[:, 0] = 0
|
||||
feats_i = x_eval[out_index, :count]
|
||||
x_st_i = SparseTensor(feats=feats_i, coords=coords_i.to(torch.int32))
|
||||
t_i = t_eval[out_index].unsqueeze(0) if t_eval.shape[0] > 1 else t_eval
|
||||
c_i = c_eval[out_index].unsqueeze(0) if c_eval.shape[0] > 1 else c_eval
|
||||
|
||||
if mode == "shape_generation":
|
||||
if is_512_run:
|
||||
sparse_out = self.img2shape_512(x_st_i, t_i, c_i)
|
||||
else:
|
||||
sparse_out = self.img2shape(x_st_i, t_i, c_i)
|
||||
else:
|
||||
slat = transformer_options.get("shape_slat")
|
||||
if slat is None:
|
||||
raise ValueError("shape_slat can't be None")
|
||||
if slat.ndim == 3:
|
||||
if slat.shape[0] != logical_batch:
|
||||
raise ValueError(
|
||||
f"shape_slat batch {slat.shape[0]} doesn't match coord_counts batch {logical_batch}"
|
||||
)
|
||||
if slat.shape[1] < count:
|
||||
raise ValueError(
|
||||
f"shape_slat tokens {slat.shape[1]} can't cover coord count {count} for batch {i}"
|
||||
)
|
||||
slat_feats = slat[i, :count].to(x_st_i.device)
|
||||
else:
|
||||
slat_feats = slat[:count].to(x_st_i.device)
|
||||
x_st_i = x_st_i.replace(feats=torch.cat([x_st_i.feats, slat_feats], dim=-1))
|
||||
sparse_out = self.shape2txt(x_st_i, t_i, c_i)
|
||||
|
||||
sparse_outs.append(sparse_out.feats)
|
||||
active_coord_counts.append(count)
|
||||
|
||||
out_channels = sparse_outs[0].shape[-1]
|
||||
sparse_out_norms = [float(feats.square().sum().detach().cpu().item()) for feats in sparse_outs]
|
||||
print(
|
||||
"TRELLIS2_SPARSE_OUT_TRACE",
|
||||
{
|
||||
"mode": mode,
|
||||
"coords_rows": int(coords.shape[0]),
|
||||
"active_coord_counts": active_coord_counts,
|
||||
"sparse_out_norms": sparse_out_norms,
|
||||
},
|
||||
)
|
||||
padded = sparse_outs[0].new_zeros((B, N, out_channels))
|
||||
for out_index, (count, feats_i) in enumerate(zip(active_coord_counts, sparse_outs)):
|
||||
padded[out_index, :count] = feats_i
|
||||
dense_out = padded.transpose(1, 2).unsqueeze(-1)
|
||||
elif coords.shape[0] == N:
|
||||
feats_flat = x_eval.reshape(-1, C)
|
||||
coords_list = []
|
||||
for i in range(B):
|
||||
c = coords.clone()
|
||||
c[:, 0] = i
|
||||
coords_list.append(c)
|
||||
batched_coords = torch.cat(coords_list, dim=0)
|
||||
elif coords.shape[0] == B * N:
|
||||
feats_flat = x_eval.reshape(-1, C)
|
||||
batched_coords = coords
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Trellis2 expected coords rows {N} or {B * N}, got {coords.shape[0]}"
|
||||
)
|
||||
else:
|
||||
batched_coords = coords
|
||||
feats_flat = x_eval
|
||||
|
||||
x_st = SparseTensor(feats=feats_flat, coords=batched_coords.to(torch.int32))
|
||||
if dense_out is None:
|
||||
x_st = SparseTensor(feats=feats_flat, coords=batched_coords.to(torch.int32))
|
||||
|
||||
if mode == "shape_generation":
|
||||
if dense_out is not None:
|
||||
out = dense_out
|
||||
elif mode == "shape_generation":
|
||||
if is_512_run:
|
||||
out = self.img2shape_512(x_st, t_eval, c_eval)
|
||||
else:
|
||||
@ -856,23 +1022,152 @@ class Trellis2(nn.Module):
|
||||
if slat is None:
|
||||
raise ValueError("shape_slat can't be None")
|
||||
|
||||
base_slat_feats = slat[:N]
|
||||
slat_feats_batched = base_slat_feats.repeat(B, 1).to(x_st.device)
|
||||
if slat.ndim == 3:
|
||||
if coord_counts is not None:
|
||||
logical_batch = coord_counts.shape[0]
|
||||
if slat.shape[0] != logical_batch:
|
||||
raise ValueError(
|
||||
f"shape_slat batch {slat.shape[0]} doesn't match coord_counts batch {logical_batch}"
|
||||
)
|
||||
if B % logical_batch != 0:
|
||||
raise ValueError(
|
||||
f"Trellis2 coord_counts batch {logical_batch} doesn't divide latent batch {B}"
|
||||
)
|
||||
repeat_factor = B // logical_batch
|
||||
slat_list = []
|
||||
for _ in range(repeat_factor):
|
||||
for i in range(logical_batch):
|
||||
count = int(coord_counts[i].item())
|
||||
if slat.shape[1] < count:
|
||||
raise ValueError(
|
||||
f"shape_slat tokens {slat.shape[1]} can't cover coord count {count} for batch {i}"
|
||||
)
|
||||
slat_list.append(slat[i, :count])
|
||||
slat_feats_batched = torch.cat(slat_list, dim=0).to(x_st.device)
|
||||
else:
|
||||
if slat.shape[0] != B:
|
||||
raise ValueError(f"shape_slat batch {slat.shape[0]} doesn't match latent batch {B}")
|
||||
if slat.shape[1] != N:
|
||||
raise ValueError(f"shape_slat tokens {slat.shape[1]} doesn't match latent tokens {N}")
|
||||
slat_feats_batched = slat.reshape(B * N, -1).to(x_st.device)
|
||||
else:
|
||||
base_slat_feats = slat[:N]
|
||||
slat_feats_batched = base_slat_feats.repeat(B, 1).to(x_st.device)
|
||||
x_st = x_st.replace(feats=torch.cat([x_st.feats, slat_feats_batched], dim=-1))
|
||||
out = self.shape2txt(x_st, t_eval, c_eval)
|
||||
else: # structure
|
||||
orig_bsz = x.shape[0]
|
||||
if shape_rule and orig_bsz > 1:
|
||||
half = orig_bsz // 2
|
||||
x = x[half:]
|
||||
timestep = timestep[half:] if timestep.shape[0] > 1 else timestep
|
||||
out = self.structure_model(x, timestep, cond if shape_rule and orig_bsz > 1 else context)
|
||||
if shape_rule and orig_bsz > 1:
|
||||
out = out.repeat(2, 1, 1, 1, 1)
|
||||
cond_or_uncond = transformer_options.get("cond_or_uncond") or []
|
||||
batch_groups = len(cond_or_uncond) if len(cond_or_uncond) > 0 and orig_bsz % len(cond_or_uncond) == 0 else 1
|
||||
logical_batch = orig_bsz // batch_groups
|
||||
print(
|
||||
"TRELLIS2_STRUCTURE_INPUT_TRACE",
|
||||
{
|
||||
"orig_bsz": int(orig_bsz),
|
||||
"batch_groups": int(batch_groups),
|
||||
"logical_batch": int(logical_batch),
|
||||
"cond_or_uncond": cond_or_uncond,
|
||||
"x_norms": [float(v) for v in x.square().sum(dim=(1, 2, 3, 4)).detach().cpu().tolist()],
|
||||
"x_sums": [float(v) for v in x.sum(dim=(1, 2, 3, 4)).detach().cpu().tolist()],
|
||||
"c_norms": [float(v) for v in context.square().sum(dim=(1, 2)).detach().cpu().tolist()],
|
||||
"c_sums": [float(v) for v in context.sum(dim=(1, 2)).detach().cpu().tolist()],
|
||||
},
|
||||
)
|
||||
|
||||
if logical_batch > 1:
|
||||
x_groups = x.reshape(batch_groups, logical_batch, *x.shape[1:])
|
||||
if timestep.shape[0] > 1:
|
||||
t_groups = timestep.reshape(batch_groups, logical_batch, *timestep.shape[1:])
|
||||
else:
|
||||
t_groups = timestep
|
||||
c_groups = context.reshape(batch_groups, logical_batch, *context.shape[1:])
|
||||
|
||||
if shape_rule and batch_groups > 1:
|
||||
selected_group_indices = [batch_groups - 1]
|
||||
else:
|
||||
selected_group_indices = list(range(batch_groups))
|
||||
|
||||
out_groups = []
|
||||
selected_x_norms = []
|
||||
selected_x_sums = []
|
||||
selected_c_norms = []
|
||||
selected_c_sums = []
|
||||
for sample_index in range(logical_batch):
|
||||
if shape_rule and batch_groups > 1:
|
||||
half = orig_bsz // 2
|
||||
x_i = x[half + sample_index].unsqueeze(0)
|
||||
if timestep.shape[0] > 1:
|
||||
t_i = timestep[half + sample_index].unsqueeze(0)
|
||||
else:
|
||||
t_i = timestep
|
||||
if cond.shape[0] > 1:
|
||||
c_i = cond[sample_index].unsqueeze(0)
|
||||
else:
|
||||
c_i = cond
|
||||
else:
|
||||
x_i = x_groups[selected_group_indices, sample_index]
|
||||
if timestep.shape[0] > 1:
|
||||
t_i = t_groups[selected_group_indices, sample_index]
|
||||
else:
|
||||
t_i = timestep
|
||||
c_i = c_groups[selected_group_indices, sample_index]
|
||||
selected_x_norms.extend(float(v) for v in x_i.square().sum(dim=(1, 2, 3, 4)).detach().cpu().tolist())
|
||||
selected_x_sums.extend(float(v) for v in x_i.sum(dim=(1, 2, 3, 4)).detach().cpu().tolist())
|
||||
selected_c_norms.extend(float(v) for v in c_i.square().sum(dim=(1, 2)).detach().cpu().tolist())
|
||||
selected_c_sums.extend(float(v) for v in c_i.sum(dim=(1, 2)).detach().cpu().tolist())
|
||||
out_groups.append(self.structure_model(x_i, t_i, c_i))
|
||||
|
||||
print(
|
||||
"TRELLIS2_STRUCTURE_SELECTED_TRACE",
|
||||
{
|
||||
"selected_group_indices": selected_group_indices,
|
||||
"selected_x_norms": selected_x_norms,
|
||||
"selected_x_sums": selected_x_sums,
|
||||
"selected_c_norms": selected_c_norms,
|
||||
"selected_c_sums": selected_c_sums,
|
||||
},
|
||||
)
|
||||
|
||||
out = out_groups[0].new_zeros((orig_bsz, *out_groups[0].shape[1:]))
|
||||
for sample_index, out_sample in enumerate(out_groups):
|
||||
if shape_rule and batch_groups > 1:
|
||||
repeated = out_sample[0]
|
||||
for group_index in range(batch_groups):
|
||||
out[group_index * logical_batch + sample_index] = repeated
|
||||
else:
|
||||
for local_group_index, group_index in enumerate(selected_group_indices):
|
||||
out[group_index * logical_batch + sample_index] = out_sample[local_group_index]
|
||||
else:
|
||||
if shape_rule and orig_bsz > 1:
|
||||
half = orig_bsz // 2
|
||||
x = x[half:]
|
||||
timestep = timestep[half:] if timestep.shape[0] > 1 else timestep
|
||||
out = self.structure_model(x, timestep, cond if shape_rule and orig_bsz > 1 else context)
|
||||
if shape_rule and orig_bsz > 1:
|
||||
out = out.repeat(2, 1, 1, 1, 1)
|
||||
|
||||
print(
|
||||
"TRELLIS2_STRUCTURE_OUTPUT_TRACE",
|
||||
{
|
||||
"out_norms": [float(v) for v in out.square().sum(dim=(1, 2, 3, 4)).detach().cpu().tolist()],
|
||||
"out_sums": [float(v) for v in out.sum(dim=(1, 2, 3, 4)).detach().cpu().tolist()],
|
||||
},
|
||||
)
|
||||
|
||||
if not_struct_mode:
|
||||
out = out.feats
|
||||
out = out.view(B, N, -1).transpose(1, 2).unsqueeze(-1)
|
||||
if rule and orig_bsz > 1:
|
||||
out = out.repeat(orig_bsz, 1, 1, 1)
|
||||
if dense_out is None:
|
||||
out = out.feats
|
||||
out = out.view(B, N, -1).transpose(1, 2).unsqueeze(-1)
|
||||
if rule and orig_bsz > B:
|
||||
out = out.repeat(orig_bsz // B, 1, 1, 1)
|
||||
print(
|
||||
"TRELLIS2_DENSE_OUT_TRACE",
|
||||
{
|
||||
"mode": mode,
|
||||
"coords_rows": int(coords.shape[0]) if coords is not None else None,
|
||||
"output_shape": list(out.shape),
|
||||
"output_norms": [float(v) for v in out.squeeze(-1).square().sum(dim=(1, 2)).detach().cpu().tolist()],
|
||||
"coord_counts": coord_counts.tolist() if coord_counts is not None else None,
|
||||
},
|
||||
)
|
||||
return out
|
||||
|
||||
@ -7,6 +7,23 @@ import logging
|
||||
import comfy.nested_tensor
|
||||
|
||||
def prepare_noise_inner(latent_image, generator, noise_inds=None):
|
||||
coord_counts = getattr(latent_image, "trellis_coord_counts", None)
|
||||
if coord_counts is not None:
|
||||
noise = torch.zeros(latent_image.size(), dtype=torch.float32, layout=latent_image.layout, device="cpu")
|
||||
base_state = generator.get_state()
|
||||
for i, count in enumerate(coord_counts.tolist()):
|
||||
local_generator = torch.Generator(device="cpu")
|
||||
local_generator.set_state(base_state.clone())
|
||||
sample_noise = torch.randn(
|
||||
[1, latent_image.size(1), int(count), latent_image.size(3)],
|
||||
dtype=torch.float32,
|
||||
layout=latent_image.layout,
|
||||
generator=local_generator,
|
||||
device="cpu",
|
||||
)
|
||||
noise[i:i + 1, :, :int(count), :] = sample_noise
|
||||
return noise.to(dtype=latent_image.dtype)
|
||||
|
||||
if noise_inds is None:
|
||||
return torch.randn(latent_image.size(), dtype=torch.float32, layout=latent_image.layout, generator=generator, device="cpu").to(dtype=latent_image.dtype)
|
||||
|
||||
|
||||
@ -96,6 +96,70 @@ def shape_norm(shape_latent, coords):
|
||||
samples = samples * std + mean
|
||||
return samples
|
||||
|
||||
|
||||
def infer_batched_coord_layout(coords):
|
||||
if coords.ndim != 2 or coords.shape[1] != 4:
|
||||
raise ValueError(f"Expected Trellis2 coords with shape [N, 4], got {tuple(coords.shape)}")
|
||||
|
||||
if coords.shape[0] == 0:
|
||||
raise ValueError("Trellis2 coords can't be empty")
|
||||
|
||||
batch_ids = coords[:, 0].to(torch.int64)
|
||||
batch_size = int(batch_ids.max().item()) + 1
|
||||
counts = torch.bincount(batch_ids, minlength=batch_size)
|
||||
|
||||
if (counts == 0).any():
|
||||
raise ValueError(f"Non-contiguous Trellis2 batch ids in coords: {batch_ids.unique(sorted=True).tolist()}")
|
||||
|
||||
max_tokens = int(counts.max().item())
|
||||
return batch_size, counts, max_tokens
|
||||
|
||||
|
||||
def flatten_batched_sparse_latent(samples, coords, coord_counts):
|
||||
samples = samples.squeeze(-1).transpose(1, 2)
|
||||
if coord_counts is None:
|
||||
return samples.reshape(-1, samples.shape[-1]), coords
|
||||
|
||||
feat_list = []
|
||||
coord_list = []
|
||||
for i in range(coord_counts.shape[0]):
|
||||
count = int(coord_counts[i].item())
|
||||
coords_i = coords[coords[:, 0] == i]
|
||||
if coords_i.shape[0] != count:
|
||||
raise ValueError(f"Trellis2 coords rows for batch {i} expected {count}, got {coords_i.shape[0]}")
|
||||
feat_list.append(samples[i, :count])
|
||||
coord_list.append(coords_i)
|
||||
|
||||
return torch.cat(feat_list, dim=0), torch.cat(coord_list, dim=0)
|
||||
|
||||
|
||||
def split_batched_sparse_latent(samples, coords, coord_counts):
|
||||
samples = samples.squeeze(-1).transpose(1, 2)
|
||||
if coord_counts is None:
|
||||
return [(samples.reshape(-1, samples.shape[-1]), coords)]
|
||||
|
||||
items = []
|
||||
for i in range(coord_counts.shape[0]):
|
||||
count = int(coord_counts[i].item())
|
||||
coords_i = coords[coords[:, 0] == i]
|
||||
if coords_i.shape[0] != count:
|
||||
raise ValueError(f"Trellis2 coords rows for batch {i} expected {count}, got {coords_i.shape[0]}")
|
||||
items.append((samples[i, :count], coords_i))
|
||||
return items
|
||||
|
||||
|
||||
def log_sparse_batch_trace(tag, items):
|
||||
feat_norms = [float(feats.square().sum().detach().cpu().item()) for feats, _ in items]
|
||||
coord_rows = [int(coords_i.shape[0]) for _, coords_i in items]
|
||||
print(
|
||||
tag,
|
||||
{
|
||||
"batch_size": len(items),
|
||||
"coord_rows": coord_rows,
|
||||
"feat_norms": feat_norms,
|
||||
},
|
||||
)
|
||||
|
||||
def paint_mesh_with_voxels(mesh, voxel_coords, voxel_colors, resolution):
|
||||
"""
|
||||
Generic function to paint a mesh using nearest-neighbor colors from a sparse voxel field.
|
||||
@ -169,12 +233,32 @@ class VaeDecodeShapeTrellis(IO.ComfyNode):
|
||||
|
||||
vae = vae.first_stage_model
|
||||
coords = samples["coords"]
|
||||
coord_counts = samples.get("coord_counts")
|
||||
|
||||
samples = samples["samples"]
|
||||
samples = samples.squeeze(-1).transpose(1, 2).reshape(-1, 32).to(device)
|
||||
samples = shape_norm(samples, coords)
|
||||
if coord_counts is None:
|
||||
samples, coords = flatten_batched_sparse_latent(samples, coords, coord_counts)
|
||||
samples = shape_norm(samples.to(device), coords.to(device))
|
||||
mesh, subs = vae.decode_shape_slat(samples, resolution)
|
||||
else:
|
||||
split_items = split_batched_sparse_latent(samples, coords, coord_counts)
|
||||
mesh = []
|
||||
subs_per_sample = []
|
||||
for feats_i, coords_i in split_items:
|
||||
coords_i = coords_i.to(device).clone()
|
||||
coords_i[:, 0] = 0
|
||||
sample_i = shape_norm(feats_i.to(device), coords_i)
|
||||
mesh_i, subs_i = vae.decode_shape_slat(sample_i, resolution)
|
||||
mesh.append(mesh_i[0])
|
||||
subs_per_sample.append(subs_i)
|
||||
|
||||
subs = []
|
||||
for stage_index in range(len(subs_per_sample[0])):
|
||||
stage_tensors = [sample_subs[stage_index] for sample_subs in subs_per_sample]
|
||||
feats_list = [stage_tensor.feats for stage_tensor in stage_tensors]
|
||||
coords_list = [stage_tensor.coords for stage_tensor in stage_tensors]
|
||||
subs.append(SparseTensor.from_tensor_list(feats_list, coords_list))
|
||||
|
||||
mesh, subs = vae.decode_shape_slat(samples, resolution)
|
||||
face_list = [m.faces for m in mesh]
|
||||
vert_list = [m.vertices for m in mesh]
|
||||
if all(v.shape == vert_list[0].shape for v in vert_list) and all(f.shape == face_list[0].shape for f in face_list):
|
||||
@ -210,12 +294,14 @@ class VaeDecodeTextureTrellis(IO.ComfyNode):
|
||||
|
||||
vae = vae.first_stage_model
|
||||
coords = samples["coords"]
|
||||
coord_counts = samples.get("coord_counts")
|
||||
|
||||
samples = samples["samples"]
|
||||
samples = samples.squeeze(-1).transpose(1, 2).reshape(-1, 32).to(device)
|
||||
samples, coords = flatten_batched_sparse_latent(samples, coords, coord_counts)
|
||||
samples = samples.to(device)
|
||||
std = tex_slat_normalization["std"].to(samples)
|
||||
mean = tex_slat_normalization["mean"].to(samples)
|
||||
samples = SparseTensor(feats = samples, coords=coords)
|
||||
samples = SparseTensor(feats = samples, coords=coords.to(device))
|
||||
samples = samples * std + mean
|
||||
|
||||
voxel = vae.decode_tex_slat(samples, shape_subs)
|
||||
@ -273,7 +359,13 @@ class VaeDecodeStructureTrellis2(IO.ComfyNode):
|
||||
decoder = decoder.to(load_device)
|
||||
samples = samples["samples"]
|
||||
samples = samples.to(load_device)
|
||||
decoded = decoder(samples)>0
|
||||
if samples.shape[0] > 1:
|
||||
decoded_items = []
|
||||
for i in range(samples.shape[0]):
|
||||
decoded_items.append(decoder(samples[i:i + 1]) > 0)
|
||||
decoded = torch.cat(decoded_items, dim=0)
|
||||
else:
|
||||
decoded = decoder(samples) > 0
|
||||
decoder.to(offload_device)
|
||||
current_res = decoded.shape[2]
|
||||
|
||||
@ -305,32 +397,102 @@ class Trellis2UpsampleCascade(IO.ComfyNode):
|
||||
device = comfy.model_management.get_torch_device()
|
||||
comfy.model_management.load_model_gpu(vae.patcher)
|
||||
|
||||
feats = shape_latent_512["samples"].squeeze(-1).transpose(1, 2).reshape(-1, 32).to(device)
|
||||
coords_512 = shape_latent_512["coords"].to(device)
|
||||
|
||||
slat = shape_norm(feats, coords_512)
|
||||
|
||||
coord_counts = shape_latent_512.get("coord_counts")
|
||||
decoder = vae.first_stage_model.shape_dec
|
||||
|
||||
slat.feats = slat.feats.to(next(decoder.parameters()).dtype)
|
||||
hr_coords = decoder.upsample(slat, upsample_times=4)
|
||||
|
||||
lr_resolution = 512
|
||||
hr_resolution = int(target_resolution)
|
||||
target_resolution = int(target_resolution)
|
||||
|
||||
while True:
|
||||
quant_coords = torch.cat([
|
||||
hr_coords[:, :1],
|
||||
((hr_coords[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(),
|
||||
], dim=1)
|
||||
final_coords = quant_coords.unique(dim=0)
|
||||
num_tokens = final_coords.shape[0]
|
||||
if coord_counts is None:
|
||||
feats, coords_512 = flatten_batched_sparse_latent(
|
||||
shape_latent_512["samples"],
|
||||
shape_latent_512["coords"],
|
||||
coord_counts,
|
||||
)
|
||||
feats = feats.to(device)
|
||||
coords_512 = coords_512.to(device)
|
||||
print(
|
||||
"TRELLIS2_UPSAMPLE_INPUT_TRACE",
|
||||
{
|
||||
"batch_size": 1,
|
||||
"coord_rows": [int(coords_512.shape[0])],
|
||||
"feat_norms": [float(feats.square().sum().detach().cpu().item())],
|
||||
},
|
||||
)
|
||||
slat = shape_norm(feats, coords_512)
|
||||
slat.feats = slat.feats.to(next(decoder.parameters()).dtype)
|
||||
hr_coords = decoder.upsample(slat, upsample_times=4)
|
||||
|
||||
if num_tokens < max_tokens or hr_resolution <= 1024:
|
||||
break
|
||||
hr_resolution -= 128
|
||||
hr_resolution = target_resolution
|
||||
while True:
|
||||
quant_coords = torch.cat([
|
||||
hr_coords[:, :1],
|
||||
((hr_coords[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(),
|
||||
], dim=1)
|
||||
final_coords = quant_coords.unique(dim=0)
|
||||
num_tokens = final_coords.shape[0]
|
||||
|
||||
return IO.NodeOutput(final_coords,)
|
||||
if num_tokens < max_tokens or hr_resolution <= 1024:
|
||||
break
|
||||
hr_resolution -= 128
|
||||
|
||||
print(
|
||||
"TRELLIS2_UPSAMPLE_OUTPUT_TRACE",
|
||||
{
|
||||
"batch_size": 1,
|
||||
"coord_rows": [int(final_coords.shape[0])],
|
||||
"hr_resolution": int(hr_resolution),
|
||||
},
|
||||
)
|
||||
return IO.NodeOutput(final_coords,)
|
||||
|
||||
final_coords_list = []
|
||||
items = split_batched_sparse_latent(
|
||||
shape_latent_512["samples"],
|
||||
shape_latent_512["coords"],
|
||||
coord_counts,
|
||||
)
|
||||
log_sparse_batch_trace("TRELLIS2_UPSAMPLE_INPUT_TRACE", items)
|
||||
decoder_dtype = next(decoder.parameters()).dtype
|
||||
|
||||
output_coord_rows = []
|
||||
output_resolutions = []
|
||||
for batch_index, (feats_i, coords_i) in enumerate(items):
|
||||
feats_i = feats_i.to(device)
|
||||
coords_i = coords_i.to(device).clone()
|
||||
coords_i[:, 0] = 0
|
||||
slat_i = shape_norm(feats_i, coords_i)
|
||||
slat_i.feats = slat_i.feats.to(decoder_dtype)
|
||||
hr_coords_i = decoder.upsample(slat_i, upsample_times=4)
|
||||
|
||||
hr_resolution = target_resolution
|
||||
while True:
|
||||
quant_coords_i = torch.cat([
|
||||
hr_coords_i[:, :1],
|
||||
((hr_coords_i[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(),
|
||||
], dim=1)
|
||||
final_coords_i = quant_coords_i.unique(dim=0)
|
||||
num_tokens = final_coords_i.shape[0]
|
||||
|
||||
if num_tokens < max_tokens or hr_resolution <= 1024:
|
||||
break
|
||||
hr_resolution -= 128
|
||||
|
||||
final_coords_i = final_coords_i.clone()
|
||||
final_coords_i[:, 0] = batch_index
|
||||
final_coords_list.append(final_coords_i)
|
||||
output_coord_rows.append(int(final_coords_i.shape[0]))
|
||||
output_resolutions.append(int(hr_resolution))
|
||||
|
||||
print(
|
||||
"TRELLIS2_UPSAMPLE_OUTPUT_TRACE",
|
||||
{
|
||||
"batch_size": len(final_coords_list),
|
||||
"coord_rows": output_coord_rows,
|
||||
"hr_resolution": output_resolutions,
|
||||
},
|
||||
)
|
||||
|
||||
return IO.NodeOutput(torch.cat(final_coords_list, dim=0),)
|
||||
|
||||
dino_mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
|
||||
dino_std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
|
||||
@ -406,6 +568,7 @@ class Trellis2Conditioning(IO.ComfyNode):
|
||||
|
||||
cond_512_list = []
|
||||
cond_1024_list = []
|
||||
composite_trace = []
|
||||
|
||||
for b in range(batch_size):
|
||||
item_image = image[b]
|
||||
@ -460,6 +623,14 @@ class Trellis2Conditioning(IO.ComfyNode):
|
||||
|
||||
# to match trellis2 code (quantize -> dequantize)
|
||||
composite_uint8 = (composite_np * 255.0).round().clip(0, 255).astype(np.uint8)
|
||||
composite_trace.append(
|
||||
{
|
||||
"sample_index": int(b),
|
||||
"shape": list(composite_uint8.shape),
|
||||
"sum": int(composite_uint8.sum(dtype=np.int64)),
|
||||
"prefix": composite_uint8[:2, :2, :].reshape(-1).tolist(),
|
||||
}
|
||||
)
|
||||
|
||||
cropped_pil = Image.fromarray(composite_uint8)
|
||||
|
||||
@ -471,6 +642,19 @@ class Trellis2Conditioning(IO.ComfyNode):
|
||||
cond_1024_batched = torch.cat(cond_1024_list, dim=0)
|
||||
neg_cond_batched = torch.zeros_like(cond_512_batched)
|
||||
neg_embeds_batched = torch.zeros_like(cond_1024_batched)
|
||||
print(
|
||||
"TRELLIS2_CONDITIONING_TRACE",
|
||||
{
|
||||
"batch_size": int(batch_size),
|
||||
"cond_512_norms": [float(v) for v in cond_512_batched.square().sum(dim=(1, 2)).detach().cpu().tolist()],
|
||||
"cond_512_sums": [float(v) for v in cond_512_batched.sum(dim=(1, 2)).detach().cpu().tolist()],
|
||||
"cond_512_prefix": cond_512_batched[:, 0, :8].detach().cpu().tolist(),
|
||||
"cond_1024_norms": [float(v) for v in cond_1024_batched.square().sum(dim=(1, 2)).detach().cpu().tolist()],
|
||||
"cond_1024_sums": [float(v) for v in cond_1024_batched.sum(dim=(1, 2)).detach().cpu().tolist()],
|
||||
"cond_1024_prefix": cond_1024_batched[:, 0, :8].detach().cpu().tolist(),
|
||||
"composite_trace": composite_trace,
|
||||
},
|
||||
)
|
||||
|
||||
positive = [[cond_512_batched, {"embeds": cond_1024_batched}]]
|
||||
negative = [[neg_cond_batched, {"embeds": neg_embeds_batched}]]
|
||||
@ -509,8 +693,32 @@ class EmptyShapeLatentTrellis2(IO.ComfyNode):
|
||||
else:
|
||||
raise ValueError(f"Invalid input to EmptyShapeLatent: {type(structure_or_coords)}")
|
||||
in_channels = 32
|
||||
# image like format
|
||||
latent = torch.randn(1, in_channels, coords.shape[0], 1)
|
||||
batch_size, coord_counts, max_tokens = infer_batched_coord_layout(coords)
|
||||
if batch_size == 1:
|
||||
coord_counts = None
|
||||
latent = torch.randn(1, in_channels, coords.shape[0], 1)
|
||||
else:
|
||||
latent = torch.zeros(batch_size, in_channels, max_tokens, 1)
|
||||
base_state = torch.random.get_rng_state()
|
||||
for i in range(batch_size):
|
||||
count = int(coord_counts[i].item())
|
||||
generator = torch.Generator(device="cpu")
|
||||
generator.set_state(base_state.clone())
|
||||
latent_i = torch.randn(1, in_channels, count, 1, generator=generator)
|
||||
latent[i, :, :count] = latent_i[0]
|
||||
if coords.shape[0] > 1000:
|
||||
norms = [float(v) for v in latent.squeeze(-1).square().sum(dim=(1, 2)).detach().cpu().tolist()]
|
||||
print(
|
||||
"TRELLIS2_EMPTY_SHAPE_TRACE",
|
||||
{
|
||||
"coords_rows": int(coords.shape[0]),
|
||||
"batch_size": int(batch_size),
|
||||
"coord_counts": coord_counts.tolist() if coord_counts is not None else None,
|
||||
"latent_norms": norms,
|
||||
},
|
||||
)
|
||||
if coord_counts is not None:
|
||||
latent.trellis_coord_counts = coord_counts.clone()
|
||||
model = model.clone()
|
||||
model.model_options = model.model_options.copy()
|
||||
if "transformer_options" in model.model_options:
|
||||
@ -519,11 +727,17 @@ class EmptyShapeLatentTrellis2(IO.ComfyNode):
|
||||
model.model_options["transformer_options"] = {}
|
||||
|
||||
model.model_options["transformer_options"]["coords"] = coords
|
||||
if coord_counts is not None:
|
||||
model.model_options["transformer_options"]["coord_counts"] = coord_counts
|
||||
if is_512_pass:
|
||||
model.model_options["transformer_options"]["generation_mode"] = "shape_generation_512"
|
||||
else:
|
||||
model.model_options["transformer_options"]["generation_mode"] = "shape_generation"
|
||||
return IO.NodeOutput({"samples": latent, "coords": coords, "type": "trellis2"}, model)
|
||||
output = {"samples": latent, "coords": coords, "type": "trellis2"}
|
||||
if coord_counts is not None:
|
||||
output["coord_counts"] = coord_counts
|
||||
output["batch_index"] = [0] * batch_size
|
||||
return IO.NodeOutput(output, model)
|
||||
|
||||
class EmptyTextureLatentTrellis2(IO.ComfyNode):
|
||||
@classmethod
|
||||
@ -553,10 +767,45 @@ class EmptyTextureLatentTrellis2(IO.ComfyNode):
|
||||
coords = structure_or_coords.int()
|
||||
|
||||
shape_latent = shape_latent["samples"]
|
||||
batch_size, coord_counts, max_tokens = infer_batched_coord_layout(coords)
|
||||
if shape_latent.ndim == 4:
|
||||
shape_latent = shape_latent.squeeze(-1).transpose(1, 2).reshape(-1, channels)
|
||||
if shape_latent.shape[0] != batch_size:
|
||||
raise ValueError(
|
||||
f"shape_latent batch {shape_latent.shape[0]} doesn't match coords batch {batch_size}"
|
||||
)
|
||||
shape_latent = shape_latent.squeeze(-1).transpose(1, 2)
|
||||
if shape_latent.shape[1] < max_tokens:
|
||||
raise ValueError(
|
||||
f"shape_latent tokens {shape_latent.shape[1]} can't cover coords max tokens {max_tokens}"
|
||||
)
|
||||
|
||||
latent = torch.randn(1, channels, coords.shape[0], 1)
|
||||
if batch_size == 1:
|
||||
coord_counts = None
|
||||
latent = torch.randn(1, channels, coords.shape[0], 1)
|
||||
else:
|
||||
latent = torch.zeros(batch_size, channels, max_tokens, 1)
|
||||
base_state = torch.random.get_rng_state()
|
||||
for i in range(batch_size):
|
||||
count = int(coord_counts[i].item())
|
||||
generator = torch.Generator(device="cpu")
|
||||
generator.set_state(base_state.clone())
|
||||
latent_i = torch.randn(1, channels, count, 1, generator=generator)
|
||||
latent[i, :, :count] = latent_i[0]
|
||||
if coords.shape[0] > 1000:
|
||||
norms = [float(v) for v in latent.squeeze(-1).square().sum(dim=(1, 2)).detach().cpu().tolist()]
|
||||
shape_norms = [float(v) for v in shape_latent.square().sum(dim=(1, 2)).detach().cpu().tolist()] if shape_latent.ndim == 3 else None
|
||||
print(
|
||||
"TRELLIS2_EMPTY_TEXTURE_TRACE",
|
||||
{
|
||||
"coords_rows": int(coords.shape[0]),
|
||||
"batch_size": int(batch_size),
|
||||
"coord_counts": coord_counts.tolist() if coord_counts is not None else None,
|
||||
"latent_norms": norms,
|
||||
"shape_latent_norms": shape_norms,
|
||||
},
|
||||
)
|
||||
if coord_counts is not None:
|
||||
latent.trellis_coord_counts = coord_counts.clone()
|
||||
model = model.clone()
|
||||
model.model_options = model.model_options.copy()
|
||||
if "transformer_options" in model.model_options:
|
||||
@ -565,9 +814,15 @@ class EmptyTextureLatentTrellis2(IO.ComfyNode):
|
||||
model.model_options["transformer_options"] = {}
|
||||
|
||||
model.model_options["transformer_options"]["coords"] = coords
|
||||
if coord_counts is not None:
|
||||
model.model_options["transformer_options"]["coord_counts"] = coord_counts
|
||||
model.model_options["transformer_options"]["generation_mode"] = "texture_generation"
|
||||
model.model_options["transformer_options"]["shape_slat"] = shape_latent
|
||||
return IO.NodeOutput({"samples": latent, "coords": coords, "type": "trellis2"}, model)
|
||||
output = {"samples": latent, "coords": coords, "type": "trellis2"}
|
||||
if coord_counts is not None:
|
||||
output["coord_counts"] = coord_counts
|
||||
output["batch_index"] = [0] * batch_size
|
||||
return IO.NodeOutput(output, model)
|
||||
|
||||
|
||||
class EmptyStructureLatentTrellis2(IO.ComfyNode):
|
||||
@ -587,8 +842,15 @@ class EmptyStructureLatentTrellis2(IO.ComfyNode):
|
||||
def execute(cls, batch_size):
|
||||
in_channels = 8
|
||||
resolution = 16
|
||||
latent = torch.randn(batch_size, in_channels, resolution, resolution, resolution)
|
||||
return IO.NodeOutput({"samples": latent, "type": "trellis2"})
|
||||
generator = torch.Generator(device="cpu")
|
||||
generator.manual_seed(11426)
|
||||
latent = torch.randn(1, in_channels, resolution, resolution, resolution, generator=generator).repeat(batch_size, 1, 1, 1, 1)
|
||||
norms = [float(v) for v in latent.square().sum(dim=(1, 2, 3, 4)).detach().cpu().tolist()]
|
||||
print("TRELLIS2_EMPTY_STRUCTURE_TRACE", {"batch_size": int(batch_size), "latent_norms": norms})
|
||||
output = {"samples": latent, "type": "trellis2"}
|
||||
if batch_size > 1:
|
||||
output["batch_index"] = [0] * batch_size
|
||||
return IO.NodeOutput(output)
|
||||
|
||||
def simplify_fn(vertices, faces, colors=None, target=100000):
|
||||
if vertices.ndim == 3:
|
||||
|
||||
Loading…
Reference in New Issue
Block a user