imagen/client.ts

import { fetchEventSource } from "@microsoft/fetch-event-source";

const byId = <T extends HTMLElement>(id: string) =>
	document.getElementById(id) as T;
const select = (id: string) => byId<HTMLSelectElement>(id);
const textarea = (id: string) => byId<HTMLTextAreaElement>(id);

const persistedFields = ["size", "prompt"] as const;
for (const f of persistedFields) {
	const el = byId<HTMLSelectElement | HTMLTextAreaElement>(f);
	const saved = localStorage.getItem("aip:" + f);
	if (saved) el.value = saved;
	const save = () => localStorage.setItem("aip:" + f, el.value);
	el.addEventListener("input", save);
	el.addEventListener("change", save);
}

const refImages: string[] = [];
const refPreview = byId<HTMLDivElement>("refPreview");

function renderRefPreview() {
	refPreview.innerHTML = "";
	refImages.forEach((src, i) => {
		const thumb = document.createElement("div");
		thumb.className = "thumb";
		const img = document.createElement("img");
		img.src = src;
		img.alt = "reference " + (i + 1);
		const btn = document.createElement("button");
		btn.type = "button";
		btn.className = "remove";
		btn.textContent = "\u00d7";
		btn.title = "Remove";
		btn.onclick = () => {
			refImages.splice(i, 1);
			renderRefPreview();
		};
		thumb.appendChild(img);
		thumb.appendChild(btn);
		refPreview.appendChild(thumb);
	});
}

byId<HTMLInputElement>("refImages").addEventListener("change", async (e) => {
	const el = e.target as HTMLInputElement;
	const files = Array.from(el.files ?? []);
	for (const file of files) {
		if (!file.type.startsWith("image/")) continue;
		const dataUrl = await new Promise<string>((resolve, reject) => {
			const r = new FileReader();
			r.onload = () => resolve(r.result as string);
			r.onerror = () => reject(r.error);
			r.readAsDataURL(file);
		});
		refImages.push(dataUrl);
	}
	renderRefPreview();
	el.value = "";
});

class FatalError extends Error {}

byId<HTMLButtonElement>("generate").addEventListener("click", async () => {
	const btn = byId<HTMLButtonElement>("generate");
	const status = byId<HTMLSpanElement>("status");
	const result = byId<HTMLDivElement>("result");

	const body = {
		size: select("size").value,
		prompt: textarea("prompt").value.trim(),
		referenceImages: refImages,
	};

	if (!body.prompt) {
		status.textContent = "Please enter a prompt.";
		status.className = "status error";
		return;
	}

	btn.disabled = true;
	status.className = "status";
	status.textContent = "Generating...";
	result.innerHTML = "";

	const preview = document.createElement("img");
	preview.alt = body.prompt;
	let appended = false;
	const showPreview = (src: string) => {
		preview.src = src;
		if (!appended) {
			result.appendChild(preview);
			appended = true;
		}
	};

	const abort = new AbortController();
	try {
		await fetchEventSource("/api/generate", {
			method: "POST",
			headers: { "Content-Type": "application/json" },
			body: JSON.stringify(body),
			signal: abort.signal,
			openWhenHidden: true,
			async onopen(res) {
				const ct = res.headers.get("content-type") ?? "";
				if (res.ok && ct.includes("event-stream")) return;
				const data = (await res.json().catch(() => ({}))) as { error?: string };
				throw new FatalError(data.error ?? "HTTP " + res.status);
			},
			onmessage(ev) {
				if (!ev.data) return;
				const payload = JSON.parse(ev.data) as {
					image?: string;
					index?: number;
					message?: string;
				};
				if (ev.event === "partial" && payload.image) {
					showPreview(payload.image);
					status.textContent =
						"Receiving preview " + ((payload.index ?? 0) + 1) + "...";
				} else if (ev.event === "final" && payload.image) {
					showPreview(payload.image);
					status.textContent = "Done.";
				} else if (ev.event === "error") {
					throw new FatalError(payload.message ?? "Unknown error");
				} else if (ev.event === "done") {
					abort.abort();
				}
			},
			onerror(err) {
				throw err;
			},
		});
	} catch (err) {
		if (!abort.signal.aborted) {
			const message = err instanceof Error ? err.message : String(err);
			status.textContent = "Error: " + message;
			status.className = "status error";
		}
	} finally {
		btn.disabled = false;
	}
});