集成即梦生图和豆包生图

2026-07-31 06:09:47 +08:00 · 2026-03-14 00:37:46 +08:00
parent c62cda4424
commit 32003da694
4 changed files with 578 additions and 11 deletions
@@ -0,0 +1,111 @@
+import type { CliArgs } from "../types";
+
+export function getDefaultModel(): string {
+  return process.env.SEEDREAM_IMAGE_MODEL || "doubao-seedream-5-0-260128";
+}
+
+function getApiKey(): string | null {
+  return process.env.ARK_API_KEY || null;
+}
+
+function getBaseUrl(): string {
+  return process.env.SEEDREAM_BASE_URL || "https://ark.cn-beijing.volces.com/api/v3";
+}
+
+/**
+ * Convert aspect ratio to Seedream size format
+ * Seedream API accepts: "2k" (default), "3k", or WIDTHxHEIGHT format
+ * Note: API uses lowercase "2k"/"3k", not "2K"/"3K"
+ */
+function getSeedreamSize(ar: string | null, quality: CliArgs["quality"], imageSize?: string | null): string {
+  // If explicit size is provided
+  if (imageSize) {
+    const upper = imageSize.toUpperCase();
+    if (upper === "2K" || upper === "3K") {
+      return upper.toLowerCase(); // API expects "2k" or "3k"
+    }
+    // For widthxheight format, pass through as-is
+    if (imageSize.includes("x")) {
+      return imageSize;
+    }
+  }
+
+  // Default to 2k (smallest option supported by API)
+  return "2k";
+}
+
+type SeedreamImageResponse = {
+  model: string;
+  created: number;
+  data: Array<{
+    url: string;
+    size: string;
+  }>;
+  usage: {
+    generated_images: number;
+    output_tokens: number;
+    total_tokens: number;
+  };
+};
+
+export async function generateImage(
+  prompt: string,
+  model: string,
+  args: CliArgs
+): Promise<Uint8Array> {
+  const apiKey = getApiKey();
+  if (!apiKey) {
+    throw new Error(
+      "ARK_API_KEY is required. " +
+        "Get your API key from https://console.volcengine.com/ark"
+    );
+  }
+
+  const baseUrl = getBaseUrl();
+  const size = getSeedreamSize(args.aspectRatio, args.quality, args.imageSize);
+
+  console.error(`Calling Seedream API (${model}) with size: ${size}`);
+
+  const requestBody = {
+    model,
+    prompt,
+    size,
+    output_format: "png",
+    watermark: false,
+  };
+
+  const response = await fetch(`${baseUrl}/images/generations`, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+      Authorization: `Bearer ${apiKey}`,
+    },
+    body: JSON.stringify(requestBody),
+  });
+
+  if (!response.ok) {
+    const err = await response.text();
+    throw new Error(`Seedream API error (${response.status}): ${err}`);
+  }
+
+  const result = (await response.json()) as SeedreamImageResponse;
+
+  if (!result.data || result.data.length === 0) {
+    throw new Error("No image data in Seedream response");
+  }
+
+  const imageUrl = result.data[0].url;
+  if (!imageUrl) {
+    throw new Error("No image URL in Seedream response");
+  }
+
+  // Download image from URL
+  console.error(`Downloading image from ${imageUrl}...`);
+  const imgResponse = await fetch(imageUrl);
+  if (!imgResponse.ok) {
+    throw new Error(`Failed to download image from ${imageUrl}`);
+  }
+
+  const buffer = await imgResponse.arrayBuffer();
+  return new Uint8Array(buffer);
+}