diff --git a/CHANGELOG.md b/CHANGELOG.md index 6bf58f0f7a..9892b9b7bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ - **feat(dashboard):** collapse quota rows and sort by remaining quota in the usage view. (thanks @j2-cuong) - **feat(dashboard):** add a settings toggle for tool-source diagnostics logging. (thanks @DuyPrX) - **feat(oauth):** import a ChatGPT/Codex connection from a raw access token (no refresh token required). (thanks @ryanngit) +- **feat(providers):** add NVIDIA NIM image generation (FLUX models). (thanks @eng2007) ### 🔧 Bug Fixes diff --git a/open-sse/config/imageRegistry.ts b/open-sse/config/imageRegistry.ts index 544f51959b..330a37c730 100644 --- a/open-sse/config/imageRegistry.ts +++ b/open-sse/config/imageRegistry.ts @@ -563,6 +563,35 @@ export const IMAGE_PROVIDERS: Record = { supportedSizes: ["1024x1024", "1024x1280", "1280x1024"], }, + // NVIDIA NIM image generation (FLUX models). Distinct from the NVIDIA *chat* entry + // (open-sse/config/providers/registry/nvidia/index.ts, host integrate.api.nvidia.com, + // OpenAI-compatible) — image generation lives on ai.api.nvidia.com/v1/genai/ + // with a native NIM body per model, so it gets a dedicated `nvidia-nim` format/handler + // (handleNvidiaNimImageGeneration) rather than reusing the OpenAI image path. + // Ported from upstream 9router#1195. + nvidia: { + id: "nvidia", + baseUrl: "https://ai.api.nvidia.com/v1/genai", + authType: "apikey", + authHeader: "bearer", + format: "nvidia-nim", + models: [ + { id: "black-forest-labs/flux.1-dev", name: "FLUX.1 Dev", inputModalities: ["text", "image"] }, + { id: "black-forest-labs/flux.1-schnell", name: "FLUX.1 Schnell" }, + { + id: "black-forest-labs/flux.1-kontext-dev", + name: "FLUX.1 Kontext Dev (Edit)", + inputModalities: ["text", "image"], + }, + { + id: "black-forest-labs/flux.2-klein-4b", + name: "FLUX.2 Klein 4B", + inputModalities: ["text", "image"], + }, + ], + supportedSizes: ["1024x1024", "768x1344", "512x512"], + }, + // SenseNova (商汤日日新) Text-to-Image on the free Token Plan. OpenAI-compatible // `/v1/images/generations`, so the generic OpenAI image handler routes it — same // SenseNova api-key/connection as the chat provider. (9router#2233) diff --git a/open-sse/handlers/imageGeneration.ts b/open-sse/handlers/imageGeneration.ts index 852d4132b5..2bfc452a0a 100644 --- a/open-sse/handlers/imageGeneration.ts +++ b/open-sse/handlers/imageGeneration.ts @@ -61,6 +61,7 @@ import { extractMarkdownImageUrls, CHATGPT_WEB_IMAGE_ID_RE, } from "./imageGeneration/providers/chatgptWeb.ts"; +import { handleNvidiaNimImageGeneration } from "./imageGeneration/providers/nvidiaNim.ts"; interface KieImageOptions { @@ -523,6 +524,17 @@ export async function handleImageGeneration({ }); } + if (providerConfig.format === "nvidia-nim") { + return handleNvidiaNimImageGeneration({ + model, + provider, + providerConfig, + body, + credentials, + log, + }); + } + return handleOpenAIImageGeneration({ model, provider, providerConfig, body, credentials, log }); } diff --git a/open-sse/handlers/imageGeneration/providers/nvidiaNim.ts b/open-sse/handlers/imageGeneration/providers/nvidiaNim.ts new file mode 100644 index 0000000000..863c20667d --- /dev/null +++ b/open-sse/handlers/imageGeneration/providers/nvidiaNim.ts @@ -0,0 +1,286 @@ +// NVIDIA NIM image generation (FLUX models) — ported from upstream 9router#1195. +// Unlike the NVIDIA *chat* entry (open-sse/config/providers/registry/nvidia/index.ts, +// host integrate.api.nvidia.com, OpenAI-compatible), NVIDIA NIM *image* generation lives +// on a different host (ai.api.nvidia.com) with a native per-model NIM body — so it gets +// its own provider handler rather than reusing the OpenAI-compatible image path. +// +// Invoke shape: POST https://ai.api.nvidia.com/v1/genai/ +// Authorization: Bearer +// where is the registered model id itself (e.g. "black-forest-labs/flux.1-dev"). +// +// Response shape varies across the NIM `genai` catalog (SDXL-style `artifacts[].base64`, +// list-of-strings `images[]`, OpenAI-style `data[].b64_json`, or single-value shorthands) +// — normalizeNvidiaNimImages() accepts every variant so a NIM response-shape change +// degrades to "no images" rather than throwing. + +import { saveCallLog } from "@/lib/usageDb"; +import { sanitizeErrorMessage } from "../../../utils/error.ts"; + +const FLUX_1_DEV = "black-forest-labs/flux.1-dev"; +const FLUX_1_KONTEXT_DEV = "black-forest-labs/flux.1-kontext-dev"; + +function numberFromInput(value: unknown): number | null { + if (value === undefined || value === null || value === "") return null; + const num = Number(value); + return Number.isFinite(num) ? num : null; +} + +function parseSizeString(size: unknown): { width: number; height: number } | null { + if (typeof size !== "string" || !size || size === "auto") return null; + const match = /^(\d+)x(\d+)$/.exec(size); + if (!match) return null; + return { width: Number(match[1]), height: Number(match[2]) }; +} + +function parseDimensions(body: Record): { width: number; height: number } | null { + const width = numberFromInput(body.width); + const height = numberFromInput(body.height); + if (width !== null && height !== null) return { width, height }; + return parseSizeString(body.size); +} + +function copyIfPresent( + target: Record, + source: Record, + key: string +): void { + if (source[key] !== undefined && source[key] !== null && source[key] !== "") { + target[key] = source[key]; + } +} + +function copyNumberIfPresent( + target: Record, + source: Record, + key: string, + options: { greaterThan?: number } = {} +): void { + if (source[key] === undefined || source[key] === null || source[key] === "") return; + const value = Number(source[key]); + if (!Number.isFinite(value)) return; + if (options.greaterThan !== undefined && !(value > options.greaterThan)) return; + target[key] = value; +} + +function normalizeImageArray(image: unknown): unknown[] { + if (Array.isArray(image)) return image.filter(Boolean); + return image ? [image] : []; +} + +// FLUX.1 Dev only accepts dimensions in the 768-1344px range, in 64px increments — +// out-of-range values are silently dropped rather than sent upstream (matches upstream +// 9router#1195 behavior, verified against build.nvidia.com model page constraints). +function isFlux1DevDimension(value: number): boolean { + return Number.isInteger(value) && value >= 768 && value <= 1344 && value % 64 === 0; +} + +/** + * Build the per-model NIM request body. Each FLUX family member on NIM accepts a + * slightly different parameter set: + * - flux.1-dev: mode (base/depth/canny) + cfg_scale (only forwarded if > 1) + strict + * 768-1344/64px-increment width/height validation; input image only sent for + * non-"base" modes (depth/canny control image) + * - flux.1-kontext-dev: image-conditioned edit — requires `image`, uses `aspect_ratio` + * instead of width/height (the model preserves/derives its own output dimensions) + * - flux.1-schnell / flux.2-klein-4b: width/height/seed/steps, optional `image` sent as + * an array when present (edit-style input) + */ +export function buildNvidiaNimRequestBody( + model: string, + body: Record +): Record { + const req: Record = { prompt: body.prompt }; + const dimensions = parseDimensions(body); + + if (dimensions && model !== FLUX_1_KONTEXT_DEV) { + if (model !== FLUX_1_DEV || (isFlux1DevDimension(dimensions.width) && isFlux1DevDimension(dimensions.height))) { + req.width = dimensions.width; + req.height = dimensions.height; + } + } + + if (model === FLUX_1_DEV) { + const mode = body.mode || "base"; + req.mode = mode; + if (mode !== "base") { + const images = normalizeImageArray(body.image); + if (images.length > 0) req.image = images[0]; + } + } else if (model === FLUX_1_KONTEXT_DEV) { + const images = normalizeImageArray(body.image); + if (images.length > 0) req.image = images[0]; + copyIfPresent(req, body, "aspect_ratio"); + } else if (body.image) { + req.image = normalizeImageArray(body.image); + } + + if (model === FLUX_1_DEV) { + copyNumberIfPresent(req, body, "cfg_scale", { greaterThan: 1 }); + } else { + copyIfPresent(req, body, "cfg_scale"); + } + copyIfPresent(req, body, "seed"); + copyIfPresent(req, body, "steps"); + return req; +} + +function imageItemFromValue(value: unknown): { b64_json?: string; url?: string; finish_reason?: string } | null { + if (!value) return null; + if (typeof value === "string") return { b64_json: value }; + if (typeof value !== "object") return null; + const obj = value as Record; + if (typeof obj.url === "string") return { url: obj.url }; + const base64 = obj.base64 || obj.b64_json || obj.image || obj.data; + if (typeof base64 !== "string") return null; + const item: { b64_json: string; finish_reason?: string } = { b64_json: base64 }; + const finishReason = obj.finishReason || obj.finish_reason; + if (typeof finishReason === "string") item.finish_reason = finishReason; + return item; +} + +/** + * Normalize the NIM response into `{ created, data: [{ b64_json | url, finish_reason? }] }`. + * Accepts every response shape seen across the NIM `genai` catalog rather than a single + * assumed format. + */ +export function normalizeNvidiaNimImages(responseBody: unknown): { + created: number; + data: Array<{ b64_json?: string; url?: string; finish_reason?: string }>; +} { + const obj = (responseBody && typeof responseBody === "object" ? responseBody : {}) as Record< + string, + unknown + >; + + // Already OpenAI-shaped — pass through. + if (typeof obj.created === "number" && Array.isArray(obj.data)) { + return obj as { created: number; data: Array<{ b64_json?: string; url?: string }> }; + } + + const candidates: unknown[] = []; + if (Array.isArray(obj.artifacts)) candidates.push(...obj.artifacts); + if (Array.isArray(obj.images)) candidates.push(...obj.images); + if (Array.isArray(obj.data)) candidates.push(...obj.data); + if (obj.artifact) candidates.push(obj.artifact); + if (obj.image) candidates.push(obj.image); + if (obj.base64) candidates.push(obj.base64); + const result = obj.result as Record | undefined; + if (result?.image) candidates.push(result.image); + if (result && Array.isArray(result.artifacts)) candidates.push(...result.artifacts); + + return { + created: Math.floor(Date.now() / 1000), + data: candidates.map(imageItemFromValue).filter((item): item is NonNullable => item !== null), + }; +} + +export async function handleNvidiaNimImageGeneration({ + model, + provider, + providerConfig, + body, + credentials, + log, +}: { + model: string; + provider: string; + providerConfig: { baseUrl: string }; + body: Record; + credentials?: { apiKey?: string; accessToken?: string } | null; + log?: { + info: (scope: string, message: string) => void; + error: (scope: string, message: string) => void; + } | null; +}) { + const startTime = Date.now(); + const token = credentials?.apiKey || credentials?.accessToken || ""; + + if (model === FLUX_1_KONTEXT_DEV && !body.image) { + return { + success: false, + status: 400, + error: "NVIDIA FLUX.1 Kontext Dev requires an input image", + }; + } + + const requestBody = buildNvidiaNimRequestBody(model, body); + + if (log) { + const promptPreview = String(body.prompt ?? "").slice(0, 60); + log.info("IMAGE", `${provider}/${model} (nvidia-nim) | prompt: "${promptPreview}..."`); + } + + const url = `${providerConfig.baseUrl.replace(/\/$/, "")}/${model}`; + + try { + const response = await fetch(url, { + method: "POST", + headers: { + "Content-Type": "application/json", + Accept: "application/json", + Authorization: `Bearer ${token}`, + }, + body: JSON.stringify(requestBody), + }); + + if (!response.ok) { + const errorText = await response.text(); + if (log) + log.error("IMAGE", `${provider} error ${response.status}: ${errorText.slice(0, 200)}`); + saveCallLog({ + method: "POST", + path: "/v1/images/generations", + status: response.status, + model: `${provider}/${model}`, + provider, + duration: Date.now() - startTime, + error: errorText.slice(0, 500), + }).catch(() => {}); + return { success: false, status: response.status, error: errorText }; + } + + const payload = await response.json(); + const normalized = normalizeNvidiaNimImages(payload); + + if (normalized.data.length === 0) { + saveCallLog({ + method: "POST", + path: "/v1/images/generations", + status: 502, + model: `${provider}/${model}`, + provider, + duration: Date.now() - startTime, + error: "No images returned from NVIDIA NIM", + }).catch(() => {}); + return { success: false, status: 502, error: "No images returned from NVIDIA NIM" }; + } + + saveCallLog({ + method: "POST", + path: "/v1/images/generations", + status: 200, + model: `${provider}/${model}`, + provider, + duration: Date.now() - startTime, + responseBody: { images_count: normalized.data.length }, + }).catch(() => {}); + + return { success: true, data: normalized }; + } catch (err) { + if (log) log.error("IMAGE", `${provider} fetch error: ${(err as Error).message}`); + saveCallLog({ + method: "POST", + path: "/v1/images/generations", + status: 502, + model: `${provider}/${model}`, + provider, + duration: Date.now() - startTime, + error: (err as Error).message, + }).catch(() => {}); + return { + success: false, + status: 502, + error: `Image provider error: ${sanitizeErrorMessage((err as Error).message || err)}`, + }; + } +} diff --git a/tests/unit/nvidia-nim-image.test.ts b/tests/unit/nvidia-nim-image.test.ts new file mode 100644 index 0000000000..2f07ae325e --- /dev/null +++ b/tests/unit/nvidia-nim-image.test.ts @@ -0,0 +1,314 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +// Ported from upstream 9router#1195: NVIDIA NIM image generation (FLUX models). +// NVIDIA already exists as a CHAT provider (integrate.api.nvidia.com, +// OpenAI-compatible) — image generation is a distinct host (ai.api.nvidia.com) +// with a native NIM body shape, so it gets its own `nvidia-nim` format/handler +// (handleNvidiaNimImageGeneration) rather than reusing the OpenAI image path. + +import { handleImageGeneration } from "../../open-sse/handlers/imageGeneration.ts"; +import { getImageProvider } from "../../open-sse/config/imageRegistry.ts"; +import { + buildNvidiaNimRequestBody, + normalizeNvidiaNimImages, +} from "../../open-sse/handlers/imageGeneration/providers/nvidiaNim.ts"; + +test("nvidia is registered as an nvidia-nim image provider with the 4 FLUX models", () => { + const cfg = getImageProvider("nvidia"); + assert.ok(cfg, "expected an IMAGE_PROVIDERS entry for nvidia"); + assert.equal(cfg.format, "nvidia-nim"); + assert.equal(cfg.baseUrl, "https://ai.api.nvidia.com/v1/genai"); + assert.equal(cfg.authType, "apikey"); + assert.equal(cfg.authHeader, "bearer"); + + const ids = cfg.models.map((m) => m.id); + assert.deepEqual(ids, [ + "black-forest-labs/flux.1-dev", + "black-forest-labs/flux.1-schnell", + "black-forest-labs/flux.1-kontext-dev", + "black-forest-labs/flux.2-klein-4b", + ]); +}); + +test("handleImageGeneration(nvidia/flux.1-schnell): URL construction + minimal body shaping", async () => { + const originalFetch = globalThis.fetch; + let capturedUrl; + let capturedOptions; + + globalThis.fetch = async (url, options) => { + capturedUrl = String(url); + capturedOptions = options; + return new Response( + JSON.stringify({ artifacts: [{ base64: "base64nvidia", finishReason: "SUCCESS" }] }), + { status: 200, headers: { "content-type": "application/json" } } + ); + }; + + try { + const result = await handleImageGeneration({ + body: { + model: "nvidia/black-forest-labs/flux.1-schnell", + prompt: "A neon city", + width: 1344, + height: 1024, + seed: 7, + steps: 4, + }, + credentials: { apiKey: "nv-token" }, + log: null, + }); + + assert.equal(result.success, true); + assert.equal(capturedUrl, "https://ai.api.nvidia.com/v1/genai/black-forest-labs/flux.1-schnell"); + assert.equal(capturedOptions.method, "POST"); + assert.equal(capturedOptions.headers.Authorization, "Bearer nv-token"); + assert.equal(capturedOptions.headers.Accept, "application/json"); + + const requestBody = JSON.parse(capturedOptions.body); + assert.deepEqual(requestBody, { + prompt: "A neon city", + width: 1344, + height: 1024, + seed: 7, + steps: 4, + }); + + assert.equal(result.data.data[0].b64_json, "base64nvidia"); + assert.equal(result.data.data[0].finish_reason, "SUCCESS"); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("handleImageGeneration(nvidia/flux.2-klein-4b): sends edit input image as an array", async () => { + const originalFetch = globalThis.fetch; + + globalThis.fetch = async () => + new Response(JSON.stringify({ artifacts: [{ base64: "base64nvidiaedit" }] }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + + try { + const result = await handleImageGeneration({ + body: { + model: "nvidia/black-forest-labs/flux.2-klein-4b", + prompt: "Make the frog wear tiny glasses", + image: "data:image/png;example_id,0", + width: 1024, + height: 1024, + seed: 0, + steps: 4, + }, + credentials: { apiKey: "nv-token" }, + log: null, + }); + + assert.equal(result.success, true); + assert.equal(result.data.data[0].b64_json, "base64nvidiaedit"); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("buildNvidiaNimRequestBody(flux.2-klein-4b): image sent as an array, width/height passthrough", () => { + const req = buildNvidiaNimRequestBody("black-forest-labs/flux.2-klein-4b", { + prompt: "Make the frog wear tiny glasses", + image: "data:image/png;example_id,0", + width: 1024, + height: 1024, + seed: 0, + steps: 4, + }); + assert.deepEqual(req, { + prompt: "Make the frog wear tiny glasses", + width: 1024, + height: 1024, + image: ["data:image/png;example_id,0"], + seed: 0, + steps: 4, + }); +}); + +test("buildNvidiaNimRequestBody(flux.1-dev): omits input image in base mode", () => { + const req = buildNvidiaNimRequestBody("black-forest-labs/flux.1-dev", { + prompt: "A simple coffee shop interior", + mode: "base", + image: "data:image/png;example_id,0", + cfg_scale: 1.1, + width: 768, + height: 1344, + seed: 0, + steps: 50, + }); + assert.deepEqual(req, { + prompt: "A simple coffee shop interior", + mode: "base", + width: 768, + height: 1344, + cfg_scale: 1.1, + seed: 0, + steps: 50, + }); +}); + +test("buildNvidiaNimRequestBody(flux.1-dev): drops out-of-range dimensions and non-positive cfg_scale", () => { + const req = buildNvidiaNimRequestBody("black-forest-labs/flux.1-dev", { + prompt: "A simple coffee shop interior", + mode: "base", + image: "data:image/png;example_id,0", + cfg_scale: 0, + width: 1792, // out of the 768-1344 range + height: 1024, + seed: 0, + steps: 50, + }); + assert.deepEqual(req, { + prompt: "A simple coffee shop interior", + mode: "base", + seed: 0, + steps: 50, + }); + assert.ok(!("width" in req) && !("height" in req) && !("cfg_scale" in req)); +}); + +test("buildNvidiaNimRequestBody(flux.1-dev): sends control image as a string for non-base modes", () => { + const req = buildNvidiaNimRequestBody("black-forest-labs/flux.1-dev", { + prompt: "A simple coffee shop interior", + mode: "depth", + image: "data:image/png;example_id,0", + cfg_scale: 3.5, + width: 1024, + height: 1024, + seed: 0, + steps: 50, + }); + assert.deepEqual(req, { + prompt: "A simple coffee shop interior", + width: 1024, + height: 1024, + mode: "depth", + image: "data:image/png;example_id,0", + cfg_scale: 3.5, + seed: 0, + steps: 50, + }); +}); + +test("buildNvidiaNimRequestBody(flux.1-kontext-dev): uses aspect_ratio instead of width/height", () => { + const req = buildNvidiaNimRequestBody("black-forest-labs/flux.1-kontext-dev", { + prompt: "Now the mouse is holding pizza instead", + image: "data:image/png;example_id,0", + aspect_ratio: "match_input_image", + width: 1024, + height: 1024, + steps: 30, + cfg_scale: 3.5, + seed: 0, + }); + assert.deepEqual(req, { + prompt: "Now the mouse is holding pizza instead", + image: "data:image/png;example_id,0", + aspect_ratio: "match_input_image", + cfg_scale: 3.5, + seed: 0, + steps: 30, + }); +}); + +test("handleImageGeneration(nvidia/flux.1-kontext-dev): requires an input image, does not fetch", async () => { + const originalFetch = globalThis.fetch; + let fetchCalled = false; + globalThis.fetch = async () => { + fetchCalled = true; + throw new Error("fetch should not be called without an input image"); + }; + + try { + const result = await handleImageGeneration({ + body: { + model: "nvidia/black-forest-labs/flux.1-kontext-dev", + prompt: "Now the mouse is holding pizza instead", + }, + credentials: { apiKey: "nv-token" }, + log: null, + }); + + assert.equal(result.success, false); + assert.equal(result.status, 400); + assert.match(result.error, /requires an input image/i); + assert.equal(fetchCalled, false); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("normalizeNvidiaNimImages: accepts artifacts[], images[], data[], and single-value shapes", () => { + assert.deepEqual(normalizeNvidiaNimImages({ artifacts: [{ base64: "a1" }] }).data, [ + { b64_json: "a1" }, + ]); + assert.deepEqual(normalizeNvidiaNimImages({ images: ["b1", "b2"] }).data, [ + { b64_json: "b1" }, + { b64_json: "b2" }, + ]); + assert.deepEqual(normalizeNvidiaNimImages({ data: [{ b64_json: "c1" }] }).data, [ + { b64_json: "c1" }, + ]); + assert.deepEqual(normalizeNvidiaNimImages({ image: "d1" }).data, [{ b64_json: "d1" }]); + assert.deepEqual(normalizeNvidiaNimImages({ result: { image: "e1" } }).data, [ + { b64_json: "e1" }, + ]); + // Already-OpenAI-shaped responses pass through untouched. + const passthrough = { created: 123, data: [{ b64_json: "f1" }] }; + assert.deepEqual(normalizeNvidiaNimImages(passthrough), passthrough); + // Unrecognized shape -> empty data, never throws. + assert.deepEqual(normalizeNvidiaNimImages({ nonsense: true }).data, []); +}); + +test("handleImageGeneration(nvidia): upstream error body never leaks a stack trace", async () => { + const originalFetch = globalThis.fetch; + globalThis.fetch = async () => { + throw new Error(`boom at /home/user/secret/path/imageGeneration.ts:123:45`); + }; + + try { + const result = await handleImageGeneration({ + body: { + model: "nvidia/black-forest-labs/flux.1-schnell", + prompt: "test", + }, + credentials: { apiKey: "nv-token" }, + log: null, + }); + + assert.equal(result.success, false); + assert.equal(result.status, 502); + assert.ok(!result.error.includes("/home/"), "error must not leak an absolute path/stack"); + assert.ok(!result.error.includes(":123:45"), "error must not leak a stack trace location"); + } finally { + globalThis.fetch = originalFetch; + } +}); + +test("handleImageGeneration(nvidia): upstream non-2xx response is surfaced with its status", async () => { + const originalFetch = globalThis.fetch; + globalThis.fetch = async () => + new Response("rate limited", { status: 429, headers: { "content-type": "text/plain" } }); + + try { + const result = await handleImageGeneration({ + body: { + model: "nvidia/black-forest-labs/flux.1-schnell", + prompt: "test", + }, + credentials: { apiKey: "nv-token" }, + log: null, + }); + + assert.equal(result.success, false); + assert.equal(result.status, 429); + } finally { + globalThis.fetch = originalFetch; + } +});