第三步补全接入
This commit is contained in:
Ukawka
@@ -0,0 +1,152 @@
|
||||
import {execFileSync} from 'node:child_process';
|
||||
import {copyFileSync, existsSync, mkdirSync, readFileSync, readdirSync, writeFileSync} from 'node:fs';
|
||||
import path from 'node:path';
|
||||
import {projectRoot, publicStage} from '../lib/config.mjs';
|
||||
import {selectPython} from '../lib/python-runtime.mjs';
|
||||
import {readImage, sha256} from '../lib/images.mjs';
|
||||
import {validateParts} from '../lib/validate-parts.mjs';
|
||||
import {prepareRun, jsonFile, finishRun} from '../lib/run-artifacts.mjs';
|
||||
import {editImage, decodeEditedImages} from '../providers/image-edits.mjs';
|
||||
import {redact} from '../providers/http.mjs';
|
||||
|
||||
const script = path.join(projectRoot, 'completion_geometry.py');
|
||||
const readJson = (file) => JSON.parse(readFileSync(file, 'utf8').replace(/^\uFEFF/u, ''));
|
||||
function pythonCall(python, args) {
|
||||
try {
|
||||
execFileSync(python, [script, ...args], {cwd: projectRoot, windowsHide: true,
|
||||
encoding: 'utf8', timeout: 120000, env: {...process.env, PYTHONUTF8: '1'}, stdio: ['ignore', 'pipe', 'pipe']});
|
||||
} catch (error) { throw new Error(`补全几何处理失败：${error.stderr || error.message}`); }
|
||||
}
|
||||
function filesIn(dir, prefix = '') {
|
||||
return readdirSync(dir, {withFileTypes: true}).flatMap((entry) => entry.isDirectory()
|
||||
? filesIn(path.join(dir, entry.name), `${prefix}${entry.name}/`) : [`${prefix}${entry.name}`]);
|
||||
}
|
||||
|
||||
export async function runCompletion({stage, inputDir, outDir, partIds = [], prepareOnly = false,
|
||||
candidatesDir, completionSize = 1024, extensionRadius = 32, python: requestedPython,
|
||||
resume = false, dryRun = false}, transport) {
|
||||
if (!inputDir) throw new Error('部件补全需要 --input-dir 指向 SAM 输出目录');
|
||||
if (!outDir && !dryRun) throw new Error('部件补全需要独立的 --out-dir');
|
||||
if (prepareOnly && candidatesDir) throw new Error('--prepare-only 与 --candidates-dir 不能同时使用');
|
||||
const size = Number(completionSize), radius = Number(extensionRadius);
|
||||
if (!Number.isInteger(size) || size < 128 || size > 4096 || !Number.isInteger(radius) || radius < 1 || radius > 256) {
|
||||
throw new Error('--completion-size 必须为 128～4096 的整数；--extension-radius 必须为 1～256 的整数');
|
||||
}
|
||||
const input = path.resolve(inputDir);
|
||||
const plan = readJson(path.join(input, 'semantic-parts.json'));
|
||||
const validation = validateParts(plan);
|
||||
if (plan.schema_version !== 2 || !validation.valid) throw new Error(`补全需要有效 v2 计划：${validation.errors.join(', ')}`);
|
||||
const names = new Set(plan.parts.map((p) => p.id));
|
||||
if (partIds.some((id) => !names.has(id)) || new Set(partIds).size !== partIds.length) throw new Error('--part 包含未知或重复部件 ID');
|
||||
const segmentation = readJson(path.join(input, 'segmentation.json'));
|
||||
if (!Array.isArray(segmentation.parts) || new Set(segmentation.parts.map((p) => p.id)).size !== segmentation.parts.length || segmentation.parts.some((p) => !names.has(p.id))) {
|
||||
throw new Error('SAM 部件 ID 必须唯一且存在于语义计划中');
|
||||
}
|
||||
const source = readImage(path.join(input, 'source.png'), {pngOnly: true});
|
||||
const inputFiles = ['source.png', 'semantic-parts.json', 'segmentation.json', ...segmentation.parts.map((p) => `masks/${p.id}.png`)];
|
||||
const inputs = Object.fromEntries(inputFiles.map((file) => [file, sha256(readFileSync(path.join(input, file)))]));
|
||||
const imported = {};
|
||||
if (candidatesDir) {
|
||||
for (const p of plan.parts.filter((p) => p.needs_completion && (!partIds.length || partIds.includes(p.id)))) {
|
||||
const file = path.resolve(candidatesDir, `${p.id}.png`);
|
||||
if (existsSync(file)) imported[p.id] = {file, sha256: sha256(readFileSync(file))};
|
||||
}
|
||||
}
|
||||
const python = selectPython(requestedPython);
|
||||
if (dryRun) {
|
||||
console.log(JSON.stringify({stage: 'completion', input, width: source.width, height: source.height,
|
||||
selected_parts: partIds, completion_size: size, extension_radius: radius,
|
||||
mode: prepareOnly ? 'prepare' : candidatesDir ? 'offline_candidates' : 'image_edit', network_requested: false}, null, 2));
|
||||
return {dryRun: true};
|
||||
}
|
||||
const codeFiles = ['stages/completion-stage.mjs', 'completion_geometry.py', 'normalize_reference.py',
|
||||
'lib/config.mjs', 'lib/python-runtime.mjs', 'lib/run-artifacts.mjs', 'lib/images.mjs',
|
||||
'lib/validate-parts.mjs', 'schemas/semantic-parts.schema.json', 'providers/image-edits.mjs', 'providers/http.mjs'];
|
||||
const identity = {schema_version: 1, stage: 'completion', output_kind: 'bounded_part_completion',
|
||||
mode: prepareOnly ? 'prepare' : candidatesDir ? 'offline_candidates' : 'image_edit',
|
||||
resolved_stage: publicStage(stage), endpoint_sha256: sha256(stage.baseUrl), prompt_sha256: sha256(stage.prompt),
|
||||
inputs, part_ids: partIds, completion_size: size, extension_radius: radius,
|
||||
imported_candidates: Object.fromEntries(Object.entries(imported).map(([id, value]) => [id, value.sha256])),
|
||||
code_sha256: sha256(codeFiles.map((file) => `${file}\n${readFileSync(path.join(projectRoot, file), 'utf8')}`).join('\n'))};
|
||||
const fingerprint = sha256(JSON.stringify(identity));
|
||||
const run = prepareRun(outDir, fingerprint, resume);
|
||||
if (run.reused) return {out: run.out, reused: true};
|
||||
const manifest = {...identity, fingerprint, status: 'running', quality_status: 'pending', started_at: new Date().toISOString(), network_requested: false};
|
||||
jsonFile(run.manifestPath, manifest, stage.apiKey);
|
||||
const result = {schema_version: 1, quality_status: 'needs_review', parts: [], skipped: []};
|
||||
try {
|
||||
for (const file of inputFiles) {
|
||||
mkdirSync(path.dirname(path.join(run.out, file)), {recursive: true});
|
||||
copyFileSync(path.join(input, file), path.join(run.out, file));
|
||||
}
|
||||
pythonCall(python, ['prepare', '--out', run.out, '--size', String(size), '--radius', String(radius), ...partIds.flatMap((id) => ['--part', id])]);
|
||||
const prepared = readJson(path.join(run.out, 'completion-plan.json'));
|
||||
result.skipped = prepared.skipped;
|
||||
result.parts = prepared.tasks.map((task) => ({id: task.id, status: 'prepared', upstream_warnings: task.warnings, semantic_verified: false}));
|
||||
for (const [index, task] of prepared.tasks.entries()) {
|
||||
const folder = path.join(run.out, 'tasks', task.id);
|
||||
const prompt = `${stage.prompt}\n\nPart: ${task.id}: ${task.semantic}\nCompletion requirements: ${JSON.stringify(task.completion)}\nCanvas: ${size} x ${size} pixels. Keep the part at its current position and native scale.\nExtend only the hidden contour and requested joint connections, by at most ${radius} pixels from the visible silhouette. Keep all other background pure white or transparent. Do not add the occluder or neighboring body parts.\n${stage.capabilities.multiple_images ? 'The second image is local context only; do not copy its other parts into the output.' : 'Only the isolated visible part is provided.'}\nReturn exactly the same canvas dimensions.\n`;
|
||||
const fullPrompt = `${prompt}\nOccluded by: ${task.occluded_by.join(', ') || 'not specified'}.\nLocal pixel bounds for proposed extensions (xyxy, exclusive upper bounds): ${JSON.stringify(task.region_bounds_xyxy)}. Null means no supported region for that requirement. These boxes bound the allowed regions; do not fill whole rectangles.\n`;
|
||||
writeFileSync(path.join(folder, 'prompt.txt'), fullPrompt);
|
||||
if (prepareOnly) continue;
|
||||
const record = result.parts[index];
|
||||
console.log(`部件补全 [${index + 1}/${prepared.tasks.length}] ${task.id}`);
|
||||
let outputs;
|
||||
if (candidatesDir) {
|
||||
if (!imported[task.id]) { record.status = 'missing_candidate'; continue; }
|
||||
outputs = [{bytes: readFileSync(imported[task.id].file), extension: 'png'}];
|
||||
} else {
|
||||
const images = [readImage(path.join(folder, 'input.png'), {pngOnly: true})];
|
||||
if (stage.capabilities.multiple_images) images.push(readImage(path.join(folder, 'context.png'), {pngOnly: true}));
|
||||
const mask = stage.capabilities.mask ? readImage(path.join(folder, 'edit-mask.png'), {pngOnly: true}) : null;
|
||||
jsonFile(path.join(folder, 'request.json'), {model: stage.model, images: images.map((im) => im.sha256),
|
||||
edit_mask_sent: Boolean(mask), allowed_mask_sha256: sha256(readFileSync(path.join(folder, 'allowed-mask.png')))}, stage.apiKey);
|
||||
manifest.network_requested = true;
|
||||
jsonFile(run.manifestPath, manifest, stage.apiKey);
|
||||
try {
|
||||
const {payload, metadata} = await editImage(stage, {images, mask, prompt: fullPrompt}, transport);
|
||||
jsonFile(path.join(folder, 'response.raw.json'), payload, stage.apiKey);
|
||||
jsonFile(path.join(folder, 'request-metadata.json'), metadata, stage.apiKey);
|
||||
outputs = await decodeEditedImages(payload, stage, transport);
|
||||
} catch (error) {
|
||||
record.status = 'request_failed';
|
||||
if (error.diagnostics) jsonFile(path.join(folder, 'response.error.json'), error.diagnostics, stage.apiKey);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
record.candidates = [];
|
||||
for (const [i, output] of outputs.entries()) {
|
||||
const filename = `candidate-${i + 1}.${output.extension}`;
|
||||
writeFileSync(path.join(folder, filename), output.bytes);
|
||||
const merged = path.join(folder, `merged-${i + 1}.png`), qaFile = path.join(folder, `qa-${i + 1}.json`);
|
||||
pythonCall(python, ['merge', '--out', folder, '--candidate', path.join(folder, filename), '--merged', merged, '--qa', qaFile]);
|
||||
const qa = readJson(qaFile);
|
||||
record.candidates.push({file: `tasks/${task.id}/${filename}`, qa});
|
||||
if (qa.accepted && !record.selected_file) {
|
||||
record.selected_file = `tasks/${task.id}/merged-${i + 1}.png`;
|
||||
record.qa = qa;
|
||||
}
|
||||
}
|
||||
record.status = record.selected_file ? 'geometry_passed_needs_review' : 'rejected';
|
||||
if (!record.qa) record.qa = {accepted: false, reasons: record.candidates.flatMap((c) => c.qa.reasons)};
|
||||
jsonFile(path.join(run.out, 'completion-result.json'), result, stage.apiKey);
|
||||
}
|
||||
manifest.quality_status = prepareOnly ? 'prepared_needs_review' : 'needs_review';
|
||||
manifest.prepared_count = result.parts.length;
|
||||
manifest.accepted_geometry_count = result.parts.filter((p) => p.selected_file).length;
|
||||
manifest.rejected_count = result.parts.filter((p) => p.status === 'rejected').length;
|
||||
manifest.skipped_count = result.skipped.length;
|
||||
manifest.completed_at = new Date().toISOString();
|
||||
jsonFile(path.join(run.out, 'completion-result.json'), result, stage.apiKey);
|
||||
pythonCall(python, ['finalize', '--out', run.out]);
|
||||
finishRun(run, manifest, filesIn(run.out).filter((file) => file !== 'manifest.json'), stage.apiKey);
|
||||
return {out: run.out, reused: false, prepared: result.parts.length, accepted: manifest.accepted_geometry_count, rejected: manifest.rejected_count};
|
||||
} catch (error) {
|
||||
manifest.status = 'failed'; manifest.error = redact(error.message, stage.apiKey);
|
||||
manifest.quality_status = 'failed';
|
||||
jsonFile(run.manifestPath, manifest, stage.apiKey);
|
||||
jsonFile(path.join(run.out, 'completion-result.json'), result, stage.apiKey);
|
||||
try { pythonCall(python, ['finalize', '--out', run.out]); } catch { /* Preserve the original error and artifacts. */ }
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user