144 lines
4.4 KiB
C
144 lines
4.4 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* pth-convert: Convert PyTorch .pth files to safetensors format
|
|
*
|
|
* Usage: pth-convert <input.pth> -o <output.safetensors> [--inspect]
|
|
*/
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
#include "pth/pth.h"
|
|
#include "safetensors/safetensors.h"
|
|
|
|
static void print_usage(const char *prog) {
|
|
fprintf(stderr, "Usage: %s <input.pth> -o <output.safetensors> [--inspect]\n", prog);
|
|
fprintf(stderr, "\nOptions:\n");
|
|
fprintf(stderr, " --inspect Print tensor table and exit (no conversion)\n");
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
const char *input = NULL;
|
|
const char *output = NULL;
|
|
int inspect = 0;
|
|
|
|
for (int i = 1; i < argc; i++) {
|
|
if (strcmp(argv[i], "-o") == 0 && i + 1 < argc) {
|
|
output = argv[++i];
|
|
} else if (strcmp(argv[i], "--inspect") == 0) {
|
|
inspect = 1;
|
|
} else if (argv[i][0] != '-') {
|
|
input = argv[i];
|
|
} else {
|
|
fprintf(stderr, "Unknown option: %s\n", argv[i]);
|
|
print_usage(argv[0]);
|
|
return 1;
|
|
}
|
|
}
|
|
|
|
if (!input) {
|
|
fprintf(stderr, "Error: no input file specified\n");
|
|
print_usage(argv[0]);
|
|
return 1;
|
|
}
|
|
|
|
if (!inspect && !output) {
|
|
fprintf(stderr, "Error: no output file specified (use -o or --inspect)\n");
|
|
print_usage(argv[0]);
|
|
return 1;
|
|
}
|
|
|
|
/* Open .pth file */
|
|
printf("Opening %s...\n", input);
|
|
pt_file *pth = pt_open(input);
|
|
if (!pth) {
|
|
fprintf(stderr, "Error: failed to open %s\n", input);
|
|
return 1;
|
|
}
|
|
|
|
int n = pt_tensor_count(pth);
|
|
printf("Found %d tensors\n\n", n);
|
|
|
|
/* Print tensor table */
|
|
if (inspect || !output) {
|
|
printf("%-50s %-8s %-20s %12s\n", "Name", "Dtype", "Shape", "Size");
|
|
printf("%-50s %-8s %-20s %12s\n",
|
|
"--------------------------------------------------",
|
|
"--------", "--------------------", "------------");
|
|
|
|
size_t total_bytes = 0;
|
|
for (int i = 0; i < n; i++) {
|
|
const pt_tensor *t = pt_tensor_at(pth, i);
|
|
if (!t) continue;
|
|
|
|
/* Format shape */
|
|
char shape_str[128] = {0};
|
|
int off = 0;
|
|
off += snprintf(shape_str + off, sizeof(shape_str) - off, "[");
|
|
for (int d = 0; d < t->ndim; d++) {
|
|
if (d > 0) off += snprintf(shape_str + off, sizeof(shape_str) - off, ", ");
|
|
off += snprintf(shape_str + off, sizeof(shape_str) - off, "%lld",
|
|
(long long)t->shape[d]);
|
|
}
|
|
snprintf(shape_str + off, sizeof(shape_str) - off, "]");
|
|
|
|
/* Format size */
|
|
char size_str[32];
|
|
if (t->data_size >= 1024 * 1024)
|
|
snprintf(size_str, sizeof(size_str), "%.1f MB", (double)t->data_size / (1024 * 1024));
|
|
else if (t->data_size >= 1024)
|
|
snprintf(size_str, sizeof(size_str), "%.1f KB", (double)t->data_size / 1024);
|
|
else
|
|
snprintf(size_str, sizeof(size_str), "%zu B", t->data_size);
|
|
|
|
printf("%-50s %-8s %-20s %12s\n",
|
|
t->name, fs_dtype_name(t->dtype), shape_str, size_str);
|
|
|
|
total_bytes += t->data_size;
|
|
}
|
|
|
|
printf("\nTotal: %d tensors, %.1f MB\n", n, (double)total_bytes / (1024 * 1024));
|
|
|
|
if (inspect) {
|
|
pt_close(pth);
|
|
return 0;
|
|
}
|
|
}
|
|
|
|
/* Convert to safetensors */
|
|
printf("Converting to %s...\n", output);
|
|
|
|
fs_writer *w = fs_writer_new();
|
|
if (!w) {
|
|
fprintf(stderr, "Error: failed to create safetensors writer\n");
|
|
pt_close(pth);
|
|
return 1;
|
|
}
|
|
|
|
int skipped = 0;
|
|
for (int i = 0; i < n; i++) {
|
|
const pt_tensor *t = pt_tensor_at(pth, i);
|
|
if (!t || !t->data || t->data_size == 0) {
|
|
fprintf(stderr, " Warning: skipping tensor '%s' (no data)\n", t ? t->name : "?");
|
|
skipped++;
|
|
continue;
|
|
}
|
|
fs_writer_add_tensor(w, t->name, t->dtype, t->shape, t->ndim, t->data);
|
|
}
|
|
|
|
if (!fs_writer_save(w, output)) {
|
|
fprintf(stderr, "Error: failed to write %s\n", output);
|
|
fs_writer_free(w);
|
|
pt_close(pth);
|
|
return 1;
|
|
}
|
|
|
|
fs_writer_free(w);
|
|
pt_close(pth);
|
|
|
|
printf("Done! Converted %d tensors (%d skipped)\n", n - skipped, skipped);
|
|
return 0;
|
|
}
|