feat(baoyu-image-gen): add Agnes AI image generation provider

2026-08-07 17:33:03 +08:00 · 2026-06-06 04:53:00 +08:00
parent ce84174bf7
commit ad7a7a646d
5 changed files with 499 additions and 9 deletions
@@ -0,0 +1,189 @@
+import { readFile } from "node:fs/promises";
+import path from "node:path";
+import type { CliArgs } from "../types";
+
+const DEFAULT_MODEL = "agnes-image-2.1-flash";
+const DEFAULT_BASE_URL = "https://apihub.agnes-ai.com/v1";
+const DEFAULT_SIZE = "1024x1024";
+
+type AgnesResponse = {
+  created?: number;
+  data: Array<{ url?: string; b64_json?: string }>;
+};
+
+export function getDefaultModel(): string {
+  return process.env.AGNES_IMAGE_MODEL || DEFAULT_MODEL;
+}
+
+function getApiKey(): string {
+  const key = process.env.AGNES_API_KEY;
+  if (!key) {
+    throw new Error("AGNES_API_KEY is required. Get one from https://apihub.agnes-ai.com.");
+  }
+  return key;
+}
+
+function getBaseUrl(): string {
+  return (process.env.AGNES_BASE_URL || DEFAULT_BASE_URL).replace(/\/+$/, "");
+}
+
+export function parseAspectRatio(ar: string): { width: number; height: number } | null {
+  const match = ar.match(/^(\d+(?:\.\d+)?):(\d+(?:\.\d+)?)$/);
+  if (!match) return null;
+  const w = parseFloat(match[1]!);
+  const h = parseFloat(match[2]!);
+  if (w <= 0 || h <= 0) return null;
+  return { width: w, height: h };
+}
+
+function gcd(a: number, b: number): number {
+  let x = Math.abs(Math.round(a));
+  let y = Math.abs(Math.round(b));
+  while (y !== 0) {
+    const next = x % y;
+    x = y;
+    y = next;
+  }
+  return x || 1;
+}
+
+export function snapDim(n: number): number {
+  return Math.max(32, Math.round(n / 32) * 32);
+}
+
+export function resolveSize(args: Pick<CliArgs, "size" | "aspectRatio">): string {
+  if (args.size) return args.size;
+
+  if (args.aspectRatio) {
+    const parsed = parseAspectRatio(args.aspectRatio);
+    if (parsed) {
+      const g = gcd(parsed.width, parsed.height);
+      const rw = Math.round(parsed.width / g);
+      const rh = Math.round(parsed.height / g);
+      if (rw === 1 && rh === 1) return "1024x1024";
+      const maxEdge = 2048;
+      const scale = Math.max(1, Math.floor(maxEdge / Math.max(rw, rh)));
+      const width = rw * scale;
+      const height = rh * scale;
+      return `${snapDim(width)}x${snapDim(height)}`;
+    }
+  }
+
+  return DEFAULT_SIZE;
+}
+
+function isRemoteUrl(refPath: string): boolean {
+  return /^https?:\/\//i.test(refPath);
+}
+
+export async function resolveReferenceImages(
+  referenceImages: string[]
+): Promise<string[]> {
+  const result: string[] = [];
+  for (const refPath of referenceImages) {
+    if (isRemoteUrl(refPath)) {
+      result.push(refPath);
+      continue;
+    }
+    const bytes = await readFile(refPath);
+    const ext = path.extname(refPath).toLowerCase();
+    let mime = "image/png";
+    if (ext === ".jpg" || ext === ".jpeg") mime = "image/jpeg";
+    else if (ext === ".webp") mime = "image/webp";
+    else if (ext === ".gif") mime = "image/gif";
+    const b64 = Buffer.from(bytes).toString("base64");
+    result.push(`data:${mime};base64,${b64}`);
+  }
+  return result;
+}
+
+export function validateArgs(_model: string, args: CliArgs): void {
+  if (args.n > 1) {
+    throw new Error("Agnes image generation currently returns a single image per request. Set --n 1 or omit --n.");
+  }
+}
+
+export function getDefaultOutputExtension(_model: string, args: CliArgs): string {
+  return args.responseFormat === "url" ? ".txt" : ".png";
+}
+
+export function buildRequestBody(
+  prompt: string,
+  model: string,
+  args: Pick<CliArgs, "size" | "aspectRatio" | "referenceImages">
+): Record<string, unknown> {
+  const body: Record<string, unknown> = {
+    model,
+    prompt,
+    size: resolveSize(args),
+  };
+
+  if (args.referenceImages.length > 0) {
+    body.image = args.referenceImages;
+  }
+
+  body.extra_body = { response_format: "url" };
+
+  return body;
+}
+
+export async function extractImageFromResponse(result: AgnesResponse): Promise<Uint8Array> {
+  const img = result.data[0];
+
+  if (img?.b64_json) {
+    return Uint8Array.from(Buffer.from(img.b64_json, "base64"));
+  }
+
+  if (img?.url) {
+    const imgRes = await fetch(img.url);
+    if (!imgRes.ok) throw new Error(`Failed to download image from Agnes: ${imgRes.status}`);
+    return new Uint8Array(await imgRes.arrayBuffer());
+  }
+
+  throw new Error("No image in Agnes response");
+}
+
+export async function generateImage(
+  prompt: string,
+  model: string,
+  args: CliArgs
+): Promise<Uint8Array> {
+  const apiKey = getApiKey();
+  const baseUrl = getBaseUrl();
+
+  const referenceImages = await resolveReferenceImages(args.referenceImages);
+
+  const body = buildRequestBody(prompt, model, { ...args, referenceImages });
+
+  const controller = new AbortController();
+  const timeout = setTimeout(() => controller.abort(), 120_000);
+
+  try {
+    const res = await fetch(`${baseUrl}/images/generations`, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        Authorization: `Bearer ${apiKey}`,
+      },
+      body: JSON.stringify(body),
+      signal: controller.signal,
+    });
+
+    if (!res.ok) {
+      const err = await res.text();
+      throw new Error(`Agnes API error (${res.status}): ${err}`);
+    }
+
+    const result = (await res.json()) as AgnesResponse;
+
+    if (args.responseFormat === "url") {
+      const url = result.data[0]?.url;
+      if (!url) throw new Error("No URL in Agnes response");
+      return new Uint8Array(Buffer.from(url, "utf-8"));
+    }
+
+    return extractImageFromResponse(result);
+  } finally {
+    clearTimeout(timeout);
+  }
+}