Files
imagen/index.ts
T

130 lines
3.2 KiB
TypeScript

import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
import { generateImage } from "ai";
import index from "./index.html";
type Size = `${number}x${number}`;
async function generateWithReference({
baseURL,
apiKey,
model,
prompt,
size,
referenceImages,
}: {
baseURL: string;
apiKey: string;
model: string;
prompt: string;
size: Size;
referenceImages: string[];
}): Promise<string[]> {
const form = new FormData();
form.append("model", model);
form.append("prompt", prompt);
form.append("size", size);
for (let i = 0; i < referenceImages.length; i++) {
const dataUrl = referenceImages[i];
if (!dataUrl) continue;
const match = dataUrl.match(/^data:([^;]+);base64,(.+)$/);
if (!match) continue;
const mime = match[1]!;
const b64 = match[2]!;
const bytes = Buffer.from(b64, "base64");
const ext = mime.split("/")[1] ?? "png";
form.append("image", new Blob([bytes], { type: mime }), `ref-${i}.${ext}`);
}
const url = `${baseURL.replace(/\/+$/, "")}/images/edits`;
const res = await fetch(url, {
method: "POST",
headers: { Authorization: `Bearer ${apiKey}` },
body: form,
});
if (!res.ok) {
const text = await res.text().catch(() => "");
throw new Error(`Upstream ${res.status}: ${text || res.statusText}`);
}
const data = (await res.json()) as {
data?: Array<{ b64_json?: string }>;
};
return (data.data ?? [])
.map((item) => (item.b64_json ? `data:image/png;base64,${item.b64_json}` : null))
.filter((s): s is string => s !== null);
}
const server = Bun.serve({
hostname: "0.0.0.0",
routes: {
"/": index,
"/api/generate": {
POST: async (req) => {
try {
const { baseURL, apiKey, model, prompt, size, referenceImages } =
(await req.json()) as {
baseURL?: string;
apiKey?: string;
model?: string;
prompt?: string;
size?: Size;
referenceImages?: string[];
};
if (!baseURL || !apiKey || !model || !prompt) {
return Response.json(
{ error: "baseURL, apiKey, model, prompt are required" },
{ status: 400 },
);
}
if (Array.isArray(referenceImages) && referenceImages.length > 0) {
const images = await generateWithReference({
baseURL,
apiKey,
model,
prompt,
size: size ?? "1024x1024",
referenceImages,
});
return Response.json({ images });
}
const provider = createOpenAICompatible({
name: "custom",
apiKey,
baseURL,
});
const { images } = await generateImage({
model: provider.imageModel(model),
prompt,
size: size || "1024x1024",
});
const out = images.map((img) => {
const mediaType = (img as { mediaType?: string }).mediaType ?? "image/png";
const base64 = (img as { base64?: string }).base64;
return base64 ? `data:${mediaType};base64,${base64}` : null;
}).filter(Boolean);
return Response.json({ images: out });
} catch (err) {
const message = err instanceof Error ? err.message : String(err);
console.error("[generate] error:", err);
return Response.json({ error: message }, { status: 500 });
}
},
},
},
development: {
hmr: true,
console: true,
},
});
console.log(`Listening on ${server.url}`);