112 lines
2.8 KiB
TypeScript
112 lines
2.8 KiB
TypeScript
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);
|
|
}
|