imagen/index.ts

import { Hono } from "hono";
import { streamSSE } from "hono/streaming";
import type { SSEStreamingApi } from "hono/streaming";
import index from "./index.html";

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

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

function requireEnv(name: string): string {
	const v = process.env[name];
	if (!v) {
		throw new Error(`Missing required env: ${name} (see .env.example).`);
	}
	return v;
}

const BASE_URL = requireEnv("BASE_URL");
const API_KEY = requireEnv("API_KEY");
const MODEL = requireEnv("MODEL");

function decodeDataUrl(
	dataUrl: string,
): { bytes: Uint8Array<ArrayBuffer>; mime: string } | null {
	const match = dataUrl.match(/^data:([^;]+);base64,(.+)$/);
	if (!match) return null;
	const mime = match[1]!;
	const binary = atob(match[2]!);
	const bytes = new Uint8Array(new ArrayBuffer(binary.length));
	for (let i = 0; i < binary.length; i++) bytes[i] = binary.charCodeAt(i);
	return { bytes, mime };
}

async function callUpstream(args: {
	prompt: string;
	size: Size;
	referenceImages: string[];
	stream: boolean;
	signal?: AbortSignal;
}): Promise<Response> {
	const { prompt, size, referenceImages, stream, signal } = args;
	const isEdit = referenceImages.length > 0;
	const url = `${BASE_URL.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");
		}
		const imageField = referenceImages.length > 1 ? "image[]" : "image";
		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(
				imageField,
				new Blob([decoded.bytes], { type: decoded.mime }),
				`ref-${i}.${ext}`,
			);
		}
		return fetch(url, {
			method: "POST",
			headers: { Authorization: `Bearer ${API_KEY}` },
			body: form,
			signal,
		});
	}

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

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 emitUpstreamBlock(
	raw: string,
	stream: SSEStreamingApi,
): Promise<void> {
	const block = parseSSEBlock(raw);
	if (!block || 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")) {
		await stream.writeSSE({
			event: "partial",
			data: JSON.stringify({
				image: `data:image/png;base64,${b64}`,
				index: parsed.partial_image_index ?? 0,
			}),
		});
	} else if (type.endsWith(".completed")) {
		await stream.writeSSE({
			event: "final",
			data: JSON.stringify({ image: `data:image/png;base64,${b64}` }),
		});
	}
}

async function forwardUpstreamSSE(
	upstream: Response,
	stream: SSEStreamingApi,
): Promise<void> {
	if (!upstream.body) throw new Error("Upstream returned no body");
	const reader = upstream.body.getReader();
	const decoder = new TextDecoder();
	let buffer = "";
	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) {
			await emitUpstreamBlock(buffer.slice(0, idx), stream);
			buffer = buffer.slice(idx + 2);
		}
	}
	if (buffer.trim().length > 0) await emitUpstreamBlock(buffer, stream);
}

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

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

const app = new Hono();

app.post("/api/generate", async (c) => {
	const body = (await c.req.json()) as GenerateRequest;
	const { prompt, size, referenceImages } = body;
	if (!prompt) {
		return c.json({ error: "prompt is required" }, 400);
	}
	const refs = Array.isArray(referenceImages) ? referenceImages : [];
	const args = {
		prompt,
		size: size ?? ("1024x1024" as Size),
		referenceImages: refs,
	};

	return streamSSE(c, async (stream) => {
		const abort = new AbortController();
		stream.onAbort(() => abort.abort());

		await stream.write(": connected\n\n");
		const keepalive = setInterval(() => {
			stream.write(": keepalive\n\n").catch(() => {});
		}, 15_000);

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

			if (!upstream.ok && upstream.status === 400) {
				const errText = await upstream.text().catch(() => "");
				if (isStreamingUnsupportedError(errText)) {
					upstream = await callUpstream({
						...args,
						stream: false,
						signal: abort.signal,
					});
				} 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, stream);
			} else {
				await forwardUpstreamJSON(upstream, stream);
			}

			await stream.writeSSE({ event: "done", data: "" });
		} catch (err) {
			if (abort.signal.aborted) return;
			const message = err instanceof Error ? err.message : String(err);
			console.error("[generate] error:", err);
			await stream.writeSSE({
				event: "error",
				data: JSON.stringify({ message }),
			});
		} finally {
			clearInterval(keepalive);
		}
	});
});

const server = Bun.serve({
	hostname: "0.0.0.0",
	idleTimeout: 255,
	routes: {
		"/": index,
	},
	fetch: app.fetch,
	development: {
		hmr: true,
		console: true,
		chromeDevToolsAutomaticWorkspaceFolders: false,
	},
});

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