98 lines
3.4 KiB
C++
98 lines
3.4 KiB
C++
// Encode an image into the TRELLIS.2 DINOv3 conditioning tensor (.dinodata),
|
|
// replacing the external dump_dinodata.py:
|
|
//
|
|
// image (PNG/JPG; solid black/white backgrounds are removed automatically)
|
|
// -> background cleanup -> alpha bbox crop, premultiply, LANCZOS 512
|
|
// -> DINOv3 ViT-L/16 -> [1, 1029, 1024] cond -> .dinodata
|
|
//
|
|
// usage: dino_encode <dino.gguf> <image> [out.dinodata] [--size N] [--pre out.png]
|
|
//
|
|
#include "trellis2.h"
|
|
|
|
#define STB_IMAGE_IMPLEMENTATION
|
|
#include "stb_image.h"
|
|
#define STB_IMAGE_WRITE_IMPLEMENTATION
|
|
#include "stb_image_write.h"
|
|
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
int main(int argc, char ** argv) {
|
|
if (argc < 3) {
|
|
std::fprintf(stderr,
|
|
"usage: %s <dino.gguf> <image> [out.dinodata] [--size N] [--pre out.png]\n",
|
|
argv[0]);
|
|
return 2;
|
|
}
|
|
const std::string gguf_path = argv[1];
|
|
const std::string img_path = argv[2];
|
|
std::string out_path = "cond.dinodata";
|
|
std::string pre_path;
|
|
int size = 512;
|
|
for (int i = 3; i < argc; ++i) {
|
|
if (std::strcmp(argv[i], "--size") == 0 && i + 1 < argc) {
|
|
size = std::atoi(argv[++i]);
|
|
} else if (std::strcmp(argv[i], "--pre") == 0 && i + 1 < argc) {
|
|
pre_path = argv[++i];
|
|
} else {
|
|
out_path = argv[i];
|
|
}
|
|
}
|
|
|
|
int w = 0, h = 0, comp = 0;
|
|
unsigned char * pixels = stbi_load(img_path.c_str(), &w, &h, &comp, 4);
|
|
if (!pixels) {
|
|
std::fprintf(stderr, "failed to decode image %s: %s\n",
|
|
img_path.c_str(), stbi_failure_reason());
|
|
return 1;
|
|
}
|
|
const int removed = trellis2_remove_solid_background_rgba(
|
|
pixels, w, h, TRELLIS2_BACKGROUND_AUTO);
|
|
std::printf("image : %s %dx%d (%d channels, background pixels changed: %d)\n",
|
|
img_path.c_str(), w, h, comp, removed);
|
|
|
|
std::string err;
|
|
std::vector<uint8_t> rgb;
|
|
if (!trellis2_preprocess_rgba(pixels, w, h, size, rgb, &err)) {
|
|
std::fprintf(stderr, "preprocess failed: %s\n", err.c_str());
|
|
stbi_image_free(pixels);
|
|
return 1;
|
|
}
|
|
stbi_image_free(pixels);
|
|
|
|
if (!pre_path.empty()) {
|
|
stbi_write_png(pre_path.c_str(), size, size, 3, rgb.data(), size * 3);
|
|
std::printf("wrote : %s (preprocessed %dx%d RGB)\n", pre_path.c_str(), size, size);
|
|
}
|
|
|
|
trellis2_dino_model * model = trellis2_dino_load(gguf_path, true, &err);
|
|
if (!model) {
|
|
std::fprintf(stderr, "model load failed: %s\n", err.c_str());
|
|
return 1;
|
|
}
|
|
std::printf("model : %s (backend %s)\n", gguf_path.c_str(),
|
|
trellis2_dino_backend_name(model));
|
|
|
|
trellis2_dino_cond cond;
|
|
if (!trellis2_dino_encode_rgb(model, rgb.data(), size, cond, &err)) {
|
|
std::fprintf(stderr, "encode failed: %s\n", err.c_str());
|
|
trellis2_dino_free(model);
|
|
return 1;
|
|
}
|
|
trellis2_dino_free(model);
|
|
|
|
const trellis2_dino_fingerprint fp = trellis2_dino_fingerprints(cond);
|
|
std::printf("cond : [1, %lld, %lld] min=%.4f max=%.4f mean=%.6f l2=%.4f\n",
|
|
(long long) cond.tokens(), (long long) cond.channels(),
|
|
fp.vmin, fp.vmax, fp.mean, fp.l2);
|
|
|
|
if (!trellis2_save_dinodata(out_path, cond, &err)) {
|
|
std::fprintf(stderr, "save failed: %s\n", err.c_str());
|
|
return 1;
|
|
}
|
|
std::printf("wrote : %s (%zu floats)\n", out_path.c_str(), cond.count());
|
|
return 0;
|
|
}
|