76 lines
3.1 KiB
C++
76 lines
3.1 KiB
C++
// ss_decode — run the stage-1 SS decoder: a sparse-structure latent z_s
|
|
// (.latent from ss_sample) + the SS decoder GGUF -> an occupancy logit grid.
|
|
// The coarse voxel scaffold is logit > 0.
|
|
//
|
|
// usage: ss_decode <ss_dec.gguf> <z_s.latent> [out.occ]
|
|
//
|
|
// Writes the occupancy logit grid (channel-major [out_channels * R_out^3]
|
|
// float32) to out.occ if given. Also prints occupancy stats and, if all output
|
|
// channels are 1, the count of occupied voxels.
|
|
|
|
#include "trellis2.h"
|
|
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <fstream>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
int main(int argc, char ** argv) {
|
|
if (argc < 3) {
|
|
std::fprintf(stderr, "usage: %s <ss_dec.gguf> <z_s.latent> [out.occ]\n", argv[0]);
|
|
return 2;
|
|
}
|
|
const std::string gguf_path = argv[1];
|
|
const std::string lat_path = argv[2];
|
|
const std::string out_path = argc > 3 ? argv[3] : "";
|
|
|
|
std::printf("trellis2.cpp %s\n", trellis2_version());
|
|
|
|
std::string err;
|
|
trellis2_ss_dec_model * m = trellis2_ss_dec_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_dec_backend_name(m));
|
|
|
|
const trellis2_ss_dec_hparams hp = trellis2_ss_dec_hparams_of(m); // copy (used after free)
|
|
const int Rin = hp.res_in();
|
|
const int Rout = hp.res_out();
|
|
const size_t n_in = (size_t) hp.latent_channels * Rin * Rin * Rin;
|
|
const size_t n_out = (size_t) hp.out_channels * Rout * Rout * Rout;
|
|
|
|
// load latent
|
|
std::ifstream f(lat_path, std::ios::binary);
|
|
if (!f) { std::fprintf(stderr, "cannot open latent %s\n", lat_path.c_str()); trellis2_ss_dec_free(m); return 1; }
|
|
std::vector<float> latent(n_in);
|
|
f.read(reinterpret_cast<char *>(latent.data()), (std::streamsize) (n_in * sizeof(float)));
|
|
if ((size_t) f.gcount() != n_in * sizeof(float)) {
|
|
std::fprintf(stderr, "latent size mismatch: got %zu floats, want %zu\n",
|
|
(size_t) f.gcount() / sizeof(float), n_in);
|
|
trellis2_ss_dec_free(m);
|
|
return 1;
|
|
}
|
|
std::printf("latent : [%d,%d,%d,%d]\n", hp.latent_channels, Rin, Rin, Rin);
|
|
|
|
std::vector<float> occ(n_out, 0.0f);
|
|
if (!trellis2_ss_dec_decode(m, latent.data(), occ.data(), &err)) {
|
|
std::fprintf(stderr, "decode error: %s\n", err.c_str());
|
|
trellis2_ss_dec_free(m);
|
|
return 1;
|
|
}
|
|
trellis2_ss_dec_free(m);
|
|
|
|
double mn = 1e30, mx = -1e30, sum = 0.0;
|
|
size_t occupied = 0;
|
|
for (float v : occ) { mn = v < mn ? v : mn; mx = v > mx ? v : mx; sum += v; if (v > 0.0f) ++occupied; }
|
|
std::printf("logits : [%d,%d,%d,%d] min=%.4f max=%.4f mean=%.5f occupied(>0)=%zu/%zu (%.2f%%)\n",
|
|
hp.out_channels, Rout, Rout, Rout, mn, mx, sum / (double) n_out,
|
|
occupied, n_out, 100.0 * (double) occupied / (double) n_out);
|
|
|
|
if (!out_path.empty()) {
|
|
std::ofstream o(out_path, std::ios::binary);
|
|
o.write(reinterpret_cast<const char *>(occ.data()), (std::streamsize) (n_out * sizeof(float)));
|
|
std::printf("wrote %s (%zu floats)\n", out_path.c_str(), n_out);
|
|
}
|
|
return 0;
|
|
}
|