Add batch image jobs runtime

src/ai.ts

@@ -1,14 +1,25 @@
+import { readFile } from "node:fs/promises";
+import { extname, isAbsolute, join, resolve } from "node:path";
+
 import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
 import { streamText } from "ai";
 
 import type { Config } from "@/config";
+import type { ImageJob } from "@/jobs";
 
 type GeneratedImage = {
 	bytes: Uint8Array;
 	mediaType: string;
 };
 
+type InputImagePart = {
+	type: "image";
+	image: Uint8Array;
+	mediaType: string;
+};
+
 const MARKDOWN_IMAGE_REGEX = /!\[[^\]]*\]\((https?:\/\/[^\s)]+)\)/;
+const INPUT_IMAGES_DIR = "images";
 
 function createProvider(config: Config) {
 	return createOpenAICompatible({
@@ -18,6 +29,40 @@ function createProvider(config: Config) {
 	});
 }
 
+function mediaTypeFromPath(filePath: string): string {
+	switch (extname(filePath).toLowerCase()) {
+		case ".png":
+			return "image/png";
+		case ".jpg":
+		case ".jpeg":
+			return "image/jpeg";
+		case ".webp":
+			return "image/webp";
+		case ".gif":
+			return "image/gif";
+		default:
+			return "image/png";
+	}
+}
+
+async function buildImageParts(
+	imagePaths: string[],
+): Promise<InputImagePart[]> {
+	return Promise.all(
+		imagePaths.map(async (imagePath) => ({
+			type: "image" as const,
+			image: new Uint8Array(
+				await readFile(
+					isAbsolute(imagePath)
+						? imagePath
+						: resolve(join(INPUT_IMAGES_DIR, imagePath)),
+				),
+			),
+			mediaType: mediaTypeFromPath(imagePath),
+		})),
+	);
+}
+
 function writeSectionTitle(title: string, alreadyPrinted: boolean): void {
 	if (alreadyPrinted) {
 		return;
@@ -51,12 +96,18 @@ async function downloadImage(url: string): Promise<GeneratedImage> {
 
 export async function streamPrompt(
 	config: Config,
-	prompt: string,
+	job: Pick<ImageJob, "prompt" | "images">,
 ): Promise<GeneratedImage> {
 	const provider = createProvider(config);
+	const imageParts = await buildImageParts(job.images ?? []);
 	const result = streamText({
 		model: provider(config.model),
-		prompt,
+		messages: [
+			{
+				role: "user",
+				content: [{ type: "text", text: job.prompt }, ...imageParts],
+			},
+		],
 	});
 
 	let reasoningStarted = false;

src/image.ts

@@ -1,7 +1,16 @@
 import { mkdir, writeFile } from "node:fs/promises";
 import { join } from "node:path";
 
-const IMAGES_DIR = "images";
+const OUTPUT_DIR = "output";
+
+function sanitizeFileNamePart(value: string): string {
+	return value
+		.trim()
+		.toLowerCase()
+		.replace(/[^a-z0-9]+/g, "-")
+		.replace(/^-+|-+$/g, "")
+		.slice(0, 60);
+}
 
 function mediaTypeToExtension(mediaType: string): string {
 	switch (mediaType) {
@@ -18,19 +27,21 @@ function mediaTypeToExtension(mediaType: string): string {
 	}
 }
 
-function createTimestampFileName(mediaType: string): string {
+function createTimestampFileName(mediaType: string, name?: string): string {
 	const extension = mediaTypeToExtension(mediaType);
 	const timestamp = new Date().toISOString().replace(/[.:]/g, "-");
-	return `${timestamp}.${extension}`;
+	const prefix = name ? `${sanitizeFileNamePart(name)}-` : "";
+	return `${prefix}${timestamp}.${extension}`;
 }
 
 export async function saveImage(
 	bytes: Uint8Array,
 	mediaType: string,
+	name?: string,
 ): Promise<string> {
-	await mkdir(IMAGES_DIR, { recursive: true });
+	await mkdir(OUTPUT_DIR, { recursive: true });
 
-	const filePath = join(IMAGES_DIR, createTimestampFileName(mediaType));
+	const filePath = join(OUTPUT_DIR, createTimestampFileName(mediaType, name));
 	await writeFile(filePath, bytes);
 
 	return filePath;

src/index.ts

@@ -1,27 +1,32 @@
 import { streamPrompt } from "@/ai";
 import { loadConfig } from "@/config";
 import { saveImage } from "@/image";
-
-const DEFAULT_PROMPT =
-	"A cute nano banano mascot in a clean sci-fi lab, bright colors, polished illustration";
-
-function buildPrompt(): string {
-	const cliPrompt = Bun.argv.slice(2).join(" ").trim();
-	return cliPrompt || DEFAULT_PROMPT;
-}
+import { jobs } from "@/jobs";
 
 async function main(): Promise<void> {
 	const config = loadConfig();
-	const prompt = buildPrompt();
+	const enabledJobs = jobs.filter((job) => job.enabled !== false);
 
-	console.log(`Prompt: ${prompt}`);
+	if (enabledJobs.length === 0) {
+		throw new Error("No enabled jobs found.");
+	}
 
-	const image = await streamPrompt(config, prompt);
+	for (const job of enabledJobs) {
+		console.log(`\n=== ${job.name} ===`);
+		console.log(`Prompt: ${job.prompt}`);
+		console.log(`Input images: ${job.images?.length ?? 0}`);
 
-	const savedImagePath = await saveImage(image.bytes, image.mediaType);
+		const image = await streamPrompt(config, job);
 
-	console.log(`\nSaved image: ${savedImagePath}`);
-	console.log(`Image media type: ${image.mediaType}`);
+		const savedImagePath = await saveImage(
+			image.bytes,
+			image.mediaType,
+			job.name,
+		);
+
+		console.log(`\nSaved image: ${savedImagePath}`);
+		console.log(`Image media type: ${image.mediaType}`);
+	}
 }
 
 main().catch((error: unknown) => {

src/jobs.ts

@@ -0,0 +1,16 @@
+export type ImageJob = {
+	name: string;
+	prompt: string;
+	images?: string[];
+	enabled?: boolean;
+};
+
+export const jobs: ImageJob[] = [
+	{
+		name: "nano-banano-test",
+		enabled: true,
+		prompt:
+			"根据角色参考图，设计并绘制一个细化过的插画。画面需包括这四位角色。画面描绘四位角色在一起睡觉的场景，删除帽子、鞋子和披肩，保留袜子。衣服凌乱，闭眼，互相紧靠在一起。注意保留角色的头身比、脸型、服饰细节和原图画风。注意表现角色身高关系：粉发170 两位白发162 棕发151",
+		images: ["Aishia1.png", "Celia1.png", "Latifa1.png", "Sara1.png"],
+	},
+];