ai-chatbot/lib/canvas/layout-replace-verifier.ts

86 lines
4.6 KiB
TypeScript
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import "server-only";
import { createHash } from "node:crypto";
import sharp from "sharp";
import { generateText } from "ai";
import { getLanguageModelById } from "@/lib/ai/dynamic-provider";
import { getMultimodalAiModelsWithProvider } from "@/lib/db/queries";
import { fetchImageBufferFromUrl } from "@/lib/ai/rmbg/service";
import { LAYOUT_REPLACE_QUALITY_PROMPT, parseLayoutReplaceQuality } from "./layout-replace-quality";
import { LAYOUT_REPLACE_PRODUCT_PROMPT, parseLayoutReplaceProductAnalysis } from "./layout-replace-product-crop";
import { LAYOUT_REPLACE_TEXT_PROMPT, parseLayoutReplaceTextBoxes } from "./layout-replace-text-lock";
async function visionInput(buffer: Buffer) {
const resized = await sharp(buffer, { failOn: "error" }).rotate()
.resize(1400, 1400, { fit: "inside", withoutEnlargement: true })
.flatten({ background: "#ffffff" }).jpeg({ quality: 90 }).toBuffer();
return new URL(`data:image/jpeg;base64,${resized.toString("base64")}`);
}
async function inspectLayoutImages<T>(buffers: Buffer[], prompt: string, parse: (text: string) => T): Promise<T> {
const candidates = (await getMultimodalAiModelsWithProvider()).filter(({ model }) =>
model.supportsVision && !model.supportsImageGen && !model.isSubModel
).sort((a, b) => Number(b.model.isDefault) - Number(a.model.isDefault) ||
Number(b.provider.type === "vip98k") - Number(a.provider.type === "vip98k"));
if (!candidates.length) throw new Error("未配置可用的多模态质量检查模型,未放行未经检查的替换结果");
const images = await Promise.all(buffers.map(visionInput));
let lastError: unknown;
for (const candidate of candidates.slice(0, 2)) {
try {
const model = await getLanguageModelById(candidate.model.id);
if (!model || typeof model === "string") continue;
const { text } = await generateText({
model: model as any, maxRetries: 0, maxOutputTokens: 1200,
abortSignal: AbortSignal.timeout(60_000),
messages: [{ role: "user", content: [
{ type: "text", text: prompt },
...images.map(image => ({ type: "image" as const, image })),
] }],
});
return parse(text);
} catch (error) {
lastError = error;
// The fallback is another read-only check, never another image generation.
console.warn("[Canvas Layout Replace] quality verifier unavailable", {
model: candidate.model.modelId, errorType: error instanceof Error ? error.name : "unknown",
});
}
}
throw new Error("替换结果质量检查暂不可用,未将未经检查的图片标记为成功", { cause: lastError });
}
export async function locateLayoutReplaceProduct(sourceBuffer: Buffer, templateBuffer: Buffer) {
return inspectLayoutImages([templateBuffer, sourceBuffer], LAYOUT_REPLACE_PRODUCT_PROMPT, parseLayoutReplaceProductAnalysis);
}
const textRegions = new Map<string, { expires: number; result: ReturnType<typeof inspectTemplateText> }>();
function inspectTemplateText(buffer: Buffer) {
return inspectLayoutImages([buffer], LAYOUT_REPLACE_TEXT_PROMPT, parseLayoutReplaceTextBoxes);
}
export async function locateLayoutReplaceTemplateText(buffer: Buffer) {
const key = createHash("sha256").update(buffer).digest("hex");
const cached = textRegions.get(key);
if (cached && cached.expires > Date.now()) return cached.result;
if (textRegions.size >= 32) textRegions.delete(textRegions.keys().next().value!);
const result = inspectTemplateText(buffer);
const entry = { expires: Date.now() + 30 * 60_000, result };
textRegions.set(key, entry);
result.catch(() => { if (textRegions.get(key) === entry) textRegions.delete(key); });
return result;
}
export async function verifyCanvasLayoutReplacement(params: {
sourceBuffer: Buffer;
templateBuffer: Buffer;
resultUrl: string;
resultBuffer?: Buffer;
sourceIdentity?: string;
productBuffer?: Buffer;
}) {
const resultBuffer = params.resultBuffer || await fetchImageBufferFromUrl(params.resultUrl);
const prompt = [LAYOUT_REPLACE_QUALITY_PROMPT, params.productBuffer
? "图4是图2新商品的放大裁剪,仅用于看清结构细节。逐项核对图4与图3中的开放格、分隔板、门板和下缘支撑,增减任何部件都不合格;图2原图仍为最终依据。" : "", params.sourceIdentity
? `生成前识别的新商品结构(仅作核对线索,若与图2矛盾以图2实际图像为准):${params.sourceIdentity.slice(0, 600)}` : ""].filter(Boolean).join("\n");
return inspectLayoutImages([params.templateBuffer, params.sourceBuffer, resultBuffer,
...(params.productBuffer ? [params.productBuffer] : [])], prompt, parseLayoutReplaceQuality);
}