77 lines
3.0 KiB
C++
77 lines
3.0 KiB
C++
// ss_sample — run stage-1 sampling: a DINOv3 .dinodata cond + the SS-flow DiT
|
|
// GGUF -> the sparse-structure latent z_s. The occupancy scaffold is z_s > 0.
|
|
//
|
|
// usage: ss_sample <ss_flow_dit.gguf> <cond.dinodata> [out.latent] [--seed N]
|
|
//
|
|
// Writes the z_s latent (channel-major [in_channels * R^3] float32) to out.latent
|
|
// if given — the input the SS decoder (next stage) will consume.
|
|
|
|
#include "trellis2.h"
|
|
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <fstream>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
int main(int argc, char ** argv) {
|
|
if (argc < 3) {
|
|
std::fprintf(stderr, "usage: %s <ss_flow_dit.gguf> <cond.dinodata> [out.latent] [--seed N]\n", argv[0]);
|
|
return 2;
|
|
}
|
|
const std::string gguf_path = argv[1];
|
|
const std::string dino_path = argv[2];
|
|
std::string out_path;
|
|
uint64_t seed = 0;
|
|
for (int i = 3; i < argc; ++i) {
|
|
if (std::strcmp(argv[i], "--seed") == 0 && i + 1 < argc) seed = std::strtoull(argv[++i], nullptr, 10);
|
|
else if (out_path.empty()) out_path = argv[i];
|
|
}
|
|
|
|
std::printf("trellis2.cpp %s\n", trellis2_version());
|
|
|
|
trellis2_dino_cond cond;
|
|
std::string err;
|
|
if (!trellis2_load_dinodata(dino_path, cond, &err)) {
|
|
std::fprintf(stderr, "cond load error: %s\n", err.c_str());
|
|
return 1;
|
|
}
|
|
std::printf("cond : %lld tokens x %lld channels\n",
|
|
(long long) cond.tokens(), (long long) cond.channels());
|
|
|
|
trellis2_ss_flow_model * m = trellis2_ss_flow_load(gguf_path, true, &err);
|
|
if (!m) { std::fprintf(stderr, "model load error: %s\n", err.c_str()); return 1; }
|
|
std::printf("backend: %s\n", trellis2_ss_flow_backend_name(m));
|
|
|
|
const trellis2_ss_flow_hparams hp = trellis2_ss_flow_hparams_of(m); // copy (used after free)
|
|
const size_t N = (size_t) hp.resolution * hp.resolution * hp.resolution;
|
|
const size_t n = (size_t) hp.in_channels * N;
|
|
|
|
trellis2_ss_sampler_params P; // pipeline defaults (12 steps, gs 7.5, ...)
|
|
P.seed = seed;
|
|
|
|
std::vector<float> latent(n, 0.0f);
|
|
if (!trellis2_ss_flow_sample(m, cond.data.data(), (int) cond.tokens(), (int) cond.channels(),
|
|
&P, /*noise*/ nullptr, latent.data(), &err)) {
|
|
std::fprintf(stderr, "sample error: %s\n", err.c_str());
|
|
trellis2_ss_flow_free(m);
|
|
return 1;
|
|
}
|
|
trellis2_ss_flow_free(m);
|
|
|
|
double mn = 1e30, mx = -1e30, sum = 0.0;
|
|
size_t occ = 0;
|
|
for (float v : latent) { mn = v < mn ? v : mn; mx = v > mx ? v : mx; sum += v; if (v > 0.0f) ++occ; }
|
|
std::printf("z_s : [%d,%d,%d,%d] min=%.4f max=%.4f mean=%.5f occupancy(>0)=%.2f%%\n",
|
|
hp.in_channels, hp.resolution, hp.resolution, hp.resolution,
|
|
mn, mx, sum / (double) n, 100.0 * (double) occ / (double) n);
|
|
|
|
if (!out_path.empty()) {
|
|
std::ofstream f(out_path, std::ios::binary);
|
|
f.write(reinterpret_cast<const char *>(latent.data()), (std::streamsize) (n * sizeof(float)));
|
|
std::printf("wrote %s (%zu floats)\n", out_path.c_str(), n);
|
|
}
|
|
return 0;
|
|
}
|