ai-chatbot/scripts/verify-canvas-layout-replace-quality.ts

import { mkdir, writeFile } from "node:fs/promises";
import { resolve } from "node:path";
import dotenv from "dotenv";
import postgres from "postgres";
import { createAtlasCloud } from "@/lib/ai/providers/atlascloud";
import { createSnapgenImage } from "@/lib/ai/providers/snapgen-image";
import { createVip98k } from "@/lib/ai/providers/vip98k";
import { resolveModelRuntimeId } from "@/lib/ai/model-runtime";
import { buildCanvasLayoutReplacePrompt } from "@/lib/canvas/layout-replace-model-selection";
import { resolveCanvasLayoutReplaceSourceProfile } from "@/lib/canvas/layout-replace-source-profile";

dotenv.config({ path: ".env.local" });
dotenv.config();

function getArgument(name: string) {
  const prefix = `--${name}=`;
  return process.argv.find((argument) => argument.startsWith(prefix))?.slice(prefix.length) || "";
}

function requireHttpArgument(name: string) {
  const value = getArgument(name).trim();
  if (!/^https?:\/\//i.test(value)) {
    throw new Error(`缺少有效参数 --${name}=https://...`);
  }
  return value;
}

function extractOriginalImageUrl(markdown: string) {
  const match = markdown.match(/!\[thumb\|\|\|([^|\]]+)(?:\|\|\|[^\]]*)?\]\([^)]+\)/);
  if (match?.[1] && /^https?:\/\//i.test(match[1])) return match[1];
  const direct = markdown.match(/!\[[^\]]*\]\((https?:\/\/[^)]+)\)/);
  return direct?.[1] || "";
}

async function main() {
  const databaseUrl = process.env.POSTGRES_URL;
  if (!databaseUrl) throw new Error("POSTGRES_URL 未配置");

  const sourceImageUrl = requireHttpArgument("source");
  const referenceImageUrl = requireHttpArgument("reference");
  const modelRef = getArgument("model-ref").trim();
  if (!modelRef) throw new Error("缺少 --model-ref");
  const outputDirectory = resolve(
    getArgument("output-dir") || "/tmp/canvas-layout-replace-quality",
  );

  const sql = postgres(databaseUrl, { max: 1 });
  try {
    const [candidate] = await sql<
      Array<{
        apiKey: string;
        baseUrl: string;
        capabilities: Record<string, unknown> | null;
        imageGenConfig: Record<string, unknown> | null;
        modelId: string;
        providerType: string;
      }>
    >`
      select
        p."apiKey" as "apiKey",
        p."baseUrl" as "baseUrl",
        p.type as "providerType",
        m."modelId" as "modelId",
        m."imageGenConfig" as "imageGenConfig",
        m.capabilities as "capabilities"
      from "AiModel" m
      join "AiProvider" p on p.id = m."providerId"
      where m.id = ${modelRef}
        and m."isEnabled" = true
        and p."isEnabled" = true
      limit 1
    `;
    if (!candidate) throw new Error("找不到启用的测试模型");
    if (!["atlascloud", "snapgen-image", "vip98k"].includes(candidate.providerType)) {
      throw new Error(`本地质量脚本不支持供应商 ${candidate.providerType}`);
    }

    const sourceProfile =
      await resolveCanvasLayoutReplaceSourceProfile(sourceImageUrl);
    const prompt = buildCanvasLayoutReplacePrompt({
      referenceHeight: 726,
      referenceWidth: 543,
      sourceProductAspectRatio: sourceProfile?.silhouetteAspectRatio,
    });
    const runtimeModelId = resolveModelRuntimeId({ modelId: candidate.modelId });
    const generationOptions = {
      messages: [
        {
          role: "user",
          content: [
            { type: "text", text: prompt },
            { type: "image_url", image_url: { url: sourceImageUrl } },
            { type: "image_url", image_url: { url: referenceImageUrl } },
          ],
        },
      ],
      imageParams: {
        aspectRatio: "3:4",
        imageCount: 1,
        imageSize: "2K",
        promptExtend: false,
        quality: "high",
      },
    };
    let result: { text?: string };
    if (candidate.providerType === "atlascloud") {
      const model = createAtlasCloud({
        apiKey: candidate.apiKey,
        baseURL: candidate.baseUrl,
      }).chat(runtimeModelId);
      result = await model.doGenerate({
        ...generationOptions,
        imageGenConfig:
          candidate.imageGenConfig as Parameters<
            typeof model.doGenerate
          >[0]["imageGenConfig"],
      });
    } else if (candidate.providerType === "snapgen-image") {
      const model = createSnapgenImage({
        apiKey: candidate.apiKey,
        baseURL: candidate.baseUrl,
      }).chat(runtimeModelId);
      result = await model.doGenerate({
        ...generationOptions,
        imageGenConfig:
          candidate.imageGenConfig as Parameters<
            typeof model.doGenerate
          >[0]["imageGenConfig"],
      });
    } else {
      const model = createVip98k({
        apiKey: candidate.apiKey,
        baseURL: candidate.baseUrl,
      }).chat(runtimeModelId, {
        capabilities: candidate.capabilities,
        imageGenConfig: candidate.imageGenConfig,
      });
      result = await model.doGenerate({
        ...generationOptions,
        imageGenConfig: candidate.imageGenConfig,
      });
    }
    const resultUrl = extractOriginalImageUrl(result.text || "");
    if (!resultUrl) throw new Error("模型返回中没有结果图片");

    const response = await fetch(resultUrl, {
      signal: AbortSignal.timeout(120_000),
    });
    if (!response.ok) throw new Error(`下载测试结果失败（HTTP ${response.status}）`);
    await mkdir(outputDirectory, { recursive: true });
    const outputPath = resolve(outputDirectory, `result-${Date.now()}.png`);
    await writeFile(outputPath, Buffer.from(await response.arrayBuffer()));
    process.stdout.write(
      `${JSON.stringify({ modelId: runtimeModelId, outputPath, resultUrl })}\n`,
    );
  } finally {
    await sql.end();
  }
}

main().catch((error) => {
  process.stderr.write(`${error instanceof Error ? error.message : String(error)}\n`);
  process.exit(1);
});