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

63 lines
3.1 KiB
TypeScript

import "server-only";
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, parseLayoutReplaceProductBox } from "./layout-replace-product-crop";
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, parseLayoutReplaceProductBox);
}
export async function verifyCanvasLayoutReplacement(params: {
sourceBuffer: Buffer;
templateBuffer: Buffer;
resultUrl: string;
resultBuffer?: Buffer;
}) {
const resultBuffer = params.resultBuffer || await fetchImageBufferFromUrl(params.resultUrl);
return inspectLayoutImages([params.templateBuffer, params.sourceBuffer, resultBuffer],
LAYOUT_REPLACE_QUALITY_PROMPT, parseLayoutReplaceQuality);
}