163 lines
8.0 KiB
TypeScript
163 lines
8.0 KiB
TypeScript
// Фоновое отслеживание незавершённых задач: результаты из ComfyUI записываются в базу,
|
||
// даже если ни одна страница не открыта.
|
||
import type {
|
||
FileRef, Job, LiveStatus, ModelArtifacts, ModelGenerator, TextureArtifacts, TextureSource,
|
||
} from '../shared/api.ts';
|
||
import { comfyJson } from './comfy.ts';
|
||
import { glbTriangles } from './glb.ts';
|
||
import { TEXTURE_VIEWS } from './workflows.ts';
|
||
import {
|
||
batchByJob, deleteJob, finishBatch, finishJob, finishProjectClean, finishProjectModel, finishProjectTexture, pendingJobs,
|
||
projectByCleanJob, projectByModelJob, projectByTextureJob, transaction,
|
||
} from './db.ts';
|
||
|
||
/** Сколько ждать задачу, которой нет ни в очереди, ни в истории ComfyUI (ComfyUI перезапускался). */
|
||
const LOST_AFTER_MS = 30_000;
|
||
|
||
const live = new Map<string, LiveStatus>();
|
||
const missingSince = new Map<string, number>();
|
||
|
||
/** Положение незавершённой задачи в очереди ComfyUI — для показа на страницах. */
|
||
export function liveStatus(id: string): LiveStatus | undefined {
|
||
return live.get(id);
|
||
}
|
||
|
||
interface HistoryEntry {
|
||
status?: { status_str?: string; completed?: boolean; messages?: [string, { exception_message?: string }][] };
|
||
outputs?: Record<string, Record<string, unknown>>;
|
||
}
|
||
|
||
/** Файлы из outputs истории: картинки (SaveImage) и GLB (SaveGLB кладёт их под другим ключом) — по расширению. */
|
||
function outputFiles(h: HistoryEntry) {
|
||
const images: FileRef[] = [];
|
||
const models: FileRef[] = [];
|
||
for (const node of Object.values(h.outputs ?? {})) {
|
||
for (const list of Object.values(node)) {
|
||
if (!Array.isArray(list)) continue;
|
||
for (const i of list as FileRef[]) {
|
||
if (typeof i?.filename !== 'string') continue;
|
||
if (/\.glb$/i.test(i.filename)) models.push(i);
|
||
else if (/\.(png|webp|jpe?g)$/i.test(i.filename)) images.push(i);
|
||
}
|
||
}
|
||
}
|
||
return { images, models };
|
||
}
|
||
|
||
type Files = { images?: FileRef[]; models?: FileRef[] };
|
||
/** Файл по метке префикса (см. workflows.ts): <prefix><tag>_NNNNN_.png. */
|
||
const tagged = (files: FileRef[] | undefined, tag: string) => files?.find(i => i.filename.includes(tag + '_')) ?? null;
|
||
|
||
/**
|
||
* Модель без текстур: у Hunyuan3D-2 mini — единственный GLB, у TRELLIS.2 — <prefix>_mesh_…glb; плюс превью.
|
||
*/
|
||
function modelArtifacts(generator: ModelGenerator, r: Files): ModelArtifacts | null {
|
||
const glb = generator === 'trellis' ? tagged(r.models, '_mesh') : r.models?.[0];
|
||
if (!glb) return null;
|
||
return { generator, glb, preview: tagged(r.images, '_preview'), faces: glbTriangles(glb), updated: Date.now() };
|
||
}
|
||
|
||
/**
|
||
* Текстурная группа: текстурированный GLB (не _mesh), base color, metallic, roughness, развёртка и виды —
|
||
* рендеры _view<угол> у TRELLIS.2 или 6 нарисованных видов _mv у Paint.
|
||
*/
|
||
function textureArtifacts(source: TextureSource, r: Files): TextureArtifacts | null {
|
||
const glb = r.models?.find(m => !m.filename.includes('_mesh_')) ?? null;
|
||
if (!glb) return null;
|
||
const images = r.images ?? [];
|
||
const views = source === 'trellis'
|
||
? TEXTURE_VIEWS.map(a => tagged(images, '_view' + a)).filter((i): i is FileRef => i != null)
|
||
: images.filter(i => i.filename.includes('_mv_')).sort((x, y) => x.filename.localeCompare(y.filename));
|
||
return {
|
||
source, glb, faces: glbTriangles(glb), baseColor: tagged(images, '_basecolor'), metallic: tagged(images, '_metallic'),
|
||
roughness: tagged(images, '_roughness'), uv: tagged(images, '_uv'), views, updated: Date.now(),
|
||
};
|
||
}
|
||
|
||
/**
|
||
* Завершение задачи. Картинки text2img прикрепляются к её батчу (Images.batchId);
|
||
* результат очистки и модель проекта заменяют прежние в проекте, а сама задача удаляется — история не хранится.
|
||
*/
|
||
function finish(job: Job, state: 'done' | 'error', result: { error?: string; images?: FileRef[]; models?: FileRef[] }) {
|
||
transaction(() => {
|
||
if (job.type === 'cleanup') {
|
||
const projectId = job.params.projectId ?? projectByCleanJob(job.id);
|
||
if (projectId != null) finishProjectClean(projectId, state === 'done' ? result.images?.[0] ?? null : null, result.error);
|
||
deleteJob(job.id);
|
||
return;
|
||
}
|
||
// модель проекта (у TRELLIS.2 — сразу с текстурами) — в проект, задача удаляется (история не хранится)
|
||
const modelProject = job.type === 'model' ? job.params.projectId ?? projectByModelJob(job.id) : null;
|
||
if (job.type === 'model' && modelProject != null) {
|
||
const generator = job.params.generator ?? 'hunyuan';
|
||
const done = state === 'done';
|
||
const model = done ? modelArtifacts(generator, result) : null;
|
||
const texture = done && generator === 'trellis' ? textureArtifacts('trellis', result) : null;
|
||
finishProjectModel(modelProject, model, texture, result.error ?? (done && !model ? 'нет GLB в результате' : undefined));
|
||
deleteJob(job.id);
|
||
return;
|
||
}
|
||
// текстура Hunyuan3D-2 Paint — в проект, задача удаляется
|
||
const textureProject = job.type === 'texture' ? job.params.projectId ?? projectByTextureJob(job.id) : null;
|
||
if (textureProject != null) {
|
||
finishProjectTexture(textureProject, state === 'done' ? textureArtifacts('paint', result) : null, result.error);
|
||
deleteJob(job.id);
|
||
return;
|
||
}
|
||
finishJob(job.id, state, result);
|
||
const batchId = job.type === 'text2img' ? job.params.batchId ?? batchByJob(job.id) : null;
|
||
if (batchId != null) finishBatch(batchId, state, result.images ?? [], result.error);
|
||
});
|
||
}
|
||
|
||
async function tick() {
|
||
const jobs = pendingJobs();
|
||
if (!jobs.length) {
|
||
live.clear();
|
||
missingSince.clear();
|
||
return;
|
||
}
|
||
// ComfyUI недоступен — задачи остаются незавершёнными до его возвращения
|
||
const q = await comfyJson<{ queue_running?: unknown[][]; queue_pending?: unknown[][] }>('/queue');
|
||
const running = new Set((q.queue_running ?? []).map(x => String(x[1])));
|
||
const queued = (q.queue_pending ?? []).map(x => String(x[1]));
|
||
|
||
for (const job of jobs) {
|
||
if (running.has(job.id)) {
|
||
live.set(job.id, { state: 'running' });
|
||
missingSince.delete(job.id);
|
||
continue;
|
||
}
|
||
const pos = queued.indexOf(job.id);
|
||
if (pos >= 0) {
|
||
live.set(job.id, { state: 'queued', position: pos + 1 });
|
||
missingSince.delete(job.id);
|
||
continue;
|
||
}
|
||
const h = (await comfyJson<Record<string, HistoryEntry>>('/history/' + encodeURIComponent(job.id)))[job.id];
|
||
if (h?.status?.status_str === 'error') {
|
||
const err = (h.status.messages ?? []).find(m => m[0] === 'execution_error' || m[0] === 'execution_interrupted');
|
||
finish(job, 'error', { error: err ? err[1].exception_message ?? 'прервано' : 'ошибка выполнения' });
|
||
} else if (h) {
|
||
finish(job, 'done', outputFiles(h));
|
||
} else {
|
||
// между постановкой в очередь и появлением в истории задачу может не быть видно — ждём
|
||
const since = missingSince.get(job.id) ?? Date.now();
|
||
missingSince.set(job.id, since);
|
||
if (Date.now() - since < LOST_AFTER_MS) continue;
|
||
finish(job, 'error', { error: 'задача не найдена в ComfyUI (перезапускался?)' });
|
||
}
|
||
live.delete(job.id);
|
||
missingSince.delete(job.id);
|
||
}
|
||
}
|
||
|
||
let busy = false;
|
||
export function startTracker(intervalMs = 1500) {
|
||
setInterval(async () => {
|
||
if (busy) return;
|
||
busy = true;
|
||
try { await tick(); } catch { /* ComfyUI недоступен — попробуем на следующем шаге */ } finally { busy = false; }
|
||
}, intervalMs);
|
||
}
|