diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py index 3b7717915..986218a58 100644 --- a/comfy/latent_formats.py +++ b/comfy/latent_formats.py @@ -762,6 +762,7 @@ class Hunyuan3Dv2_1(LatentFormat): class Trellis2(LatentFormat): # TODO latent_channels = 32 + trellis3d_preview = True # routes the sampler preview to Trellis3DPreviewer class Hunyuan3Dv2mini(LatentFormat): latent_channels = 64 latent_dimensions = 1 diff --git a/comfy/ldm/trellis2/model.py b/comfy/ldm/trellis2/model.py index 7e263169a..99c930c8b 100644 --- a/comfy/ldm/trellis2/model.py +++ b/comfy/ldm/trellis2/model.py @@ -8,7 +8,7 @@ from comfy.ldm.trellis2.attention import ( ) from comfy.ldm.genmo.joint_model.layers import TimestepEmbedder from comfy.ldm.flux.math import apply_rope, apply_rope1 - +from comfy.ldm.trellis2 import sampling_preview class SparseGELU(nn.GELU): def forward(self, input: VarLenTensor) -> VarLenTensor: @@ -1100,6 +1100,8 @@ class Trellis2(nn.Module): # Pre-computed per-stage back-projected features proj_feats = kwargs.get("trellis2_proj_feats") + sampling_preview.set_context(mode=mode, coords=coords, coord_counts=coord_counts) + is_first_shape_pass = False if mode == "shape_generation_512": is_first_shape_pass = True diff --git a/comfy/ldm/trellis2/sampling_preview.py b/comfy/ldm/trellis2/sampling_preview.py new file mode 100644 index 000000000..76b2ef2c4 --- /dev/null +++ b/comfy/ldm/trellis2/sampling_preview.py @@ -0,0 +1,40 @@ +"""Side-channel for per-step sampling previews of the Trellis2/Pixal3D cascade. + +The sampler callback only receives the denoised latent `x0`, but a 3D preview of +the texture stage also needs the sparse voxel coords (the latent alone is just an +unordered [B, C, N, 1] feature stack). The diffusion model's forward writes the +current stage context here each step; the latent previewer reads it. Single prompt +samples one stage at a time, so a module-level holder is safe. + +Intentionally dependency-free to avoid import cycles with comfy.ldm.trellis2.model. +""" + +_context = {} + +# Fitted texture latent -> base-color factors: (W [C, 3], b [3]) on CPU, or None. +# Calibrated by VaeDecodeTextureTrellis from real decoded albedo and read by the +# latent previewer for a faithful texture-stage color preview. +_tex_rgb = None + + +def set_context(mode=None, coords=None, coord_counts=None): + _context["mode"] = mode + _context["coords"] = coords + _context["coord_counts"] = coord_counts + + +def get_context(): + return _context + + +def clear(): + _context.clear() + + +def set_tex_rgb(W, b): + global _tex_rgb + _tex_rgb = (W, b) + + +def get_tex_rgb(): + return _tex_rgb diff --git a/latent_preview.py b/latent_preview.py index a9d777661..40c9e0fe3 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -75,10 +75,151 @@ class Latent2RGBPreviewer(LatentPreviewer): return preview_to_image(latent_image) +class Trellis3DPreviewer(LatentPreviewer): + """Per-step preview for the Trellis2/Pixal3D cascade. + + Structure stage: x0 is a dense [B, 32, 16, 16, 16] grid — project the per-cell + activation norm orthographically to a 2D occupancy heatmap (no decode, no coords). + Texture stage: x0 is sparse [B, 32, N, 1] — splat the first 3 latent channels as + pseudo-color onto the fixed voxel coords (read from the sampling side-channel). + Shape stage adds no visible motion (coords are fixed, only sub-voxel detail + evolves) and a full decode per step is too costly, so it's skipped. + + Both stages render through one orthographic point splatter (static view). + """ + _SIZE = 128 + _FILL = 0.9 # fraction of frame the texture splat spans (leaves a border) + _STRUCTURE_ZOOM = 0.66 # <1 pulls the SS camera back, leaving margin around the blob + + def _splat(self, points, colors, rad): + # points: [K, 3] voxel-index coords. colors: [K, 3] in [0, 1]. + # Center + isotropic-normalize, project orthographically front-on + # (x->horizontal, y->up, z->depth), then splat a square footprint per point + # with one global far->near sort (painter's). + S = self._SIZE + dev = points.device # keep every tensor here + p = points.float() + p = p - (p.amax(0) + p.amin(0)) * 0.5 + p = p / p.abs().amax().clamp(min=1e-8) + x, y, z = p[:, 0], p[:, 1], p[:, 2] + depth = z # into-screen + m = self._FILL + u = ((x * m * 0.5 + 0.5) * (S - 1)).long().clamp(0, S - 1) + v = (((-y) * m * 0.5 + 0.5) * (S - 1)).long().clamp(0, S - 1) # image up = +y + cols = colors.to(dev) + us, vs, ds, cs = [], [], [], [] + for dv in range(-rad, rad + 1): + for du in range(-rad, rad + 1): + us.append((u + du).clamp(0, S - 1)) + vs.append((v + dv).clamp(0, S - 1)) + ds.append(depth) + cs.append(cols) + order = torch.cat(ds).argsort() + img = torch.zeros(S, S, 3, device=dev) + img[torch.cat(vs)[order], torch.cat(us)[order]] = torch.cat(cs)[order] + return preview_to_image(img, do_scale=False) + + @staticmethod + def _turbo(x): + # Anton Mikhailov polynomial approximation of the turbo colormap. x: any shape + # in [0, 1] -> (..., 3) RGB. + x = x.clamp(0.0, 1.0) + x2 = x * x; x3 = x2 * x; x4 = x2 * x2; x5 = x4 * x + r = 0.13572138 + 4.61539260*x - 42.66032258*x2 + 132.13108234*x3 - 152.94239396*x4 + 59.28637943*x5 + g = 0.09140261 + 2.19418839*x + 4.84296658*x2 - 14.18503333*x3 + 4.27729857*x4 + 2.82956604*x5 + b = 0.10667330 + 12.64194608*x - 60.58204836*x2 + 110.36276771*x3 - 89.90310912*x4 + 27.34824973*x5 + return torch.stack([r, g, b], dim=-1).clamp(0.0, 1.0) + + def _structure(self, x0): + # x0: [B, 32, D, H, W]; the model only consumes the first 8 channels. + # Dense orthographic max-projection -> filled occupancy heatmap (turbo-colored, + # intensity-weighted so empty space stays black). + act = x0[0, :min(8, x0.shape[1])].float().norm(dim=0) # [D, H, W] + proj = act.amax(dim=2) # project along one axis + proj = (proj - proj.amin()) / (proj.amax() - proj.amin() + 1e-8) + inner = max(1, int(round(self._SIZE * self._STRUCTURE_ZOOM))) + img = torch.nn.functional.interpolate(proj[None, None], size=(inner, inner), mode="nearest") + pad = self._SIZE - inner + pl, pt = pad // 2, pad // 2 + gray = torch.nn.functional.pad(img, (pl, pad - pl, pt, pad - pt))[0, 0] # [S, S], zero margin + rgb = self._turbo(gray) * gray.unsqueeze(-1) # [S, S, 3], black where empty + return preview_to_image(rgb, do_scale=False) + + @staticmethod + def _latent_color(latent): + # Prefer the calibrated latent->base_color map (fit from real decoded + # albedo by VaeDecodeTextureTrellis); fall back to PCA pseudo-color until + # a texture decode has trained it. + try: + from comfy.ldm.trellis2 import sampling_preview + factors = sampling_preview.get_tex_rgb() + except Exception: + factors = None + if factors is not None: + W, b = factors + rgb = latent @ W.to(latent) + b.to(latent) + return rgb.clamp(0, 1) + return Trellis3DPreviewer._pca_color(latent) + + @staticmethod + def _pca_color(latent): + # latent: [n, C]. Map the 3 directions of maximum variance to RGB. + # Higher contrast and more coherent than picking 3 fixed channels. + X = latent - latent.mean(dim=0, keepdim=True) + cov = (X.transpose(0, 1) @ X) / max(X.shape[0] - 1, 1) # [C, C] + _, evecs = torch.linalg.eigh(cov) # ascending eigenvalues + pcs = evecs[:, -3:] # [C, 3] top-3 components + # Deterministic sign per component (largest-magnitude entry positive) to + # stop the preview's hues from flickering as the latent rotates each step. + sign = torch.sign(pcs[pcs.abs().argmax(dim=0), torch.arange(3, device=pcs.device)]) + pcs = pcs * sign.clamp(min=-1.0) + proj = X @ pcs # [n, 3] + pmin = proj.amin(dim=0, keepdim=True) + pmax = proj.amax(dim=0, keepdim=True) + return ((proj - pmin) / (pmax - pmin + 1e-8)).clamp(0, 1) + + def _texture(self, x0, coords): + if coords.shape[-1] == 4: + b0 = coords[:, 0] == 0 + spatial = coords[b0][:, 1:4].float() + else: + spatial = coords[:, :3].float() + n0 = spatial.shape[0] + if n0 == 0: + return None + latent = x0[0, :, :n0, 0].float().transpose(0, 1) # [n0, C] + colors = self._latent_color(latent) # [n0, 3] + res = float(spatial.max().item()) + 1.0 + rad = max(1, int(round(self._SIZE * self._FILL / max(res, 1) / 2))) + return self._splat(spatial, colors, rad) + + def decode_latent_to_preview(self, x0): + try: + from comfy.ldm.trellis2 import sampling_preview + ctx = sampling_preview.get_context() + if x0.ndim == 5: + return self._structure(x0) + mode = ctx.get("mode") + coords = ctx.get("coords") + if mode == "texture_generation" and coords is not None: + return self._texture(x0, coords) + except Exception as e: + logging.debug(f"Trellis3DPreviewer: skipping preview ({e})") + return None + + def decode_latent_to_preview_image(self, preview_format, x0): + preview_image = self.decode_latent_to_preview(x0) + if preview_image is None: + return None + return ("JPEG", preview_image, MAX_PREVIEW_RESOLUTION) + + def get_previewer(device, latent_format): previewer = None method = args.preview_method if method != LatentPreviewMethod.NoPreviews: + if getattr(latent_format, "trellis3d_preview", False): + return Trellis3DPreviewer() # TODO previewer methods taesd_decoder_path = None if latent_format.taesd_decoder_name is not None: