100 lines
3.5 KiB
Python
100 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Reference for the stage-1 SS decoder, to validate the C++ trellis2_ss_dec_decode
|
|
against the real SparseStructureDecoder.
|
|
|
|
Builds SparseStructureDecoder in float32 (lossless upcast of the fp16
|
|
checkpoint), decodes a sparse-structure latent z_s, and writes a
|
|
self-describing binary `ss_dec_ref.bin`:
|
|
|
|
magic : 8 bytes "SSDEC001"
|
|
int32 : latent_channels, res_in, out_channels, res_out
|
|
float32 : latent[latent_channels * res_in^3] channel-major
|
|
float32 : logits[out_channels * res_out^3] channel-major
|
|
|
|
The latent comes from --latent (a .latent produced by ss_sample / ref_ss_sample)
|
|
if given; otherwise a fixed-seed standard-normal z_s is used so the C++ side can
|
|
feed the identical input.
|
|
|
|
Usage:
|
|
python ref_ss_dec.py [--device mps|cpu] [--latent z_s.latent] [--seed N]
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import struct
|
|
import sys
|
|
|
|
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
|
|
|
|
SHIV = os.environ.get("TRELLIS2_PY", "/trellis2")
|
|
sys.path.insert(0, SHIV)
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
DEFAULT_CKPT = os.environ.get("TRELLIS2_CKPT", os.path.join(
|
|
os.path.dirname(__file__), "..", "models", "TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16"))
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--ckpt", default=DEFAULT_CKPT)
|
|
ap.add_argument("--device", default="cpu", choices=["mps", "cpu", "cuda"])
|
|
ap.add_argument("--latent", default=None, help="optional z_s .latent (channel-major float32)")
|
|
ap.add_argument("--seed", type=int, default=1234)
|
|
ap.add_argument("--out", default=os.path.join(os.path.dirname(__file__), "ss_dec_ref.bin"))
|
|
args = ap.parse_args()
|
|
|
|
from safetensors.torch import load_file
|
|
from trellis2.models.sparse_structure_vae import SparseStructureDecoder
|
|
|
|
dev = torch.device(args.device)
|
|
with open(args.ckpt + ".json") as f:
|
|
cfg = json.load(f)["args"]
|
|
|
|
model = SparseStructureDecoder(**cfg)
|
|
model.convert_to_fp32() # set self.dtype=f32 AND convert torso modules
|
|
model.eval()
|
|
sd = {k: v.float() for k, v in load_file(args.ckpt + ".safetensors").items()}
|
|
missing, unexpected = model.load_state_dict(sd, strict=False)
|
|
assert not missing and not unexpected, (missing, unexpected)
|
|
model.to(dev)
|
|
|
|
Cin = cfg["latent_channels"]
|
|
Rin = 16
|
|
Oc = cfg["out_channels"]
|
|
|
|
if args.latent:
|
|
z_cm = np.fromfile(args.latent, dtype="<f4")
|
|
assert z_cm.size == Cin * Rin**3, (z_cm.size, Cin * Rin**3)
|
|
z = torch.from_numpy(z_cm.reshape(1, Cin, Rin, Rin, Rin).copy()).float().to(dev)
|
|
print(f"latent : loaded {args.latent}")
|
|
else:
|
|
g = torch.Generator().manual_seed(args.seed)
|
|
z = torch.randn(1, Cin, Rin, Rin, Rin, generator=g).to(dev)
|
|
print(f"latent : random seed={args.seed}")
|
|
|
|
with torch.no_grad():
|
|
logits = model(z) # [1, Oc, Rout, Rout, Rout]
|
|
|
|
Rout = logits.shape[-1]
|
|
lg = logits.detach().cpu().numpy().astype(np.float32)
|
|
print(f"logits : [{Oc},{Rout},{Rout},{Rout}] min={lg.min():.5f} max={lg.max():.5f} "
|
|
f"mean={lg.mean():.6f} occupied(>0)={(lg > 0).mean() * 100:.2f}%")
|
|
|
|
z_out = z.detach().cpu().numpy().astype(np.float32).reshape(Cin, -1).reshape(-1)
|
|
lg_out = lg.reshape(Oc, -1).reshape(-1)
|
|
|
|
with open(args.out, "wb") as f:
|
|
f.write(b"SSDEC001")
|
|
f.write(struct.pack("<4i", Cin, Rin, Oc, Rout))
|
|
f.write(z_out.astype("<f4").tobytes())
|
|
f.write(lg_out.astype("<f4").tobytes())
|
|
print(f"wrote {args.out} ({os.path.getsize(args.out):,} bytes)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|