diff --git a/skills/baoyu-image-gen/scripts/providers/openrouter.test.ts b/skills/baoyu-image-gen/scripts/providers/openrouter.test.ts new file mode 100644 index 0000000..eda5437 --- /dev/null +++ b/skills/baoyu-image-gen/scripts/providers/openrouter.test.ts @@ -0,0 +1,131 @@ +import assert from "node:assert/strict"; +import test from "node:test"; + +import type { CliArgs } from "../types.ts"; +import { + buildContent, + buildRequestBody, + extractImageFromResponse, + getAspectRatio, + getImageSize, + validateArgs, +} from "./openrouter.ts"; + +const GEMINI_MODEL = "google/gemini-3.1-flash-image-preview"; +const GEMINI_25_MODEL = "google/gemini-2.5-flash-image-preview"; +const FLUX_MODEL = "black-forest-labs/flux.2-pro"; + +function makeArgs(overrides: Partial = {}): CliArgs { + return { + prompt: null, + promptFiles: [], + imagePath: null, + provider: null, + model: null, + aspectRatio: null, + size: null, + quality: null, + imageSize: null, + referenceImages: [], + n: 1, + batchFile: null, + jobs: null, + json: false, + help: false, + ...overrides, + }; +} + +test("OpenRouter request body uses image_config and string content for text-only prompts", () => { + const args = makeArgs({ aspectRatio: "16:9", quality: "2k" }); + const body = buildRequestBody("hello", GEMINI_MODEL, args, []); + + assert.deepEqual(body.image_config, { + image_size: "2K", + aspect_ratio: "16:9", + }); + assert.deepEqual(body.provider, { + require_parameters: true, + }); + assert.deepEqual(body.modalities, ["image", "text"]); + assert.equal(body.stream, false); + assert.equal(body.messages[0].content, "hello"); +}); + +test("OpenRouter request body omits image_config and provider guard when no image options are requested", () => { + const body = buildRequestBody("hello", FLUX_MODEL, makeArgs(), []); + + assert.equal(body.image_config, undefined); + assert.equal(body.provider, undefined); + assert.deepEqual(body.modalities, ["image"]); + assert.equal(body.stream, false); + assert.equal(body.messages[0].content, "hello"); +}); + +test("OpenRouter request body keeps multimodal array content when references are provided", () => { + const content = buildContent("hello", ["data:image/png;base64,abc"]); + assert.ok(Array.isArray(content)); + assert.deepEqual(content[0], { type: "text", text: "hello" }); + assert.deepEqual(content[1], { + type: "image_url", + image_url: { url: "data:image/png;base64,abc" }, + }); +}); + +test("OpenRouter size and aspect helpers infer expected defaults", () => { + assert.equal(getImageSize(makeArgs()), null); + assert.equal(getImageSize(makeArgs({ quality: "normal" })), "1K"); + assert.equal(getImageSize(makeArgs({ size: "2048x1024" })), "2K"); + assert.equal(getAspectRatio(GEMINI_MODEL, makeArgs({ size: "2048x1024" })), null); + assert.equal(getAspectRatio(GEMINI_MODEL, makeArgs({ size: "1600x900" })), "16:9"); + assert.equal(getAspectRatio(GEMINI_MODEL, makeArgs({ size: "1024x4096" })), "1:4"); + assert.equal(getAspectRatio(GEMINI_25_MODEL, makeArgs({ size: "1024x4096" })), null); + assert.equal(getAspectRatio(FLUX_MODEL, makeArgs({ size: "1024x4096" })), null); +}); + +test("OpenRouter validates explicit aspect ratios against model support", () => { + assert.doesNotThrow(() => + validateArgs(GEMINI_MODEL, makeArgs({ aspectRatio: "1:4" })), + ); + assert.throws( + () => validateArgs(GEMINI_25_MODEL, makeArgs({ aspectRatio: "1:4" })), + /does not support aspect ratio 1:4/, + ); + assert.throws( + () => validateArgs(FLUX_MODEL, makeArgs({ aspectRatio: "1:4" })), + /does not support aspect ratio 1:4/, + ); +}); + +test("OpenRouter response extraction supports inline image data and finish_reason errors", async () => { + const bytes = await extractImageFromResponse({ + choices: [ + { + message: { + images: [ + { + image_url: { + url: `data:image/png;base64,${Buffer.from("hello").toString("base64")}`, + }, + }, + ], + }, + }, + ], + }); + assert.equal(Buffer.from(bytes).toString("utf8"), "hello"); + + await assert.rejects( + () => + extractImageFromResponse({ + choices: [ + { + finish_reason: "error", + native_finish_reason: "MALFORMED_FUNCTION_CALL", + message: { content: null }, + }, + ], + }), + /finish_reason=MALFORMED_FUNCTION_CALL/, + ); +}); diff --git a/skills/baoyu-image-gen/scripts/providers/openrouter.ts b/skills/baoyu-image-gen/scripts/providers/openrouter.ts index 4e55e06..86e0db7 100644 --- a/skills/baoyu-image-gen/scripts/providers/openrouter.ts +++ b/skills/baoyu-image-gen/scripts/providers/openrouter.ts @@ -3,6 +3,19 @@ import { readFile } from "node:fs/promises"; import type { CliArgs } from "../types"; const DEFAULT_MODEL = "google/gemini-3.1-flash-image-preview"; +const COMMON_ASPECT_RATIOS = [ + "1:1", + "2:3", + "3:2", + "3:4", + "4:3", + "4:5", + "5:4", + "9:16", + "16:9", + "21:9", +]; +const GEMINI_EXTENDED_ASPECT_RATIOS = ["1:4", "4:1", "1:8", "8:1"]; type OpenRouterImageEntry = { image_url?: string | { url?: string | null } | null; @@ -18,9 +31,11 @@ type OpenRouterMessagePart = { type OpenRouterResponse = { choices?: Array<{ + finish_reason?: string | null; + native_finish_reason?: string | null; message?: { images?: OpenRouterImageEntry[]; - content?: string | OpenRouterMessagePart[]; + content?: string | OpenRouterMessagePart[] | null; }; }>; }; @@ -29,6 +44,24 @@ export function getDefaultModel(): string { return process.env.OPENROUTER_IMAGE_MODEL || DEFAULT_MODEL; } +function normalizeModelId(model: string): string { + return model.trim().toLowerCase().split(":")[0]!; +} + +export function isGeminiImageModel(model: string): boolean { + const normalized = normalizeModelId(model); + return normalized.startsWith("google/gemini-") && normalized.includes("image-preview"); +} + +function getSupportedAspectRatios(model: string): Set { + const normalized = normalizeModelId(model); + if (normalized !== "google/gemini-3.1-flash-image-preview") { + return new Set(COMMON_ASPECT_RATIOS); + } + + return new Set([...COMMON_ASPECT_RATIOS, ...GEMINI_EXTENDED_ASPECT_RATIOS]); +} + function getApiKey(): string | null { return process.env.OPENROUTER_API_KEY || null; } @@ -103,17 +136,45 @@ function inferImageSize(size: string | null): "1K" | "2K" | "4K" | null { return "4K"; } -function getImageSize(args: CliArgs): "1K" | "2K" | "4K" { +export function getImageSize(args: CliArgs): "1K" | "2K" | "4K" | null { if (args.imageSize) return args.imageSize as "1K" | "2K" | "4K"; const inferredFromSize = inferImageSize(args.size); if (inferredFromSize) return inferredFromSize; - return args.quality === "normal" ? "1K" : "2K"; + if (args.quality === "normal") return "1K"; + if (args.quality === "2k") return "2K"; + return null; } -function getAspectRatio(args: CliArgs): string | null { - return args.aspectRatio || inferAspectRatio(args.size); +export function getAspectRatio(model: string, args: CliArgs): string | null { + if (args.aspectRatio) return args.aspectRatio; + + const inferred = inferAspectRatio(args.size); + if (!inferred || !getSupportedAspectRatios(model).has(inferred)) { + return null; + } + + return inferred; +} + +function getModalities(model: string): string[] { + return isGeminiImageModel(model) ? ["image", "text"] : ["image"]; +} + +export function validateArgs(model: string, args: CliArgs): void { + if (!args.aspectRatio) { + return; + } + + const supported = getSupportedAspectRatios(model); + if (supported.has(args.aspectRatio)) { + return; + } + + throw new Error( + `OpenRouter model ${model} does not support aspect ratio ${args.aspectRatio}. Supported values: ${Array.from(supported).join(", ")}` + ); } function getMimeType(filename: string): string { @@ -129,7 +190,14 @@ async function readImageAsDataUrl(filePath: string): Promise { return `data:${getMimeType(filePath)};base64,${bytes.toString("base64")}`; } -function buildContent(prompt: string, referenceImages: string[]): Array> { +export function buildContent( + prompt: string, + referenceImages: string[], +): string | Array> { + if (referenceImages.length === 0) { + return prompt; + } + const content: Array> = [{ type: "text", text: prompt }]; for (const imageUrl of referenceImages) { @@ -171,8 +239,9 @@ async function downloadImage(value: string): Promise { return Uint8Array.from(Buffer.from(value, "base64")); } -async function extractImageFromResponse(result: OpenRouterResponse): Promise { - const message = result.choices?.[0]?.message; +export async function extractImageFromResponse(result: OpenRouterResponse): Promise { + const choice = result.choices?.[0]; + const message = choice?.message; for (const image of message?.images ?? []) { const imageUrl = extractImageUrl(image); @@ -194,7 +263,50 @@ async function extractImageFromResponse(result: OpenRouterResponse): Promise { + const imageConfig: Record = {}; + + const imageSize = getImageSize(args); + if (imageSize) { + imageConfig.image_size = imageSize; + } + + const aspectRatio = getAspectRatio(model, args); + if (aspectRatio) { + imageConfig.aspect_ratio = aspectRatio; + } + + const body: Record = { + messages: [ + { + role: "user", + content: buildContent(prompt, referenceImages), + }, + ], + modalities: getModalities(model), + stream: false, + }; + + if (Object.keys(imageConfig).length > 0) { + body.image_config = imageConfig; + body.provider = { + require_parameters: true, + }; + } + + return body; } export async function generateImage( @@ -212,32 +324,15 @@ export async function generateImage( referenceImages.push(await readImageAsDataUrl(refPath)); } - const imageGenerationOptions: Record = { - size: getImageSize(args), - }; - - const aspectRatio = getAspectRatio(args); - if (aspectRatio) { - imageGenerationOptions.aspect_ratio = aspectRatio; - } - const body = { model, - messages: [ - { - role: "user", - content: buildContent(prompt, referenceImages), - }, - ], - modalities: ["image", "text"], - max_tokens: 256, - imageGenerationOptions, - providerPreferences: { - require_parameters: true, - }, + ...buildRequestBody(prompt, model, args, referenceImages), }; - console.log(`Generating image with OpenRouter (${model})...`, imageGenerationOptions); + console.log( + `Generating image with OpenRouter (${model})...`, + (body.image_config as Record), + ); const response = await fetch(`${getBaseUrl()}/chat/completions`, { method: "POST",