115 lines
4.7 KiB
Python
115 lines
4.7 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Generate a reference forward pass of the TRELLIS.2 SS-flow DiT in float32, so the
|
|
C++ implementation can be validated against it bit-for-bit (modulo fp rounding).
|
|
|
|
Loads SparseStructureFlowModel directly from the checkpoint (no full pipeline),
|
|
runs forward(x, t, cond) on CPU in float32 with a fixed seed, and writes a
|
|
self-describing binary `ss_flow_ref.bin`:
|
|
|
|
magic : 8 bytes "SSFREF01"
|
|
int32 : resolution, in_channels, out_channels, cond_tokens, cond_channels
|
|
float32: t
|
|
float32: x [in_channels * resolution^3] channel-major (x[c*R^3 + n])
|
|
float32: cond[cond_tokens * cond_channels] token-major (the .dinodata layout)
|
|
float32: out[out_channels * resolution^3] channel-major (reference output)
|
|
|
|
Usage:
|
|
python ref_ss_flow.py --dinodata /path/MushroomBoy.dinodata [--t 500] [--out ss_flow_ref.bin]
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import struct
|
|
import sys
|
|
|
|
os.environ.setdefault("ATTN_BACKEND", "sdpa")
|
|
os.environ.setdefault("SPARSE_ATTN_BACKEND", "sdpa")
|
|
os.environ.setdefault("SPARSE_CONV_BACKEND", "none")
|
|
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.2-4B/ckpts/ss_flow_img_dit_1_3B_64_bf16"))
|
|
|
|
|
|
def load_dinodata(path):
|
|
with open(path, "rb") as f:
|
|
assert f.read(8) == b"DINOCOND", "bad magic"
|
|
version, dtype, 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 # [1, tokens, channels]
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--dinodata", default=os.path.join(os.path.dirname(__file__), "..", "dumps", "fixture.dinodata"))
|
|
ap.add_argument("--ckpt", default=DEFAULT_CKPT, help="checkpoint stem (no extension)")
|
|
ap.add_argument("--t", type=float, default=500.0)
|
|
ap.add_argument("--seed", type=int, default=1234)
|
|
ap.add_argument("--out", default=os.path.join(os.path.dirname(__file__), "ss_flow_ref.bin"))
|
|
args = ap.parse_args()
|
|
|
|
from safetensors.torch import load_file
|
|
from trellis2.models.sparse_structure_flow import SparseStructureFlowModel
|
|
|
|
with open(args.ckpt + ".json") as f:
|
|
cfg = json.load(f)["args"]
|
|
print(f"building SparseStructureFlowModel: {cfg['num_blocks']} blocks, d={cfg['model_channels']}")
|
|
|
|
torch.manual_seed(args.seed)
|
|
model = SparseStructureFlowModel(**cfg)
|
|
model.convert_to(torch.float32) # blocks -> f32 AND sets self.dtype=f32 (else
|
|
# forward's manual_cast downcasts activations)
|
|
model.eval()
|
|
sd = {k: v.float() for k, v in load_file(args.ckpt + ".safetensors").items()}
|
|
# rope_phases is a computed buffer (not stored in the checkpoint).
|
|
missing, unexpected = model.load_state_dict(sd, strict=False)
|
|
assert not unexpected, f"unexpected keys: {unexpected}"
|
|
assert missing == ["rope_phases"], f"unexpected missing keys: {missing}"
|
|
|
|
R = cfg["resolution"]
|
|
Cin = cfg["in_channels"]
|
|
Cout = cfg["out_channels"]
|
|
|
|
cond_np = load_dinodata(args.dinodata) # [1, Lkv, Cctx]
|
|
Lkv, Cctx = cond_np.shape[1], cond_np.shape[2]
|
|
assert Cctx == cfg["cond_channels"], f"cond channels {Cctx} != {cfg['cond_channels']}"
|
|
cond = torch.from_numpy(cond_np.copy()).float()
|
|
|
|
rng = np.random.default_rng(args.seed)
|
|
x_np = rng.standard_normal((1, Cin, R, R, R)).astype(np.float32)
|
|
x = torch.from_numpy(x_np)
|
|
t = torch.tensor([args.t], dtype=torch.float32)
|
|
|
|
with torch.no_grad():
|
|
out = model(x, t, cond) # [1, Cout, R, R, R]
|
|
out_np = out.detach().cpu().numpy().astype(np.float32)
|
|
print(f"forward done. x{tuple(x_np.shape)} t={args.t} cond{tuple(cond_np.shape)} -> out{tuple(out_np.shape)}")
|
|
print(f"out: min={out_np.min():.5f} max={out_np.max():.5f} mean={out_np.mean():.6f} l2={np.linalg.norm(out_np):.5f}")
|
|
|
|
# Flatten to the C++ layouts.
|
|
x_cm = x_np.reshape(Cin, -1).reshape(-1) # channel-major [Cin * R^3]
|
|
out_cm = out_np.reshape(Cout, -1).reshape(-1) # channel-major [Cout * R^3]
|
|
cond_tm = cond_np.reshape(-1) # token-major [Lkv * Cctx]
|
|
|
|
with open(args.out, "wb") as f:
|
|
f.write(b"SSFREF01")
|
|
f.write(struct.pack("<5i", R, Cin, Cout, Lkv, Cctx))
|
|
f.write(struct.pack("<f", args.t))
|
|
f.write(x_cm.astype("<f4").tobytes())
|
|
f.write(cond_tm.astype("<f4").tobytes())
|
|
f.write(out_cm.astype("<f4").tobytes())
|
|
print(f"wrote {args.out} ({os.path.getsize(args.out):,} bytes)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|