81 lines
3.9 KiB
TypeScript
81 lines
3.9 KiB
TypeScript
// Однократный перенос старых прогонов text2img в ImageBatches.
|
|
// ComfyUI (SaveImage) записывает в каждый PNG граф запуска (текстовый чанк "prompt") — из него
|
|
// восстанавливается пользовательский ввод; картинки одного батча — подряд идущие файлы с одинаковым графом.
|
|
import fs from 'node:fs';
|
|
import path from 'node:path';
|
|
import zlib from 'node:zlib';
|
|
import { COMFY_OUTPUT } from './config.ts';
|
|
import { attachImages, batchOfImage, getMeta, insertTextBatch, setMeta, transaction } from './db.ts';
|
|
import { inputFromGraph } from './workflows.ts';
|
|
|
|
const KEY = 'migration.imageBatches.v1';
|
|
const SRC_DIR = 'img2glb_src';
|
|
/** Картинки одного батча сохраняются почти одновременно. */
|
|
const SAME_BATCH_MS = 15_000;
|
|
|
|
/** Текстовые чанки PNG (tEXt, iTXt, zTXt). */
|
|
function pngTexts(file: string): Record<string, string> {
|
|
const b = fs.readFileSync(file);
|
|
const out: Record<string, string> = {};
|
|
let o = 8;
|
|
while (o + 8 <= b.length) {
|
|
const len = b.readUInt32BE(o);
|
|
const type = b.toString('latin1', o + 4, o + 8);
|
|
const d = b.subarray(o + 8, o + 8 + len);
|
|
const nul = d.indexOf(0);
|
|
if (type === 'tEXt') {
|
|
out[d.toString('latin1', 0, nul)] = d.toString('latin1', nul + 1);
|
|
} else if (type === 'zTXt') {
|
|
out[d.toString('latin1', 0, nul)] = zlib.inflateSync(d.subarray(nul + 2)).toString('latin1');
|
|
} else if (type === 'iTXt') {
|
|
const compressed = d[nul + 1] === 1;
|
|
let p = d.indexOf(0, nul + 3) + 1; // после языка
|
|
p = d.indexOf(0, p) + 1; // после переведённого ключа
|
|
const text = d.subarray(p);
|
|
out[d.toString('latin1', 0, nul)] = (compressed ? zlib.inflateSync(text) : text).toString('utf8');
|
|
} else if (type === 'IEND') {
|
|
break;
|
|
}
|
|
o += 12 + len;
|
|
}
|
|
return out;
|
|
}
|
|
|
|
export function migrateOldRuns() {
|
|
if (getMeta(KEY)) return;
|
|
const dir = path.join(COMFY_OUTPUT, SRC_DIR);
|
|
let names: string[] = [];
|
|
try { names = fs.readdirSync(dir).filter(n => /\.png$/i.test(n)); } catch { /* папки нет — переносить нечего */ }
|
|
|
|
const files = names
|
|
.map(name => ({ name, mtime: fs.statSync(path.join(dir, name)).mtimeMs }))
|
|
.filter(f => batchOfImage({ filename: f.name, subfolder: SRC_DIR, type: 'output' }) == null)
|
|
.sort((a, b) => a.mtime - b.mtime || a.name.localeCompare(b.name));
|
|
|
|
// группировка: одинаковый граф, подряд, в пределах SAME_BATCH_MS от первой картинки
|
|
const groups: { graph: string; first: number; names: string[] }[] = [];
|
|
let skipped = 0;
|
|
for (const f of files) {
|
|
let graph: string | undefined;
|
|
try { graph = pngTexts(path.join(dir, f.name)).prompt; } catch { /* битый PNG */ }
|
|
if (!graph) { skipped++; continue; }
|
|
const last = groups[groups.length - 1];
|
|
if (last && last.graph === graph && f.mtime - last.first <= SAME_BATCH_MS) last.names.push(f.name);
|
|
else groups.push({ graph, first: f.mtime, names: [f.name] });
|
|
}
|
|
|
|
let created = 0;
|
|
transaction(() => {
|
|
for (const g of groups) {
|
|
let parsed;
|
|
try { parsed = inputFromGraph(JSON.parse(g.graph)); } catch { parsed = null; }
|
|
if (!parsed) { skipped += g.names.length; continue; }
|
|
const id = insertTextBatch(parsed.input, { created: Math.round(g.first), state: 'done', seedUsed: parsed.seed, prompt: parsed.prompt });
|
|
attachImages(id, g.names.map(filename => ({ filename, subfolder: SRC_DIR, type: 'output' })));
|
|
created++;
|
|
}
|
|
setMeta(KEY, new Date().toISOString());
|
|
});
|
|
console.log(`Перенос старых прогонов: батчей ${created}, картинок ${files.length - skipped}` + (skipped ? `, пропущено ${skipped}` : ''));
|
|
}
|