diff --git a/open-sse/config/imageRegistry.ts b/open-sse/config/imageRegistry.ts index d40e3e1023..2315608f02 100644 --- a/open-sse/config/imageRegistry.ts +++ b/open-sse/config/imageRegistry.ts @@ -234,6 +234,7 @@ export const IMAGE_PROVIDERS: Record = { kie: { id: "kie", baseUrl: "https://api.kie.ai/api/v1/gpt4o-image/generate", + statusUrl: "https://api.kie.ai/api/v1/gpt4o-image/record-info", authType: "apikey", authHeader: "bearer", format: "kie-image", diff --git a/open-sse/config/musicRegistry.ts b/open-sse/config/musicRegistry.ts index 72f9681023..6cce3b22ee 100644 --- a/open-sse/config/musicRegistry.ts +++ b/open-sse/config/musicRegistry.ts @@ -15,6 +15,7 @@ interface MusicModel { interface MusicProvider { id: string; baseUrl: string; + statusUrl?: string; authType: string; authHeader: string; format: string; @@ -25,6 +26,7 @@ export const MUSIC_PROVIDERS: Record = { kie: { id: "kie", baseUrl: "https://api.kie.ai", + statusUrl: "https://api.kie.ai/api/v1/generate/record-info", authType: "apikey", authHeader: "bearer", format: "kie-music", diff --git a/open-sse/config/videoRegistry.ts b/open-sse/config/videoRegistry.ts index a7b4439c1a..2778c3e0b8 100644 --- a/open-sse/config/videoRegistry.ts +++ b/open-sse/config/videoRegistry.ts @@ -15,6 +15,7 @@ interface VideoModel { interface VideoProvider { id: string; baseUrl: string; + statusUrl?: string; authType: string; authHeader: string; format: string; @@ -25,6 +26,7 @@ export const VIDEO_PROVIDERS: Record = { kie: { id: "kie", baseUrl: "https://api.kie.ai", + statusUrl: "https://api.kie.ai/api/v1/jobs/recordInfo", authType: "apikey", authHeader: "bearer", format: "kie-video", diff --git a/open-sse/handlers/imageGeneration.ts b/open-sse/handlers/imageGeneration.ts index 070022ca55..fb14b6540c 100644 --- a/open-sse/handlers/imageGeneration.ts +++ b/open-sse/handlers/imageGeneration.ts @@ -19,6 +19,7 @@ import { randomUUID } from "crypto"; import { getImageProvider, parseImageModel } from "../config/imageRegistry.ts"; import { mapImageSize } from "../translator/image/sizeMapper.ts"; import { saveCallLog } from "@/lib/usageDb"; +import { sleep } from "../utils/sleep.ts"; import { submitComfyWorkflow, pollComfyResult, @@ -26,6 +27,15 @@ import { extractComfyOutputFiles, } from "../utils/comfyuiClient.ts"; +interface KieImageOptions { + model: string; + provider: string; + providerConfig: any; + body: any; + credentials: any; + log: any; +} + const OPENAI_IMAGE_TO_IMAGE_MODELS = new Set([ "black-forest-labs/FLUX.1-redux", "black-forest-labs/FLUX.1-depth", @@ -296,23 +306,24 @@ async function handleKieImageGeneration({ body, credentials, log, -}) { +}: KieImageOptions) { const startTime = Date.now(); const token = credentials?.apiKey || credentials?.accessToken; - const timeoutMs = normalizePositiveNumber(body.timeout_ms, 180000); + const timeoutMs = normalizePositiveNumber(body.timeout_ms, 300000); const pollIntervalMs = normalizePositiveNumber(body.poll_interval_ms, 2500); - const payload: Record = { + + const payload = { prompt: body.prompt, - size: typeof body.size === "string" ? body.size : "1:1", - nVariants: Number(body.n) > 0 ? Number(body.n) : 1, + image_size: mapImageSize(body.size, "1:1"), + num_images: body.n || 1, }; - try { - if (log) { - const promptPreview = String(body.prompt ?? "").slice(0, 60); - log.info("IMAGE", `${provider}/${model} (kie-image) | prompt: "${promptPreview}..."`); - } + if (log) { + const promptPreview = String(body.prompt ?? "").slice(0, 60); + log.info("IMAGE", `${provider}/${model} (kie-image) | prompt: "${promptPreview}..."`); + } + try { const createRes = await fetch(providerConfig.baseUrl, { method: "POST", headers: { @@ -347,9 +358,13 @@ async function handleKieImageGeneration({ }); } - const statusUrl = providerConfig.baseUrl - .replace(/\/generate$/, "/record-info") - .replace("/api/v1/gpt4o-image/generate", "/api/v1/gpt4o-image/record-info"); + // Use statusUrl from providerConfig if available, fallback to dynamic derivation + const statusUrl = + providerConfig.statusUrl || + providerConfig.baseUrl + .replace(/\/generate$/, "/record-info") + .replace("/api/v1/gpt4o-image/generate", "/api/v1/gpt4o-image/record-info"); + const deadline = Date.now() + timeoutMs; while (Date.now() < deadline) { const pollUrl = new URL(statusUrl); @@ -442,12 +457,10 @@ async function handleKieImageGeneration({ model, status: 502, startTime, - error: `Image provider error: ${err.message}`, - requestBody: payload, + error: `Image provider error: ${err instanceof Error ? err.message : String(err)}`, }); } } - /** * Handle Gemini-format image generation (Antigravity / Nano Banana) * Uses Gemini's generateContent API with responseModalities: ["TEXT", "IMAGE"] @@ -2039,10 +2052,6 @@ function normalizePositiveNumber(value, fallback) { return Math.floor(n); } -function sleep(ms) { - return new Promise((resolve) => setTimeout(resolve, ms)); -} - /** * Handle SD WebUI image generation (local, no auth) * POST {baseUrl} with { prompt, negative_prompt, width, height, steps } diff --git a/open-sse/handlers/musicGeneration.ts b/open-sse/handlers/musicGeneration.ts index d05494b073..a8092d88c3 100644 --- a/open-sse/handlers/musicGeneration.ts +++ b/open-sse/handlers/musicGeneration.ts @@ -22,6 +22,7 @@ import { extractComfyOutputFiles, } from "../utils/comfyuiClient.ts"; import { saveCallLog } from "@/lib/usageDb"; +import { sleep } from "../utils/sleep.ts"; /** * Handle music generation request @@ -176,6 +177,13 @@ async function handleKieMusicGeneration({ body, credentials, log, +}: { + model: string; + provider: string; + providerConfig: any; + body: any; + credentials: any; + log: any; }) { const startTime = Date.now(); const timeoutMs = Number(body.timeout_ms) > 0 ? Number(body.timeout_ms) : 300000; @@ -216,10 +224,10 @@ async function handleKieMusicGeneration({ } const deadline = Date.now() + timeoutMs; - const statusBaseUrl = `${baseUrl}/api/v1/generate/record-info`; + const statusUrl = providerConfig.statusUrl || `${baseUrl}/api/v1/generate/record-info`; while (Date.now() < deadline) { - const pollUrl = new URL(statusBaseUrl); + const pollUrl = new URL(statusUrl); pollUrl.searchParams.set("taskId", String(taskId)); const recordRes = await fetch(pollUrl.toString(), { @@ -290,10 +298,10 @@ async function handleKieMusicGeneration({ error: `KIE music polling timed out after ${timeoutMs}ms`, }; } catch (err) { - return { success: false, status: 502, error: `Music provider error: ${err.message}` }; + return { + success: false, + status: 502, + error: `Music provider error: ${err instanceof Error ? err.message : String(err)}`, + }; } } - -function sleep(ms: number) { - return new Promise((resolve) => setTimeout(resolve, ms)); -} diff --git a/open-sse/handlers/videoGeneration.ts b/open-sse/handlers/videoGeneration.ts index 15d6c570ac..98b05e9208 100644 --- a/open-sse/handlers/videoGeneration.ts +++ b/open-sse/handlers/videoGeneration.ts @@ -23,6 +23,7 @@ import { extractComfyOutputFiles, } from "../utils/comfyuiClient.ts"; import { saveCallLog } from "@/lib/usageDb"; +import { sleep } from "../utils/sleep.ts"; /** * Handle video generation request @@ -267,6 +268,28 @@ async function handleSDWebUIVideoGeneration({ model, provider, providerConfig, b } } +function normalizeKieVideoResult(recordData: any): string[] { + let resultJson: Record = {}; + try { + resultJson = + typeof recordData?.data?.resultJson === "string" + ? JSON.parse(recordData.data.resultJson) + : recordData?.data?.resultJson || {}; + } catch { + resultJson = {}; + } + + const urls = Array.isArray(resultJson?.resultUrls) + ? (resultJson.resultUrls as string[]) + : Array.isArray(resultJson?.videoUrls) + ? (resultJson.videoUrls as string[]) + : Array.isArray(recordData?.data?.response?.resultUrls) + ? (recordData.data.response.resultUrls as string[]) + : []; + + return urls.filter((url: unknown) => typeof url === "string" && url.length > 0); +} + async function handleKieVideoGeneration({ model, provider, @@ -274,6 +297,13 @@ async function handleKieVideoGeneration({ body, credentials, log, +}: { + model: string; + provider: string; + providerConfig: any; + body: any; + credentials: any; + log: any; }) { const startTime = Date.now(); const timeoutMs = Number(body.timeout_ms) > 0 ? Number(body.timeout_ms) : 300000; @@ -318,10 +348,10 @@ async function handleKieVideoGeneration({ } const deadline = Date.now() + timeoutMs; - const statusBaseUrl = `${baseUrl}/api/v1/jobs/recordInfo`; + const statusUrl = providerConfig.statusUrl || `${baseUrl}/api/v1/jobs/recordInfo`; while (Date.now() < deadline) { - const pollUrl = new URL(statusBaseUrl); + const pollUrl = new URL(statusUrl); pollUrl.searchParams.set("taskId", String(taskId)); const recordRes = await fetch(pollUrl.toString(), { @@ -342,25 +372,8 @@ async function handleKieVideoGeneration({ ).toLowerCase(); if (state === "success" || state === "1" || state === "finished") { - let resultJson: Record = {}; - try { - resultJson = - typeof recordData?.data?.resultJson === "string" - ? JSON.parse(recordData.data.resultJson) - : recordData?.data?.resultJson || {}; - } catch { - resultJson = {}; - } - const urls = Array.isArray(resultJson?.resultUrls) - ? resultJson.resultUrls - : Array.isArray(resultJson?.videoUrls) - ? resultJson.videoUrls - : Array.isArray(recordData?.data?.response?.resultUrls) - ? recordData.data.response.resultUrls - : []; - const videos = urls - .filter((url: unknown) => typeof url === "string" && url.length > 0) - .map((url: unknown) => ({ url: url as string, format: "mp4" })); + const videoUrls = normalizeKieVideoResult(recordData); + const videos = videoUrls.map((url) => ({ url, format: "mp4" })); saveCallLog({ method: "POST", @@ -405,10 +418,10 @@ async function handleKieVideoGeneration({ error: `KIE video polling timed out after ${timeoutMs}ms`, }; } catch (err) { - return { success: false, status: 502, error: `Video provider error: ${err.message}` }; + return { + success: false, + status: 502, + error: `Video provider error: ${err instanceof Error ? err.message : String(err)}`, + }; } } - -function sleep(ms: number) { - return new Promise((resolve) => setTimeout(resolve, ms)); -} diff --git a/open-sse/utils/sleep.ts b/open-sse/utils/sleep.ts new file mode 100644 index 0000000000..3216842dee --- /dev/null +++ b/open-sse/utils/sleep.ts @@ -0,0 +1,9 @@ +/** + * Shared sleep utility to pause execution for a given number of milliseconds. + * + * @param ms - Number of milliseconds to sleep + * @returns Promise that resolves after the specified time + */ +export function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +}