117 lines
4.8 KiB
Python
117 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Reference for the full stage-1 flow-Euler sampling loop, to validate the C++
|
|
trellis2_ss_flow_sample against the real FlowEulerGuidanceIntervalSampler.
|
|
|
|
Builds SparseStructureFlowModel in float32, loads the DINOv3 cond from a
|
|
.dinodata file, draws fixed noise, runs the sampler, and writes a
|
|
self-describing binary `ss_sample_ref.bin`:
|
|
|
|
magic : 8 bytes "SSSAMP01"
|
|
int32 : resolution, in_channels, cond_tokens, cond_channels, steps
|
|
float32 : guidance_strength, guidance_rescale, gi_min, gi_max, rescale_t, sigma_min
|
|
float32 : noise [in_channels * R^3] channel-major
|
|
float32 : cond [cond_tokens * cond_channels] token-major
|
|
float32 : latent[in_channels * R^3] channel-major (reference z_s)
|
|
|
|
Usage:
|
|
python ref_ss_sample.py [--device mps|cpu] [--dinodata .../MushroomBoy.dinodata]
|
|
"""
|
|
|
|
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"
|
|
_v, _d, 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)
|
|
ap.add_argument("--device", default="cpu", choices=["mps", "cpu", "cuda"])
|
|
ap.add_argument("--seed", type=int, default=1234)
|
|
ap.add_argument("--out", default=os.path.join(os.path.dirname(__file__), "ss_sample_ref.bin"))
|
|
args = ap.parse_args()
|
|
|
|
from safetensors.torch import load_file
|
|
from trellis2.models.sparse_structure_flow import SparseStructureFlowModel
|
|
from trellis2.pipelines import samplers
|
|
|
|
dev = torch.device(args.device)
|
|
with open(args.ckpt + ".json") as f:
|
|
cfg = json.load(f)["args"]
|
|
|
|
model = SparseStructureFlowModel(**cfg)
|
|
model.convert_to(torch.float32)
|
|
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 unexpected and missing == ["rope_phases"], (missing, unexpected)
|
|
model.to(dev)
|
|
|
|
R, Cin = cfg["resolution"], cfg["in_channels"]
|
|
cond_np = load_dinodata(args.dinodata) # [1, Lkv, Cctx]
|
|
Lkv, Cctx = cond_np.shape[1], cond_np.shape[2]
|
|
cond = torch.from_numpy(cond_np.copy()).float().to(dev)
|
|
neg_cond = torch.zeros_like(cond)
|
|
|
|
# Fixed noise (CPU generator for reproducibility, then move to device).
|
|
g = torch.Generator().manual_seed(args.seed)
|
|
noise = torch.randn(1, Cin, R, R, R, generator=g).to(dev)
|
|
|
|
sampler = samplers.FlowEulerGuidanceIntervalSampler(sigma_min=1e-5)
|
|
params = dict(steps=12, rescale_t=5.0, guidance_strength=7.5,
|
|
guidance_interval=[0.6, 1.0], guidance_rescale=0.7)
|
|
print(f"sampling on {dev} steps={params['steps']} gs={params['guidance_strength']} "
|
|
f"rescale={params['guidance_rescale']} interval={params['guidance_interval']}")
|
|
with torch.no_grad():
|
|
z_s = sampler.sample(model, noise, cond, neg_cond, verbose=True, **params).samples
|
|
|
|
z_np = z_s.detach().cpu().numpy().astype(np.float32) # [1, Cin, R,R,R]
|
|
noise_np = noise.detach().cpu().numpy().astype(np.float32)
|
|
print(f"z_s: min={z_np.min():.5f} max={z_np.max():.5f} mean={z_np.mean():.6f} "
|
|
f"l2={np.linalg.norm(z_np):.5f} (occupancy>0: {(z_np>0).mean()*100:.2f}%)")
|
|
|
|
noise_cm = noise_np.reshape(Cin, -1).reshape(-1)
|
|
z_cm = z_np.reshape(Cin, -1).reshape(-1)
|
|
cond_tm = cond_np.reshape(-1)
|
|
|
|
with open(args.out, "wb") as f:
|
|
f.write(b"SSSAMP01")
|
|
f.write(struct.pack("<5i", R, Cin, Lkv, Cctx, params["steps"]))
|
|
f.write(struct.pack("<6f", params["guidance_strength"], params["guidance_rescale"],
|
|
params["guidance_interval"][0], params["guidance_interval"][1],
|
|
params["rescale_t"], 1e-5))
|
|
f.write(noise_cm.astype("<f4").tobytes())
|
|
f.write(cond_tm.astype("<f4").tobytes())
|
|
f.write(z_cm.astype("<f4").tobytes())
|
|
print(f"wrote {args.out} ({os.path.getsize(args.out):,} bytes)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|