imagen/index.ts

import index from "./index.html";

type Size = `${number}x${number}`;

type GenerateRequest = {
	baseURL?: string;
	apiKey?: string;
	model?: string;
	prompt?: string;
	size?: Size;
	referenceImages?: string[];
};

type SSEController = ReadableStreamDefaultController<Uint8Array>;

const encoder = new TextEncoder();

function sseEvent(controller: SSEController, event: string, data: unknown): void {
	controller.enqueue(
		encoder.encode(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`),
	);
}

function sseComment(controller: SSEController, text: string): void {
	controller.enqueue(encoder.encode(`: ${text}\n\n`));
}

function decodeDataUrl(dataUrl: string): { bytes: Buffer; mime: string } | null {
	const match = dataUrl.match(/^data:([^;]+);base64,(.+)$/);
	if (!match) return null;
	const mime = match[1]!;
	const b64 = match[2]!;
	return { bytes: Buffer.from(b64, "base64"), mime };
}

async function callUpstream(args: {
	baseURL: string;
	apiKey: string;
	model: string;
	prompt: string;
	size: Size;
	referenceImages: string[];
	stream: boolean;
}): Promise<Response> {
	const { baseURL, apiKey, model, prompt, size, referenceImages, stream } = args;
	const isEdit = referenceImages.length > 0;
	const url = `${baseURL.replace(/\/+$/, "")}/images/${isEdit ? "edits" : "generations"}`;

	if (isEdit) {
		const form = new FormData();
		form.append("model", model);
		form.append("prompt", prompt);
		form.append("size", size);
		if (stream) {
			form.append("stream", "true");
			form.append("partial_images", "2");
		}
		for (let i = 0; i < referenceImages.length; i++) {
			const dataUrl = referenceImages[i];
			if (!dataUrl) continue;
			const decoded = decodeDataUrl(dataUrl);
			if (!decoded) continue;
			const ext = decoded.mime.split("/")[1] ?? "png";
			form.append(
				"image",
				new Blob([decoded.bytes], { type: decoded.mime }),
				`ref-${i}.${ext}`,
			);
		}
		return fetch(url, {
			method: "POST",
			headers: { Authorization: `Bearer ${apiKey}` },
			body: form,
		});
	}

	const body: Record<string, unknown> = { model, prompt, size };
	if (stream) {
		body.stream = true;
		body.partial_images = 2;
	}
	return fetch(url, {
		method: "POST",
		headers: {
			Authorization: `Bearer ${apiKey}`,
			"Content-Type": "application/json",
		},
		body: JSON.stringify(body),
	});
}

function parseSSEBlock(raw: string): { event: string; data: string } | null {
	let eventName = "message";
	const dataLines: string[] = [];
	for (const line of raw.split("\n")) {
		if (line.startsWith(":")) continue;
		if (line.startsWith("event:")) eventName = line.slice(6).trim();
		else if (line.startsWith("data:")) dataLines.push(line.slice(5).trim());
	}
	if (dataLines.length === 0) return null;
	return { event: eventName, data: dataLines.join("\n") };
}

async function forwardUpstreamSSE(
	upstream: Response,
	controller: SSEController,
): Promise<void> {
	if (!upstream.body) throw new Error("Upstream returned no body");
	const reader = upstream.body.getReader();
	const decoder = new TextDecoder();
	let buffer = "";

	const handle = (raw: string) => {
		const block = parseSSEBlock(raw);
		if (!block) return;
		if (block.data === "[DONE]") return;
		let parsed: {
			type?: string;
			b64_json?: string;
			partial_image_index?: number;
		};
		try {
			parsed = JSON.parse(block.data);
		} catch {
			return;
		}
		const type = parsed.type ?? block.event;
		const b64 = parsed.b64_json;
		if (!b64) return;
		if (type.endsWith(".partial_image")) {
			sseEvent(controller, "partial", {
				image: `data:image/png;base64,${b64}`,
				index: parsed.partial_image_index ?? 0,
			});
		} else if (type.endsWith(".completed")) {
			sseEvent(controller, "final", {
				image: `data:image/png;base64,${b64}`,
			});
		}
	};

	while (true) {
		const { value, done } = await reader.read();
		if (done) break;
		buffer += decoder.decode(value, { stream: true });
		let idx: number;
		while ((idx = buffer.indexOf("\n\n")) !== -1) {
			handle(buffer.slice(0, idx));
			buffer = buffer.slice(idx + 2);
		}
	}
	if (buffer.trim().length > 0) handle(buffer);
}

async function forwardUpstreamJSON(
	upstream: Response,
	controller: SSEController,
): Promise<void> {
	const data = (await upstream.json()) as {
		data?: Array<{ b64_json?: string }>;
	};
	for (const item of data.data ?? []) {
		if (!item.b64_json) continue;
		sseEvent(controller, "final", {
			image: `data:image/png;base64,${item.b64_json}`,
		});
	}
}

function isStreamingUnsupportedError(errText: string): boolean {
	return /\b(stream|partial_images)\b/i.test(errText);
}

const server = Bun.serve({
	hostname: "0.0.0.0",
	routes: {
		"/": index,
		"/api/generate": {
			POST: async (req) => {
				const body = (await req.json()) as GenerateRequest;
				const { baseURL, apiKey, model, prompt, size, referenceImages } = body;
				if (!baseURL || !apiKey || !model || !prompt) {
					return Response.json(
						{ error: "baseURL, apiKey, model, prompt are required" },
						{ status: 400 },
					);
				}
				const refs = Array.isArray(referenceImages) ? referenceImages : [];
				const args = {
					baseURL,
					apiKey,
					model,
					prompt,
					size: size ?? ("1024x1024" as Size),
					referenceImages: refs,
				};

				const stream = new ReadableStream<Uint8Array>({
					async start(controller) {
						const keepalive = setInterval(() => {
							try {
								sseComment(controller, "keepalive");
							} catch {}
						}, 20_000);

						try {
							let upstream = await callUpstream({ ...args, stream: true });

							if (!upstream.ok && upstream.status === 400) {
								const errText = await upstream.text().catch(() => "");
								if (isStreamingUnsupportedError(errText)) {
									upstream = await callUpstream({ ...args, stream: false });
								} else {
									throw new Error(`Upstream 400: ${errText || upstream.statusText}`);
								}
							}

							if (!upstream.ok) {
								const errText = await upstream.text().catch(() => "");
								throw new Error(
									`Upstream ${upstream.status}: ${errText || upstream.statusText}`,
								);
							}

							const contentType = upstream.headers.get("content-type") ?? "";
							if (contentType.includes("event-stream")) {
								await forwardUpstreamSSE(upstream, controller);
							} else {
								await forwardUpstreamJSON(upstream, controller);
							}

							sseEvent(controller, "done", {});
						} catch (err) {
							const message = err instanceof Error ? err.message : String(err);
							console.error("[generate] error:", err);
							try {
								sseEvent(controller, "error", { message });
							} catch {}
						} finally {
							clearInterval(keepalive);
							try {
								controller.close();
							} catch {}
						}
					},
				});

				return new Response(stream, {
					headers: {
						"Content-Type": "text/event-stream",
						"Cache-Control": "no-cache, no-transform",
						"X-Accel-Buffering": "no",
						Connection: "keep-alive",
					},
				});
			},
		},
	},
	development: {
		hmr: true,
		console: true,
	},
});

console.log(`Listening on ${server.url}`);