From 1ba57577fb31de59be03823b5bb039dc40b83893 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 5 Sep 2026 14:52:37 +0000 Subject: [PATCH] Add Workers AI model catalogue and multi-shape inference adapter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit runWorkersAi assumed one contract — {prompt, seed} in, binary image bytes out — and threw `unexpected Workers AI response shape` on anything else. Workers AI models do not share that contract, and the differences don't follow vendor or task: - flux-1-schnell takes `steps` and returns JSON {image: }, so the fastest model in the catalogue failed at render time, after the renderer had already been debited. - The Stable Diffusion family takes `num_steps` and returns a ReadableStream. - flux-2-* take multipart form-data, even for a prompt-only call. - Leonardo's two models disagree with each other: phoenix-1.0 streams bytes, lucid-origin returns base64 JSON. - Vision and instruct models return text, which had no output kind. Adds src/ai/workersAiCatalogue.ts: one row per allowlisted model recording its request shape, response shape, content type, defaults and parameter allowlist, each checked against that model's own docs page. runWorkersAi now dispatches on the row rather than guessing, and img2img/inpaint/vision gain the binary inputs they always needed (image as a uint8 array, mask for inpainting). The same table closes a separate hole: provider_model_id was any free-form string forwarded straight to env.AI.run(), so a typo or an unpriced model published successfully and only failed once a user had been charged. Registration now rejects ids outside the catalogue and names what is allowed. InferenceOutput gains a `text` kind for the companion lane. ModelKind is deliberately unchanged — companion models are an internal assist for prompt expansion and reverse-prompting, not something a creator publishes; widening the published-model schema is a separate change. Stored render objects were also always named .png regardless of what the model emitted; the extension now follows the content type. Verification: 15 new unit tests over the pure helpers (all four response shapes, per-family request building, param allowlisting, catalogue membership), 48/48 unit tests, 146/146 smoke checks including three new end-to-end registration checks, typecheck clean, dry-run deploy clean on both environments. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01QuSa99HABSrWRarT9BQ5Tz --- workers/api/scripts/smoke_api.mjs | 41 +++ workers/api/src/ai/inference.ts | 255 +++++++++++++++-- workers/api/src/ai/workersAiCatalogue.ts | 298 ++++++++++++++++++++ workers/api/src/render/handlers.ts | 12 +- workers/api/test/workersai.adapter.test.mjs | 251 +++++++++++++++++ 5 files changed, 833 insertions(+), 24 deletions(-) create mode 100644 workers/api/src/ai/workersAiCatalogue.ts create mode 100644 workers/api/test/workersai.adapter.test.mjs diff --git a/workers/api/scripts/smoke_api.mjs b/workers/api/scripts/smoke_api.mjs index 4d2d826..ad38bb7 100644 --- a/workers/api/scripts/smoke_api.mjs +++ b/workers/api/scripts/smoke_api.mjs @@ -316,6 +316,47 @@ let modelSlug = null; JSON.stringify(r.data).slice(0, 200), ); } +// Workers AI model ids are an allowlist, not free text. Before the +// catalogue existed any string was forwarded to env.AI.run(), so a typo +// published fine and failed only after the renderer had been debited. +let waModelSlug = null; +{ + const r = await api("POST", "/v1/models", { + title: `Smoke WA Model ${RUN}`, + kind: "image", + provider: "workers_ai", + visibility: "private", + }); + waModelSlug = r.data?.model?.slug ?? null; + record( + "POST /v1/models (workers_ai image)", + r.status === 201 && !!waModelSlug, + JSON.stringify(r.data).slice(0, 160), + ); +} +{ + const r = await api("POST", `/v1/models/${waModelSlug}/versions`, { + // One character short of the real id. + provider_model_id: "@cf/black-forest-labs/flux-1-schnel", + price_tokens: 10, + }); + record( + "workers_ai version rejects an unlisted model id", + r.status === 400 && r.data?.error === "model_not_in_catalogue", + JSON.stringify(r.data).slice(0, 160), + ); +} +{ + const r = await api("POST", `/v1/models/${waModelSlug}/versions`, { + provider_model_id: "@cf/black-forest-labs/flux-1-schnell", + price_tokens: 10, + }); + record( + "workers_ai version accepts a catalogued model id", + r.status === 201 && r.data?.version?.version === 1, + JSON.stringify(r.data).slice(0, 160), + ); +} { const r = await api("GET", `/v1/models/${modelSlug}`); record( diff --git a/workers/api/src/ai/inference.ts b/workers/api/src/ai/inference.ts index 1fd29e4..dc968ae 100644 --- a/workers/api/src/ai/inference.ts +++ b/workers/api/src/ai/inference.ts @@ -1,6 +1,7 @@ import Anthropic from "@anthropic-ai/sdk"; import type { Env } from "../types"; import type { ModelKind, ModelProvider } from "../db/render"; +import { getCatalogueEntry, type CatalogueEntry } from "./workersAiCatalogue"; /** * Model-execution layer for the render-token service. @@ -28,11 +29,23 @@ export interface InferenceInput { * provider. */ weightsRef?: string | null; + /** + * Source image for the img2img / inpaint / vision lanes. Raw bytes — + * the adapter converts to whatever encoding the target model wants. + */ + sourceImage?: Uint8Array | null; + /** Inpainting mask. Required by the inpaint family, ignored elsewhere. */ + maskImage?: Uint8Array | null; } export interface InferenceOutput { - outputKind: "code" | "image"; - /** Present for `code` jobs — the generated source, inline. */ + /** + * `text` covers the companion lane (vision reverse-prompting, prompt + * expansion). Deliberately NOT added to `ModelKind` — companion + * models are an internal assist, not something a creator publishes. + */ + outputKind: "code" | "image" | "text"; + /** Present for `code` and `text` jobs — the generated text, inline. */ text: string | null; /** Present for `image` jobs — raw bytes for the caller to persist to R2. */ bytes: Uint8Array | null; @@ -72,6 +85,21 @@ export function isRenderMock(env: Env): boolean { * so scripts/smoke_api.mjs and local `wrangler dev` runs are hermetic. */ async function runMock(input: InferenceInput): Promise { + // Companion (vision / instruct) models return text, and are reached by + // provider model id rather than by ModelKind — so the mock has to look + // them up the same way the real path does, or mock runs would claim an + // image where production returns prose. + const mockEntry = + input.provider === "workers_ai" + ? getCatalogueEntry(input.providerModelId) + : undefined; + if (mockEntry?.outputKind === "text") { + const text = + `mock ${mockEntry.family} output — model=${mockEntry.id} ` + + `seed=${input.seed}: ${input.prompt.slice(0, 120)}`; + const outputHash = await sha256Hex(new TextEncoder().encode(text)); + return { outputKind: "text", text, bytes: null, contentType: null, outputHash }; + } if (input.kind === "code") { const text = `// mock render — provider=${input.provider} model=${input.providerModelId}\n` + @@ -159,11 +187,180 @@ async function runAnthropic( } /** - * Cloudflare Workers AI — image generation (reference art / textures). - * Output feeds the studio's working material only; it is never eligible - * for the frozen-artifact bundle (enforced in the freeze pipeline, not - * here — this function just runs the model). + * Cloudflare Workers AI. + * + * Dispatches on the catalogue entry rather than assuming one shape: + * Workers AI models disagree on both request and response format, and + * the disagreement does not follow vendor or task (see the header of + * ./workersAiCatalogue.ts). An unlisted model is refused here rather + * than being forwarded to `env.AI.run()` on trust. + * + * Image output feeds the studio's working material only; it is never + * eligible for the frozen-artifact bundle (enforced in the freeze + * pipeline, not here — this function just runs the model). + */ + +/** Decode a base64 string into bytes. */ +function base64ToBytes(b64: string): Uint8Array { + const binary = atob(b64); + return Uint8Array.from(binary, (ch) => ch.charCodeAt(0)); +} + +/** + * Caller params, narrowed to the keys this model accepts. Unknown keys + * are dropped: some models reject them and others ignore them silently, + * and a caller should not be able to discover which by experiment. + */ +export function filterParams( + entry: CatalogueEntry, + params: Record | null, +): Record { + const out: Record = { ...entry.defaults }; + if (!params) return out; + for (const key of entry.paramAllowlist) { + if (params[key] !== undefined) out[key] = params[key]; + } + return out; +} + +/** Numeric seed where the seed is numeric; undefined otherwise. */ +function numericSeed(seed: string): number | undefined { + return /^\d+$/.test(seed) ? parseInt(seed, 10) : undefined; +} + +/** + * Build the JSON body for a model. Families differ in more than + * parameters: the SD lanes want pixel arrays, vision wants an image + * alongside a prompt, and instruct wants a chat transcript. + */ +export function buildJsonInput( + entry: CatalogueEntry, + input: InferenceInput, +): Record { + const params = filterParams(entry, input.params); + const seed = numericSeed(input.seed); + + if (entry.family === "instruct") { + const messages: { role: string; content: string }[] = []; + if (input.systemPrompt) { + messages.push({ role: "system", content: input.systemPrompt }); + } + messages.push({ role: "user", content: input.prompt }); + return { ...params, messages }; + } + + if (entry.family === "vision") { + if (!input.sourceImage) { + throw new InferenceError( + `${entry.label} requires a source image`, + "missing_source_image", + ); + } + return { ...params, prompt: input.prompt, image: Array.from(input.sourceImage) }; + } + + const body: Record = { ...params, prompt: input.prompt }; + if (seed !== undefined) body.seed = seed; + + if (entry.family === "img2img" || entry.family === "inpaint") { + if (!input.sourceImage) { + throw new InferenceError( + `${entry.label} requires a source image`, + "missing_source_image", + ); + } + // Documented as an array of 8-bit unsigned integers, not raw bytes. + body.image = Array.from(input.sourceImage); + } + if (entry.family === "inpaint") { + if (!input.maskImage) { + throw new InferenceError( + `${entry.label} requires a mask`, + "missing_mask", + ); + } + body.mask = Array.from(input.maskImage); + } + return body; +} + +/** + * FLUX.2 models take multipart form-data even for a prompt-only call. + * FormData exposes neither its serialized body nor its boundary, so it + * is round-tripped through a Response to get both. + */ +function buildMultipartInput( + entry: CatalogueEntry, + input: InferenceInput, +): { multipart: { body: unknown; contentType: string } } { + const params = filterParams(entry, input.params); + const seed = numericSeed(input.seed); + const form = new FormData(); + form.append("prompt", input.prompt); + for (const [key, value] of Object.entries(params)) { + if (value !== undefined && value !== null) form.append(key, String(value)); + } + if (seed !== undefined) form.append("seed", String(seed)); + if (input.sourceImage) { + form.append( + "input_image_0", + new Blob([input.sourceImage], { type: "image/png" }), + ); + } + const serialized = new Response(form); + const contentType = serialized.headers.get("content-type"); + if (!contentType) { + throw new InferenceError("could not serialize multipart body", "provider_error"); + } + return { multipart: { body: serialized.body, contentType } }; +} + +/** + * Decode whatever the model returned into bytes or text, per the + * catalogue's declared response shape. */ +export async function decodeOutput( + entry: CatalogueEntry, + result: unknown, +): Promise<{ bytes: Uint8Array | null; text: string | null }> { + if (entry.responseShape === "binary") { + if (result instanceof ReadableStream) { + const buf = await new Response(result as ReadableStream).arrayBuffer(); + return { bytes: new Uint8Array(buf), text: null }; + } + if (result instanceof Uint8Array) return { bytes: result, text: null }; + if (result instanceof ArrayBuffer) { + return { bytes: new Uint8Array(result), text: null }; + } + throw new InferenceError( + `${entry.id} did not return binary output`, + "provider_error", + ); + } + + if (entry.responseShape === "json_base64_image") { + const image = (result as { image?: unknown } | null)?.image; + if (typeof image !== "string" || !image) { + throw new InferenceError( + `${entry.id} did not return a base64 image`, + "provider_error", + ); + } + return { bytes: base64ToBytes(image), text: null }; + } + + // json_text + const field = entry.textField ?? "response"; + const value = (result as Record | null)?.[field]; + if (typeof value !== "string" || !value.trim()) { + throw new InferenceError( + `${entry.id} did not return text in "${field}"`, + "provider_error", + ); + } + return { bytes: null, text: value.trim() }; +} + async function runWorkersAi( env: Env, input: InferenceInput, @@ -174,37 +371,49 @@ async function runWorkersAi( "provider_unconfigured", ); } + const entry = getCatalogueEntry(input.providerModelId); + if (!entry) { + throw new InferenceError( + `${input.providerModelId} is not an allowlisted Workers AI model`, + "model_not_in_catalogue", + ); + } + + const body = + entry.requestShape === "multipart" + ? buildMultipartInput(entry, input) + : buildJsonInput(entry, input); + let result: unknown; try { - result = await env.AI.run(input.providerModelId, { - prompt: input.prompt, - seed: /^\d+$/.test(input.seed) ? parseInt(input.seed, 10) : undefined, - }); + // The binding's per-model overloads can't be satisfied by a model id + // only known at runtime; the catalogue is what guarantees the body + // matches the model. + result = await (env.AI as unknown as { + run(model: string, body: unknown): Promise; + }).run(entry.id, body); } catch (e) { throw new InferenceError( e instanceof Error ? e.message : "workers ai request failed", "provider_error", ); } - // Workers AI image models return a ReadableStream (binary) response. - let bytes: Uint8Array; - if (result instanceof ReadableStream) { - const buf = await new Response(result).arrayBuffer(); - bytes = new Uint8Array(buf); - } else if (result instanceof Uint8Array) { - bytes = result; - } else { - throw new InferenceError( - "unexpected Workers AI response shape", - "provider_error", - ); + + const { bytes, text } = await decodeOutput(entry, result); + + if (entry.outputKind === "text") { + const outputHash = await sha256Hex(new TextEncoder().encode(text ?? "")); + return { outputKind: "text", text, bytes: null, contentType: null, outputHash }; + } + if (!bytes) { + throw new InferenceError(`${entry.id} returned no image bytes`, "provider_error"); } const outputHash = await sha256Hex(bytes); return { outputKind: "image", text: null, bytes, - contentType: "image/png", + contentType: entry.contentType, outputHash, }; } diff --git a/workers/api/src/ai/workersAiCatalogue.ts b/workers/api/src/ai/workersAiCatalogue.ts new file mode 100644 index 0000000..e1f24cc --- /dev/null +++ b/workers/api/src/ai/workersAiCatalogue.ts @@ -0,0 +1,298 @@ +/** + * Allowlisted Cloudflare Workers AI models, and the shape of each one. + * + * Workers AI models do NOT share a request or response contract, and the + * differences are not inferable from vendor or task: + * + * - `flux-1-schnell` takes `steps` and returns JSON `{ image: }`. + * - The Stable-Diffusion family takes `num_steps` and returns a raw + * `ReadableStream`. + * - `flux-2-*` take multipart form-data, even for a prompt-only call. + * - Leonardo's two models disagree with each other: `phoenix-1.0` + * streams bytes, `lucid-origin` returns base64 JSON. + * + * So the shape is recorded per model rather than guessed at call time. + * Every row here was checked against + * https://developers.cloudflare.com/workers-ai/models// — when + * adding a model, check its page rather than copying a neighbouring row. + * + * This table is also the registration allowlist: `provider_model_id` was + * previously any free-form string passed straight to `env.AI.run()`, so a + * typo or an unpriced model only failed after a user had been charged. + */ + +export type WorkersAiFamily = + | "text2img" + | "img2img" + | "inpaint" + | "vision" + | "instruct"; + +/** How the request body is built. */ +export type RequestShape = "json" | "multipart"; + +/** How the response is decoded. */ +export type ResponseShape = "binary" | "json_base64_image" | "json_text"; + +export interface CatalogueEntry { + id: string; + label: string; + family: WorkersAiFamily; + requestShape: RequestShape; + responseShape: ResponseShape; + outputKind: "image" | "text"; + /** Null for text models — they produce no stored artifact. */ + contentType: string | null; + /** Applied before caller params; caller may override allowlisted keys. */ + defaults: Readonly>; + /** + * Caller-supplied params forwarded to the model. Anything else is + * dropped — an unknown key is rejected by some models and silently + * ignored by others, and neither is a good failure mode. + */ + paramAllowlist: readonly string[]; + /** Which JSON field carries the text, for `json_text` responses. */ + textField?: "response" | "description"; + /** + * RELATIVE cost weight, cheapest = 1. A deliberate placeholder: real + * per-model pricing is published per model page (and is variously + * per-step, per-512px-tile, or both), so these are an ordering for the + * billing layer to reason about, NOT dollar amounts. Substitute real + * figures from https://developers.cloudflare.com/workers-ai/platform/pricing/ + * before charging against them. + */ + costUnits: number; +} + +/** Stable-Diffusion-family params — shared by every `@cf` SD derivative. */ +const SD_PARAMS = [ + "negative_prompt", + "height", + "width", + "num_steps", + "guidance", + "seed", +] as const; + +const ENTRIES: readonly CatalogueEntry[] = [ + // ---- Draft / real-time ------------------------------------------------- + { + id: "@cf/black-forest-labs/flux-1-schnell", + label: "FLUX.1 [schnell]", + family: "text2img", + requestShape: "json", + // Returns { image: "" } — NOT a stream. This is the case + // the previous single-path implementation threw on. + responseShape: "json_base64_image", + outputKind: "image", + contentType: "image/jpeg", + defaults: { steps: 4 }, + paramAllowlist: ["steps", "seed"], + costUnits: 1, + }, + { + id: "@cf/bytedance/stable-diffusion-xl-lightning", + label: "SDXL Lightning", + family: "text2img", + requestShape: "json", + responseShape: "binary", + outputKind: "image", + contentType: "image/png", + defaults: { num_steps: 4 }, + paramAllowlist: SD_PARAMS, + costUnits: 1, + }, + { + id: "@cf/lykon/dreamshaper-8-lcm", + label: "DreamShaper 8 LCM", + family: "text2img", + requestShape: "json", + responseShape: "binary", + outputKind: "image", + contentType: "image/png", + defaults: { num_steps: 8 }, + paramAllowlist: SD_PARAMS, + costUnits: 1, + }, + { + id: "@cf/black-forest-labs/flux-2-klein-4b", + label: "FLUX.2 [klein] 4B", + family: "text2img", + // Multipart form-data even for a prompt-only call. + requestShape: "multipart", + responseShape: "json_base64_image", + outputKind: "image", + contentType: "image/jpeg", + // Distilled: `steps` is fixed at 4 upstream and cannot be adjusted. + defaults: { width: 1024, height: 1024 }, + paramAllowlist: ["width", "height", "guidance", "seed"], + costUnits: 2, + }, + + // ---- Final / high fidelity -------------------------------------------- + { + id: "@cf/black-forest-labs/flux-2-klein-9b", + label: "FLUX.2 [klein] 9B", + family: "text2img", + requestShape: "multipart", + responseShape: "json_base64_image", + outputKind: "image", + contentType: "image/jpeg", + defaults: { width: 1024, height: 1024 }, + paramAllowlist: ["width", "height", "guidance", "seed"], + costUnits: 4, + }, + { + id: "@cf/black-forest-labs/flux-2-dev", + label: "FLUX.2 [dev]", + family: "text2img", + requestShape: "multipart", + responseShape: "json_base64_image", + outputKind: "image", + contentType: "image/jpeg", + defaults: { width: 1024, height: 768 }, + paramAllowlist: ["width", "height", "guidance", "seed", "steps"], + costUnits: 8, + }, + { + id: "@cf/leonardo/phoenix-1.0", + label: "Leonardo Phoenix 1.0", + family: "text2img", + requestShape: "json", + // Leonardo's two models differ from each other — phoenix streams. + responseShape: "binary", + outputKind: "image", + contentType: "image/png", + defaults: { num_steps: 25, guidance: 2, width: 1024, height: 1024 }, + paramAllowlist: ["num_steps", "guidance", "seed", "width", "height", "negative_prompt"], + costUnits: 6, + }, + { + id: "@cf/leonardo/lucid-origin", + label: "Leonardo Lucid Origin", + family: "text2img", + requestShape: "json", + // ...while lucid-origin returns base64 JSON. + responseShape: "json_base64_image", + outputKind: "image", + contentType: "image/jpeg", + defaults: { num_steps: 25, guidance: 4.5, width: 1120, height: 1120 }, + paramAllowlist: ["num_steps", "steps", "guidance", "seed", "width", "height"], + costUnits: 6, + }, + { + id: "@cf/stabilityai/stable-diffusion-xl-base-1.0", + label: "Stable Diffusion XL 1.0", + family: "text2img", + requestShape: "json", + responseShape: "binary", + outputKind: "image", + contentType: "image/png", + defaults: { num_steps: 20 }, + paramAllowlist: SD_PARAMS, + costUnits: 3, + }, + + // ---- Edit / remix ------------------------------------------------------ + { + id: "@cf/runwayml/stable-diffusion-v1-5-img2img", + label: "SD 1.5 img2img", + family: "img2img", + requestShape: "json", + responseShape: "binary", + outputKind: "image", + contentType: "image/png", + defaults: { num_steps: 20, strength: 1 }, + paramAllowlist: [...SD_PARAMS, "strength"], + costUnits: 2, + }, + { + id: "@cf/runwayml/stable-diffusion-v1-5-inpainting", + label: "SD 1.5 inpainting", + family: "inpaint", + requestShape: "json", + responseShape: "binary", + outputKind: "image", + contentType: "image/png", + defaults: { num_steps: 20, strength: 1 }, + paramAllowlist: [...SD_PARAMS, "strength"], + costUnits: 2, + }, + + // ---- Companion: vision (reverse-prompt, critique) ---------------------- + { + id: "@cf/meta/llama-3.2-11b-vision-instruct", + label: "Llama 3.2 11B Vision", + family: "vision", + requestShape: "json", + responseShape: "json_text", + outputKind: "text", + contentType: null, + defaults: { max_tokens: 512 }, + paramAllowlist: ["max_tokens", "temperature"], + textField: "response", + costUnits: 1, + }, + { + id: "@cf/llava-hf/llava-1.5-7b-hf", + label: "LLaVA 1.5 7B", + family: "vision", + requestShape: "json", + responseShape: "json_text", + outputKind: "text", + contentType: null, + defaults: { max_tokens: 512 }, + paramAllowlist: ["max_tokens", "temperature"], + // LLaVA returns `description` rather than `response`. + textField: "description", + costUnits: 1, + }, + + // ---- Companion: instruct (prompt expansion) ---------------------------- + { + id: "@cf/meta/llama-3.3-70b-instruct-fp8-fast", + label: "Llama 3.3 70B Instruct (fp8)", + family: "instruct", + requestShape: "json", + responseShape: "json_text", + outputKind: "text", + contentType: null, + defaults: { max_tokens: 1024 }, + paramAllowlist: ["max_tokens", "temperature"], + textField: "response", + costUnits: 2, + }, + { + id: "@cf/qwen/qwen2.5-coder-32b-instruct", + label: "Qwen2.5 Coder 32B", + family: "instruct", + requestShape: "json", + responseShape: "json_text", + outputKind: "text", + contentType: null, + defaults: { max_tokens: 1024 }, + paramAllowlist: ["max_tokens", "temperature"], + textField: "response", + costUnits: 2, + }, +] as const; + +const BY_ID = new Map(ENTRIES.map((e) => [e.id, e])); + +export function getCatalogueEntry(id: string): CatalogueEntry | undefined { + return BY_ID.get(id); +} + +export function isCatalogueModel(id: string): boolean { + return BY_ID.has(id); +} + +/** Every allowlisted id — used in the registration error to say what IS allowed. */ +export function catalogueModelIds(): string[] { + return ENTRIES.map((e) => e.id); +} + +/** Models a creator may publish: the image lanes, not the internal companions. */ +export function publishableModelIds(): string[] { + return ENTRIES.filter((e) => e.outputKind === "image").map((e) => e.id); +} diff --git a/workers/api/src/render/handlers.ts b/workers/api/src/render/handlers.ts index 9db4ab8..4fd765c 100644 --- a/workers/api/src/render/handlers.ts +++ b/workers/api/src/render/handlers.ts @@ -7,6 +7,7 @@ import { checkRateLimit } from "../lib/rateLimit"; import { getProjectById } from "../db/projects"; import { applyLedgerEntry } from "../db/tokens"; import { runInference, InferenceError } from "../ai/inference"; +import { isCatalogueModel, publishableModelIds } from "../ai/workersAiCatalogue"; import { isModelKind, isModelProvider, @@ -350,6 +351,14 @@ export async function publishVersionHandler( if (model.provider === "fal_custom" && !weightsRef) { return badRequest(c, "weights_ref_required_for_fal_custom"); } + // Workers AI ids were previously forwarded to env.AI.run() unchecked, so + // a typo or an unpriced model published fine and only failed once a + // renderer had been debited. The catalogue is the allowlist. + if (model.provider === "workers_ai" && !isCatalogueModel(providerModelId)) { + return badRequest(c, "model_not_in_catalogue", { + allowed: publishableModelIds(), + }); + } const version = await insertModelVersion(c.env.DB, { modelId: model.id, @@ -528,7 +537,8 @@ export async function renderHandler( if (!c.env.CAPTURES) { throw new InferenceError("R2 CAPTURES bucket is not bound", "storage_unconfigured"); } - outputKey = `renders/${session.uid}/${job.id}-${seed}.png`; + const ext = result.contentType === "image/jpeg" ? "jpg" : "png"; + outputKey = `renders/${session.uid}/${job.id}-${seed}.${ext}`; await c.env.CAPTURES.put(outputKey, result.bytes, { httpMetadata: { contentType: result.contentType ?? "image/png" }, customMetadata: { diff --git a/workers/api/test/workersai.adapter.test.mjs b/workers/api/test/workersai.adapter.test.mjs new file mode 100644 index 0000000..5b4618d --- /dev/null +++ b/workers/api/test/workersai.adapter.test.mjs @@ -0,0 +1,251 @@ +// Workers AI adapter: request building and response decoding. +// +// These are the parts where Workers AI's per-model inconsistency bites. +// The models genuinely disagree — flux-1-schnell takes `steps` and +// returns JSON base64, the Stable Diffusion family takes `num_steps` and +// returns a ReadableStream, and Leonardo's two models disagree with each +// other — so the catalogue's per-model shape is the thing under test. +// +// The previous single-path implementation accepted only a stream or a +// Uint8Array, which meant flux-1-schnell (the fastest, cheapest model in +// the catalogue) threw `unexpected Workers AI response shape` at render +// time, after the renderer had already been debited. The +// `json_base64_image` cases below are the regression guard for that. + +import { build } from "esbuild"; +import { mkdtempSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join, resolve, dirname } from "node:path"; +import { pathToFileURL, fileURLToPath } from "node:url"; +import assert from "node:assert/strict"; +import { test } from "node:test"; + +const srcDir = resolve(dirname(fileURLToPath(import.meta.url)), "..", "src"); +const tmp = mkdtempSync(join(tmpdir(), "ga-waiadapter-")); +const outFile = join(tmp, "adapter.mjs"); +const entry = join(tmp, "entry.ts"); + +writeFileSync( + entry, + `export { buildJsonInput, decodeOutput, filterParams, InferenceError } from ${JSON.stringify(join(srcDir, "ai", "inference"))}; +export { getCatalogueEntry, isCatalogueModel, publishableModelIds, catalogueModelIds } from ${JSON.stringify(join(srcDir, "ai", "workersAiCatalogue"))}; +`, +); + +await build({ + entryPoints: [entry], + bundle: true, + format: "esm", + platform: "node", + target: "node18", + // Bundle everything: the output lands in a temp dir outside the + // project, so anything left external would not resolve at import time. + outfile: outFile, +}); + +const { + buildJsonInput, + decodeOutput, + filterParams, + getCatalogueEntry, + isCatalogueModel, + publishableModelIds, +} = await import(pathToFileURL(outFile).href); + +const FLUX = "@cf/black-forest-labs/flux-1-schnell"; +const SDXL = "@cf/bytedance/stable-diffusion-xl-lightning"; +const IMG2IMG = "@cf/runwayml/stable-diffusion-v1-5-img2img"; +const INPAINT = "@cf/runwayml/stable-diffusion-v1-5-inpainting"; +const LLAVA = "@cf/llava-hf/llava-1.5-7b-hf"; +const QWEN = "@cf/qwen/qwen2.5-coder-32b-instruct"; +const PHOENIX = "@cf/leonardo/phoenix-1.0"; +const LUCID = "@cf/leonardo/lucid-origin"; + +const baseInput = (providerModelId, over = {}) => ({ + kind: "image", + provider: "workers_ai", + providerModelId, + systemPrompt: null, + prompt: "a cyberpunk lizard", + params: null, + seed: "12345", + ...over, +}); + +// ---- response decoding --------------------------------------------------- + +test("decodes json_base64_image — the shape that used to throw", async () => { + const entryFlux = getCatalogueEntry(FLUX); + assert.equal(entryFlux.responseShape, "json_base64_image"); + const original = new Uint8Array([137, 80, 78, 71, 13, 10, 26, 10]); + const b64 = Buffer.from(original).toString("base64"); + const { bytes, text } = await decodeOutput(entryFlux, { image: b64 }); + assert.equal(text, null); + assert.deepEqual([...bytes], [...original]); +}); + +test("decodes a binary ReadableStream", async () => { + const entrySd = getCatalogueEntry(SDXL); + assert.equal(entrySd.responseShape, "binary"); + const payload = new Uint8Array([1, 2, 3, 4, 5]); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(payload); + controller.close(); + }, + }); + const { bytes } = await decodeOutput(entrySd, stream); + assert.deepEqual([...bytes], [...payload]); +}); + +test("decodes a bare Uint8Array on a binary model", async () => { + const payload = new Uint8Array([9, 9, 9]); + const { bytes } = await decodeOutput(getCatalogueEntry(SDXL), payload); + assert.deepEqual([...bytes], [...payload]); +}); + +test("decodes json_text from the per-model field", async () => { + // LLaVA answers in `description`; the instruct models use `response`. + const llava = await decodeOutput(getCatalogueEntry(LLAVA), { + description: " a red square ", + }); + assert.equal(llava.text, "a red square"); + assert.equal(llava.bytes, null); + + const qwen = await decodeOutput(getCatalogueEntry(QWEN), { + response: "expanded prompt", + }); + assert.equal(qwen.text, "expanded prompt"); +}); + +test("the two Leonardo models are decoded differently from each other", () => { + assert.equal(getCatalogueEntry(PHOENIX).responseShape, "binary"); + assert.equal(getCatalogueEntry(LUCID).responseShape, "json_base64_image"); +}); + +test("a mismatched response is rejected, not silently accepted", async () => { + await assert.rejects( + () => decodeOutput(getCatalogueEntry(FLUX), { nope: true }), + /did not return a base64 image/, + ); + await assert.rejects( + () => decodeOutput(getCatalogueEntry(SDXL), { image: "abc" }), + /did not return binary output/, + ); + await assert.rejects( + () => decodeOutput(getCatalogueEntry(QWEN), { response: " " }), + /did not return text/, + ); +}); + +// ---- request building ---------------------------------------------------- + +test("step parameter is named per model, not globally", () => { + // Getting this wrong silently drops the caller's step count. + const flux = buildJsonInput(getCatalogueEntry(FLUX), baseInput(FLUX)); + assert.equal(flux.steps, 4); + assert.equal(flux.num_steps, undefined); + + const sd = buildJsonInput(getCatalogueEntry(SDXL), baseInput(SDXL)); + assert.equal(sd.num_steps, 4); + assert.equal(sd.steps, undefined); +}); + +test("numeric seeds are forwarded, non-numeric ones are omitted", () => { + const withSeed = buildJsonInput(getCatalogueEntry(FLUX), baseInput(FLUX)); + assert.equal(withSeed.seed, 12345); + const hexSeed = buildJsonInput( + getCatalogueEntry(FLUX), + baseInput(FLUX, { seed: "0xdeadbeef" }), + ); + assert.equal(hexSeed.seed, undefined); +}); + +test("img2img sends the source as a uint8 array", () => { + const src = new Uint8Array([10, 20, 30]); + const body = buildJsonInput( + getCatalogueEntry(IMG2IMG), + baseInput(IMG2IMG, { sourceImage: src }), + ); + assert.deepEqual(body.image, [10, 20, 30]); + assert.equal(body.mask, undefined); +}); + +test("inpainting sends both image and mask", () => { + const body = buildJsonInput( + getCatalogueEntry(INPAINT), + baseInput(INPAINT, { + sourceImage: new Uint8Array([1, 2]), + maskImage: new Uint8Array([255, 0]), + }), + ); + assert.deepEqual(body.image, [1, 2]); + assert.deepEqual(body.mask, [255, 0]); +}); + +test("edit lanes refuse to run without their binary inputs", () => { + // Better to fail here than to be charged for a render that cannot work. + assert.throws( + () => buildJsonInput(getCatalogueEntry(IMG2IMG), baseInput(IMG2IMG)), + /requires a source image/, + ); + assert.throws( + () => + buildJsonInput( + getCatalogueEntry(INPAINT), + baseInput(INPAINT, { sourceImage: new Uint8Array([1]) }), + ), + /requires a mask/, + ); + assert.throws( + () => buildJsonInput(getCatalogueEntry(LLAVA), baseInput(LLAVA)), + /requires a source image/, + ); +}); + +test("instruct models get a chat transcript, with the system prompt when set", () => { + const body = buildJsonInput( + getCatalogueEntry(QWEN), + baseInput(QWEN, { systemPrompt: "you expand prompts" }), + ); + assert.deepEqual(body.messages, [ + { role: "system", content: "you expand prompts" }, + { role: "user", content: "a cyberpunk lizard" }, + ]); + assert.equal(body.prompt, undefined); + + const noSystem = buildJsonInput(getCatalogueEntry(QWEN), baseInput(QWEN)); + assert.equal(noSystem.messages.length, 1); + assert.equal(noSystem.messages[0].role, "user"); +}); + +// ---- params + allowlist -------------------------------------------------- + +test("caller params override defaults but unknown keys are dropped", () => { + const params = filterParams(getCatalogueEntry(SDXL), { + num_steps: 12, + guidance: 9, + // Not on SDXL's allowlist — some models reject unknown keys and + // others ignore them, so neither is forwarded. + definitely_not_a_param: "x", + steps: 99, + }); + assert.equal(params.num_steps, 12); + assert.equal(params.guidance, 9); + assert.equal(params.definitely_not_a_param, undefined); + assert.equal(params.steps, undefined); +}); + +test("catalogue membership gates registration", () => { + assert.ok(isCatalogueModel(FLUX)); + assert.ok(!isCatalogueModel("@cf/black-forest-labs/flux-1-schnel")); // typo + assert.ok(!isCatalogueModel("../../etc/passwd")); + assert.ok(!isCatalogueModel("")); +}); + +test("companion text models are not offered as publishable", () => { + const publishable = publishableModelIds(); + assert.ok(publishable.includes(FLUX)); + assert.ok(!publishable.includes(LLAVA)); + assert.ok(!publishable.includes(QWEN)); +});