Initial release
This commit is contained in:
civ
@@ -0,0 +1,239 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Reference dumps for the 1024_cascade high-resolution geometry stage.
|
||||
|
||||
Produces dumps/reference_cascade.gguf with the full HR chain so the C++ port can
|
||||
be validated stage by stage:
|
||||
|
||||
cond_512, cond_1024 the two DINOv3 conds (embedded → self-contained test)
|
||||
coords32 [L,4] 32^3 scaffold (from the SS reference latent)
|
||||
lr_noise [L,32] LR sampling noise (seed 4321)
|
||||
lr_slat [L,32] 512-model sampler output, denormalized
|
||||
up_coords [Nup,4] decoder.upsample(lr_slat, upsample_times=4) → 512^3
|
||||
hr_coords [Lhr,4] quantized+unique → 64^3 (the HR flow scaffold)
|
||||
hr_noise [Lhr,32] HR sampling noise (seed 5678)
|
||||
hr_flow_t500_out [Lhr,32] 1024-model flow forward at t=500 on hr_coords
|
||||
hr_slat [Lhr,32] 1024-model sampler output, denormalized
|
||||
lvl{i}.*, out7, out_coords per-level decode taps at resolution 1024 (→ 1024^3)
|
||||
|
||||
The reference is generated with TF32 disabled (ref_common.setup) so the golden
|
||||
values are true fp32. Run inside the container (see scripts/refgen.sh).
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import struct
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import ref_common # noqa: E402
|
||||
|
||||
ref_common.setup() # TF32 off + sdpa sparse attention + pure-torch sparse conv
|
||||
|
||||
import numpy as np # noqa: E402
|
||||
import torch # noqa: E402
|
||||
import torch.nn.functional as F # noqa: E402
|
||||
|
||||
|
||||
def read_dinodata(path):
|
||||
with open(path, "rb") as f:
|
||||
assert f.read(8) == b"DINOCOND", "bad magic"
|
||||
_, _, ndim = struct.unpack("<III", f.read(12))
|
||||
shape = struct.unpack("<%dI" % ndim, f.read(4 * ndim))
|
||||
arr = np.frombuffer(f.read(), dtype="<f4").reshape(shape)
|
||||
return arr
|
||||
|
||||
|
||||
def load_ss_sample_latent(path):
|
||||
with open(path, "rb") as f:
|
||||
assert f.read(8) == b"SSSAMP01", "bad magic"
|
||||
R, cin, lkv, cctx, steps = struct.unpack("<5i", f.read(20))
|
||||
n = cin * R * R * R
|
||||
f.seek(-(n * 4), os.SEEK_END)
|
||||
z = np.frombuffer(f.read(n * 4), dtype="<f4").reshape(1, cin, R, R, R)
|
||||
return torch.from_numpy(z.copy())
|
||||
|
||||
|
||||
def load_flow(stem, dev):
|
||||
from safetensors.torch import load_file
|
||||
from trellis2.models.structured_latent_flow import SLatFlowModel
|
||||
with open(stem + ".json") as f:
|
||||
cfg = json.load(f)["args"]
|
||||
cfg.pop("initialization", None)
|
||||
cfg.pop("dtype", None)
|
||||
m = SLatFlowModel(**cfg, dtype="float32")
|
||||
sd = {k: v.float() for k, v in load_file(stem + ".safetensors").items()}
|
||||
missing, unexpected = m.load_state_dict(sd, strict=False)
|
||||
assert not unexpected, unexpected
|
||||
m.convert_to(torch.float32)
|
||||
m.eval().to(dev)
|
||||
return m, cfg
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--models", default=os.path.join(ref_common.MODELS, "TRELLIS.2-4B", "ckpts"))
|
||||
ap.add_argument("--ss-dec", default=os.path.join(ref_common.MODELS, "TRELLIS-image-large",
|
||||
"ckpts", "ss_dec_conv3d_16l8_fp16"))
|
||||
ap.add_argument("--cond-512", default=os.path.join(ref_common.DUMPS, "fixture.dinodata"))
|
||||
ap.add_argument("--cond-1024", default=os.path.join(ref_common.DUMPS, "fixture_1024.dinodata"))
|
||||
ap.add_argument("--ss-latent", default=os.path.join(ref_common.REPO, "tests", "ss_sample_ref.bin"))
|
||||
ap.add_argument("--pipeline-json", default=os.path.join(ref_common.MODELS, "TRELLIS.2-4B", "pipeline.json"))
|
||||
ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
||||
ap.add_argument("--t", type=float, default=500.0)
|
||||
ap.add_argument("--lr-seed", type=int, default=4321)
|
||||
ap.add_argument("--hr-seed", type=int, default=5678)
|
||||
ap.add_argument("--lr-resolution", type=int, default=512)
|
||||
ap.add_argument("--resolution", type=int, default=1024)
|
||||
ap.add_argument("--out", default=os.path.join(ref_common.DUMPS, "reference_cascade.gguf"))
|
||||
args = ap.parse_args()
|
||||
|
||||
from safetensors.torch import load_file
|
||||
from trellis2.models.sparse_structure_vae import SparseStructureDecoder
|
||||
from trellis2.models.sc_vaes.fdg_vae import FlexiDualGridVaeDecoder
|
||||
from trellis2.pipelines.samplers import FlowEulerGuidanceIntervalSampler
|
||||
from trellis2.modules import sparse as sp
|
||||
|
||||
dev = torch.device(args.device)
|
||||
caps = {}
|
||||
|
||||
cond_512 = torch.from_numpy(read_dinodata(args.cond_512).copy()).float().to(dev)
|
||||
cond_1024 = torch.from_numpy(read_dinodata(args.cond_1024).copy()).float().to(dev)
|
||||
caps["cond_512"] = cond_512[0]
|
||||
caps["cond_1024"] = cond_1024[0]
|
||||
|
||||
with open(args.pipeline_json) as f:
|
||||
pj = json.load(f)["args"]
|
||||
norm = pj["shape_slat_normalization"]
|
||||
sampler_params = pj["shape_slat_sampler"]["params"]
|
||||
mean = torch.tensor(norm["mean"], device=dev)[None]
|
||||
std = torch.tensor(norm["std"], device=dev)[None]
|
||||
sampler = FlowEulerGuidanceIntervalSampler(sigma_min=1e-5)
|
||||
|
||||
# ── 32^3 scaffold from the SS reference latent ───────────────────────────
|
||||
with open(args.ss_dec + ".json") as f:
|
||||
ss_cfg = json.load(f)["args"]
|
||||
ss_cfg["use_fp16"] = False
|
||||
ss_dec = SparseStructureDecoder(**ss_cfg)
|
||||
ss_dec.load_state_dict({k: v.float() for k, v in load_file(args.ss_dec + ".safetensors").items()})
|
||||
ss_dec.dtype = torch.float32
|
||||
ss_dec.eval().float().to(dev)
|
||||
z_s = load_ss_sample_latent(args.ss_latent).to(dev)
|
||||
with torch.no_grad():
|
||||
occ = ss_dec(z_s) > 0
|
||||
occ = (F.max_pool3d(occ.float(), 2, 2, 0) > 0.5)
|
||||
coords = torch.argwhere(occ)[:, [0, 2, 3, 4]].int().contiguous()
|
||||
L = coords.shape[0]
|
||||
caps["coords32"] = coords.float()
|
||||
print(f"scaffold: {L} voxels at 32^3")
|
||||
del ss_dec
|
||||
|
||||
# ── LR flow: sample with the 512 model + cond_512, denormalize ───────────
|
||||
flow_lr, cfg_lr = load_flow(os.path.join(args.models, "slat_flow_img2shape_dit_1_3B_512_bf16"), dev)
|
||||
g = torch.Generator().manual_seed(args.lr_seed)
|
||||
lr_noise = torch.randn(L, cfg_lr["in_channels"], generator=g).to(dev)
|
||||
caps["lr_noise"] = lr_noise
|
||||
x0 = sp.SparseTensor(feats=lr_noise.clone(), coords=coords.to(dev))
|
||||
with torch.no_grad():
|
||||
lr_slat = sampler.sample(flow_lr, x0, cond=cond_512, neg_cond=torch.zeros_like(cond_512),
|
||||
**sampler_params, verbose=True).samples
|
||||
lr_slat = lr_slat * std + mean
|
||||
caps["lr_slat"] = lr_slat.feats
|
||||
print(f"lr_slat: {lr_slat.feats.shape} mean={lr_slat.feats.mean().item():.5f}")
|
||||
del flow_lr
|
||||
if dev.type == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# ── shape decoder: upsample(×4) → 512^3 candidate coords ─────────────────
|
||||
dstem = os.path.join(args.models, "shape_dec_next_dc_f16c32_fp16")
|
||||
with open(dstem + ".json") as f:
|
||||
dcfg = json.load(f)["args"]
|
||||
dcfg.pop("use_fp16", None)
|
||||
dcfg.pop("resolution", None)
|
||||
dec = FlexiDualGridVaeDecoder(resolution=args.resolution, use_fp16=False, **dcfg)
|
||||
dec.load_state_dict({k: v.float() for k, v in load_file(dstem + ".safetensors").items()})
|
||||
# The decoder runs on CPU: the 1024^3 expansion materializes millions of
|
||||
# voxels through the pure-torch sparse conv and would OOM the 16 GB GPU (and
|
||||
# it is exactly what the C++ port runs on CPU). Flows stay on GPU.
|
||||
dec.eval().float().cpu()
|
||||
|
||||
with torch.no_grad():
|
||||
up_coords = dec.upsample(lr_slat.cpu(), upsample_times=4) # [Nup, 4] at 512^3
|
||||
caps["up_coords"] = up_coords.float()
|
||||
print(f"upsample: {up_coords.shape[0]} candidate coords at {args.lr_resolution}^3")
|
||||
|
||||
# ── quantize + unique → 64^3 HR scaffold (verbatim pipeline formula) ─────
|
||||
hr_res = args.resolution
|
||||
quant_coords = torch.cat([
|
||||
up_coords[:, :1],
|
||||
((up_coords[:, 1:] + 0.5) / args.lr_resolution * (hr_res // 16)).int(),
|
||||
], dim=1)
|
||||
hr_coords = quant_coords.unique(dim=0)
|
||||
Lhr = hr_coords.shape[0]
|
||||
caps["hr_coords"] = hr_coords.float()
|
||||
print(f"hr scaffold: {Lhr} voxels at {hr_res // 16}^3")
|
||||
|
||||
# ── HR flow: forward @ t=500 + full sampler with 1024 model + cond_1024 ──
|
||||
flow_hr, cfg_hr = load_flow(os.path.join(args.models, "slat_flow_img2shape_dit_1_3B_1024_bf16"), dev)
|
||||
g = torch.Generator().manual_seed(args.hr_seed)
|
||||
hr_noise = torch.randn(Lhr, cfg_hr["in_channels"], generator=g).to(dev)
|
||||
caps["hr_noise"] = hr_noise
|
||||
xh = sp.SparseTensor(feats=hr_noise.clone(), coords=hr_coords.to(dev))
|
||||
with torch.no_grad():
|
||||
out = flow_hr(xh, torch.tensor([args.t], device=dev), cond_1024)
|
||||
caps["hr_flow_t500_out"] = out.feats
|
||||
print(f"hr flow t={args.t}: l2={out.feats.norm().item():.4f}")
|
||||
|
||||
xh0 = sp.SparseTensor(feats=hr_noise.clone(), coords=hr_coords.to(dev))
|
||||
with torch.no_grad():
|
||||
hr_slat = sampler.sample(flow_hr, xh0, cond=cond_1024, neg_cond=torch.zeros_like(cond_1024),
|
||||
**sampler_params, verbose=True).samples
|
||||
hr_slat = hr_slat * std + mean
|
||||
caps["hr_slat"] = hr_slat.feats
|
||||
print(f"hr_slat: {hr_slat.feats.shape} mean={hr_slat.feats.mean().item():.5f}")
|
||||
del flow_hr
|
||||
if dev.type == "cuda":
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# ── final decode of the HR slat at resolution 1024, level by level ───────
|
||||
# Only the final 7-channel output is kept: the per-level intermediates are
|
||||
# multiple GB at 1024^3 (and the decoder's level logic is already validated
|
||||
# exactly at the 512 tier). We run the base SparseUnetVaeDecoder forward
|
||||
# manually to get the raw 7 channels — dec(...) would run the FDG mesh
|
||||
# conversion (stubbed o_voxel). Running the decode on CPU is host-RAM heavy.
|
||||
dec.set_resolution(args.resolution)
|
||||
hr_slat_cpu = hr_slat.cpu()
|
||||
with torch.no_grad():
|
||||
h = dec.from_latent(hr_slat_cpu.float())
|
||||
for i, res in enumerate(dec.blocks):
|
||||
for j, block in enumerate(res):
|
||||
if i < len(dec.blocks) - 1 and j == len(res) - 1:
|
||||
h, _sub = block(h)
|
||||
else:
|
||||
h = block(h)
|
||||
print(f"decode level {i}: {h.feats.shape[0]} voxels x {h.feats.shape[1]} ch")
|
||||
hn = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
|
||||
out7 = dec.output_layer(hn)
|
||||
caps["out7"] = out7.feats
|
||||
caps["out_coords"] = out7.coords.float()
|
||||
print(f"out7: {out7.feats.shape}")
|
||||
|
||||
import gguf
|
||||
writer = gguf.GGUFWriter(args.out, "reference")
|
||||
manifest = {"shapes": {}, "atol": 2e-3, "rtol": 2e-3,
|
||||
"lr_resolution": args.lr_resolution, "resolution": args.resolution}
|
||||
for name, t in caps.items():
|
||||
a = t.detach().cpu().float().numpy()
|
||||
manifest["shapes"][name] = list(a.shape)
|
||||
writer.add_tensor(name, np.ascontiguousarray(a.reshape(-1), dtype=np.float32))
|
||||
writer.write_header_to_file()
|
||||
writer.write_kv_data_to_file()
|
||||
writer.write_tensors_to_file()
|
||||
writer.close()
|
||||
with open(os.path.join(ref_common.DUMPS, "manifest_cascade.json"), "w") as f:
|
||||
json.dump(manifest, f, indent=1)
|
||||
print(f"wrote {args.out} ({os.path.getsize(args.out):,} bytes), {len(caps)} tensors")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user