mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-03 21:20:49 +08:00
42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
"""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, model_frame=None):
|
|
_context["mode"] = mode
|
|
_context["coords"] = coords
|
|
_context["coord_counts"] = coord_counts
|
|
_context["model_frame"] = model_frame
|
|
|
|
|
|
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
|