diff --git a/apps/api/src/app.ts b/apps/api/src/app.ts index 2a9e5c1..d3ff4e7 100644 --- a/apps/api/src/app.ts +++ b/apps/api/src/app.ts @@ -15,6 +15,7 @@ import type { ComposioProvider, ConnectorRegistry, DestinationEmulator, + MediaOperationRegistration, RemoteConnectorDependencies, } from "@rakazo/adapters"; import { @@ -143,6 +144,8 @@ export async function createApp( email?: TransactionalEmailProvider; remoteConnectors?: RemoteConnectorDependencies; logger?: Logger; + /** Test seam: media operations instead of the configured ones. */ + mediaOperations?: MediaOperationRegistration[]; } = {}, ): Promise { const { @@ -155,6 +158,7 @@ export async function createApp( email: emailOverride, remoteConnectors, logger: loggerOverride, + mediaOperations, ...envOverrides } = overrides; const env = { ...loadEnv(process.env), ...envOverrides }; @@ -384,9 +388,9 @@ export async function createApp( events, jobs, artifacts, - registrations: createMediaOperations({ - loadSettings: () => loadMediaSettings({ prisma, secrets }), - }), + registrations: + mediaOperations ?? + createMediaOperations({ loadSettings: () => loadMediaSettings({ prisma, secrets }) }), }); const executor = createRunExecutor({ prisma, @@ -468,6 +472,7 @@ export async function createApp( const router = createRouter({ cloudAgent, + media, prisma, events, auth, diff --git a/apps/api/src/media-jobs.ts b/apps/api/src/media-jobs.ts new file mode 100644 index 0000000..4a289e8 --- /dev/null +++ b/apps/api/src/media-jobs.ts @@ -0,0 +1,118 @@ +import { ORPCError } from "@orpc/server"; +import type { MediaRuntime } from "@rakazo/adapter-kit"; +import { + type EncryptedSecretStore, + loadMediaSettings, + MediaJobRejectedError, + MediaSettingsInvalidError, + mediaSettingsView, + saveMediaSettings, +} from "@rakazo/adapters"; +import type { Actor, MediaJobStartInput, MediaSettingsUpdate } from "@rakazo/contracts"; +import type { PrismaClient } from "@rakazo/db"; +import { resolveThreadTarget } from "./thread-target.js"; + +export interface MediaRouteDeps { + prisma: PrismaClient; + secrets: EncryptedSecretStore; + media?: MediaRuntime; +} + +/** Media backends run on the deployment owner's own server, so only they configure them. */ +function requireDeploymentOwner(actor: Actor) { + if (!actor.isDeploymentOwner) throw new ORPCError("FORBIDDEN"); +} + +function requireMedia(deps: MediaRouteDeps): MediaRuntime { + if (!deps.media) throw new ORPCError("NOT_IMPLEMENTED"); + return deps.media; +} + +export async function getMediaSettings(deps: MediaRouteDeps, actor: Actor) { + requireDeploymentOwner(actor); + return mediaSettingsView(await loadMediaSettings(deps)); +} + +export async function updateMediaSettings( + deps: MediaRouteDeps, + actor: Actor, + input: MediaSettingsUpdate, +) { + requireDeploymentOwner(actor); + try { + return await saveMediaSettings(deps, input); + } catch (error) { + if (error instanceof MediaSettingsInvalidError) { + throw new ORPCError("BAD_REQUEST", { message: error.message }); + } + throw error; + } +} + +/** + * Starts an operation a user picked on a card (export, transcribe, edit) in the same conversation. + * The thread is authorized like every thread procedure; the runtime checks that the operation's + * source artifacts belong to the actor, and the job can only ever read the actor's artifacts. + */ +export async function startMediaJob(deps: MediaRouteDeps, actor: Actor, input: MediaJobStartInput) { + const media = requireMedia(deps); + const target = await resolveThreadTarget(deps.prisma, actor, input); + const source = await deps.prisma.message.findFirst({ + where: { id: input.messageId, threadId: target.threadId }, + select: { botId: true }, + }); + if (!source) throw new ORPCError("NOT_FOUND"); + const bots = + target.kind === "bot" + ? [target.botId] + : [ + ...(target.leadBotId && target.memberBotIds.includes(target.leadBotId) + ? [target.leadBotId] + : []), + ...target.memberBotIds, + ]; + const botId = source.botId && bots.includes(source.botId) ? source.botId : bots[0]; + if (!botId) throw new ORPCError("NOT_FOUND"); + try { + return await media.start({ + operation: input.operation, + input: input.input, + spaceId: actor.spaceId, + userId: actor.userId, + botId, + groupId: target.kind === "group" ? target.groupId : null, + threadId: target.threadId, + runId: null, + }); + } catch (error) { + if (error instanceof MediaJobRejectedError) { + throw new ORPCError("BAD_REQUEST", { message: error.message }); + } + throw error; + } +} + +async function ownedJobId(deps: MediaRouteDeps, actor: Actor, jobId: string): Promise { + const job = await deps.prisma.mediaJob.findFirst({ + where: { id: jobId, spaceId: actor.spaceId, userId: actor.userId }, + select: { id: true }, + }); + if (!job) throw new ORPCError("NOT_FOUND"); + return job.id; +} + +export async function getMediaJob(deps: MediaRouteDeps, actor: Actor, jobId: string) { + const media = requireMedia(deps); + const snapshot = await media.get(await ownedJobId(deps, actor, jobId)); + if (!snapshot) throw new ORPCError("NOT_FOUND"); + return snapshot; +} + +export async function cancelMediaJob(deps: MediaRouteDeps, actor: Actor, jobId: string) { + const media = requireMedia(deps); + const id = await ownedJobId(deps, actor, jobId); + await media.cancel(id); + const snapshot = await media.get(id); + if (!snapshot) throw new ORPCError("NOT_FOUND"); + return snapshot; +} diff --git a/apps/api/src/router.ts b/apps/api/src/router.ts index 7756f51..e1447e8 100644 --- a/apps/api/src/router.ts +++ b/apps/api/src/router.ts @@ -6,6 +6,7 @@ import type { ArtifactStore, ConnectorCatalogItem, JobPublisher, + MediaRuntime, MemoryStore, SandboxProvider, } from "@rakazo/adapter-kit"; @@ -166,6 +167,13 @@ import { renameMainConversation, } from "./main-conversation.js"; import { buildMcpUpdateMaterial } from "./mcp-material.js"; +import { + cancelMediaJob, + getMediaJob, + getMediaSettings, + startMediaJob, + updateMediaSettings, +} from "./media-jobs.js"; import { disconnectMemoryProvider, getMemoryProviderConfig, @@ -447,6 +455,7 @@ function mcpAssignmentDto(row: { export interface RouterDeps { cloudAgent?: CloudAgentConnection | null; + media?: MediaRuntime; prisma: PrismaClient; events: ThreadEvents; auth: Auth; @@ -3482,6 +3491,27 @@ export function createRouter(deps: RouterDeps) { return { ok: true as const }; }), }, + media: { + settings: { + get: authed.media.settings.get.handler(({ context }) => + getMediaSettings(deps, context.actor), + ), + update: authed.media.settings.update.handler(({ context, input }) => + updateMediaSettings(deps, context.actor, input), + ), + }, + jobs: { + start: authed.media.jobs.start.handler(({ context, input }) => + startMediaJob(deps, context.actor, input), + ), + get: authed.media.jobs.get.handler(({ context, input }) => + getMediaJob(deps, context.actor, input.jobId), + ), + cancel: authed.media.jobs.cancel.handler(({ context, input }) => + cancelMediaJob(deps, context.actor, input.jobId), + ), + }, + }, connections: { catalog: authed.connections.catalog.handler(async ({ context, input }) => { const adapterContext = connectionContext( diff --git a/packages/adapters/src/builtin-tools.ts b/packages/adapters/src/builtin-tools.ts index d8fc9b3..5d25634 100644 --- a/packages/adapters/src/builtin-tools.ts +++ b/packages/adapters/src/builtin-tools.ts @@ -2,6 +2,9 @@ import type { ConnectorTool } from "@rakazo/adapter-kit"; import { BotSecretDestination, BotSecretName, + EditImageInputSchema, + GenerateImageInputSchema, + GenerateVideoInputSchema, MODEL_TIERS, SecretAskPurpose, SecretHttpRequest, @@ -11,7 +14,14 @@ import { import { z } from "zod"; import { mediaToolInputSchema } from "./media-tools.js"; -type MediaToolSchema = Parameters[0]; +/** A media tool's parameters from its operation's input schema, plus `continue_after`. */ +function mediaParameters(schema: z.ZodType) { + const json = z.toJSONSchema(schema, { io: "input" }); + return mediaToolInputSchema({ + properties: (json.properties ?? {}) as Record, + required: json.required, + }); +} export const SUBAGENT_MODEL_TIER_DESCRIPTION = "Optional model tier: fast for quick simple work, code for programming, standard otherwise. Uses the model this space maps to the tier, or your own model when it maps none. Ignored when model_provider and model_id are set."; @@ -330,11 +340,29 @@ export const builtinAgentTools: ConnectorTool[] = [ }, }, }, + { + name: "generate_image", + description: + "Generate up to 4 images from a prompt on the self-hosted image server, or rework a source image (image-to-image).", + inputSchema: mediaParameters(GenerateImageInputSchema), + }, + { + name: "generate_video", + description: + "Generate a short video clip from a prompt on the self-hosted image server, optionally starting from an image.", + inputSchema: mediaParameters(GenerateVideoInputSchema), + }, + { + name: "edit_image", + description: + "Edit an image artifact into a new one: crop, rotate, flip, resize (crop, pad, or blur to fit a format), add borders, convert the format, or repaint it from a prompt (only when generate_image is available). Edits apply in that order.", + inputSchema: mediaParameters(EditImageInputSchema), + }, { name: "transcribe", description: "Transcribe speech in an audio or video artifact into SRT and VTT subtitle files, shown on a media card. Set continue_after to get the transcript back, e.g. to burn the .srt in with edit_video.", - inputSchema: mediaToolInputSchema(z.toJSONSchema(TranscribeInputSchema) as MediaToolSchema), + inputSchema: mediaParameters(TranscribeInputSchema), }, { name: "add_mcp_server", diff --git a/packages/adapters/src/comfyui-provider.test.ts b/packages/adapters/src/comfyui-provider.test.ts new file mode 100644 index 0000000..2773c1f --- /dev/null +++ b/packages/adapters/src/comfyui-provider.test.ts @@ -0,0 +1,227 @@ +import sharp from "sharp"; +import { describe, expect, it } from "vitest"; +import type { ComfyUiProviderConfig } from "./comfyui-provider.js"; +import { + createComfyUiProvider, + describeExecutionError, + parseGpuStatus, +} from "./comfyui-provider.js"; +import { ComfyUiEmulator } from "./testing/comfyui-emulator.js"; + +const BASE_URL = "http://127.0.0.1:8188"; +const IMAGE = { prompt: "a lighthouse", width: 512, height: 512 }; + +function setup(config: Partial = {}, emulator = new ComfyUiEmulator()) { + const provider = createComfyUiProvider({ + baseUrl: BASE_URL, + fetch: emulator.fetch, + pollIntervalMs: 1, + ...config, + }); + return { provider, emulator }; +} + +function submitted(emulator: ComfyUiEmulator, index = 0) { + const graph = emulator.prompts[index]?.graph ?? {}; + const byClass = (classType: string) => + Object.values(graph).find((node) => node.class_type === classType)?.inputs ?? {}; + return { graph, byClass }; +} + +describe("createComfyUiProvider", () => { + it("refuses metadata hosts, public hosts and keys over public http", () => { + expect(() => createComfyUiProvider({ baseUrl: "http://169.254.169.254" })).toThrow( + "blocked metadata or link-local host", + ); + expect(() => createComfyUiProvider({ baseUrl: "http://comfy.example.com" })).toThrow( + "Public model endpoints are blocked", + ); + expect(() => createComfyUiProvider({ baseUrl: "ftp://127.0.0.1" })).toThrow("http or https"); + }); + + it("validates custom workflows up front", () => { + expect(() => createComfyUiProvider({ baseUrl: BASE_URL, imageWorkflow: "{" })).toThrow( + "Workflow is not valid JSON", + ); + expect(() => + createComfyUiProvider({ baseUrl: BASE_URL, videoWorkflow: '{"nodes": [], "links": []}' }), + ).toThrow("API format"); + }); + + it("keeps a reverse-proxy path prefix and sends the API key", async () => { + const emulator = new ComfyUiEmulator({ apiKey: "comfy-key", basePath: "/comfy" }); + const { provider } = setup({ baseUrl: `${BASE_URL}/comfy/`, apiKey: "comfy-key" }, emulator); + await provider.generateImage(IMAGE); + expect(emulator.requests.map((request) => request.path)).toContain("prompt"); + expect(emulator.requests.every((request) => request.authorization === "Bearer comfy-key")).toBe( + true, + ); + const { provider: keyless } = setup({ baseUrl: `${BASE_URL}/comfy` }, emulator); + await expect(keyless.generateImage(IMAGE)).rejects.toThrow("ComfyUI returned 401"); + }); + + it("fills the default text-to-image workflow", async () => { + const { provider, emulator } = setup(); + await provider.generateImage({ ...IMAGE, width: 1001, negativePrompt: "fog", seed: 42 }); + const { byClass } = submitted(emulator); + expect(byClass("CheckpointLoaderSimple").ckpt_name).toBe("sd_xl_base_1.0.safetensors"); + expect(byClass("EmptyLatentImage")).toMatchObject({ width: 1000, height: 512 }); + expect(byClass("KSampler")).toMatchObject({ seed: 42, denoise: 1 }); + const texts = emulator.prompts[0]!.graph; + expect(Object.values(texts).map((node) => node.inputs.text)).toEqual( + expect.arrayContaining(["a lighthouse", "fog"]), + ); + }); + + it("uses the configured model, else the first checkpoint the video workflows do not use", async () => { + const emulator = new ComfyUiEmulator(); + emulator.checkpoints = ["ltx-video-2b-v0.9.5.safetensors", "v1-5-pruned.safetensors"]; + emulator.checkpointListFormat = "legacy"; + await setup({}, emulator).provider.generateImage(IMAGE); + expect(submitted(emulator, 0).byClass("CheckpointLoaderSimple").ckpt_name).toBe( + "v1-5-pruned.safetensors", + ); + emulator.checkpoints.push("dreamshaper.safetensors"); + await setup({ imageModel: "dreamshaper.safetensors" }, emulator).provider.generateImage(IMAGE); + expect(submitted(emulator, 1).byClass("CheckpointLoaderSimple").ckpt_name).toBe( + "dreamshaper.safetensors", + ); + emulator.checkpoints = ["ltx-video-2b-v0.9.5.safetensors"]; + await expect(setup({}, emulator).provider.generateImage(IMAGE)).rejects.toThrow( + "ComfyUI has no image checkpoint installed", + ); + }); + + it("uploads the source image for image-to-image", async () => { + const { provider, emulator } = setup(); + const bytes = new Uint8Array(Buffer.from((await sourcePng()).bytes)); + await provider.generateImage({ + ...IMAGE, + source: { bytes, mimeType: "image/png" }, + strength: 0.3, + }); + const { byClass } = submitted(emulator); + const uploaded = String(byClass("LoadImage").image); + expect(uploaded).toMatch(/^evoo-.+\.png$/); + expect(Buffer.from(emulator.uploads.get(uploaded)!).equals(Buffer.from(bytes))).toBe(true); + expect(byClass("KSampler").denoise).toBe(0.3); + }); + + it("lets a custom workflow replace only the default it matches", async () => { + const custom = { + "1": { + class_type: "CustomTextToImage", + inputs: { text: "style: {{prompt}}", seed: "{{seed}}" }, + }, + "2": { class_type: "SaveImage", inputs: { images: ["1", 0], filename_prefix: "x" } }, + }; + const { provider, emulator } = setup({ imageWorkflow: JSON.stringify(custom) }); + await provider.generateImage(IMAGE); + expect(submitted(emulator, 0).byClass("CustomTextToImage").text).toBe("style: a lighthouse"); + await provider.generateImage({ ...IMAGE, source: await sourcePng() }); + expect(submitted(emulator, 1).byClass("LoadImage").image).toBeDefined(); + }); + + it("fills the LTX-Video workflow with an 8n+1 frame count and reports the clip", async () => { + const { provider, emulator } = setup(); + const [clip] = await provider.generateVideo({ + prompt: "a paper boat", + width: 470, + height: 836, + durationSeconds: 5, + }); + const { byClass } = submitted(emulator); + expect(byClass("EmptyLTXVLatentVideo")).toMatchObject({ width: 480, height: 832, length: 121 }); + expect(byClass("CreateVideo").fps).toBe(24); + expect(clip).toMatchObject({ + mimeType: "video/mp4", + width: 480, + height: 832, + durationMs: 5042, + }); + }); + + it("reads video from a custom workflow's gifs output and skips frame previews", async () => { + const custom = { + "1": { class_type: "TextToFrames", inputs: { text: "{{prompt}}" } }, + "2": { class_type: "PreviewImage", inputs: { images: ["1", 0] } }, + "3": { class_type: "VHS_VideoCombine", inputs: { images: ["1", 0], frame_rate: 16 } }, + }; + const { provider } = setup({ videoWorkflow: JSON.stringify(custom) }); + const clips = await provider.generateVideo({ ...IMAGE, durationSeconds: 2 }); + expect(clips.map((clip) => [clip.mimeType, clip.durationMs])).toEqual([ + ["video/mp4", undefined], + ]); + }); + + it("explains rejected workflows and failed runs", async () => { + const emulator = new ComfyUiEmulator(); + const { provider } = setup({ imageModel: "missing.safetensors" }, emulator); + await expect(provider.generateImage(IMAGE)).rejects.toThrow( + "ComfyUI rejected the workflow: CheckpointLoaderSimple: Value not in list: ckpt_name: 'missing.safetensors' not in list", + ); + const working = setup({}, emulator).provider; + emulator.failNextExecution = "CUDA out of memory. Tried to allocate 2.00 GiB\nTraceback..."; + await expect(working.generateImage(IMAGE)).rejects.toThrow( + /^ComfyUI failed in KSampler: CUDA out of memory\. Tried to allocate 2\.00 GiB$/, + ); + emulator.failNextStatus = 500; + await expect(working.generateImage(IMAGE)).rejects.toThrow("ComfyUI returned 500"); + }); + + it("reports an unreachable server", async () => { + const provider = createComfyUiProvider({ + baseUrl: BASE_URL, + fetch: async () => { + throw new TypeError("fetch failed"); + }, + }); + await expect(provider.generateImage(IMAGE)).rejects.toThrow("ComfyUI is unreachable"); + await expect(provider.gpuStatus()).rejects.toThrow("ComfyUI is unreachable"); + }); + + it("interrupts its own running prompt when the caller cancels", async () => { + const emulator = new ComfyUiEmulator(); + emulator.pendingPolls = 1_000_000; + const { provider } = setup({ pollIntervalMs: 5 }, emulator); + const controller = new AbortController(); + const run = provider.generateImage({ ...IMAGE, signal: controller.signal }); + while (!emulator.requests.some((request) => request.path.startsWith("history/"))) { + await new Promise((resolve) => setTimeout(resolve, 1)); + } + controller.abort(new Error("cancelled by user")); + await expect(run).rejects.toThrow("cancelled by user"); + expect(emulator.deleted).toEqual(["prompt-1"]); + expect(emulator.interrupted).toEqual(["prompt-1"]); + }); + + it("reports GPU memory other processes leave to ComfyUI", async () => { + const emulator = new ComfyUiEmulator(); + emulator.gpu = { totalMb: 12_288, freeMb: 3_000, torchReservedMb: 7_000, torchFreeMb: 1_000 }; + await expect(setup({}, emulator).provider.gpuStatus()).resolves.toEqual({ + vramFreeMb: 9_000, + vramTotalMb: 12_288, + }); + emulator.gpu = null; + await expect(setup({}, emulator).provider.gpuStatus()).resolves.toBeUndefined(); + expect(parseGpuStatus({ devices: [{ type: "cuda", vram_total: 0, vram_free: 0 }] })).toBe( + undefined, + ); + }); + + it("names interrupted runs", () => { + expect(describeExecutionError([["execution_interrupted", {}]])).toBe( + "ComfyUI interrupted the run", + ); + expect(describeExecutionError(undefined)).toBe("ComfyUI failed"); + }); +}); + +async function sourcePng() { + const png = await sharp({ + create: { width: 32, height: 32, channels: 3, background: { r: 1, g: 2, b: 3 } }, + }) + .png() + .toBuffer(); + return { bytes: new Uint8Array(png), mimeType: "image/png" }; +} diff --git a/packages/adapters/src/comfyui-provider.ts b/packages/adapters/src/comfyui-provider.ts new file mode 100644 index 0000000..e2410a1 --- /dev/null +++ b/packages/adapters/src/comfyui-provider.ts @@ -0,0 +1,593 @@ +import { randomInt, randomUUID } from "node:crypto"; +import { setTimeout as sleep } from "node:timers/promises"; +import type { + AdapterDescriptor, + GeneratedMedia, + GpuStatus, + ImageGenerationCapabilities, + ImageGenerationProvider, + ImageGenerationRequest, + MediaFile, + MediaProgressListener, + VideoGenerationCapabilities, + VideoGenerationProvider, + VideoGenerationRequest, +} from "@rakazo/adapter-kit"; +import type { AttachmentKind } from "@rakazo/contracts"; +import { + ARTIFACT_MAX_BYTES_BY_KIND, + attachmentKindOf, + MEDIA_VIDEO_MAX_SECONDS, +} from "@rakazo/contracts"; +import { attachmentExtensionForMimeType, inferAttachmentMimeType } from "@rakazo/core"; +import type { ComfyWorkflow, WorkflowTemplate, WorkflowValues } from "@rakazo/core/media/workflows"; +import { + DEFAULT_WORKFLOWS, + fillWorkflow, + parseWorkflowTemplate, +} from "@rakazo/core/media/workflows"; +import { + assertAllowedOpenAiCompatibleRequestUrl, + assertHttpsForKeyedOpenAiCompatibleUrl, +} from "./openai-compatible-url.js"; +import { createOpenAiCompatibleFetch } from "./pi-openai-compatible-provider.js"; +import { readBodyCapped } from "./web-ssrf.js"; + +export const COMFYUI_PROVIDER_ID = "comfyui"; +/** Frame rate of generated video; frame counts follow LTX-Video's 8n+1 rule. */ +export const COMFYUI_VIDEO_FPS = 24; +const MAX_IMAGE_SIDE = 2048; +const REQUEST_TIMEOUT_MS = 30_000; +const TRANSFER_TIMEOUT_MS = 120_000; +const CANCEL_TIMEOUT_MS = 5_000; +const IMAGE_PROMPT_TIMEOUT_MS = 10 * 60_000; +const VIDEO_PROMPT_TIMEOUT_MS = 30 * 60_000; +const POLL_INTERVAL_MS = 1_000; +const MAX_JSON_BYTES = 8 * 1024 * 1024; +const MAX_ERROR_CHARS = 400; +const MIB = 1024 * 1024; +const IMAGE_STEPS = 25; +const VIDEO_STEPS = 30; +const DEFAULT_STRENGTH = 0.6; +const DEFAULT_NEGATIVE_PROMPTS = { + image: "blurry, low quality, distorted, watermark, text", + video: "worst quality, inconsistent motion, blurry, jittery, distorted", +}; +const JSON_HEADERS = { Accept: "application/json", "Content-Type": "application/json" }; + +export interface ComfyUiProviderConfig { + baseUrl: string; + apiKey?: string; + /** Checkpoint for `{{model}}`; the first installed image checkpoint when omitted. */ + imageModel?: string; + /** API-format workflow JSON that replaces the matching default image workflow. */ + imageWorkflow?: string; + /** API-format workflow JSON that replaces the matching default video workflow. */ + videoWorkflow?: string; + fetch?: typeof globalThis.fetch; + /** How long to wait between checks for a finished prompt; tests shorten it. */ + pollIntervalMs?: number; +} + +export interface ComfyUiProvider extends ImageGenerationProvider, VideoGenerationProvider { + describe(): AdapterDescriptor; + gpuStatus(signal?: AbortSignal): Promise; +} + +type MediaKind = "image" | "video"; +type ReadBody = (response: Response, signal: AbortSignal) => Promise; + +interface OutputFile { + filename: string; + subfolder: string; + type: string; + mimeType: string; + kind: AttachmentKind; +} + +interface GenerationJob { + kind: MediaKind; + template: WorkflowTemplate; + prompt: string; + negativePrompt?: string; + width: number; + height: number; + seed?: number; + count: number; + source?: MediaFile; + strength?: number; + frames?: number; + signal?: AbortSignal; + onProgress?: MediaProgressListener; +} + +/** + * Image and video generation on a self-hosted ComfyUI server through its HTTP API: submit an + * API-format workflow, poll its history, download the outputs. Requests go through the same + * address-checked fetch as OpenAI-compatible model servers. + */ +export function createComfyUiProvider(config: ComfyUiProviderConfig): ComfyUiProvider { + return new ComfyUiClient(config); +} + +class ComfyUiClient implements ComfyUiProvider { + private readonly baseUrl: URL; + private readonly apiKey: string | undefined; + private readonly imageModel: string | undefined; + private readonly imageWorkflow: WorkflowTemplate | undefined; + private readonly videoWorkflow: WorkflowTemplate | undefined; + private readonly fetch: typeof globalThis.fetch; + private readonly pollIntervalMs: number; + private readonly clientId = randomUUID(); + + constructor(config: ComfyUiProviderConfig) { + this.baseUrl = comfyUiBaseUrl(config.baseUrl); + this.apiKey = config.apiKey?.trim() || undefined; + assertHttpsForKeyedOpenAiCompatibleUrl(this.baseUrl, this.apiKey); + this.imageModel = config.imageModel?.trim() || undefined; + this.imageWorkflow = config.imageWorkflow?.trim() + ? parseWorkflowTemplate(config.imageWorkflow) + : undefined; + this.videoWorkflow = config.videoWorkflow?.trim() + ? parseWorkflowTemplate(config.videoWorkflow) + : undefined; + this.fetch = createOpenAiCompatibleFetch(config.fetch); + this.pollIntervalMs = config.pollIntervalMs ?? POLL_INTERVAL_MS; + } + + describe() { + return { + id: COMFYUI_PROVIDER_ID, + contractVersion: "1", + adapterVersion: "0.1.0", + capabilities: { + textToImage: true, + imageToImage: true, + inpaint: false, + maxWidth: MAX_IMAGE_SIDE, + maxHeight: MAX_IMAGE_SIDE, + textToVideo: true, + imageToVideo: true, + maxDurationSeconds: MEDIA_VIDEO_MAX_SECONDS, + }, + }; + } + + async gpuStatus(signal?: AbortSignal): Promise { + return parseGpuStatus(await this.getJson("system_stats", signal)); + } + + async generateImage(request: ImageGenerationRequest): Promise { + if (request.mask) throw new Error("ComfyUI inpainting is not supported"); + return this.generate({ + ...request, + kind: "image", + template: pickTemplate( + this.imageWorkflow, + Boolean(request.source), + DEFAULT_WORKFLOWS.textToImage, + DEFAULT_WORKFLOWS.imageToImage, + ), + width: roundSide(request.width, 8), + height: roundSide(request.height, 8), + count: request.count ?? 1, + }); + } + + async generateVideo(request: VideoGenerationRequest): Promise { + const seconds = Math.min(Math.max(request.durationSeconds, 1), MEDIA_VIDEO_MAX_SECONDS); + return this.generate({ + ...request, + kind: "video", + template: pickTemplate( + this.videoWorkflow, + Boolean(request.source), + DEFAULT_WORKFLOWS.textToVideo, + DEFAULT_WORKFLOWS.imageToVideo, + ), + width: roundSide(request.width, 32), + height: roundSide(request.height, 32), + frames: 8 * Math.round((seconds * COMFYUI_VIDEO_FPS) / 8) + 1, + count: 1, + }); + } + + private async generate(job: GenerationJob): Promise { + const { template, signal } = job; + const values: WorkflowValues = { + prompt: job.prompt, + negative_prompt: job.negativePrompt ?? DEFAULT_NEGATIVE_PROMPTS[job.kind], + width: job.width, + height: job.height, + steps: job.kind === "image" ? IMAGE_STEPS : VIDEO_STEPS, + strength: job.strength ?? DEFAULT_STRENGTH, + frames: job.frames ?? 1, + fps: COMFYUI_VIDEO_FPS, + }; + if (template.placeholders.has("model")) values.model = await this.imageCheckpoint(signal); + if (template.placeholders.has("image")) { + if (!job.source) throw new Error("ComfyUI workflow needs an input image"); + values.image = await this.upload(job.source, signal); + } + const sized = template.placeholders.has("width") && template.placeholders.has("height"); + const timed = template.placeholders.has("frames") && template.placeholders.has("fps"); + const firstSeed = job.seed ?? randomInt(2 ** 48 - 1); + const results: GeneratedMedia[] = []; + for (let index = 0; index < job.count; index += 1) { + const graph = fillWorkflow(template, { ...values, seed: firstSeed + index }); + const outputs = await this.runPrompt(graph, job.kind, signal); + const files = outputFiles(outputs, job.kind); + if (files.length === 0) throw new Error(`ComfyUI produced no ${job.kind}`); + for (const file of files) { + results.push({ + bytes: await this.download(file, signal), + mimeType: file.mimeType, + name: file.filename, + ...(sized && file.kind === job.kind ? { width: job.width, height: job.height } : {}), + ...(timed && file.kind === "video" + ? { durationMs: Math.round(((values.frames ?? 1) / COMFYUI_VIDEO_FPS) * 1000) } + : {}), + }); + } + job.onProgress?.((index + 1) / job.count); + } + return results; + } + + /** Submits a prompt and waits for its outputs; a cancelled or timed-out prompt is dequeued. */ + private async runPrompt( + graph: ComfyWorkflow, + kind: MediaKind, + signal: AbortSignal | undefined, + ): Promise> { + const deadline = AbortSignal.timeout( + kind === "image" ? IMAGE_PROMPT_TIMEOUT_MS : VIDEO_PROMPT_TIMEOUT_MS, + ); + const runSignal = signal ? AbortSignal.any([signal, deadline]) : deadline; + const promptId = await this.submit(graph, runSignal); + try { + for (;;) { + const outputs = await this.finishedOutputs(promptId, runSignal); + if (outputs) return outputs; + await sleep(this.pollIntervalMs, undefined, { signal: runSignal }); + } + } catch (error) { + if (!runSignal.aborted) throw error; + await this.cancelPrompt(promptId); + if (signal?.aborted) throw signal.reason ?? error; + throw new Error("ComfyUI timed out"); + } + } + + private async submit(graph: ComfyWorkflow, signal: AbortSignal): Promise { + const payload = { prompt: graph, client_id: this.clientId }; + return this.postJson("prompt", payload, signal, async (response, readSignal) => { + if (response.status === 400) { + const rejection = await readJson(response, readSignal).catch(() => undefined); + throw new Error(describeWorkflowRejection(rejection)); + } + const result = await readOkJson(response, readSignal); + const promptId = isRecord(result) ? result.prompt_id : undefined; + if (typeof promptId !== "string" || !promptId) { + throw new Error("ComfyUI did not accept the workflow"); + } + return promptId; + }); + } + + private async finishedOutputs( + promptId: string, + signal: AbortSignal, + ): Promise | undefined> { + const history = await this.getJson(`history/${encodeURIComponent(promptId)}`, signal); + const entry = isRecord(history) ? history[promptId] : undefined; + if (!isRecord(entry)) return undefined; + const status = isRecord(entry.status) ? entry.status : {}; + if (status.status_str === "error") throw new Error(describeExecutionError(status.messages)); + if (status.completed === false) return undefined; + return isRecord(entry.outputs) ? entry.outputs : {}; + } + + /** Best effort: drops the prompt from the queue, or interrupts it when it is ours and running. */ + private async cancelPrompt(promptId: string): Promise { + try { + await this.postJson( + "queue", + { delete: [promptId] }, + undefined, + discardBody, + CANCEL_TIMEOUT_MS, + ); + const queue = await this.getJson("queue", undefined, CANCEL_TIMEOUT_MS); + const running = + isRecord(queue) && Array.isArray(queue.queue_running) ? queue.queue_running : []; + if (running.some((entry) => Array.isArray(entry) && entry[1] === promptId)) { + const interrupt = { prompt_id: promptId }; + await this.postJson("interrupt", interrupt, undefined, discardBody, CANCEL_TIMEOUT_MS); + } + } catch { + // The prompt then finishes on its own; nothing waits for it. + } + } + + private async imageCheckpoint(signal: AbortSignal | undefined): Promise { + if (this.imageModel) return this.imageModel; + const info = await this.getJson("object_info/CheckpointLoaderSimple", signal); + const videoCheckpoints = new Set( + [DEFAULT_WORKFLOWS.textToVideo, DEFAULT_WORKFLOWS.imageToVideo, this.videoWorkflow].flatMap( + (template) => (template ? checkpointNames(template.graph) : []), + ), + ); + const checkpoint = parseCheckpointList(info).find((name) => !videoCheckpoints.has(name)); + if (!checkpoint) throw new Error("ComfyUI has no image checkpoint installed"); + return checkpoint; + } + + private async upload(file: MediaFile, signal: AbortSignal | undefined): Promise { + const form = new FormData(); + const extension = attachmentExtensionForMimeType(file.mimeType) || ".png"; + form.append( + "image", + new Blob([Buffer.from(file.bytes)], { type: file.mimeType }), + `evoo-${randomUUID()}${extension}`, + ); + form.append("type", "input"); + form.append("overwrite", "true"); + // Serialize here: the address-checked fetch runs on the undici package, whose FormData + // class differs from the global one. + const encoded = new Response(form); + const headers = { + Accept: "application/json", + "Content-Type": encoded.headers.get("content-type") ?? "", + }; + const body = new Uint8Array(await encoded.arrayBuffer()); + return this.send("upload/image", { method: "POST", headers, body }, signal, { + timeoutMs: TRANSFER_TIMEOUT_MS, + read: async (response, readSignal) => { + const result = await readOkJson(response, readSignal); + const name = isRecord(result) && typeof result.name === "string" ? result.name : ""; + if (!name) throw new Error("ComfyUI did not accept the image"); + const subfolder = + isRecord(result) && typeof result.subfolder === "string" ? result.subfolder : ""; + return subfolder ? `${subfolder}/${name}` : name; + }, + }); + } + + private download(file: OutputFile, signal: AbortSignal | undefined): Promise { + const query = new URLSearchParams({ + filename: file.filename, + subfolder: file.subfolder, + type: file.type, + }); + return this.send(`view?${query}`, {}, signal, { + timeoutMs: TRANSFER_TIMEOUT_MS, + read: async (response, readSignal) => { + await assertOk(response); + return readBodyCapped(response, ARTIFACT_MAX_BYTES_BY_KIND[file.kind], readSignal); + }, + }); + } + + private getJson(path: string, signal: AbortSignal | undefined, timeoutMs?: number) { + const init = { headers: { Accept: "application/json" } }; + return this.send(path, init, signal, { timeoutMs, read: readOkJson }); + } + + private postJson( + path: string, + payload: unknown, + signal: AbortSignal | undefined, + read: ReadBody, + timeoutMs?: number, + ): Promise { + const init = { method: "POST", headers: JSON_HEADERS, body: JSON.stringify(payload) }; + return this.send(path, init, signal, { timeoutMs, read }); + } + + /** One request, with the timeout covering the response body too. */ + private async send( + path: string, + init: { + method?: string; + headers?: Record; + body?: string | Uint8Array; + }, + signal: AbortSignal | undefined, + options: { timeoutMs?: number; read: ReadBody }, + ): Promise { + const timeout = AbortSignal.timeout(options.timeoutMs ?? REQUEST_TIMEOUT_MS); + const requestSignal = signal ? AbortSignal.any([signal, timeout]) : timeout; + const headers: Record = { ...init.headers }; + if (this.apiKey) headers.Authorization = `Bearer ${this.apiKey}`; + try { + const response = await this.fetch(new URL(path, this.baseUrl).href, { + method: init.method ?? "GET", + headers, + body: init.body, + redirect: "error", + signal: requestSignal, + }); + return await options.read(response, requestSignal); + } catch (error) { + if (signal?.aborted) throw error; + if (timeout.aborted) throw new Error("ComfyUI timed out"); + // undici reports refused or reset connections as a bare "fetch failed". + if (error instanceof TypeError && error.message === "fetch failed") { + throw new Error("ComfyUI is unreachable"); + } + throw error; + } + } +} + +function comfyUiBaseUrl(raw: string): URL { + const url = assertAllowedOpenAiCompatibleRequestUrl(raw.trim()); + url.search = ""; + url.hash = ""; + url.pathname = `${url.pathname.replace(/\/+$/, "")}/`; + return url; +} + +function isRecord(value: unknown): value is Record { + return Boolean(value) && typeof value === "object" && !Array.isArray(value); +} + +function truncate(text: string): string { + return text.length > MAX_ERROR_CHARS ? `${text.slice(0, MAX_ERROR_CHARS - 1)}…` : text; +} + +function roundSide(side: number, multiple: number): number { + const rounded = Math.round(side / multiple) * multiple; + return Math.min(Math.max(rounded, multiple * Math.ceil(64 / multiple)), MAX_IMAGE_SIDE); +} + +/** A custom workflow replaces the default that matches whether it takes an input image. */ +function pickTemplate( + custom: WorkflowTemplate | undefined, + withImage: boolean, + textDefault: WorkflowTemplate, + imageDefault: WorkflowTemplate, +): WorkflowTemplate { + if (custom && custom.placeholders.has("image") === withImage) return custom; + return withImage ? imageDefault : textDefault; +} + +function checkpointNames(graph: ComfyWorkflow): string[] { + return Object.values(graph).flatMap((node) => + typeof node.inputs.ckpt_name === "string" ? [node.inputs.ckpt_name] : [], + ); +} + +async function discardBody(response: Response): Promise { + await response.body?.cancel().catch(() => undefined); +} + +async function assertOk(response: Response): Promise { + if (response.ok) return; + await discardBody(response); + throw new Error(`ComfyUI returned ${response.status}`); +} + +async function readJson(response: Response, signal: AbortSignal): Promise { + const bytes = await readBodyCapped(response, MAX_JSON_BYTES, signal); + try { + return JSON.parse(new TextDecoder().decode(bytes)); + } catch { + throw new Error("ComfyUI returned invalid JSON"); + } +} + +async function readOkJson(response: Response, signal: AbortSignal): Promise { + await assertOk(response); + return readJson(response, signal); +} + +/** + * ComfyUI can evict its own models, so memory its PyTorch process holds counts as available; + * only memory other processes use (another model server on the same GPU) does not. + */ +export function parseGpuStatus(body: unknown): GpuStatus | undefined { + const devices = isRecord(body) && Array.isArray(body.devices) ? body.devices : []; + for (const device of devices) { + if (!isRecord(device) || device.type === "cpu") continue; + const total = finiteNumber(device.vram_total); + const free = finiteNumber(device.vram_free); + if (!total || free === undefined) continue; + const reserved = finiteNumber(device.torch_vram_total) ?? 0; + const reservedFree = finiteNumber(device.torch_vram_free) ?? 0; + const available = Math.min(total, Math.max(0, free - reservedFree + reserved)); + return { vramFreeMb: Math.floor(available / MIB), vramTotalMb: Math.floor(total / MIB) }; + } + return undefined; +} + +function finiteNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) ? value : undefined; +} + +/** Older ComfyUI lists combo values as `[[...names], {...}]`, newer as `["COMBO", {options}]`. */ +export function parseCheckpointList(body: unknown): string[] { + const loader = isRecord(body) ? body.CheckpointLoaderSimple : undefined; + const input = isRecord(loader) ? loader.input : undefined; + const required = isRecord(input) ? input.required : undefined; + const spec = isRecord(required) ? required.ckpt_name : undefined; + if (!Array.isArray(spec)) return []; + const names = Array.isArray(spec[0]) + ? spec[0] + : isRecord(spec[1]) && Array.isArray(spec[1].options) + ? spec[1].options + : []; + return names.filter((name): name is string => typeof name === "string" && name.length > 0); +} + +/** Saved files win over previews; video runs return clips (or animations), not frames. */ +function outputFiles(outputs: Record, kind: MediaKind): OutputFile[] { + const files: OutputFile[] = []; + for (const node of Object.values(outputs)) { + if (!isRecord(node)) continue; + for (const key of ["images", "gifs", "videos"]) { + const entries = node[key]; + if (!Array.isArray(entries)) continue; + for (const entry of entries) { + if (!isRecord(entry) || typeof entry.filename !== "string" || !entry.filename) continue; + const mimeType = inferAttachmentMimeType(entry.filename); + if (!mimeType) continue; + files.push({ + filename: entry.filename, + subfolder: typeof entry.subfolder === "string" ? entry.subfolder : "", + type: typeof entry.type === "string" ? entry.type : "output", + mimeType, + kind: attachmentKindOf(mimeType), + }); + } + } + } + const videos = files.filter((file) => file.kind === "video"); + const wanted = + kind === "image" + ? files.filter((file) => file.kind === "image") + : videos.length > 0 + ? videos + : files.filter((file) => file.mimeType === "image/gif" || file.mimeType === "image/webp"); + const saved = wanted.filter((file) => file.type === "output"); + return saved.length > 0 ? saved : wanted; +} + +export function describeWorkflowRejection(body: unknown): string { + const problems: string[] = []; + const nodeErrors = isRecord(body) && isRecord(body.node_errors) ? body.node_errors : {}; + for (const node of Object.values(nodeErrors)) { + if (!isRecord(node) || !Array.isArray(node.errors)) continue; + for (const error of node.errors) { + if (!isRecord(error)) continue; + const parts = [node.class_type, error.message, error.details]; + problems.push(parts.filter((part) => typeof part === "string" && part).join(": ")); + } + } + if (problems.length === 0 && isRecord(body) && isRecord(body.error)) { + if (typeof body.error.message === "string") problems.push(body.error.message); + } + return truncate( + problems.length > 0 + ? `ComfyUI rejected the workflow: ${problems.join("; ")}` + : "ComfyUI rejected the workflow", + ); +} + +export function describeExecutionError(messages: unknown): string { + const events = Array.isArray(messages) ? messages : []; + for (const event of events) { + if (!Array.isArray(event)) continue; + const [name, data] = event; + if (name === "execution_interrupted") return "ComfyUI interrupted the run"; + if (name !== "execution_error" || !isRecord(data)) continue; + const node = typeof data.node_type === "string" ? data.node_type : "a node"; + const message = + typeof data.exception_message === "string" + ? data.exception_message.trim().split("\n")[0] + : ""; + return truncate(`ComfyUI failed in ${node}${message ? `: ${message}` : ""}`); + } + return "ComfyUI failed"; +} diff --git a/packages/adapters/src/executor-tools.ts b/packages/adapters/src/executor-tools.ts index 2e819b8..1045aa2 100644 --- a/packages/adapters/src/executor-tools.ts +++ b/packages/adapters/src/executor-tools.ts @@ -1,6 +1,7 @@ -import type { ConnectorTool } from "@rakazo/adapter-kit"; +import type { ConnectorTool, MediaCapabilities } from "@rakazo/adapter-kit"; import { agentConnectionTools, builtinAgentTools } from "./builtin-tools.js"; import { selectCloudAgentTools } from "./cloud-agent-tools-select.js"; +import { selectMediaTools } from "./media-tools.js"; import { selectMemoryTools } from "./memory-tools.js"; import { filterImageReturningComputerTools } from "./model-vision.js"; import { @@ -81,6 +82,8 @@ export function selectBuiltinToolsForRun(options: { trigger: string; semanticMemoryEnabled: boolean; cloudAgentEnabled?: boolean; + /** Media tools appear only for operations the space can run. */ + mediaCapabilities?: MediaCapabilities; messagingChannelRun: boolean; /** Cross-owner agent connections only exist for chat-linked bots. */ agentConnectionsEnabled?: boolean; @@ -88,21 +91,24 @@ export function selectBuiltinToolsForRun(options: { foreignAgentRun?: boolean; bot: BuiltinToolPolicy; }) { - const selected = selectCloudAgentTools( - selectMemoryTools( - filterBuiltinToolsForRun( - filterBuiltinToolsForThread( - filterPageBrowserTools( - filterImageReturningComputerTools(builtinAgentTools, options.graphicalToolsAllowed), - options.pageBrowserAllowed ?? options.graphicalToolsAllowed, + const selected = selectMediaTools( + selectCloudAgentTools( + selectMemoryTools( + filterBuiltinToolsForRun( + filterBuiltinToolsForThread( + filterPageBrowserTools( + filterImageReturningComputerTools(builtinAgentTools, options.graphicalToolsAllowed), + options.pageBrowserAllowed ?? options.graphicalToolsAllowed, + ), + options.groupId, ), - options.groupId, + options.trigger, ), - options.trigger, + options.semanticMemoryEnabled, ), - options.semanticMemoryEnabled, + Boolean(options.cloudAgentEnabled), ), - Boolean(options.cloudAgentEnabled), + options.mediaCapabilities, ).filter( (tool) => !(options.foreignAgentRun && tool.name === "update_user_profile") && diff --git a/packages/adapters/src/executor.ts b/packages/adapters/src/executor.ts index 9d2c8bd..4cca822 100644 --- a/packages/adapters/src/executor.ts +++ b/packages/adapters/src/executor.ts @@ -283,6 +283,7 @@ import { needsOAuthProbe, parseMcpServerToolArgs, } from "./mcp-server-tool.js"; +import { executeMediaTool, isMediaToolName } from "./media-tools.js"; import { loadAgentMemoryContext } from "./memory-context.js"; import type { MemoryProviderResolver } from "./memory-provider-factory.js"; import { selectMemoryTools } from "./memory-tools.js"; @@ -1478,6 +1479,7 @@ export function createRunExecutor(deps: ExecutorDeps) { trigger: run.trigger, semanticMemoryEnabled, cloudAgentEnabled: cloudAgentsEnabled(cloudAgent, run.spaceId), + mediaCapabilities: await deps.media?.capabilities(run.spaceId).catch(() => undefined), messagingChannelRun, agentConnectionsEnabled: hasMessagingIdentity, foreignAgentRun, @@ -2599,6 +2601,24 @@ export function createRunExecutor(deps: ExecutorDeps) { return computerScreenToolResult(() => tool(browser, computer, context, args), finish); } + if (isMediaToolName(name)) { + return finish( + await executeMediaTool( + deps.media, + { + spaceId: run.spaceId, + userId: run.userId, + botId: bot.id, + groupId: thread.groupId, + threadId: run.threadId, + runId: run.id, + operationKey: effectKey, + }, + name, + args, + ), + ); + } if (name.startsWith("cloud_agent_")) { return finish( await executeCloudAgentTool( diff --git a/packages/adapters/src/image-generation-conformance.test.ts b/packages/adapters/src/image-generation-conformance.test.ts new file mode 100644 index 0000000..92883a5 --- /dev/null +++ b/packages/adapters/src/image-generation-conformance.test.ts @@ -0,0 +1,111 @@ +import type { + ImageGenerationProvider, + MediaFile, + VideoGenerationProvider, +} from "@rakazo/adapter-kit"; +import { attachmentKindOf } from "@rakazo/contracts"; +import sharp from "sharp"; +import { describe, expect, it } from "vitest"; +import { createComfyUiProvider } from "./comfyui-provider.js"; +import { ComfyUiEmulator } from "./testing/comfyui-emulator.js"; + +type MediaGenerationProvider = ImageGenerationProvider & VideoGenerationProvider; + +const factories: Record MediaGenerationProvider> = { + comfyui: () => + createComfyUiProvider({ + baseUrl: "http://127.0.0.1:8188", + fetch: new ComfyUiEmulator().fetch, + pollIntervalMs: 1, + }), +}; + +async function sourceImage(): Promise { + const png = await sharp({ + create: { width: 96, height: 64, channels: 3, background: { r: 10, g: 20, b: 30 } }, + }) + .png() + .toBuffer(); + return { bytes: new Uint8Array(png), mimeType: "image/png", name: "source.png" }; +} + +/** + * Offline conformance: every image/video generation provider describes its limits, returns + * the requested number of decodable images, is reproducible by seed, takes a source image, + * reports progress, returns video, and honors cancellation. + */ +for (const [name, create] of Object.entries(factories)) { + describe(`${name} media generation conformance (offline)`, () => { + it("describes the contract and its limits", () => { + const image: ImageGenerationProvider = create(); + const video: VideoGenerationProvider = create(); + expect(image.describe().contractVersion).toBe("1"); + expect(image.describe().capabilities.maxWidth).toBeGreaterThanOrEqual(1024); + expect(image.describe().capabilities.maxHeight).toBeGreaterThanOrEqual(1024); + expect(video.describe().contractVersion).toBe("1"); + expect(video.describe().capabilities.maxDurationSeconds).toBeGreaterThan(0); + }); + + it("returns the requested number of images, reproducible by seed", async () => { + const provider = create(); + const progress: number[] = []; + const request = { prompt: "a red kite", width: 256, height: 192, count: 2, seed: 7 }; + const images = await provider.generateImage({ + ...request, + onProgress: (fraction) => progress.push(fraction), + }); + expect(images).toHaveLength(2); + for (const image of images) { + expect(attachmentKindOf(image.mimeType)).toBe("image"); + const metadata = await sharp(image.bytes).metadata(); + expect([metadata.width, metadata.height]).toEqual([256, 192]); + } + expect(Buffer.from(images[0]!.bytes).equals(Buffer.from(images[1]!.bytes))).toBe(false); + expect(progress.at(-1)).toBe(1); + const again = await provider.generateImage({ ...request, count: 1 }); + expect(Buffer.from(again[0]!.bytes).equals(Buffer.from(images[0]!.bytes))).toBe(true); + }); + + it("generates from a source image", async () => { + const images = await create().generateImage({ + prompt: "the same scene at night", + width: 96, + height: 64, + source: await sourceImage(), + strength: 0.5, + }); + expect(images).toHaveLength(1); + expect(attachmentKindOf(images[0]!.mimeType)).toBe("image"); + }); + + it("returns video from text and from a start frame", async () => { + const provider = create(); + for (const source of [undefined, await sourceImage()]) { + const [clip, ...rest] = await provider.generateVideo({ + prompt: "waves rolling onto a beach", + width: 768, + height: 512, + durationSeconds: 2, + source, + }); + expect(rest).toHaveLength(0); + expect(attachmentKindOf(clip!.mimeType)).toBe("video"); + expect(clip!.bytes.byteLength).toBeGreaterThan(0); + } + }); + + it("reports GPU memory with free not above total", async () => { + const status = await create().gpuStatus?.(); + if (!status) return; + expect(status.vramTotalMb).toBeGreaterThan(0); + expect(status.vramFreeMb).toBeGreaterThanOrEqual(0); + expect(status.vramFreeMb).toBeLessThanOrEqual(status.vramTotalMb); + }); + + it("rejects when the caller cancels", async () => { + const request = { prompt: "x", width: 64, height: 64, signal: AbortSignal.abort() }; + await expect(create().generateImage(request)).rejects.toThrow(); + await expect(create().generateVideo({ ...request, durationSeconds: 1 })).rejects.toThrow(); + }); + }); +} diff --git a/packages/adapters/src/index.ts b/packages/adapters/src/index.ts index 22c3355..24adde8 100644 --- a/packages/adapters/src/index.ts +++ b/packages/adapters/src/index.ts @@ -24,6 +24,7 @@ export * from "./cloud-agent-provider-env.js"; export * from "./cloud-agent-service.js"; export * from "./cloud-agent-tools.js"; export * from "./cloud-agent-tools-select.js"; +export * from "./comfyui-provider.js"; export * from "./composio-catalog-cache.js"; export * from "./composio-connector.js"; export * from "./composio-emulator.js"; diff --git a/packages/adapters/src/media-operations.ts b/packages/adapters/src/media-operations.ts index 3905c49..839a488 100644 --- a/packages/adapters/src/media-operations.ts +++ b/packages/adapters/src/media-operations.ts @@ -1,4 +1,5 @@ import type { MediaSettings } from "@rakazo/contracts"; +import { createImageOperations } from "./media/image-operations.js"; import { createTranscriptionOperations } from "./media/transcription.js"; import type { MediaOperationRegistration } from "./media-runtime.js"; @@ -12,5 +13,8 @@ export interface MediaOperationDeps { * from this one list, so both processes run the same operations. */ export function createMediaOperations(deps: MediaOperationDeps): MediaOperationRegistration[] { - return [...createTranscriptionOperations({ loadSettings: deps.loadSettings })]; + return [ + ...createImageOperations(deps), + ...createTranscriptionOperations({ loadSettings: deps.loadSettings }), + ]; } diff --git a/packages/adapters/src/media-runtime.postgres.test.ts b/packages/adapters/src/media-runtime.postgres.test.ts new file mode 100644 index 0000000..8fdf406 --- /dev/null +++ b/packages/adapters/src/media-runtime.postgres.test.ts @@ -0,0 +1,615 @@ +import { randomUUID } from "node:crypto"; +import { existsSync } from "node:fs"; +import { mkdtemp, readdir, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import path from "node:path"; +import type { + BackgroundJob, + JobPublisher, + MediaJobContext, + MediaOperationResult, +} from "@rakazo/adapter-kit"; +import { MediaJobDeferredError, mediaRunJobKey } from "@rakazo/adapter-kit"; +import type { MediaJobBlock, MessageBlock } from "@rakazo/contracts"; +import { clearThread, createDb, type PrismaClient } from "@rakazo/db"; +import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; +import { z } from "zod"; +import { storeArtifactFromFile } from "./artifact-files.js"; +import { LocalArtifactStore } from "./artifacts.js"; +import { + createMediaRuntime, + MEDIA_WAKE_TRIGGER, + MediaJobRejectedError, + type MediaOperationRegistration, + type MediaRuntimeOptions, +} from "./media-runtime.js"; + +const describePostgres = + process.env.VERIFY_DATABASE && process.env.DATABASE_URL ? describe.sequential : describe.skip; + +const PNG = new Uint8Array([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 1, 2, 3, 4]); + +type FakeInput = { prompt: string; source?: string }; +type Run = (context: MediaJobContext) => Promise; + +function deferred() { + let resolve!: (value: T) => void; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +describePostgres("media job runtime (PostgreSQL, fake operations)", () => { + let prisma: PrismaClient; + let db: ReturnType; + let root: string; + const users: string[] = []; + + beforeAll(async () => { + db = createDb(process.env.DATABASE_URL!); + prisma = db.prisma; + root = await mkdtemp(path.join(tmpdir(), "media-runtime-test-")); + }); + afterAll(async () => { + if (!db) return; + await prisma.organization.deleteMany({ where: { id: { in: users } } }); + await prisma.user.deleteMany({ where: { id: { in: users } } }); + await prisma.$disconnect(); + await db.pool.end(); + await rm(root, { recursive: true, force: true }); + }); + + async function setup( + options: { + run?: Run; + available?: () => Promise; + gpu?: MediaOperationRegistration["gpu"]; + runtime?: MediaRuntimeOptions; + } = {}, + ) { + const id = randomUUID(); + users.push(id); + await prisma.user.create({ + data: { id, name: "Media test", email: `${id}@example.test`, emailVerified: true }, + }); + await prisma.organization.create({ + data: { id, name: "Media test", slug: id, createdAt: new Date() }, + }); + await prisma.space.create({ + data: { id, organizationId: id, name: "Media test", isDefault: true }, + }); + await prisma.member.create({ + data: { id, organizationId: id, userId: id, role: "owner", createdAt: new Date() }, + }); + const bot = await prisma.bot.create({ + data: { spaceId: id, userId: id, name: "Studio", color: "test-color" }, + }); + const thread = await prisma.thread.create({ data: { spaceId: id, userId: id, botId: bot.id } }); + const task = await prisma.task.create({ + data: { + spaceId: id, + userId: id, + botId: bot.id, + threadId: thread.id, + prompt: "Make a picture", + status: "running", + }, + }); + const run = await prisma.run.create({ + data: { + spaceId: id, + userId: id, + botId: bot.id, + threadId: thread.id, + taskId: task.id, + status: "running", + trigger: "user", + }, + }); + const artifacts = new LocalArtifactStore(path.join(root, id)); + const enqueued: BackgroundJob[] = []; + const cancelled: string[] = []; + const jobs: JobPublisher = { + enqueue: async (job) => { + enqueued.push(job); + }, + cancel: async (key) => { + cancelled.push(key); + }, + close: async () => undefined, + }; + const notify = vi.fn(async () => undefined); + const handlerCalls: string[] = []; + const defaultRun: Run = async (context) => { + const input = context.input as FakeInput; + if (input.source) await context.fetchArtifact(input.source); + await context.progress(0.5); + const file = path.join(context.tempDir, "out.png"); + const poster = path.join(context.tempDir, "poster.jpg"); + await writeFile(file, PNG); + await writeFile(poster, PNG); + await context.addOutput({ + path: file, + name: "picture.png", + mimeType: "image/png", + width: 64, + height: 32, + posterPath: poster, + parentArtifactId: input.source, + }); + return { summary: `Drew ${input.prompt}.` }; + }; + const registration: MediaOperationRegistration = { + operation: "generate_image", + inputSchema: z.object({ prompt: z.string().min(1), source: z.string().optional() }), + title: (input) => input.prompt, + available: options.available ?? (async () => true), + artifactIds: (input) => (input.source ? [input.source] : []), + handler: { + operation: "generate_image", + run: async (context) => { + handlerCalls.push(context.jobId); + return (options.run ?? defaultRun)(context); + }, + }, + ...(options.gpu ? { gpu: options.gpu } : {}), + }; + const tempRoot = path.join(root, id, "tmp"); + const runtimeDeps = { + prisma, + events: { notify }, + jobs, + artifacts, + registrations: [registration as MediaOperationRegistration], + }; + const runtimeOptions: MediaRuntimeOptions = { + tempRoot, + heartbeatMs: 50, + progressIntervalMs: 0, + ...options.runtime, + }; + const runtime = createMediaRuntime(runtimeDeps, runtimeOptions); + const start = (input: Record, extra: Record = {}) => + runtime.start({ + operation: "generate_image", + input, + spaceId: id, + userId: id, + botId: bot.id, + groupId: null, + threadId: thread.id, + runId: run.id, + ...extra, + }); + const job = (jobId: string) => prisma.mediaJob.findUniqueOrThrow({ where: { id: jobId } }); + const card = async (jobId: string): Promise => { + const row = await job(jobId); + const message = await prisma.message.findUniqueOrThrow({ where: { id: row.messageId! } }); + const block = (message.blocks as MessageBlock[]).find( + (entry): entry is MediaJobBlock => entry.kind === "media", + ); + if (!block) throw new Error("media card missing"); + return block; + }; + const sourceArtifact = async () => { + const file = path.join(root, `${randomUUID()}.png`); + await writeFile(file, PNG); + return storeArtifactFromFile( + { prisma, artifacts }, + { + path: file, + name: "source.png", + mimeType: "image/png", + spaceId: id, + userId: id, + botId: bot.id, + groupId: null, + runId: null, + }, + ); + }; + return { + id, + bot, + thread, + run, + runtime, + runtimeDeps, + runtimeOptions, + enqueued, + cancelled, + notify, + handlerCalls, + tempRoot, + start, + job, + card, + sourceArtifact, + }; + } + + it("runs a job to done, stores outputs with posters and versions, and cleans its temp dir", async () => { + const ctx = await setup(); + const source = await ctx.sourceArtifact(); + const started = await ctx.start({ prompt: "a red fox", source: source.id }); + + const queued = await ctx.card(started.jobId); + expect(queued).toMatchObject({ status: "queued", title: "a red fox", outputs: [] }); + expect(ctx.enqueued.map((job) => job.name)).toEqual(["media.run"]); + expect(ctx.enqueued[0]).toMatchObject({ payload: { jobId: started.jobId }, queue: "media" }); + const message = await prisma.message.findUniqueOrThrow({ where: { id: started.messageId } }); + expect(message).toMatchObject({ role: "bot", botId: ctx.bot.id, runId: ctx.run.id }); + + await ctx.runtime.runJob(started.jobId); + + const done = await ctx.job(started.jobId); + expect(done).toMatchObject({ status: "done", progress: 1, attempts: 1, leaseToken: null }); + expect(done.finishedAt).toBeInstanceOf(Date); + const block = await ctx.card(started.jobId); + expect(block.status).toBe("done"); + expect(block.outputs).toHaveLength(1); + const [output] = block.outputs; + expect(output).toMatchObject({ + kind: "image", + mimeType: "image/png", + name: "picture.png", + size: PNG.byteLength, + width: 64, + height: 32, + parentArtifactId: source.id, + }); + expect(done.outputArtifactIds).toEqual([output!.posterArtifactId, output!.artifactId]); + const stored = await prisma.artifact.findUniqueOrThrow({ where: { id: output!.artifactId } }); + expect(stored).toMatchObject({ + spaceId: ctx.id, + userId: ctx.id, + botId: ctx.bot.id, + runId: ctx.run.id, + parentArtifactId: source.id, + width: 64, + height: 32, + }); + const events = await prisma.event.findMany({ + where: { threadId: ctx.thread.id, type: "thread.message.updated" }, + orderBy: { seq: "asc" }, + }); + const statuses = events.map( + (event) => ((event.payload as { blocks: MessageBlock[] }).blocks[0] as MediaJobBlock).status, + ); + expect(statuses[0]).toBe("running"); + expect(statuses.at(-1)).toBe("done"); + expect(ctx.notify).toHaveBeenCalled(); + expect(existsSync(path.join(ctx.tempRoot, started.jobId))).toBe(false); + // No wake without continue_after. + expect( + await prisma.run.count({ where: { trigger: MEDIA_WAKE_TRIGGER, spaceId: ctx.id } }), + ).toBe(0); + expect(await ctx.runtime.get(started.jobId)).toMatchObject({ + id: started.jobId, + status: "done", + outputs: block.outputs, + }); + }); + + it("validates input, availability and readable artifacts before creating a job", async () => { + const ctx = await setup(); + await expect(ctx.start({ prompt: "" })).rejects.toBeInstanceOf(MediaJobRejectedError); + await expect(ctx.start({ prompt: "x", source: "not-an-artifact" })).rejects.toThrow( + "Unknown artifact.", + ); + const other = await setup(); + const foreign = await other.sourceArtifact(); + await expect(ctx.start({ prompt: "x", source: foreign.id })).rejects.toThrow( + "Unknown artifact.", + ); + await expect(ctx.start({ prompt: "x" }, { threadId: other.thread.id })).rejects.toThrow( + "This conversation is no longer available.", + ); + const unavailable = await setup({ available: async () => false }); + await expect(unavailable.start({ prompt: "x" })).rejects.toThrow("not available"); + expect( + await prisma.mediaJob.count({ where: { spaceId: { in: [ctx.id, unavailable.id] } } }), + ).toBe(0); + }); + + it("returns the first job for a replayed tool call", async () => { + const ctx = await setup(); + const first = await ctx.start({ prompt: "fox" }, { operationKey: `effect:${ctx.id}` }); + const replay = await ctx.start({ prompt: "fox" }, { operationKey: `effect:${ctx.id}` }); + expect(replay).toEqual(first); + expect(await prisma.mediaJob.count({ where: { spaceId: ctx.id } })).toBe(1); + }); + + it("fails a job with the operation's error and still cleans up", async () => { + let tempDir = ""; + const ctx = await setup({ + run: async (context) => { + tempDir = context.tempDir; + await writeFile(path.join(context.tempDir, "partial.bin"), PNG); + throw new Error(`ComfyUI rejected the workflow (${context.tempDir}/x)`); + }, + }); + const { jobId } = await ctx.start({ prompt: "fox" }); + await ctx.runtime.runJob(jobId); + const failed = await ctx.job(jobId); + expect(failed.status).toBe("failed"); + expect(failed.error).toBe("ComfyUI rejected the workflow (/x)"); + expect((await ctx.card(jobId)).error).toBe(failed.error); + expect(existsSync(tempDir)).toBe(false); + }); + + it("refuses outputs outside the temp dir and sources of another user", async () => { + const other = await setup(); + const foreign = await other.sourceArtifact(); + const ctx = await setup({ + run: async (context) => { + await context.fetchArtifact(foreign.id); + return { summary: "" }; + }, + }); + const { jobId } = await ctx.start({ prompt: "fox" }); + await ctx.runtime.runJob(jobId); + expect(await ctx.job(jobId)).toMatchObject({ + status: "failed", + error: "A source file is no longer available.", + }); + + const outside = await setup({ + run: async (context) => { + await context.addOutput({ path: "/etc/hostname", name: "x.png", mimeType: "image/png" }); + return { summary: "" }; + }, + }); + const escaped = await outside.start({ prompt: "fox" }); + await outside.runtime.runJob(escaped.jobId); + expect((await outside.job(escaped.jobId)).status).toBe("failed"); + expect(await prisma.artifact.count({ where: { spaceId: outside.id } })).toBe(0); + }); + + it("defers with backoff, then fails once the deferral budget is spent", async () => { + const ctx = await setup({ + run: async () => { + throw new MediaJobDeferredError("ComfyUI queue is full", 2_000); + }, + runtime: { maxDeferrals: 2 }, + }); + const { jobId } = await ctx.start({ prompt: "fox" }); + const before = Date.now(); + await ctx.runtime.runJob(jobId); + const waiting = await ctx.job(jobId); + expect(waiting).toMatchObject({ status: "waiting", deferrals: 1, leaseToken: null }); + expect(waiting.nextAttemptAt!.getTime()).toBeGreaterThanOrEqual(before + 2_000); + expect((await ctx.card(jobId)).status).toBe("waiting"); + const retry = ctx.enqueued.at(-1)!; + expect(retry).toMatchObject({ name: "media.run", payload: { jobId }, queue: "media" }); + expect(retry.availableAt?.getTime()).toBe(waiting.nextAttemptAt!.getTime()); + + // Too early: the delayed job owns the next attempt. + await ctx.runtime.runJob(jobId); + expect(ctx.handlerCalls).toHaveLength(1); + + await prisma.mediaJob.update({ where: { id: jobId }, data: { nextAttemptAt: new Date() } }); + await ctx.runtime.runJob(jobId); + const second = await ctx.job(jobId); + expect(second).toMatchObject({ status: "waiting", deferrals: 2 }); + // Backoff doubles. + expect(second.nextAttemptAt!.getTime() - Date.now()).toBeGreaterThan(3_000); + + await prisma.mediaJob.update({ where: { id: jobId }, data: { nextAttemptAt: new Date() } }); + await ctx.runtime.runJob(jobId); + expect(await ctx.job(jobId)).toMatchObject({ + status: "failed", + error: "The media server stayed busy. Try again later.", + }); + }); + + it("waits for free GPU memory before running the handler", async () => { + let free = 1_000; + const ctx = await setup({ + gpu: { + minVramMb: 6_000, + status: async () => ({ vramFreeMb: free, vramTotalMb: 24_000 }), + }, + }); + const { jobId } = await ctx.start({ prompt: "fox" }); + await ctx.runtime.runJob(jobId); + expect(ctx.handlerCalls).toHaveLength(0); + expect(await ctx.job(jobId)).toMatchObject({ status: "waiting", deferrals: 1 }); + + free = 20_000; + await prisma.mediaJob.update({ where: { id: jobId }, data: { nextAttemptAt: new Date() } }); + await ctx.runtime.runJob(jobId); + expect(ctx.handlerCalls).toHaveLength(1); + expect((await ctx.job(jobId)).status).toBe("done"); + }); + + it("cancels a queued job at once and a running job through the lease heartbeat", async () => { + const ctx = await setup({ + run: async (context) => { + await new Promise((resolve) => { + context.signal.addEventListener("abort", () => resolve(), { once: true }); + }); + context.signal.throwIfAborted(); + return { summary: "" }; + }, + }); + const queued = await ctx.start({ prompt: "queued" }); + await ctx.runtime.cancel(queued.jobId); + expect(await ctx.job(queued.jobId)).toMatchObject({ status: "cancelled" }); + expect((await ctx.card(queued.jobId)).status).toBe("cancelled"); + expect(ctx.cancelled).toEqual([mediaRunJobKey(queued.jobId)]); + await ctx.runtime.runJob(queued.jobId); + expect(ctx.handlerCalls).toHaveLength(0); + + // Another process (a second runtime) asks to cancel a job this one runs. + const running = await ctx.start({ prompt: "running" }); + const runningJob = ctx.runtime.runJob(running.jobId); + await vi.waitFor(() => expect(ctx.handlerCalls).toEqual([running.jobId])); + const otherProcess = createMediaRuntime(ctx.runtimeDeps, ctx.runtimeOptions); + await otherProcess.cancel(running.jobId); + expect((await ctx.job(running.jobId)).status).toBe("running"); + await runningJob; + expect(await ctx.job(running.jobId)).toMatchObject({ status: "cancelled", leaseToken: null }); + expect((await ctx.card(running.jobId)).status).toBe("cancelled"); + }); + + it("never runs one job twice and keeps one job on the GPU at a time", async () => { + const release = deferred(); + const entered = deferred(); + const ctx = await setup({ + run: async () => { + entered.resolve(); + await release.promise; + return { summary: "ok" }; + }, + }); + const otherProcess = createMediaRuntime(ctx.runtimeDeps, ctx.runtimeOptions); + const first = await ctx.start({ prompt: "one" }); + const second = await ctx.start({ prompt: "two" }); + + const a = ctx.runtime.runJob(first.jobId); + await entered.promise; + await otherProcess.runJob(first.jobId); + await otherProcess.runJob(second.jobId); + expect(ctx.handlerCalls).toEqual([first.jobId]); + expect(await ctx.job(second.jobId)).toMatchObject({ status: "waiting", attempts: 0 }); + + release.resolve(); + await a; + await prisma.mediaJob.update({ + where: { id: second.jobId }, + data: { nextAttemptAt: new Date() }, + }); + await otherProcess.runJob(second.jobId); + expect(ctx.handlerCalls).toEqual([first.jobId, second.jobId]); + expect((await ctx.job(second.jobId)).status).toBe("done"); + }); + + it("recovers a job whose runner died and fails it after too many restarts", async () => { + const ctx = await setup(); + const { jobId } = await ctx.start({ prompt: "fox" }); + const crashedAt = new Date(Date.now() - 1_000); + await prisma.mediaJob.update({ + where: { id: jobId }, + data: { + status: "running", + attempts: 1, + leaseToken: "dead-runner", + leaseExpiresAt: crashedAt, + updatedAt: crashedAt, + }, + }); + ctx.enqueued.length = 0; + await ctx.runtime.reconcile(); + expect(ctx.enqueued).toEqual([ + expect.objectContaining({ name: "media.run", payload: { jobId } }), + ]); + await ctx.runtime.runJob(jobId); + expect(await ctx.job(jobId)).toMatchObject({ status: "done", attempts: 2 }); + + const doomed = await ctx.start({ prompt: "crashes" }); + await prisma.mediaJob.update({ + where: { id: doomed.jobId }, + data: { + status: "running", + attempts: 3, + leaseToken: "dead-runner", + leaseExpiresAt: crashedAt, + }, + }); + await ctx.runtime.runJob(doomed.jobId); + expect(await ctx.job(doomed.jobId)).toMatchObject({ + status: "failed", + error: "The job was interrupted. Try again.", + }); + expect(ctx.handlerCalls).toEqual([jobId]); + }); + + it("re-enqueues lost and overdue jobs, outside the queue once they look stuck", async () => { + const ctx = await setup(); + const fresh = await ctx.start({ prompt: "fresh" }); + const lost = await ctx.start({ prompt: "lost" }); + const stuck = await ctx.start({ prompt: "stuck" }); + const minuteAgo = new Date(Date.now() - 60_000); + const hourAgo = new Date(Date.now() - 3_600_000); + await prisma.mediaJob.update({ + where: { id: lost.jobId }, + data: { updatedAt: minuteAgo, createdAt: minuteAgo }, + }); + await prisma.mediaJob.update({ + where: { id: stuck.jobId }, + data: { updatedAt: hourAgo, createdAt: hourAgo }, + }); + ctx.enqueued.length = 0; + await ctx.runtime.reconcile(); + const byJob = new Map( + ctx.enqueued.map((job) => [(job.payload as { jobId: string }).jobId, job]), + ); + expect(byJob.has(fresh.jobId)).toBe(false); + expect(byJob.get(lost.jobId)).toMatchObject({ queue: "media", availableAt: minuteAgo }); + expect(byJob.get(stuck.jobId)?.queue).toBeUndefined(); + }); + + it("wakes the starting bot with the result when continue_after is set", async () => { + const ctx = await setup(); + const { jobId } = await ctx.start({ prompt: "fox" }, { continueAfter: true }); + ctx.enqueued.length = 0; + await ctx.runtime.runJob(jobId); + const wake = await prisma.run.findFirstOrThrow({ + where: { spaceId: ctx.id, trigger: MEDIA_WAKE_TRIGGER }, + include: { task: true }, + }); + expect(wake).toMatchObject({ botId: ctx.bot.id, threadId: ctx.thread.id, status: "queued" }); + expect(wake.task.prompt).toContain(`Media job ${jobId}`); + expect(wake.task.prompt).toContain("Drew fox."); + expect(wake.task.prompt).toContain("picture.png"); + expect(ctx.enqueued).toEqual([ + expect.objectContaining({ name: "run.continue", payload: { runId: wake.id } }), + ]); + }); + + it("cancels a job whose conversation was cleared without running it", async () => { + const ctx = await setup(); + const { jobId } = await ctx.start({ prompt: "fox" }, { continueAfter: true }); + await prisma.run.update({ where: { id: ctx.run.id }, data: { status: "completed" } }); + await clearThread(prisma, { spaceId: ctx.id, botId: ctx.bot.id, threadId: ctx.thread.id }); + await ctx.runtime.runJob(jobId); + expect(ctx.handlerCalls).toHaveLength(0); + expect(await ctx.job(jobId)).toMatchObject({ status: "cancelled" }); + expect( + await prisma.run.count({ where: { spaceId: ctx.id, trigger: MEDIA_WAKE_TRIGGER } }), + ).toBe(0); + }); + + it("fails a job that outlives its run limit", async () => { + const ctx = await setup({ + run: async (context) => { + await new Promise((resolve) => setTimeout(resolve, 5_000)); + context.signal.throwIfAborted(); + return { summary: "" }; + }, + runtime: { maxRunMs: 50, abortGraceMs: 20 }, + }); + const { jobId } = await ctx.start({ prompt: "slow" }); + await ctx.runtime.runJob(jobId); + expect(await ctx.job(jobId)).toMatchObject({ + status: "failed", + error: "The job took too long.", + }); + expect(await readdir(ctx.tempRoot)).toEqual([]); + }); + + it("reports which operations a space can run", async () => { + const ctx = await setup(); + const capabilities = await ctx.runtime.capabilities(ctx.id); + expect(capabilities.generate_image).toBe(true); + expect(capabilities.transcribe).toBe(false); + const broken = await setup({ + available: async () => { + throw new Error("probe failed"); + }, + }); + expect((await broken.runtime.capabilities(broken.id)).generate_image).toBe(false); + }); +}); diff --git a/packages/adapters/src/media-runtime.ts b/packages/adapters/src/media-runtime.ts index 681f2e9..928f08f 100644 --- a/packages/adapters/src/media-runtime.ts +++ b/packages/adapters/src/media-runtime.ts @@ -1,7 +1,11 @@ -import { createHash, randomUUID } from "node:crypto"; -import { mkdir, readFile, rm, stat, writeFile } from "node:fs/promises"; +import { randomUUID } from "node:crypto"; +import { constants, createWriteStream } from "node:fs"; +import { copyFile, mkdir, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import path from "node:path"; +import { Readable } from "node:stream"; +import { pipeline } from "node:stream/promises"; +import type { ReadableStream as NodeReadableStream } from "node:stream/web"; import type { AdapterContext, ArtifactStore, @@ -17,13 +21,13 @@ import type { StartMediaJobInput, } from "@rakazo/adapter-kit"; import { + ArtifactTooLargeError, MediaJobDeferredError, mediaRunJob, mediaRunJobKey, runContinueJob, } from "@rakazo/adapter-kit"; import { - ARTIFACT_MAX_BYTES_BY_KIND, attachmentKindOf, MEDIA_JOB_ACTIVE_STATUSES, type MediaJobBlock, @@ -33,11 +37,7 @@ import { type MediaOutput, type MessageBlock, } from "@rakazo/contracts"; -import { - attachmentExtensionForMimeType, - inferAttachmentMimeType, - validateAttachmentMimeType, -} from "@rakazo/core"; +import { attachmentExtensionForMimeType, inferAttachmentMimeType } from "@rakazo/core"; import { appendEventInTransaction, createThreadMessageInTransaction, @@ -49,6 +49,7 @@ import { } from "@rakazo/db"; import { getLogger } from "@rakazo/logging"; import type { z } from "zod"; +import { storeArtifactFromFile } from "./artifact-files.js"; /** Run trigger of the turn that wakes a bot when its media job ends (`continue_after`). */ export const MEDIA_WAKE_TRIGGER = "media"; @@ -505,6 +506,7 @@ export function createMediaRuntime( const storeArtifactFile = async ( job: MediaJob, + signal: AbortSignal, file: { path: string; name: string; @@ -515,38 +517,37 @@ export function createMediaRuntime( parentArtifactId?: string; }, ) => { - // Swap point for the streaming artifact store: until it lands, files go through put(). - const bytes = new Uint8Array(await readFile(file.path)); - const context = artifactContext(job, `media-output:${job.id}`); - const stored = await deps.artifacts.put( - { name: file.name, mimeType: file.mimeType, bytes }, - context, - ); try { - return await deps.prisma.artifact.create({ - data: { - spaceId: job.spaceId, - userId: job.userId, - botId: job.botId, - groupId: job.groupId, - runId: job.runId, - name: file.name, - mimeType: file.mimeType, - size: bytes.byteLength, - hash: createHash("sha256").update(bytes).digest("hex"), - storageKey: stored.id, - width: file.width, - height: file.height, - durationMs: file.durationMs, - parentArtifactId: file.parentArtifactId, - }, + return await storeArtifactFromFile(deps, { + ...file, + spaceId: job.spaceId, + userId: job.userId, + botId: job.botId, + groupId: job.groupId, + runId: job.runId, + signal, }); } catch (error) { - await deps.artifacts.remove(stored.id, context).catch(() => undefined); + if (error instanceof ArtifactTooLargeError) throw new MediaJobRejectedError(TOO_LARGE_ERROR); throw error; } }; + /** Copies a stored artifact to `file`: a clone of the local file when possible, else a stream. */ + const copyArtifact = async (job: MediaJob, storageKey: string, file: string) => { + const context = artifactContext(job, `media-input:${job.id}`); + const local = await deps.artifacts.localPath?.(storageKey, context).catch(() => undefined); + if (local) { + await copyFile(local, file, constants.COPYFILE_EXCL | constants.COPYFILE_FICLONE); + return; + } + const stream = await deps.artifacts.createReadStream(storageKey, context); + await pipeline( + Readable.fromWeb(stream as NodeReadableStream), + createWriteStream(file, { flags: "wx", mode: 0o600 }), + ); + }; + const removeArtifacts = async (job: MediaJob, ids: string[]) => { const rows = await deps.prisma.artifact.findMany({ where: { id: { in: ids }, spaceId: job.spaceId }, @@ -657,12 +658,7 @@ export function createMediaRuntime( attachmentExtensionForMimeType(row.mimeType) || (/^\.[a-z0-9]{1,8}$/i.test(path.extname(row.name)) ? path.extname(row.name) : ""); const file = path.join(tempDir, `input-${row.id}${extension.toLowerCase()}`); - // Swap point for the streaming artifact store (localPath / createReadStream). - const bytes = await deps.artifacts.get( - row.storageKey, - artifactContext(job, `media-input:${job.id}`), - ); - await writeFile(file, bytes, { mode: 0o600, flag: "wx" }); + await copyArtifact(job, row.storageKey, file); const input: MediaJobInputFile = { artifactId: row.id, path: file, @@ -679,12 +675,8 @@ export function createMediaRuntime( }, addOutput: async (file: MediaJobOutputFile) => { signal.throwIfAborted(); - validateAttachmentMimeType(file.mimeType); const kind = attachmentKindOf(file.mimeType); const filePath = insideDir(tempDir, file.path); - const size = (await stat(filePath)).size; - if (size > ARTIFACT_MAX_BYTES_BY_KIND[kind]) - throw new MediaJobRejectedError(TOO_LARGE_ERROR); const parent = file.parentArtifactId ? await deps.prisma.artifact.findFirst({ where: { id: file.parentArtifactId, spaceId: job.spaceId, userId: job.userId }, @@ -702,7 +694,7 @@ export function createMediaRuntime( const posterName = `${path.parse(name).name}-poster${path.extname(posterPath).toLowerCase()}`; const posterMime = inferAttachmentMimeType(posterName); if (posterMime && attachmentKindOf(posterMime) === "image") { - const poster = await storeArtifactFile(job, { + const poster = await storeArtifactFile(job, signal, { path: posterPath, name: posterName, mimeType: posterMime, @@ -716,7 +708,7 @@ export function createMediaRuntime( const width = positiveInt(file.width); const height = positiveInt(file.height); const durationMs = nonNegativeInt(file.durationMs); - const artifact = await storeArtifactFile(job, { + const artifact = await storeArtifactFile(job, signal, { path: filePath, name, mimeType: file.mimeType, diff --git a/packages/adapters/src/media-settings.ts b/packages/adapters/src/media-settings.ts index d237a13..c241c47 100644 --- a/packages/adapters/src/media-settings.ts +++ b/packages/adapters/src/media-settings.ts @@ -1,10 +1,6 @@ -import { - type MediaSettings, - MediaSettingsSchema, - type MediaSettingsUpdate, - MediaSettingsUpdateSchema, - type MediaSettingsView, -} from "@rakazo/contracts"; +import type { MediaSettings, MediaSettingsUpdate, MediaSettingsView } from "@rakazo/contracts"; +import { MediaSettingsSchema, MediaSettingsUpdateSchema } from "@rakazo/contracts"; +import { parseWorkflowTemplate } from "@rakazo/core/media/workflows"; import type { PrismaClient } from "@rakazo/db"; import { assertAllowedOpenAiCompatibleRequestUrl, @@ -82,31 +78,59 @@ function keptKey( return new URL(previous.baseUrl).origin === new URL(baseUrl).origin ? previous.apiKey : undefined; } +/** A media setting the deployment owner has to correct; the message says what is wrong. */ +export class MediaSettingsInvalidError extends Error { + constructor(message: string) { + super(message); + this.name = "MediaSettingsInvalidError"; + } +} + +function assertValidWorkflow(label: string, workflow: string | undefined): void { + if (!workflow) return; + try { + parseWorkflowTemplate(workflow); + } catch (error) { + throw new Error(`${label}: ${error instanceof Error ? error.message : String(error)}`); + } +} + +function resolveSettings(previous: MediaSettings, update: MediaSettingsUpdate): MediaSettings { + const parsed = MediaSettingsUpdateSchema.parse(update); + try { + assertValidWorkflow("Image workflow", parsed.comfyui?.imageWorkflow); + assertValidWorkflow("Video workflow", parsed.comfyui?.videoWorkflow); + let comfyui: MediaSettings["comfyui"] = null; + if (parsed.comfyui) { + const baseUrl = normalizeComfyUiBaseUrl(parsed.comfyui.baseUrl); + const apiKey = keptKey(parsed.comfyui.apiKey, previous.comfyui, baseUrl); + assertHttpsForKeyedOpenAiCompatibleUrl(new URL(baseUrl), apiKey); + comfyui = { ...parsed.comfyui, baseUrl, apiKey }; + } + let transcription: MediaSettings["transcription"] = null; + if (parsed.transcription) { + const baseUrl = normalizeTranscriptionBaseUrl(parsed.transcription.baseUrl); + const apiKey = keptKey(parsed.transcription.apiKey, previous.transcription, baseUrl); + assertHttpsForKeyedOpenAiCompatibleUrl(new URL(baseUrl), apiKey); + transcription = { ...parsed.transcription, baseUrl, apiKey }; + } + return MediaSettingsSchema.parse({ comfyui, transcription }); + } catch (error) { + throw new MediaSettingsInvalidError( + error instanceof Error ? error.message : "These media settings are not valid", + ); + } +} + /** * Validates and stores the media backends (deployment owner only; the caller authorizes). - * Throws an Error with a user-facing message when a URL is not allowed. + * Throws MediaSettingsInvalidError when a URL is not allowed. */ export async function saveMediaSettings( deps: MediaSettingsDeps, update: MediaSettingsUpdate, ): Promise { - const parsed = MediaSettingsUpdateSchema.parse(update); - const previous = await loadMediaSettings(deps); - let comfyui: MediaSettings["comfyui"] = null; - if (parsed.comfyui) { - const baseUrl = normalizeComfyUiBaseUrl(parsed.comfyui.baseUrl); - const apiKey = keptKey(parsed.comfyui.apiKey, previous.comfyui, baseUrl); - assertHttpsForKeyedOpenAiCompatibleUrl(new URL(baseUrl), apiKey); - comfyui = { ...parsed.comfyui, baseUrl, apiKey }; - } - let transcription: MediaSettings["transcription"] = null; - if (parsed.transcription) { - const baseUrl = normalizeTranscriptionBaseUrl(parsed.transcription.baseUrl); - const apiKey = keptKey(parsed.transcription.apiKey, previous.transcription, baseUrl); - assertHttpsForKeyedOpenAiCompatibleUrl(new URL(baseUrl), apiKey); - transcription = { ...parsed.transcription, baseUrl, apiKey }; - } - const settings = MediaSettingsSchema.parse({ comfyui, transcription }); + const settings = resolveSettings(await loadMediaSettings(deps), update); if (!settings.comfyui && !settings.transcription) { await deps.prisma.integrationProviderConfig.deleteMany({ where: { id: MEDIA_SETTINGS_ID } }); return mediaSettingsView(settings); diff --git a/packages/adapters/src/media-tools.test.ts b/packages/adapters/src/media-tools.test.ts new file mode 100644 index 0000000..9af7370 --- /dev/null +++ b/packages/adapters/src/media-tools.test.ts @@ -0,0 +1,125 @@ +import type { MediaCapabilities, MediaRuntime, StartMediaJobInput } from "@rakazo/adapter-kit"; +import { describe, expect, it, vi } from "vitest"; +import { selectBuiltinToolsForRun } from "./executor-tools.js"; +import { MediaJobRejectedError } from "./media-runtime.js"; +import { + executeMediaTool, + isMediaToolName, + MEDIA_CONTINUE_AFTER_PROPERTY, + mediaToolInputSchema, + selectMediaTools, +} from "./media-tools.js"; + +const none: MediaCapabilities = { + generate_image: false, + generate_video: false, + edit_image: false, + edit_video: false, + export_video: false, + transcribe: false, +}; + +const scope = { + spaceId: "space-1", + userId: "user-1", + botId: "bot-1", + groupId: null, + threadId: "thread-1", + runId: "run-1", + operationKey: "effect-1", +}; + +function runtime(start: MediaRuntime["start"]): MediaRuntime { + return { + capabilities: async () => none, + start, + get: async () => null, + cancel: async () => undefined, + }; +} + +describe("media tools", () => { + it("names tools after their operations", () => { + expect(isMediaToolName("generate_image")).toBe(true); + expect(isMediaToolName("transcribe")).toBe(true); + expect(isMediaToolName("render_plot")).toBe(false); + }); + + it("shows a media tool only when the space can run its operation", () => { + const tools = [{ name: "generate_image" }, { name: "transcribe" }, { name: "web_search" }]; + expect(selectMediaTools(tools, undefined).map((tool) => tool.name)).toEqual(["web_search"]); + expect(selectMediaTools(tools, { ...none, transcribe: true }).map((tool) => tool.name)).toEqual( + ["transcribe", "web_search"], + ); + }); + + it("leaves non-media tools of a run untouched", () => { + const base = { + graphicalToolsAllowed: true, + groupId: null, + trigger: "user", + semanticMemoryEnabled: false, + messagingChannelRun: false, + bot: { kind: "specialist", restrictBuiltinTools: false, allowedBuiltinTools: [] }, + }; + const names = (mediaCapabilities?: MediaCapabilities) => + selectBuiltinToolsForRun({ ...base, mediaCapabilities }) + .map((tool) => tool.name) + .filter((name) => !isMediaToolName(name)); + expect(names({ ...none, generate_image: true, transcribe: true })).toEqual(names(undefined)); + expect(selectBuiltinToolsForRun(base).some((tool) => isMediaToolName(tool.name))).toBe(false); + }); + + it("adds continue_after to every media tool schema", () => { + expect( + mediaToolInputSchema({ + properties: { prompt: { type: "string" } }, + required: ["prompt"], + }), + ).toEqual({ + type: "object", + properties: { prompt: { type: "string" }, ...MEDIA_CONTINUE_AFTER_PROPERTY }, + required: ["prompt"], + }); + }); + + it("starts the job with the tool scope and returns at once", async () => { + const start = vi.fn(async (_input: StartMediaJobInput) => ({ + jobId: "job-1", + messageId: "message-1", + })); + const result = await executeMediaTool(runtime(start), scope, "generate_image", { + prompt: "a fox", + continue_after: true, + }); + expect(start).toHaveBeenCalledWith({ + operation: "generate_image", + input: { prompt: "a fox" }, + ...scope, + continueAfter: true, + }); + expect(result).toMatchObject({ jobId: "job-1", status: "queued" }); + expect((result as { note: string }).note).toContain("woken"); + + await executeMediaTool(runtime(start), scope, "generate_image", { prompt: "a fox" }); + expect(start.mock.calls[1]?.[0]).toMatchObject({ continueAfter: false }); + }); + + it("returns refusals to the model and rethrows other failures", async () => { + const refused = runtime(async () => { + throw new MediaJobRejectedError("Invalid generate_image input: prompt: Required"); + }); + expect(await executeMediaTool(refused, scope, "generate_image", {})).toEqual({ + error: "Invalid generate_image input: prompt: Required", + }); + expect(await executeMediaTool(undefined, scope, "transcribe", {})).toEqual({ + error: "transcribe is not available.", + }); + const broken = runtime(async () => { + throw new Error("database down"); + }); + await expect(executeMediaTool(broken, scope, "transcribe", {})).rejects.toThrow( + "database down", + ); + }); +}); diff --git a/packages/adapters/src/media/comfyui-settings.test.ts b/packages/adapters/src/media/comfyui-settings.test.ts new file mode 100644 index 0000000..5026881 --- /dev/null +++ b/packages/adapters/src/media/comfyui-settings.test.ts @@ -0,0 +1,28 @@ +import { describe, expect, it } from "vitest"; +import { saveMediaSettings } from "../media-settings.js"; + +const comfyui = (workflows: { imageWorkflow?: string; videoWorkflow?: string }) => ({ + comfyui: { baseUrl: "http://127.0.0.1:8188", ...workflows }, + transcription: null, +}); + +describe("ComfyUI settings", () => { + it("refuses custom workflows the provider could not fill", async () => { + const deps = {} as Parameters[0]; + await expect( + saveMediaSettings( + deps, + comfyui({ imageWorkflow: '{"1": {"class_type": "X", "inputs": {}}}' }), + ), + ).rejects.toThrow("Image workflow: Workflow needs a {{prompt}} placeholder"); + await expect( + saveMediaSettings( + deps, + comfyui({ + videoWorkflow: + '{"1": {"class_type": "X", "inputs": {"t": "{{prompt}}", "c": "{{cfg}}"}}}', + }), + ), + ).rejects.toThrow("Video workflow: Unknown workflow placeholders: {{cfg}}"); + }); +}); diff --git a/packages/adapters/src/media/image-edit.test.ts b/packages/adapters/src/media/image-edit.test.ts new file mode 100644 index 0000000..638c7b1 --- /dev/null +++ b/packages/adapters/src/media/image-edit.test.ts @@ -0,0 +1,110 @@ +import sharp from "sharp"; +import { describe, expect, it } from "vitest"; +import { editImageGeometry, encodeImage, imageSize } from "./image-edit.js"; + +/** Left half red, right half blue. */ +async function halves(width = 40, height = 20): Promise { + const half = Math.floor(width / 2); + const blue = await sharp({ + create: { width: width - half, height, channels: 3, background: { r: 0, g: 0, b: 255 } }, + }) + .png() + .toBuffer(); + return sharp({ create: { width, height, channels: 3, background: { r: 255, g: 0, b: 0 } } }) + .composite([{ input: blue, left: half, top: 0 }]) + .png() + .toBuffer(); +} + +async function pixel(image: Uint8Array, x: number, y: number) { + const { data, info } = await sharp(image).raw().toBuffer({ resolveWithObject: true }); + const offset = (y * info.width + x) * info.channels; + return [...data.subarray(offset, offset + info.channels)]; +} + +async function size(image: Uint8Array) { + const { width, height } = await sharp(image).metadata(); + return [width, height]; +} + +describe("editImageGeometry", () => { + it("applies crop before rotate, in source coordinates", async () => { + const result = await editImageGeometry(await halves(), { + crop: { left: 20, top: 0, width: 20, height: 10 }, + rotate: 90, + }); + expect(await size(result)).toEqual([10, 20]); + expect((await pixel(result, 5, 5)).slice(0, 3)).toEqual([0, 0, 255]); + }); + + it("flips horizontally and vertically", async () => { + const flopped = await editImageGeometry(await halves(), { flip: "horizontal" }); + expect((await pixel(flopped, 0, 0)).slice(0, 3)).toEqual([0, 0, 255]); + const flipped = await editImageGeometry(await halves(), { flip: "vertical" }); + expect((await pixel(flipped, 0, 0)).slice(0, 3)).toEqual([255, 0, 0]); + }); + + it("resizes to a preset with crop, pad or blur, and keeps the aspect for one side", async () => { + const source = await halves(160, 90); + for (const fit of ["crop", "pad", "blur"] as const) { + const result = await editImageGeometry(source, { resize: { aspect: "square", fit } }); + expect(await size(result)).toEqual([1080, 1080]); + } + const padded = await editImageGeometry(source, { resize: { aspect: "vertical", fit: "pad" } }); + expect((await pixel(padded, 0, 0)).at(-1)).toBe(0); + const scaled = await editImageGeometry(source, { resize: { width: 80, fit: "crop" } }); + expect(await size(scaled)).toEqual([80, 45]); + }); + + it("extends with a color or transparency", async () => { + const colored = await editImageGeometry(await halves(), { + extend: { top: 5, right: 0, bottom: 0, left: 3, color: "#00ff00" }, + }); + expect(await size(colored)).toEqual([43, 25]); + expect((await pixel(colored, 0, 0)).slice(0, 3)).toEqual([0, 255, 0]); + const clear = await editImageGeometry(await halves(), { + extend: { top: 2, right: 2, bottom: 2, left: 2 }, + }); + expect((await pixel(clear, 0, 0)).at(-1)).toBe(0); + }); + + it("applies EXIF orientation first", async () => { + const rotatedJpeg = await sharp(await halves()) + .jpeg() + .withMetadata({ orientation: 6 }) + .toBuffer(); + expect(await imageSize(rotatedJpeg)).toEqual({ width: 20, height: 40 }); + expect(await size(await editImageGeometry(rotatedJpeg, { flip: "vertical" }))).toEqual([ + 20, 40, + ]); + }); + + it("refuses crops outside the image and oversized results", async () => { + await expect( + editImageGeometry(await halves(), { crop: { left: 30, top: 0, width: 20, height: 10 } }), + ).rejects.toThrow("The crop area is outside the 40×20 image"); + await expect( + editImageGeometry(await halves(), { + extend: { top: 9000, right: 0, bottom: 0, left: 9000 }, + }), + ).rejects.toThrow("The edited image would be too large"); + await expect( + editImageGeometry(await halves(), { resize: { width: 8192, height: 8193, fit: "crop" } }), + ).rejects.toThrow("The edited image would be too large"); + }); +}); + +describe("encodeImage", () => { + it("encodes each format and flattens transparency for JPEG", async () => { + const clear = await editImageGeometry(await halves(), { + extend: { top: 2, right: 0, bottom: 0, left: 0 }, + }); + for (const format of ["png", "jpeg", "webp"] as const) { + const encoded = await encodeImage(clear, format); + expect((await sharp(encoded.bytes).metadata()).format).toBe(format); + expect([encoded.width, encoded.height]).toEqual([40, 22]); + } + const jpeg = await encodeImage(clear, "jpeg"); + expect((await pixel(jpeg.bytes, 0, 0)).every((channel) => channel < 16)).toBe(true); + }); +}); diff --git a/packages/adapters/src/media/image-edit.ts b/packages/adapters/src/media/image-edit.ts new file mode 100644 index 0000000..e250821 --- /dev/null +++ b/packages/adapters/src/media/image-edit.ts @@ -0,0 +1,156 @@ +import type { EditImageInput, ImageFormat } from "@rakazo/contracts"; +import { MEDIA_PRESETS } from "@rakazo/contracts"; +import type sharpFactory from "sharp"; +import type { Sharp } from "sharp"; + +type SharpFactory = typeof sharpFactory; +export type ImageGeometryEdits = Pick< + EditImageInput, + "crop" | "rotate" | "flip" | "resize" | "extend" +>; + +/** Edited images stay below this many pixels (8192²). */ +const MAX_EDIT_PIXELS = 8192 * 8192; +const TRANSPARENT = { r: 0, g: 0, b: 0, alpha: 0 }; +const BLACK = { r: 0, g: 0, b: 0 }; + +let sharpModule: Promise | undefined; + +/** sharp is a native module; load it on first use so a platform without it only loses edits. */ +export function loadSharp(): Promise { + sharpModule ??= import("sharp").then((module) => module.default); + return sharpModule; +} + +export async function sharpAvailable(): Promise { + return loadSharp().then( + () => true, + () => false, + ); +} + +/** Displayed size, with EXIF orientation applied. */ +export async function imageSize( + input: string | Uint8Array, +): Promise<{ width: number; height: number }> { + const sharp = await loadSharp(); + const { autoOrient } = await sharp(input).metadata(); + return { width: autoOrient.width, height: autoOrient.height }; +} + +/** Covers exactly width×height (center crop) as PNG: generation inputs at the size asked for. */ +export async function coverImage( + input: string | Uint8Array, + width: number, + height: number, +): Promise { + const sharp = await loadSharp(); + const png = await sharp(input, { autoOrient: true }) + .resize(width, height, { fit: "cover" }) + .png() + .toBuffer(); + return new Uint8Array(png); +} + +function parseColor(hex: string | undefined) { + if (!hex) return TRANSPARENT; + const value = Number.parseInt(hex.slice(1), 16); + return { r: (value >> 16) & 255, g: (value >> 8) & 255, b: value & 255, alpha: 1 }; +} + +function assertSize(width: number, height: number): void { + if (width * height > MAX_EDIT_PIXELS) throw new Error("The edited image would be too large"); +} + +/** Runs each step on decoded pixels, so steps apply in the order given, not sharp's own order. */ +async function materialize(sharp: SharpFactory, image: Sharp) { + const { data, info } = await image.raw().toBuffer({ resolveWithObject: true }); + const { width, height, channels } = info; + return { image: sharp(data, { raw: { width, height, channels } }), width, height }; +} + +async function resizeTo( + sharp: SharpFactory, + image: Sharp, + width: number, + height: number, + fit: "crop" | "pad" | "blur", +): Promise { + if (fit === "crop") return image.resize(width, height, { fit: "cover" }); + if (fit === "pad") + return image.resize(width, height, { fit: "contain", background: TRANSPARENT }); + const source = await image.png().toBuffer(); + const background = await sharp(source) + .resize(width, height, { fit: "cover" }) + .blur(Math.max(width, height) / 40) + .toBuffer(); + const foreground = await sharp(source).resize(width, height, { fit: "inside" }).toBuffer(); + return sharp(background).composite([{ input: foreground, gravity: "centre" }]); +} + +/** + * Applies crop, rotate, flip, resize and extend in that order and returns a lossless PNG. + * EXIF orientation is applied first, so coordinates match what viewers show. + */ +export async function editImageGeometry( + input: string | Uint8Array, + edits: ImageGeometryEdits, +): Promise { + const sharp = await loadSharp(); + let current = await materialize(sharp, sharp(input, { autoOrient: true })); + const step = async (apply: (image: Sharp) => Sharp | Promise) => { + current = await materialize(sharp, await apply(current.image)); + }; + const { crop, rotate, flip, resize, extend } = edits; + if (crop) { + if (crop.left + crop.width > current.width || crop.top + crop.height > current.height) { + throw new Error(`The crop area is outside the ${current.width}×${current.height} image`); + } + await step((image) => image.extract(crop)); + } + if (rotate) await step((image) => image.rotate(rotate)); + if (flip) await step((image) => (flip === "vertical" ? image.flip() : image.flop())); + if (resize) { + const preset = resize.aspect ? MEDIA_PRESETS[resize.aspect] : undefined; + const width = preset?.width ?? resize.width; + const height = preset?.height ?? resize.height; + const scale = width ? width / current.width : height! / current.height; + assertSize(width ?? current.width * scale, height ?? current.height * scale); + await step((image) => + width && height + ? resizeTo(sharp, image, width, height, resize.fit) + : image.resize(width ?? null, height ?? null), + ); + } + if (extend) { + const { top, right, bottom, left, color } = extend; + assertSize(current.width + left + right, current.height + top + bottom); + await step((image) => + image.extend({ top, right, bottom, left, background: parseColor(color) }), + ); + } + return current.image.png().toBuffer(); +} + +export const IMAGE_FORMAT_MIME_TYPES: Record = { + png: "image/png", + jpeg: "image/jpeg", + webp: "image/webp", +}; + +/** Encodes the final image; JPEG has no transparency, so transparent areas turn black. */ +export async function encodeImage( + input: Uint8Array, + format: ImageFormat, +): Promise<{ bytes: Buffer; width: number; height: number }> { + const sharp = await loadSharp(); + const image = sharp(input); + const encoded = + format === "jpeg" + ? image.flatten({ background: BLACK }).jpeg({ quality: 90, mozjpeg: true }) + : format === "webp" + ? image.webp({ quality: 90 }) + : image.png(); + const { data, info } = await encoded.toBuffer({ resolveWithObject: true }); + return { bytes: data, width: info.width, height: info.height }; +} diff --git a/packages/adapters/src/media/image-operations.postgres.test.ts b/packages/adapters/src/media/image-operations.postgres.test.ts new file mode 100644 index 0000000..2809d4d --- /dev/null +++ b/packages/adapters/src/media/image-operations.postgres.test.ts @@ -0,0 +1,172 @@ +import { randomUUID } from "node:crypto"; +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { JobPublisher } from "@rakazo/adapter-kit"; +import type { MediaJobBlock, MessageBlock } from "@rakazo/contracts"; +import { ComfyUiSettingsSchema } from "@rakazo/contracts"; +import type { PrismaClient } from "@rakazo/db"; +import { createDb } from "@rakazo/db"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { LocalArtifactStore } from "../artifacts.js"; +import { createMediaRuntime, MediaJobRejectedError } from "../media-runtime.js"; +import { ComfyUiEmulator } from "../testing/comfyui-emulator.js"; +import { createImageOperations } from "./image-operations.js"; + +const describePostgres = + process.env.VERIFY_DATABASE && process.env.DATABASE_URL ? describe.sequential : describe.skip; + +describePostgres("ComfyUI operations in the media runtime (PostgreSQL + ComfyUI emulator)", () => { + let db: ReturnType; + let prisma: PrismaClient; + let root: string; + const users: string[] = []; + + beforeAll(async () => { + db = createDb(process.env.DATABASE_URL!); + prisma = db.prisma; + root = await mkdtemp(join(tmpdir(), "image-operations-pg-")); + }); + + afterAll(async () => { + if (!db) return; + await prisma.mediaJob.deleteMany({ where: { userId: { in: users } } }); + await prisma.artifact.deleteMany({ where: { userId: { in: users } } }); + await prisma.organization.deleteMany({ where: { id: { in: users } } }); + await prisma.user.deleteMany({ where: { id: { in: users } } }); + await prisma.$disconnect(); + await db.pool.end(); + await rm(root, { recursive: true, force: true }); + }); + + async function setup() { + const id = randomUUID(); + users.push(id); + await prisma.user.create({ + data: { id, name: "Media test", email: `${id}@example.test`, emailVerified: true }, + }); + await prisma.organization.create({ + data: { id, name: "Media test", slug: id, createdAt: new Date() }, + }); + await prisma.space.create({ data: { id, organizationId: id, name: "Media", isDefault: true } }); + await prisma.member.create({ + data: { id, organizationId: id, userId: id, role: "owner", createdAt: new Date() }, + }); + const bot = await prisma.bot.create({ + data: { spaceId: id, userId: id, name: "Studio", color: "test-color" }, + }); + const thread = await prisma.thread.create({ data: { spaceId: id, userId: id, botId: bot.id } }); + const emulator = new ComfyUiEmulator(); + emulator.pendingPolls = 0; + const jobs: JobPublisher = { + enqueue: async () => undefined, + cancel: async () => undefined, + close: async () => undefined, + }; + const runtime = createMediaRuntime( + { + prisma, + events: { notify: async () => undefined }, + jobs, + artifacts: new LocalArtifactStore(join(root, id, "artifacts")), + registrations: createImageOperations({ + loadSettings: async () => ({ + comfyui: ComfyUiSettingsSchema.parse({ baseUrl: "http://127.0.0.1:8188" }), + transcription: null, + }), + fetch: emulator.fetch, + }), + }, + { tempRoot: join(root, id, "jobs") }, + ); + const scope = { + spaceId: id, + userId: id, + botId: bot.id, + groupId: null, + threadId: thread.id, + runId: null, + }; + return { emulator, runtime, scope }; + } + + async function card(messageId: string, jobId: string): Promise { + const message = await prisma.message.findUniqueOrThrow({ where: { id: messageId } }); + return (message.blocks as MessageBlock[]).find( + (block): block is MediaJobBlock => block.kind === "media" && block.jobId === jobId, + ); + } + + it("generates images onto the card and edits one into a linked version", async () => { + const { runtime, scope } = await setup(); + const generated = await runtime.start({ + ...scope, + operation: "generate_image", + input: { prompt: "a lighthouse at dusk", aspect: "landscape", count: 2, seed: 3 }, + }); + await runtime.runJob(generated.jobId); + const job = await runtime.get(generated.jobId); + expect(job).toMatchObject({ status: "done", title: "a lighthouse at dusk" }); + expect(job?.outputs.map((output) => [output.mimeType, output.width, output.height])).toEqual([ + ["image/png", 1344, 768], + ["image/png", 1344, 768], + ]); + expect(await card(generated.messageId, generated.jobId)).toMatchObject({ + status: "done", + outputs: job?.outputs, + }); + + const source = job!.outputs[0]!.artifactId; + const edited = await runtime.start({ + ...scope, + operation: "edit_image", + input: { artifactId: source, resize: { aspect: "vertical", fit: "blur" }, format: "jpeg" }, + }); + await runtime.runJob(edited.jobId); + const edit = await runtime.get(edited.jobId); + expect(edit).toMatchObject({ status: "done", title: "9:16 · JPEG" }); + const [output] = edit!.outputs; + expect(output).toMatchObject({ mimeType: "image/jpeg", width: 1080, height: 1920 }); + const artifact = await prisma.artifact.findUniqueOrThrow({ where: { id: output!.artifactId } }); + expect(artifact).toMatchObject({ parentArtifactId: source, width: 1080, height: 1920 }); + }); + + it("refuses source artifacts the user cannot read", async () => { + const { runtime, scope } = await setup(); + const other = await setup(); + const generated = await other.runtime.start({ + ...other.scope, + operation: "generate_image", + input: { prompt: "private" }, + }); + await other.runtime.runJob(generated.jobId); + const foreign = (await other.runtime.get(generated.jobId))!.outputs[0]!.artifactId; + await expect( + runtime.start({ + ...scope, + operation: "generate_video", + input: { prompt: "animate it", sourceArtifactId: foreign }, + }), + ).rejects.toThrow(MediaJobRejectedError); + await expect( + runtime.start({ + ...scope, + operation: "edit_image", + input: { artifactId: foreign, rotate: 90 }, + }), + ).rejects.toThrow(MediaJobRejectedError); + }); + + it("waits while the GPU lacks free memory", async () => { + const { emulator, runtime, scope } = await setup(); + emulator.gpu = { totalMb: 24_576, freeMb: 4_000, torchReservedMb: 0, torchFreeMb: 0 }; + const video = await runtime.start({ + ...scope, + operation: "generate_video", + input: { prompt: "rain on a window" }, + }); + await runtime.runJob(video.jobId); + expect(await runtime.get(video.jobId)).toMatchObject({ status: "waiting", outputs: [] }); + expect(emulator.prompts).toHaveLength(0); + }); +}); diff --git a/packages/adapters/src/media/image-operations.test.ts b/packages/adapters/src/media/image-operations.test.ts new file mode 100644 index 0000000..e028552 --- /dev/null +++ b/packages/adapters/src/media/image-operations.test.ts @@ -0,0 +1,274 @@ +import { copyFile, mkdtemp, rm, stat, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { MediaJobContext, MediaJobOutputFile } from "@rakazo/adapter-kit"; +import { MediaJobDeferredError } from "@rakazo/adapter-kit"; +import type { ComfyUiSettings, MediaOutput } from "@rakazo/contracts"; +import { attachmentKindOf, ComfyUiSettingsSchema } from "@rakazo/contracts"; +import sharp from "sharp"; +import { afterEach, describe, expect, it } from "vitest"; +import { ComfyUiEmulator } from "../testing/comfyui-emulator.js"; +import { + createImageOperations, + editImageTitle, + fileStem, + generationSize, + IMAGE_GENERATION_MIN_VRAM_MB, + VIDEO_GENERATION_MIN_VRAM_MB, +} from "./image-operations.js"; + +const dirs: string[] = []; +afterEach(async () => { + await Promise.all(dirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true }))); +}); + +async function tempDir(): Promise { + const dir = await mkdtemp(join(tmpdir(), "image-operations-")); + dirs.push(dir); + return dir; +} + +async function fixture(name: string, width: number, height: number): Promise { + const path = join(await tempDir(), name); + const image = sharp({ + create: { width, height, channels: 3, background: { r: 200, g: 120, b: 40 } }, + }); + await writeFile(path, await (name.endsWith(".jpg") ? image.jpeg() : image.png()).toBuffer()); + return path; +} + +function setup( + settings: ComfyUiSettings | null = ComfyUiSettingsSchema.parse({ + baseUrl: "http://127.0.0.1:8188", + }), +) { + const emulator = new ComfyUiEmulator(); + emulator.pendingPolls = 0; + const operations = createImageOperations({ + loadSettings: async () => ({ comfyui: settings, transcription: null }), + fetch: emulator.fetch, + }); + const operation = (name: string) => operations.find((entry) => entry.operation === name)!; + return { emulator, operations, operation }; +} + +async function run( + operation: ReturnType["operation"], + name: string, + input: unknown, + artifacts: Record = {}, +) { + const dir = await tempDir(); + const outputs: MediaJobOutputFile[] = []; + const progress: number[] = []; + const context: MediaJobContext = { + jobId: "job-1", + spaceId: "space-1", + userId: "user-1", + botId: null, + input, + tempDir: dir, + signal: new AbortController().signal, + progress: async (fraction) => { + progress.push(fraction); + }, + fetchArtifact: async (artifactId) => { + const source = artifacts[artifactId]; + if (!source) throw new Error("Artifact not found"); + const name = source.split("/").at(-1)!; + const path = join(dir, `input-${name}`); + await copyFile(source, path); + const mimeType = name.endsWith(".jpg") + ? "image/jpeg" + : name.endsWith(".mp4") + ? "video/mp4" + : "image/png"; + return { + artifactId, + path, + name, + mimeType, + kind: attachmentKindOf(mimeType), + size: (await stat(path)).size, + }; + }, + addOutput: async (file) => { + outputs.push(file); + const output: MediaOutput = { + artifactId: `output-${outputs.length}`, + kind: attachmentKindOf(file.mimeType), + mimeType: file.mimeType, + name: file.name, + size: (await stat(file.path)).size, + }; + return output; + }, + }; + const result = await operation(name).handler.run(context); + return { result, outputs, progress }; +} + +describe("image operations", () => { + it("are available by configuration and guard the GPU", async () => { + const configured = setup(); + expect(configured.operations.map((entry) => entry.operation)).toEqual([ + "generate_image", + "generate_video", + "edit_image", + ]); + for (const entry of configured.operations) { + await expect(entry.available({ spaceId: "space-1" })).resolves.toBe(true); + } + const generateImage = configured.operation("generate_image"); + expect(generateImage.gpu?.minVramMb).toBe(IMAGE_GENERATION_MIN_VRAM_MB); + expect(configured.operation("generate_video").gpu?.minVramMb).toBe( + VIDEO_GENERATION_MIN_VRAM_MB, + ); + await expect(generateImage.gpu?.status()).resolves.toEqual({ + vramFreeMb: 20_000, + vramTotalMb: 24_576, + }); + const unconfigured = setup(null); + await expect( + unconfigured.operation("generate_image").available({ spaceId: "s" }), + ).resolves.toBe(false); + await expect(unconfigured.operation("edit_image").available({ spaceId: "s" })).resolves.toBe( + true, + ); + await expect(run(unconfigured.operation, "generate_image", { prompt: "x" })).rejects.toThrow( + "ComfyUI is not configured", + ); + }); + + it("generates images at the aspect asked for", async () => { + const { operation } = setup(); + const { result, outputs, progress } = await run(operation, "generate_image", { + prompt: "A red kite, über the sea!", + aspect: "vertical", + count: 2, + seed: 11, + continue_after: true, + }); + expect( + outputs.map(({ name, mimeType, width, height }) => [name, mimeType, width, height]), + ).toEqual([ + ["a-red-kite-uber-the-sea-1.png", "image/png", 768, 1344], + ["a-red-kite-uber-the-sea-2.png", "image/png", 768, 1344], + ]); + expect(progress.at(-1)).toBe(1); + expect(result.summary).toBe("Generated 2 images with seed 11."); + }); + + it("starts from a source image at the source's aspect and links the result to it", async () => { + const { operation, emulator } = setup(); + const photo = await fixture("photo.jpg", 300, 200); + const { outputs } = await run( + operation, + "generate_image", + { prompt: "night", sourceArtifactId: "photo-1" }, + { "photo-1": photo }, + ); + const [upload] = emulator.uploads.values(); + expect((await sharp(upload!).metadata()).width).toBe(1280); + expect(outputs[0]).toMatchObject({ width: 1280, height: 832, parentArtifactId: "photo-1" }); + const clip = await fixture("clip.mp4", 8, 8); + await expect( + run(operation, "generate_image", { prompt: "x", sourceArtifactId: "c" }, { c: clip }), + ).rejects.toThrow("The source must be an image"); + }); + + it("generates video in the preset's aspect", async () => { + const { operation, emulator } = setup(); + const { outputs, result } = await run(operation, "generate_video", { + prompt: "paper boat on a river", + aspect: "vertical", + }); + expect(outputs).toEqual([ + expect.objectContaining({ + name: "paper-boat-on-a-river.mp4", + mimeType: "video/mp4", + width: 480, + height: 832, + durationMs: 5042, + }), + ]); + expect(emulator.prompts[0]?.graph["6"]?.inputs).toMatchObject({ width: 480, height: 832 }); + expect(result.summary).toMatch(/^Generated a 5 s video with seed \d+\.$/); + }); + + it("edits locally in order and keeps the source as parent", async () => { + const { operation, emulator } = setup(null); + const photo = await fixture("photo.jpg", 400, 300); + const { outputs, result } = await run( + operation, + "edit_image", + { + artifactId: "photo-1", + crop: { left: 0, top: 0, width: 300, height: 300 }, + rotate: 90, + resize: { aspect: "vertical", fit: "blur" }, + format: "webp", + }, + { "photo-1": photo }, + ); + expect(outputs).toEqual([ + expect.objectContaining({ + name: "photo-edit.webp", + mimeType: "image/webp", + width: 1080, + height: 1920, + parentArtifactId: "photo-1", + }), + ]); + expect((await sharp(outputs[0]!.path).metadata()).format).toBe("webp"); + expect(result.summary).toBe("Edited photo.jpg into photo-edit.webp (1080×1920)."); + expect(emulator.requests).toHaveLength(0); + await expect( + run(operation, "edit_image", { artifactId: "photo-1", prompt: "snow" }, { "photo-1": photo }), + ).rejects.toThrow("ComfyUI is not configured"); + }); + + it("repaints an edit on ComfyUI at the edited size, deferring while the GPU is busy", async () => { + const { operation, emulator } = setup(); + const photo = await fixture("photo.png", 300, 200); + const input = { + artifactId: "photo-1", + crop: { left: 0, top: 0, width: 250, height: 200 }, + prompt: "snow", + }; + const { outputs } = await run(operation, "edit_image", input, { "photo-1": photo }); + expect(outputs[0]).toMatchObject({ name: "photo-edit.png", width: 250, height: 200 }); + const [upload] = emulator.uploads.values(); + expect(await sharp(upload!).metadata()).toMatchObject({ width: 1152, height: 896 }); + emulator.gpu = { totalMb: 24_576, freeMb: 2_000, torchReservedMb: 0, torchFreeMb: 0 }; + await expect(run(operation, "edit_image", input, { "photo-1": photo })).rejects.toBeInstanceOf( + MediaJobDeferredError, + ); + }); +}); + +describe("helpers", () => { + it("sizes generations near the pixel budget", () => { + expect(generationSize(1, 1024 * 1024, 64)).toEqual({ width: 1024, height: 1024 }); + expect(generationSize(16 / 9, 768 * 512, 32)).toEqual({ width: 832, height: 480 }); + expect(generationSize(21 / 9, 1024 * 1024, 64)).toEqual({ width: 1536, height: 640 }); + }); + + it("names files and titles edits briefly", () => { + expect(fileStem(" Ünïcode — Straße 42 at night, very long prompt here ", "image")).toBe( + "unicode-strasse-42-at-night-very", + ); + expect(fileStem("夜の海", "video")).toBe("video"); + expect( + editImageTitle({ + artifactId: "a", + crop: { left: 0, top: 0, width: 1, height: 1 }, + rotate: 90, + resize: { aspect: "vertical", fit: "crop" }, + format: "png", + }), + ).toBe("crop · 90° · 9:16 · PNG"); + expect(editImageTitle({ artifactId: "a", resize: { width: 800, fit: "crop" } })).toBe("800×…"); + expect(editImageTitle({ artifactId: "a", prompt: "make it snow" })).toBe("make it snow"); + }); +}); diff --git a/packages/adapters/src/media/image-operations.ts b/packages/adapters/src/media/image-operations.ts new file mode 100644 index 0000000..68c08aa --- /dev/null +++ b/packages/adapters/src/media/image-operations.ts @@ -0,0 +1,315 @@ +import { randomInt } from "node:crypto"; +import { writeFile } from "node:fs/promises"; +import { join } from "node:path"; +import type { + GeneratedMedia, + GpuStatus, + MediaJobContext, + MediaJobInputFile, + MediaOperationResult, +} from "@rakazo/adapter-kit"; +import { MediaJobDeferredError } from "@rakazo/adapter-kit"; +import type { + EditImageInput, + GenerateImageInput, + GenerateVideoInput, + ImageFormat, + MediaPresetId, +} from "@rakazo/contracts"; +import { + attachmentKindOf, + EditImageInputSchema, + GenerateImageInputSchema, + GenerateVideoInputSchema, + MEDIA_PRESETS, +} from "@rakazo/contracts"; +import { attachmentExtensionForMimeType } from "@rakazo/core"; +import type { ComfyUiProvider } from "../comfyui-provider.js"; +import { createComfyUiProvider } from "../comfyui-provider.js"; +import type { MediaOperationDeps } from "../media-operations.js"; +import type { MediaOperationRegistration } from "../media-runtime.js"; +import { + coverImage, + editImageGeometry, + encodeImage, + IMAGE_FORMAT_MIME_TYPES, + imageSize, + sharpAvailable, +} from "./image-edit.js"; + +/** Free GPU memory a generation needs before it starts; fixed, so a custom workflow cannot lower it. */ +export const IMAGE_GENERATION_MIN_VRAM_MB = 6_000; +export const VIDEO_GENERATION_MIN_VRAM_MB = 12_000; +/** Pixel budgets the default workflows are tuned for (SDXL-class images, LTX-Video clips). */ +const IMAGE_PIXELS = 1024 * 1024; +const VIDEO_PIXELS = 768 * 512; +const GPU_RETRY_MS = 60_000; + +export interface ImageOperationsDeps extends MediaOperationDeps { + /** Tests point the provider at an emulator. */ + fetch?: typeof globalThis.fetch; +} + +type ProviderSource = () => Promise; + +/** A size close to `pixels` with this aspect ratio, in multiples of `multiple`. */ +export function generationSize(aspect: number, pixels: number, multiple: number) { + return { + width: Math.max(multiple, Math.round(Math.sqrt(pixels * aspect) / multiple) * multiple), + height: Math.max(multiple, Math.round(Math.sqrt(pixels / aspect) / multiple) * multiple), + }; +} + +function presetAspect(preset: MediaPresetId): number { + return MEDIA_PRESETS[preset].width / MEDIA_PRESETS[preset].height; +} + +/** File name stem from a prompt: its first words in lowercase ASCII. */ +export function fileStem(text: string, fallback: string): string { + const words = text + .normalize("NFKD") + .replace(/\p{M}/gu, "") + .toLowerCase() + .replace(/ß/g, "ss") + .match(/[a-z0-9]+/g); + return words?.slice(0, 6).join("-").slice(0, 48).replace(/-+$/, "") || fallback; +} + +/** Short card title for an edit without a prompt, such as "crop · 9:16 · PNG". */ +export function editImageTitle(input: EditImageInput): string { + if (input.prompt) return input.prompt; + const { resize } = input; + const parts = [ + input.crop && "crop", + input.rotate && `${input.rotate}°`, + input.flip && "flip", + resize && + (resize.aspect + ? MEDIA_PRESETS[resize.aspect].aspect + : `${resize.width ?? "…"}×${resize.height ?? "…"}`), + input.extend && "extend", + input.format?.toUpperCase(), + ]; + return parts.filter(Boolean).join(" · "); +} + +function reportProgress(context: MediaJobContext) { + return (fraction: number) => { + context.progress(fraction).catch(() => undefined); + }; +} + +async function requireProvider(source: ProviderSource): Promise { + const provider = await source(); + if (!provider) throw new Error("ComfyUI is not configured"); + return provider; +} + +/** Defers instead of starting a generation the GPU has no room for right now. */ +async function assertGpuReady( + provider: ComfyUiProvider, + minVramMb: number, + signal: AbortSignal, +): Promise { + const status: GpuStatus | undefined = await provider.gpuStatus(signal); + if (status && status.vramFreeMb < Math.min(minVramMb, status.vramTotalMb)) { + throw new MediaJobDeferredError("Waiting for free GPU memory", GPU_RETRY_MS); + } +} + +async function fetchImage( + context: MediaJobContext, + artifactId: string, +): Promise { + const file = await context.fetchArtifact(artifactId); + if (file.kind !== "image") throw new Error("The source must be an image"); + return { ...file, ...(await imageSize(file.path)) }; +} + +async function addGeneratedOutputs( + context: MediaJobContext, + results: GeneratedMedia[], + stem: string, + parentArtifactId: string | undefined, +): Promise { + for (const [index, result] of results.entries()) { + const extension = attachmentExtensionForMimeType(result.mimeType); + const name = `${stem}${results.length > 1 ? `-${index + 1}` : ""}${extension}`; + const path = join(context.tempDir, name); + await writeFile(path, result.bytes); + const size = + attachmentKindOf(result.mimeType) === "image" + ? await imageSize(path) + : { width: result.width, height: result.height }; + await context.addOutput({ + path, + name, + mimeType: result.mimeType, + ...size, + ...(result.durationMs !== undefined ? { durationMs: result.durationMs } : {}), + ...(parentArtifactId ? { parentArtifactId } : {}), + }); + } +} + +async function generateImage( + context: MediaJobContext, + providers: ProviderSource, +): Promise { + const input = GenerateImageInputSchema.parse(context.input); + const provider = await requireProvider(providers); + const source = input.sourceArtifactId + ? await fetchImage(context, input.sourceArtifactId) + : undefined; + const aspect = input.aspect + ? presetAspect(input.aspect) + : source + ? source.width / source.height + : 1; + const size = + input.width && input.height + ? { width: input.width, height: input.height } + : generationSize(aspect, IMAGE_PIXELS, 64); + const seed = input.seed ?? randomInt(2 ** 48 - 1); + const images = await provider.generateImage({ + prompt: input.prompt, + negativePrompt: input.negativePrompt, + ...size, + count: input.count, + seed, + source: source && { + bytes: await coverImage(source.path, size.width, size.height), + mimeType: "image/png", + }, + strength: input.strength, + signal: context.signal, + onProgress: reportProgress(context), + }); + await addGeneratedOutputs(context, images, fileStem(input.prompt, "image"), source?.artifactId); + const count = images.length === 1 ? "1 image" : `${images.length} images`; + return { summary: `Generated ${count} with seed ${seed}.` }; +} + +async function generateVideo( + context: MediaJobContext, + providers: ProviderSource, +): Promise { + const input = GenerateVideoInputSchema.parse(context.input); + const provider = await requireProvider(providers); + const source = input.sourceArtifactId + ? await fetchImage(context, input.sourceArtifactId) + : undefined; + const aspect = input.aspect + ? presetAspect(input.aspect) + : source + ? source.width / source.height + : presetAspect("landscape"); + const size = generationSize(aspect, VIDEO_PIXELS, 32); + const seed = input.seed ?? randomInt(2 ** 48 - 1); + const clips = await provider.generateVideo({ + prompt: input.prompt, + negativePrompt: input.negativePrompt, + ...size, + durationSeconds: input.durationSeconds, + seed, + source: source && { + bytes: await coverImage(source.path, size.width, size.height), + mimeType: "image/png", + }, + signal: context.signal, + onProgress: reportProgress(context), + }); + await addGeneratedOutputs(context, clips, fileStem(input.prompt, "video"), source?.artifactId); + return { summary: `Generated a ${input.durationSeconds} s video with seed ${seed}.` }; +} + +function defaultFormat(mimeType: string): ImageFormat { + if (mimeType === "image/jpeg") return "jpeg"; + if (mimeType === "image/webp") return "webp"; + return "png"; +} + +async function editImage( + context: MediaJobContext, + providers: ProviderSource, +): Promise { + const input = EditImageInputSchema.parse(context.input); + const source = await fetchImage(context, input.artifactId); + let png: Uint8Array = await editImageGeometry(source.path, input); + if (input.prompt) { + const provider = await requireProvider(providers); + await assertGpuReady(provider, IMAGE_GENERATION_MIN_VRAM_MB, context.signal); + const { width, height } = await imageSize(png); + // Repaint near the model's native size, then return to the edited size. + const size = generationSize(width / height, IMAGE_PIXELS, 64); + const [repainted] = await provider.generateImage({ + prompt: input.prompt, + ...size, + source: { bytes: await coverImage(png, size.width, size.height), mimeType: "image/png" }, + strength: input.strength, + signal: context.signal, + onProgress: reportProgress(context), + }); + if (!repainted) throw new Error("ComfyUI produced no image"); + png = await coverImage(repainted.bytes, width, height); + } + const format = input.format ?? defaultFormat(source.mimeType); + const encoded = await encodeImage(png, format); + const mimeType = IMAGE_FORMAT_MIME_TYPES[format]; + const stem = source.name.replace(/\.[^.]+$/, "") || "image"; + const name = `${stem}-edit${attachmentExtensionForMimeType(mimeType)}`; + const path = join(context.tempDir, name); + await writeFile(path, encoded.bytes); + await context.addOutput({ + path, + name, + mimeType, + width: encoded.width, + height: encoded.height, + parentArtifactId: source.artifactId, + }); + return { summary: `Edited ${source.name} into ${name} (${encoded.width}×${encoded.height}).` }; +} + +/** + * Media operations for image and video generation on ComfyUI and for image edits with sharp. + * Generation is available when ComfyUI is configured; edits whenever sharp loads (repainting + * an edit also needs ComfyUI). + */ +export function createImageOperations(deps: ImageOperationsDeps): MediaOperationRegistration[] { + const providers: ProviderSource = async () => { + const { comfyui } = await deps.loadSettings(); + return comfyui ? createComfyUiProvider({ ...comfyui, fetch: deps.fetch }) : null; + }; + const comfyUiConfigured = async () => Boolean((await deps.loadSettings()).comfyui); + const status = async (signal?: AbortSignal) => (await providers())?.gpuStatus(signal); + const sourceIds = (input: { sourceArtifactId?: string }) => + input.sourceArtifactId ? [input.sourceArtifactId] : []; + const image: MediaOperationRegistration = { + operation: "generate_image", + inputSchema: GenerateImageInputSchema, + title: (input) => input.prompt, + available: comfyUiConfigured, + handler: { operation: "generate_image", run: (context) => generateImage(context, providers) }, + gpu: { minVramMb: IMAGE_GENERATION_MIN_VRAM_MB, status }, + artifactIds: sourceIds, + }; + const video: MediaOperationRegistration = { + operation: "generate_video", + inputSchema: GenerateVideoInputSchema, + title: (input) => input.prompt, + available: comfyUiConfigured, + handler: { operation: "generate_video", run: (context) => generateVideo(context, providers) }, + gpu: { minVramMb: VIDEO_GENERATION_MIN_VRAM_MB, status }, + artifactIds: sourceIds, + }; + const edit: MediaOperationRegistration = { + operation: "edit_image", + inputSchema: EditImageInputSchema, + title: editImageTitle, + available: sharpAvailable, + handler: { operation: "edit_image", run: (context) => editImage(context, providers) }, + artifactIds: (input) => [input.artifactId], + }; + return [image, video, edit] as MediaOperationRegistration[]; +} diff --git a/packages/adapters/src/testing/comfyui-emulator.ts b/packages/adapters/src/testing/comfyui-emulator.ts new file mode 100644 index 0000000..4da22a1 --- /dev/null +++ b/packages/adapters/src/testing/comfyui-emulator.ts @@ -0,0 +1,393 @@ +import { createHash } from "node:crypto"; + +/** A 1 s, 32×32 H.264 MP4 (ffmpeg testsrc), small enough to inline. */ +const TINY_MP4_BASE64 = [ + "AAAAIGZ0eXBpc29tAAACAGlzb21pc28yYXZjMW1wNDEAAAM+bW9vdgAAAGxtdmhkAAAAAAAAAAAAAAAAAAAD6AAAA+gAAQAAAQAA", + "AAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgAA", + "Ao10cmFrAAAAXHRraGQAAAADAAAAAAAAAAAAAAABAAAAAAAAA+gAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAA", + "AAAAAAAAAAAAAABAAAAAACAAAAAgAAAAAAAkZWR0cwAAABxlbHN0AAAAAAAAAAEAAAPoAAAQAAABAAAAAAIFbWRpYQAAACBtZGhk", + "AAAAAAAAAAAAAAAAAABAAAAAQABVxAAAAAAALWhkbHIAAAAAAAAAAHZpZGUAAAAAAAAAAAAAAABWaWRlb0hhbmRsZXIAAAABsG1p", + "bmYAAAAUdm1oZAAAAAEAAAAAAAAAAAAAACRkaW5mAAAAHGRyZWYAAAAAAAAAAQAAAAx1cmwgAAAAAQAAAXBzdGJsAAAAwHN0c2QA", + "AAAAAAAAAQAAALBhdmMxAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAACAAIABIAAAASAAAAAAAAAABDExhdmMgbGlieDI2NAAAAAAA", + "AAAAAAAAAAAAAAAAAAAAGP//AAAANmF2Y0MBZAAK/+EAGGdkAAqscgRJbARAAAADAEAAAAQDxIlhGAEAB2joQ4OSyLD9+PgAAAAA", + "EHBhc3AAAAABAAAAAQAAABRidHJ0AAAAAAAADAgAAAwIAAAAGHN0dHMAAAAAAAAAAQAAAAgAAAgAAAAAFHN0c3MAAAAAAAAAAQAA", + "AAEAAAAYY3R0cwAAAAAAAAABAAAACAAAEAAAAAAcc3RzYwAAAAAAAAABAAAAAQAAAAgAAAABAAAANHN0c3oAAAAAAAAAAAAAAAgA", + "AADTAAAADAAAABgAAAAdAAAAGgAAAB0AAAAfAAAAFwAAABRzdGNvAAAAAAAAAAEAAANuAAAAPXVkdGEAAAA1bWV0YQAAAAAAAAAh", + "aGRscgAAAAAAAAAAbWRpcmFwcGwAAAAAAAAAAAAAAAAIaWxzdAAAAAhmcmVlAAABiW1kYXQAAADPZYiBABP/115QL0x5O3/uan23", + "okcn80QwiyZelIq0iY0T1naTwGHVgXgsN7YBwlEikXvo+AyEjKk/afikvNFJbO7NjVpKr2f7e+xmO5rxkaelhrWKVT83gKOM5wKs", + "QoFpCKDo56fFf6IYs3U1LAHAxZY3IIAUCFca5RyJ/ibKbmkxBCXXiBuUCCuL9FhMjcqhW2misGdpwuLgknHPr/37MowWA2zcOn9B", + "OHR5mfgyt4/0x3o95hpnI+2axZFkXHhYwK55h134UUvr+r3B9b/xAAAACEGaCC2JP6KAAAAAFEGaEEvwIGTKYTe7qicJ6jU1YvCB", + "AAAAGUGaGGk8CB5MpgTfu1UOB2sZPRDVFTCp+IAAAAAWQZogiXwIH5MpgER/wf8hVV5nVt8U2QAAABlBmiipvAgf5MpgBFfFjke1", + "iXbU9gFuIojBAAAAG0GaMMn8CB/5MpgARn/KNfso3F7XaeLq3d3VjQAAABNBmjjojwIH/5MpgAEIf9P6j9Ph", +].join(""); + +export const COMFYUI_EMULATOR_MP4 = new Uint8Array(Buffer.from(TINY_MP4_BASE64, "base64")); + +const MIB = 1024 * 1024; + +type EmulatedNode = { class_type: string; inputs: Record }; +type EmulatedGraph = Record; +type FileRef = { filename: string; subfolder: string; type: string }; + +interface EmulatedPrompt { + id: string; + graph: EmulatedGraph; + polls: number; + failure?: string; + interrupted?: boolean; + /** Dropped from the queue before it ran. */ + removed?: boolean; + /** Reported as finished by the history. */ + done?: boolean; + outputs?: Record>; +} + +export interface EmulatedComfyUiRequest { + method: string; + path: string; + authorization: string | null; + body?: unknown; +} + +function isRecord(value: unknown): value is Record { + return Boolean(value) && typeof value === "object" && !Array.isArray(value); +} + +function json(body: unknown, status = 200): Response { + return Response.json(body, { status }); +} + +/** + * In-process substitute for a ComfyUI server's HTTP API, so the real provider runs in offline + * tests: it validates submitted workflows like ComfyUI does, "runs" them after a few history + * polls, and returns small generated PNGs or an MP4 fixture. + */ +export class ComfyUiEmulator { + readonly requests: EmulatedComfyUiRequest[] = []; + readonly prompts: EmulatedPrompt[] = []; + readonly uploads = new Map(); + readonly deleted: string[] = []; + readonly interrupted: string[] = []; + checkpoints = ["sd_xl_base_1.0.safetensors", "ltx-video-2b-v0.9.5.safetensors"]; + /** `options` answers like current ComfyUI (`["COMBO", {options}]`), `legacy` like older ones. */ + checkpointListFormat: "options" | "legacy" = "options"; + /** null reports a CPU-only server. */ + gpu: { totalMb: number; freeMb: number; torchReservedMb: number; torchFreeMb: number } | null = { + totalMb: 24_576, + freeMb: 20_000, + torchReservedMb: 0, + torchFreeMb: 0, + }; + /** History polls that still report a submitted prompt as running. */ + pendingPolls = 1; + /** Answer the next request with this HTTP status. */ + failNextStatus: number | undefined; + /** Fail the next submitted prompt while it runs, with this exception message. */ + failNextExecution: string | undefined; + private readonly files = new Map(); + private nextPrompt = 1; + + /** `basePath` serves the API under a prefix, like ComfyUI behind a reverse proxy. */ + constructor(private readonly options: { apiKey?: string; basePath?: string } = {}) {} + + readonly fetch: typeof fetch = async (input, init) => { + init?.signal?.throwIfAborted(); + const url = new URL(input instanceof Request ? input.url : String(input)); + const method = init?.method ?? "GET"; + const headers = new Headers(init?.headers); + const authorization = headers.get("authorization"); + const prefix = `/${(this.options.basePath ?? "").replace(/^\/+|\/+$/g, "")}/`.replace( + "//", + "/", + ); + if (!url.pathname.startsWith(prefix)) return json({ error: "Not found" }, 404); + const path = url.pathname.slice(prefix.length); + const body = + typeof init?.body === "string" && headers.get("content-type")?.includes("json") + ? JSON.parse(init.body) + : undefined; + this.requests.push({ method, path, authorization, body }); + if (this.failNextStatus) { + const status = this.failNextStatus; + this.failNextStatus = undefined; + return new Response("Upstream failure", { status }); + } + if (this.options.apiKey && authorization !== `Bearer ${this.options.apiKey}`) { + return json({ error: "Unauthorized" }, 401); + } + const route = `${method} ${path.split("/")[0]}`; + if (route === "GET system_stats") return this.systemStats(); + if (route === "GET object_info" && path === "object_info/CheckpointLoaderSimple") { + return this.checkpointInfo(); + } + if (route === "POST prompt") return this.submit(body); + if (route === "GET history") return this.history(decodeURIComponent(path.slice(8))); + if (route === "GET view") return this.view(url.searchParams); + if (route === "POST upload") return this.upload(input, init); + if (route === "GET queue") return this.queue(); + if (route === "POST queue") { + for (const id of Array.isArray(body?.delete) ? body.delete : []) { + this.deleted.push(id); + // Like ComfyUI, deleting only drops pending prompts; a running one needs an interrupt. + const prompt = this.prompts.find((entry) => entry.id === id); + if (prompt && prompt !== this.unfinished()[0]) prompt.removed = true; + } + return json({}); + } + if (route === "POST interrupt") { + const id = String(body?.prompt_id ?? ""); + this.interrupted.push(id); + const running = this.unfinished()[0]; + if (running && (!id || running.id === id)) running.interrupted = true; + return json({}); + } + return json({ error: "Not found" }, 404); + }; + + private systemStats(): Response { + const gpu = this.gpu; + return json({ + system: { os: "posix", comfyui_version: "0.3.60", embedded_python: false }, + devices: [ + gpu + ? { + name: "cuda:0 Emulated GPU", + type: "cuda", + index: 0, + vram_total: gpu.totalMb * MIB, + vram_free: gpu.freeMb * MIB, + torch_vram_total: gpu.torchReservedMb * MIB, + torch_vram_free: gpu.torchFreeMb * MIB, + } + : { name: "cpu", type: "cpu", index: null, vram_total: 64 * 1024 * MIB, vram_free: 0 }, + ], + }); + } + + private checkpointInfo(): Response { + const spec = + this.checkpointListFormat === "legacy" + ? [this.checkpoints, { tooltip: "Checkpoint" }] + : ["COMBO", { options: this.checkpoints, tooltip: "Checkpoint" }]; + return json({ + CheckpointLoaderSimple: { + input: { required: { ckpt_name: spec } }, + output: ["MODEL", "CLIP", "VAE"], + name: "CheckpointLoaderSimple", + }, + }); + } + + private submit(body: unknown): Response { + const graph = isRecord(body) ? body.prompt : undefined; + if (!isRecord(graph) || typeof (body as { client_id?: unknown }).client_id !== "string") { + return json( + { error: { type: "invalid_prompt", message: "Cannot execute because prompt is invalid" } }, + 400, + ); + } + const nodeErrors: Record = {}; + const reject = (id: string, node: unknown, message: string, details: string) => { + nodeErrors[id] = { + errors: [{ type: "invalid", message, details, extra_info: {} }], + dependent_outputs: [], + class_type: isRecord(node) ? node.class_type : undefined, + }; + }; + for (const [id, node] of Object.entries(graph)) { + if (!isRecord(node) || typeof node.class_type !== "string" || !isRecord(node.inputs)) { + reject(id, node, "Invalid node", "class_type and inputs are required"); + continue; + } + for (const [name, value] of Object.entries(node.inputs)) { + if (typeof value === "string" && value.includes("{{")) { + reject(id, node, "Unfilled placeholder", `${name}: ${value}`); + } + if (Array.isArray(value) && !isRecord(graph[String(value[0])])) { + reject(id, node, "Missing linked node", `${name}: ${String(value[0])}`); + } + } + const checkpoint = node.inputs.ckpt_name; + if (typeof checkpoint === "string" && !this.checkpoints.includes(checkpoint)) { + reject(id, node, "Value not in list", `ckpt_name: '${checkpoint}' not in list`); + } + if (node.class_type === "LoadImage" && !this.uploads.has(String(node.inputs.image))) { + reject(id, node, "Invalid image file", `image - ${String(node.inputs.image)}`); + } + } + if (Object.keys(nodeErrors).length > 0) { + return json( + { + error: { + type: "prompt_outputs_failed_validation", + message: "Prompt outputs failed validation", + }, + node_errors: nodeErrors, + }, + 400, + ); + } + const id = `prompt-${this.nextPrompt++}`; + this.prompts.push({ + id, + graph: graph as EmulatedGraph, + polls: 0, + failure: this.failNextExecution, + }); + this.failNextExecution = undefined; + return json({ prompt_id: id, number: this.nextPrompt, node_errors: {} }); + } + + private async history(id: string): Promise { + const prompt = this.prompts.find((entry) => entry.id === id); + if (!prompt || prompt.removed) return json({}); + if (!prompt.done && !prompt.interrupted && prompt.polls < this.pendingPolls) { + prompt.polls += 1; + return json({}); + } + prompt.done = true; + if (prompt.interrupted) { + return json({ + [id]: { + prompt: [0, id, prompt.graph, {}, []], + outputs: {}, + status: { + status_str: "error", + completed: false, + messages: [["execution_interrupted", { prompt_id: id, node_id: "5" }]], + }, + }, + }); + } + if (prompt.failure) { + return json({ + [id]: { + prompt: [0, id, prompt.graph, {}, []], + outputs: {}, + status: { + status_str: "error", + completed: false, + messages: [ + ["execution_start", { prompt_id: id }], + [ + "execution_error", + { + prompt_id: id, + node_id: "5", + node_type: "KSampler", + exception_message: prompt.failure, + exception_type: "RuntimeError", + traceback: ["..."], + }, + ], + ], + }, + }, + }); + } + prompt.outputs ??= await this.execute(prompt); + return json({ + [id]: { + prompt: [0, id, prompt.graph, {}, Object.keys(prompt.outputs)], + outputs: prompt.outputs, + status: { status_str: "success", completed: true, messages: [] }, + }, + }); + } + + private async execute(prompt: EmulatedPrompt): Promise>> { + const outputs: Record> = {}; + let counter = 0; + const save = (bytes: Uint8Array, extension: string, type: string): FileRef => { + counter += 1; + const ref = { + filename: `${prompt.id}_${String(counter).padStart(5, "0")}_.${extension}`, + subfolder: type === "output" ? "evoo" : "", + type, + }; + this.files.set(fileKey(ref), bytes); + return ref; + }; + for (const [id, node] of Object.entries(prompt.graph)) { + if (node.class_type === "SaveImage" || node.class_type === "PreviewImage") { + const type = node.class_type === "SaveImage" ? "output" : "temp"; + outputs[id] = { images: [save(await this.renderImage(prompt.graph), "png", type)] }; + } else if (node.class_type === "SaveVideo") { + outputs[id] = { images: [save(COMFYUI_EMULATOR_MP4, "mp4", "output")], animated: [true] }; + } else if (node.class_type === "VHS_VideoCombine") { + outputs[id] = { gifs: [save(COMFYUI_EMULATOR_MP4, "mp4", "output")] }; + } + } + return outputs; + } + + /** A flat color derived from the seed, at the workflow's size (or its input image's). */ + private async renderImage(graph: EmulatedGraph): Promise { + const { default: sharp } = await import("sharp"); + const nodes = Object.values(graph); + const input = (name: string) => + nodes.map((node) => node.inputs[name]).find((value) => typeof value === "number"); + let width = Number(input("width") ?? 0); + let height = Number(input("height") ?? 0); + const loaded = nodes.find((node) => node.class_type === "LoadImage"); + if (loaded && !(width && height)) { + const metadata = await sharp(this.uploads.get(String(loaded.inputs.image))).metadata(); + width = metadata.width ?? 64; + height = metadata.height ?? 64; + } + const seed = String(input("seed") ?? input("noise_seed") ?? 0); + const [r = 0, g = 0, b = 0] = createHash("sha256").update(seed).digest(); + const png = await sharp({ + create: { width: width || 64, height: height || 64, channels: 3, background: { r, g, b } }, + }) + .png() + .toBuffer(); + return new Uint8Array(png); + } + + private view(params: URLSearchParams): Response { + const bytes = this.files.get( + fileKey({ + filename: params.get("filename") ?? "", + subfolder: params.get("subfolder") ?? "", + type: params.get("type") ?? "", + }), + ); + if (!bytes) return new Response("Not found", { status: 404 }); + return new Response(Buffer.from(bytes), { + headers: { "Content-Type": "application/octet-stream" }, + }); + } + + private async upload(input: RequestInfo | URL, init?: RequestInit): Promise { + const form = await new Request(String(input), { + method: "POST", + headers: init?.headers, + body: init?.body, + }).formData(); + const image = form.get("image"); + if (!(image instanceof File) || form.get("type") !== "input") { + return json({ error: "Invalid upload" }, 400); + } + this.uploads.set(image.name, new Uint8Array(await image.arrayBuffer())); + return json({ name: image.name, subfolder: "", type: "input" }); + } + + private unfinished(): EmulatedPrompt[] { + return this.prompts.filter((prompt) => !prompt.done && !prompt.removed); + } + + private queue(): Response { + const running = this.unfinished(); + return json({ + queue_running: running.slice(0, 1).map((prompt) => [0, prompt.id, prompt.graph, {}, []]), + queue_pending: running + .slice(1) + .map((prompt, index) => [index + 1, prompt.id, prompt.graph, {}, []]), + }); + } +} + +function fileKey(ref: FileRef): string { + return `${ref.type}/${ref.subfolder}/${ref.filename}`; +} diff --git a/packages/contracts/src/index.ts b/packages/contracts/src/index.ts index 72cd955..3b62947 100644 --- a/packages/contracts/src/index.ts +++ b/packages/contracts/src/index.ts @@ -10,6 +10,8 @@ export * from "./integration-settings.js"; export * from "./local-settings.js"; export * from "./mcp.js"; export * from "./media.js"; +export * from "./media-generation.js"; +export * from "./media-jobs.js"; export * from "./media-settings.js"; export * from "./media-transcription.js"; export * from "./model-tiers.js"; diff --git a/packages/contracts/src/media-generation.test.ts b/packages/contracts/src/media-generation.test.ts new file mode 100644 index 0000000..32f9b78 --- /dev/null +++ b/packages/contracts/src/media-generation.test.ts @@ -0,0 +1,42 @@ +import { describe, expect, it } from "vitest"; +import { + EditImageInputSchema, + GenerateImageInputSchema, + GenerateVideoInputSchema, +} from "./media-generation.js"; + +describe("media generation inputs", () => { + it("fills defaults and drops unknown fields such as continue_after", () => { + expect(GenerateImageInputSchema.parse({ prompt: " a fox ", continue_after: true })).toEqual({ + prompt: "a fox", + count: 1, + }); + expect(GenerateVideoInputSchema.parse({ prompt: "waves" })).toEqual({ + prompt: "waves", + durationSeconds: 5, + }); + }); + + it("bounds image count, size and video duration", () => { + expect(GenerateImageInputSchema.safeParse({ prompt: "x", count: 5 }).success).toBe(false); + expect(GenerateImageInputSchema.safeParse({ prompt: "x", width: 512 }).success).toBe(false); + expect( + GenerateImageInputSchema.safeParse({ prompt: "x", width: 512, height: 4096 }).success, + ).toBe(false); + expect(GenerateVideoInputSchema.safeParse({ prompt: "x", durationSeconds: 11 }).success).toBe( + false, + ); + }); + + it("needs at least one edit and valid edit values", () => { + expect(EditImageInputSchema.safeParse({ artifactId: "a" }).success).toBe(false); + expect(EditImageInputSchema.safeParse({ artifactId: "a", rotate: 45 }).success).toBe(false); + expect(EditImageInputSchema.safeParse({ artifactId: "a", resize: {} }).success).toBe(false); + expect( + EditImageInputSchema.safeParse({ artifactId: "a", extend: { top: 4, color: "red" } }).success, + ).toBe(false); + expect( + EditImageInputSchema.parse({ artifactId: "a", resize: { aspect: "vertical" }, rotate: 90 }), + ).toEqual({ artifactId: "a", resize: { aspect: "vertical", fit: "crop" }, rotate: 90 }); + }); +}); diff --git a/packages/contracts/src/media-generation.ts b/packages/contracts/src/media-generation.ts new file mode 100644 index 0000000..2e61c42 --- /dev/null +++ b/packages/contracts/src/media-generation.ts @@ -0,0 +1,107 @@ +import * as z from "zod"; +import { Id } from "./ids.js"; +import { MediaFitSchema, MediaPresetIdSchema } from "./media.js"; + +export const MEDIA_IMAGE_MAX_COUNT = 4; +export const MEDIA_VIDEO_MAX_SECONDS = 10; + +const Prompt = z.string().trim().min(1).max(4000); +const NegativePrompt = z.string().trim().max(4000).optional().describe("What to avoid."); +const Seed = z + .number() + .int() + .min(0) + .max(Number.MAX_SAFE_INTEGER) + .optional() + .describe("Reuse a seed to reproduce a result."); +const Strength = z.number().min(0).max(1).optional(); +const ASPECTS = "vertical 9:16, landscape 16:9, square 1:1, portrait 4:5, classic 4:3, cinema 21:9"; +const aspect = (fallback: string) => MediaPresetIdSchema.describe(`${ASPECTS}. ${fallback}`); +const Side = z.number().int().min(64).max(2048); +const Pixels = z.number().int().min(0).max(16_384); + +export const GenerateImageInputSchema = z + .object({ + prompt: Prompt.describe("What the image shows."), + negativePrompt: NegativePrompt, + aspect: aspect("Default square, or the source image's aspect ratio.").optional(), + width: Side.optional().describe("Exact width in pixels, with height; overrides aspect."), + height: Side.optional(), + count: z.number().int().min(1).max(MEDIA_IMAGE_MAX_COUNT).default(1), + seed: Seed, + sourceArtifactId: Id.optional().describe("Image artifact to start from (image-to-image)."), + strength: Strength.describe("With a source: how far to move away from it (default 0.6)."), + }) + .refine((input) => (input.width === undefined) === (input.height === undefined), { + message: "Set width and height together", + path: ["width"], + }); +export type GenerateImageInput = z.infer; + +export const GenerateVideoInputSchema = z.object({ + prompt: Prompt.describe("What happens in the clip: subject, motion, camera."), + negativePrompt: NegativePrompt, + aspect: aspect("Default landscape, or the start frame's aspect ratio.").optional(), + durationSeconds: z.number().min(1).max(MEDIA_VIDEO_MAX_SECONDS).default(5), + seed: Seed, + sourceArtifactId: Id.optional().describe("Image artifact to use as the first frame."), +}); +export type GenerateVideoInput = z.infer; + +export const ImageFormatSchema = z.enum(["png", "jpeg", "webp"]); +export type ImageFormat = z.infer; + +/** Edits apply in this order: crop, rotate, flip, resize, extend, repaint, format. */ +export const EditImageInputSchema = z + .object({ + artifactId: Id.describe("Image artifact to edit."), + crop: z + .object({ left: Pixels, top: Pixels, width: Pixels.min(1), height: Pixels.min(1) }) + .optional() + .describe("Area to keep, in source pixels."), + rotate: z.literal([90, 180, 270]).optional().describe("Clockwise, in degrees."), + flip: z.enum(["horizontal", "vertical"]).optional(), + resize: z + .object({ + aspect: aspect("Resizes to the format's exact size.").optional(), + width: z.number().int().min(1).max(8192).optional(), + height: z.number().int().min(1).max(8192).optional(), + fit: MediaFitSchema.default("crop").describe( + "How a different aspect ratio fills the size: crop, pad, or blur (blurred background).", + ), + }) + .refine((resize) => resize.aspect || resize.width || resize.height, { + message: "Set aspect, width or height", + }) + .optional() + .describe("Width or height alone keeps the aspect ratio."), + extend: z + .object({ + top: Pixels.default(0), + right: Pixels.default(0), + bottom: Pixels.default(0), + left: Pixels.default(0), + color: z + .string() + .regex(/^#[0-9a-fA-F]{6}$/) + .optional() + .describe("#RRGGBB; transparent when omitted and the format allows it."), + }) + .optional() + .describe("Borders to add, in pixels."), + prompt: Prompt.optional().describe("Repaint the image with this prompt."), + strength: Strength.describe("How far the repaint moves away from the image (default 0.6)."), + format: ImageFormatSchema.optional().describe("Default keeps the source format."), + }) + .refine( + (input) => + input.crop || + input.rotate || + input.flip || + input.resize || + input.extend || + input.prompt || + input.format, + { message: "Nothing to edit" }, + ); +export type EditImageInput = z.infer; diff --git a/packages/contracts/src/media-jobs.ts b/packages/contracts/src/media-jobs.ts new file mode 100644 index 0000000..477a6fe --- /dev/null +++ b/packages/contracts/src/media-jobs.ts @@ -0,0 +1,36 @@ +import * as z from "zod"; +import { Id } from "./ids.js"; +import { MediaJobStatusSchema, MediaOperationSchema, MediaOutputSchema } from "./media.js"; + +/** Operations a user starts from a media card (export, transcribe, edit); generation goes through a bot. */ +export const MediaCardOperationSchema = z.enum(["edit_image", "export_video", "transcribe"]); +export type MediaCardOperation = z.infer; + +export const MediaJobStartInputSchema = z + .object({ + /** The conversation, like other thread procedures: a bot chat or a group. */ + botId: Id.optional(), + groupId: Id.optional(), + /** The message the action came from; its bot posts the new card. */ + messageId: Id, + operation: MediaCardOperationSchema, + /** The operation's input; validated by the operation. */ + input: z.record(z.string(), z.unknown()), + }) + .refine((input) => Boolean(input.botId) !== Boolean(input.groupId), { + message: "Pass botId or groupId", + }); +export type MediaJobStartInput = z.infer; + +export const MediaJobSnapshotSchema = z.object({ + id: Id, + operation: MediaOperationSchema, + status: MediaJobStatusSchema, + title: z.string(), + progress: z.number().min(0).max(1).nullable(), + error: z.string().nullable(), + outputs: z.array(MediaOutputSchema), +}); +export type MediaJobSnapshotDto = z.infer; + +export const MediaJobStartOutputSchema = z.object({ jobId: Id, messageId: Id }); diff --git a/packages/contracts/src/rpc.ts b/packages/contracts/src/rpc.ts index 6640d4c..30de55a 100644 --- a/packages/contracts/src/rpc.ts +++ b/packages/contracts/src/rpc.ts @@ -82,6 +82,12 @@ import { IntegrationProviderConfigSchema, IntegrationSetupStateSchema, } from "./integration-settings.js"; +import { + MediaJobSnapshotSchema, + MediaJobStartInputSchema, + MediaJobStartOutputSchema, +} from "./media-jobs.js"; +import { MediaSettingsUpdateSchema, MediaSettingsViewSchema } from "./media-settings.js"; import { ModelTierMappingSchema, SetModelTierInput } from "./model-tiers.js"; import { CreateProjectGoalInput, @@ -690,6 +696,18 @@ export const appContract = { get: oc.output(IntegrationSetupStateSchema), save: oc.input(IntegrationProviderConfigSchema).output(z.object({ ok: z.literal(true) })), }, + media: { + /** Deployment-wide media backends; deployment owner only. */ + settings: { + get: oc.output(MediaSettingsViewSchema), + update: oc.input(MediaSettingsUpdateSchema).output(MediaSettingsViewSchema), + }, + jobs: { + start: oc.input(MediaJobStartInputSchema).output(MediaJobStartOutputSchema), + get: oc.input(z.object({ jobId: Id })).output(MediaJobSnapshotSchema), + cancel: oc.input(z.object({ jobId: Id })).output(MediaJobSnapshotSchema), + }, + }, connections: { catalog: oc .input(z.object({ query: z.string().optional(), connectorId: z.string().optional() })) diff --git a/packages/core/package.json b/packages/core/package.json index bcac6e5..562dfd6 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -34,6 +34,11 @@ "development": "./src/plot/plot-spec.ts", "default": "./src/plot/plot-spec.ts" }, + "./media/workflows": { + "types": "./src/media/workflow-template.ts", + "development": "./src/media/workflow-template.ts", + "default": "./src/media/workflow-template.ts" + }, "./node/desktop-runtime": { "types": "./src/node/desktop-runtime.ts", "development": "./src/node/desktop-runtime.ts", diff --git a/packages/core/src/attachments.ts b/packages/core/src/attachments.ts index b4a6759..cace390 100644 --- a/packages/core/src/attachments.ts +++ b/packages/core/src/attachments.ts @@ -1,12 +1,10 @@ +import type { AttachmentKind, AttachmentMimeType, MessageBlock } from "@rakazo/contracts"; import { ATTACHMENT_MAX_BASE64_LENGTH, ATTACHMENT_MAX_BYTES, - type AttachmentKind, - type AttachmentMimeType, attachmentKindOf, isAllowedAttachmentMimeType, isAttachmentImageMimeType, - type MessageBlock, } from "@rakazo/contracts"; export class AttachmentValidationError extends Error { diff --git a/packages/core/src/media/workflow-template.test.ts b/packages/core/src/media/workflow-template.test.ts new file mode 100644 index 0000000..7865463 --- /dev/null +++ b/packages/core/src/media/workflow-template.test.ts @@ -0,0 +1,113 @@ +import { describe, expect, it } from "vitest"; +import { + DEFAULT_WORKFLOWS, + fillWorkflow, + parseWorkflowTemplate, + WorkflowTemplateError, +} from "./workflow-template.js"; + +const minimal = { + "1": { class_type: "CLIPTextEncode", inputs: { text: "{{prompt}}", clip: ["2", 1] } }, + "2": { class_type: "EmptyLatentImage", inputs: { width: "{{width}}", height: 512 } }, + "3": { class_type: "SaveImage", inputs: { filename_prefix: "run/{{seed}}-{{prompt}}" } }, +}; + +describe("parseWorkflowTemplate", () => { + it("accepts JSON text or parsed graphs and lists the placeholders", () => { + for (const source of [minimal, JSON.stringify(minimal)]) { + const template = parseWorkflowTemplate(source); + expect([...template.placeholders].sort()).toEqual(["prompt", "seed", "width"]); + } + }); + + it.each([ + ["not json", "Workflow is not valid JSON"], + ["[]", "Workflow must be a JSON object"], + ["{}", "Workflow must have 1 to 1000 nodes"], + [JSON.stringify({ nodes: [], links: [] }), "Workflow must be exported in ComfyUI's API format"], + [JSON.stringify({ "1": { inputs: {} } }), "Workflow node 1 needs a class_type and inputs"], + [ + JSON.stringify({ "1": { class_type: "SaveImage", inputs: { text: "plain" } } }), + "Workflow needs a {{prompt}} placeholder", + ], + [ + JSON.stringify({ + "1": { class_type: "X", inputs: { a: "{{prompt}}", b: "{{cfg}}" } }, + }), + "Unknown workflow placeholders: {{cfg}}", + ], + ])("rejects %s", (source, message) => { + expect(() => parseWorkflowTemplate(source)).toThrow(new WorkflowTemplateError(message)); + }); + + it("rejects placeholder names that only exist on Object.prototype", () => { + const source = { + "1": { + class_type: "X", + inputs: { a: "{{prompt}}", b: "{{constructor}}", c: "{{__proto__}}" }, + }, + }; + expect(() => parseWorkflowTemplate(source)).toThrow( + "Unknown workflow placeholders: {{constructor}}, {{__proto__}}", + ); + }); + + it("rejects oversized workflow text", () => { + expect(() => parseWorkflowTemplate(" ".repeat(1_000_001))).toThrow("Workflow is too large"); + }); +}); + +describe("fillWorkflow", () => { + it("replaces whole values with typed values and embedded ones with text, without mutating", () => { + const template = parseWorkflowTemplate(minimal); + const graph = fillWorkflow(template, { prompt: "a $& cat", width: 768, seed: 7 }); + expect(graph["1"]?.inputs).toEqual({ text: "a $& cat", clip: ["2", 1] }); + expect(graph["2"]?.inputs).toEqual({ width: 768, height: 512 }); + expect(graph["3"]?.inputs.filename_prefix).toBe("run/7-a $& cat"); + expect(minimal["2"].inputs.width).toBe("{{width}}"); + }); + + it("refuses missing or mistyped values", () => { + const template = parseWorkflowTemplate(minimal); + expect(() => fillWorkflow(template, { prompt: "cat", seed: 1 })).toThrow( + "Workflow value for {{width}} is missing", + ); + expect(() => fillWorkflow(template, { prompt: "cat", width: 1.5, seed: 1 })).toThrow( + "{{width}}", + ); + }); +}); + +describe("DEFAULT_WORKFLOWS", () => { + it("only image-to-* defaults need an input image", () => { + expect(DEFAULT_WORKFLOWS.textToImage.placeholders.has("image")).toBe(false); + expect(DEFAULT_WORKFLOWS.textToVideo.placeholders.has("image")).toBe(false); + expect(DEFAULT_WORKFLOWS.imageToImage.placeholders.has("image")).toBe(true); + expect(DEFAULT_WORKFLOWS.imageToVideo.placeholders.has("image")).toBe(true); + }); + + it("fill into graphs whose links point at existing nodes", () => { + const values = { + prompt: "a lighthouse", + negative_prompt: "", + model: "model.safetensors", + image: "input.png", + width: 512, + height: 512, + seed: 1, + steps: 20, + strength: 0.6, + frames: 25, + fps: 24, + }; + for (const template of Object.values(DEFAULT_WORKFLOWS)) { + const graph = fillWorkflow(template, values); + expect(JSON.stringify(graph)).not.toContain("{{"); + for (const node of Object.values(graph)) { + for (const input of Object.values(node.inputs)) { + if (Array.isArray(input)) expect(graph[String(input[0])]).toBeDefined(); + } + } + } + }); +}); diff --git a/packages/core/src/media/workflow-template.ts b/packages/core/src/media/workflow-template.ts new file mode 100644 index 0000000..a7c74e6 --- /dev/null +++ b/packages/core/src/media/workflow-template.ts @@ -0,0 +1,159 @@ +import imageToImage from "./workflows/image-to-image.json" with { type: "json" }; +import imageToVideo from "./workflows/image-to-video.json" with { type: "json" }; +import textToImage from "./workflows/text-to-image.json" with { type: "json" }; +import textToVideo from "./workflows/text-to-video.json" with { type: "json" }; + +/** + * Values a ComfyUI workflow template can ask for. A string value that is exactly `{{name}}` + * becomes the typed value (so `"{{width}}"` turns into the number 1024); `{{name}}` inside a + * longer string is replaced by its text. Templates are data: nothing is ever evaluated. + */ +export const WORKFLOW_PLACEHOLDERS = { + prompt: "string", + negative_prompt: "string", + /** Checkpoint file name: the configured image model or the first installed checkpoint. */ + model: "string", + /** Uploaded input image (image-to-image, image-to-video). */ + image: "string", + width: "integer", + height: "integer", + seed: "integer", + steps: "integer", + /** 0–1 denoise for image-to-image. */ + strength: "number", + frames: "integer", + fps: "number", +} as const satisfies Record; + +export type WorkflowPlaceholder = keyof typeof WORKFLOW_PLACEHOLDERS; +export type WorkflowValues = { + [Name in WorkflowPlaceholder]?: (typeof WORKFLOW_PLACEHOLDERS)[Name] extends "string" + ? string + : number; +}; + +/** One node of a workflow in ComfyUI's API format ("Export (API)" in ComfyUI). */ +export interface ComfyWorkflowNode { + class_type: string; + inputs: Record; +} +export type ComfyWorkflow = Record; + +export interface WorkflowTemplate { + readonly graph: ComfyWorkflow; + /** The placeholders the graph uses; `image` marks a template that needs an input image. */ + readonly placeholders: ReadonlySet; +} + +export class WorkflowTemplateError extends Error { + constructor(message: string) { + super(message); + this.name = "WorkflowTemplateError"; + } +} + +const PLACEHOLDER_PATTERN = /\{\{([a-z_]+)\}\}/g; +const WHOLE_PLACEHOLDER = /^\{\{([a-z_]+)\}\}$/; +const MAX_WORKFLOW_CHARS = 1_000_000; +const MAX_WORKFLOW_NODES = 1_000; + +function isRecord(value: unknown): value is Record { + return Boolean(value) && typeof value === "object" && !Array.isArray(value); +} + +function isPlaceholder(name: string): name is WorkflowPlaceholder { + return Object.hasOwn(WORKFLOW_PLACEHOLDERS, name); +} + +function visitStrings(value: unknown, visit: (text: string) => void): void { + if (typeof value === "string") visit(value); + else if (Array.isArray(value)) for (const entry of value) visitStrings(entry, visit); + else if (isRecord(value)) for (const entry of Object.values(value)) visitStrings(entry, visit); +} + +/** Validates an API-format workflow (JSON text or parsed) and lists its placeholders. */ +export function parseWorkflowTemplate(source: unknown): WorkflowTemplate { + let value = source; + if (typeof source === "string") { + if (source.length > MAX_WORKFLOW_CHARS) { + throw new WorkflowTemplateError("Workflow is too large"); + } + try { + value = JSON.parse(source); + } catch { + throw new WorkflowTemplateError("Workflow is not valid JSON"); + } + } + if (!isRecord(value)) throw new WorkflowTemplateError("Workflow must be a JSON object"); + if (Array.isArray(value.nodes)) { + throw new WorkflowTemplateError("Workflow must be exported in ComfyUI's API format"); + } + const nodes = Object.entries(value); + if (nodes.length === 0 || nodes.length > MAX_WORKFLOW_NODES) { + throw new WorkflowTemplateError(`Workflow must have 1 to ${MAX_WORKFLOW_NODES} nodes`); + } + for (const [id, node] of nodes) { + if (!isRecord(node) || typeof node.class_type !== "string" || !isRecord(node.inputs)) { + throw new WorkflowTemplateError(`Workflow node ${id} needs a class_type and inputs`); + } + } + const placeholders = new Set(); + const unknown = new Set(); + visitStrings(value, (text) => { + for (const [, name = ""] of text.matchAll(PLACEHOLDER_PATTERN)) { + if (isPlaceholder(name)) placeholders.add(name); + else unknown.add(`{{${name}}}`); + } + }); + if (unknown.size > 0) { + throw new WorkflowTemplateError(`Unknown workflow placeholders: ${[...unknown].join(", ")}`); + } + if (!placeholders.has("prompt")) { + throw new WorkflowTemplateError("Workflow needs a {{prompt}} placeholder"); + } + return { graph: value as unknown as ComfyWorkflow, placeholders }; +} + +function replacePlaceholders(value: unknown, values: WorkflowValues): unknown { + if (typeof value === "string") { + const whole = WHOLE_PLACEHOLDER.exec(value)?.[1]; + if (whole) return values[whole as WorkflowPlaceholder]; + return value.replace(PLACEHOLDER_PATTERN, (_, name: string) => + String(values[name as WorkflowPlaceholder]), + ); + } + if (Array.isArray(value)) return value.map((entry) => replacePlaceholders(entry, values)); + if (isRecord(value)) { + return Object.fromEntries( + Object.entries(value).map(([key, entry]) => [key, replacePlaceholders(entry, values)]), + ); + } + return value; +} + +/** Returns a copy of the template's graph with every placeholder replaced by its typed value. */ +export function fillWorkflow(template: WorkflowTemplate, values: WorkflowValues): ComfyWorkflow { + for (const name of template.placeholders) { + const value = values[name]; + const type = WORKFLOW_PLACEHOLDERS[name]; + const valid = + type === "string" + ? typeof value === "string" + : typeof value === "number" && + Number.isFinite(value) && + (type === "number" || Number.isInteger(value)); + if (!valid) throw new WorkflowTemplateError(`Workflow value for {{${name}}} is missing`); + } + return replacePlaceholders(template.graph, values) as ComfyWorkflow; +} + +/** + * Built-in workflows. Images use core nodes only and run on any SD-family checkpoint; video + * uses ComfyUI's core LTX-Video nodes (model files in workflows/README.md). + */ +export const DEFAULT_WORKFLOWS = { + textToImage: parseWorkflowTemplate(textToImage), + imageToImage: parseWorkflowTemplate(imageToImage), + textToVideo: parseWorkflowTemplate(textToVideo), + imageToVideo: parseWorkflowTemplate(imageToVideo), +} as const; diff --git a/packages/core/src/media/workflows/README.md b/packages/core/src/media/workflows/README.md new file mode 100644 index 0000000..af97d1b --- /dev/null +++ b/packages/core/src/media/workflows/README.md @@ -0,0 +1,35 @@ +# ComfyUI workflows + +Default workflows for `generate_image`, `generate_video` and repainting in `edit_image`, in +ComfyUI's API format. A custom workflow in the Media settings replaces the matching default: +one that uses `{{image}}` replaces the image-to-image (or image-to-video) default, one without +it replaces the text-to-image (or text-to-video) default. + +## Placeholders + +A string value that is exactly `{{name}}` becomes the typed value; `{{name}}` inside a longer +string is replaced by its text. Every workflow needs `{{prompt}}`; other names are optional. + +| Placeholder | Value | +|---|---| +| `{{prompt}}`, `{{negative_prompt}}` | text | +| `{{model}}` | the configured image model, or the first installed checkpoint | +| `{{image}}` | the uploaded input image (for `LoadImage`) | +| `{{width}}`, `{{height}}` | output size in pixels | +| `{{seed}}`, `{{steps}}` | integers | +| `{{strength}}` | 0–1, image-to-image denoise | +| `{{frames}}`, `{{fps}}` | video length in frames (8n+1) and frame rate | + +To write your own, build it in ComfyUI, replace the values above with placeholders and export it +with **Export (API)**. + +## Model files + +- Images: any SD 1.5 or SDXL checkpoint in `models/checkpoints/`. +- Video (LTX-Video, core nodes, about 12 GB of VRAM): + - `models/checkpoints/ltx-video-2b-v0.9.5.safetensors` from + [Lightricks/LTX-Video](https://huggingface.co/Lightricks/LTX-Video) + - `models/text_encoders/t5xxl_fp16.safetensors` from + [comfyanonymous/flux_text_encoders](https://huggingface.co/comfyanonymous/flux_text_encoders) + +The video workflows need a ComfyUI release with the `CreateVideo` and `SaveVideo` nodes. diff --git a/packages/core/src/media/workflows/image-to-image.json b/packages/core/src/media/workflows/image-to-image.json new file mode 100644 index 0000000..e762d24 --- /dev/null +++ b/packages/core/src/media/workflows/image-to-image.json @@ -0,0 +1,45 @@ +{ + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": { "ckpt_name": "{{model}}" } + }, + "2": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{prompt}}", "clip": ["1", 1] } + }, + "3": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{negative_prompt}}", "clip": ["1", 1] } + }, + "4": { + "class_type": "LoadImage", + "inputs": { "image": "{{image}}" } + }, + "5": { + "class_type": "VAEEncode", + "inputs": { "pixels": ["4", 0], "vae": ["1", 2] } + }, + "6": { + "class_type": "KSampler", + "inputs": { + "model": ["1", 0], + "positive": ["2", 0], + "negative": ["3", 0], + "latent_image": ["5", 0], + "seed": "{{seed}}", + "steps": "{{steps}}", + "cfg": 7, + "sampler_name": "dpmpp_2m", + "scheduler": "karras", + "denoise": "{{strength}}" + } + }, + "7": { + "class_type": "VAEDecode", + "inputs": { "samples": ["6", 0], "vae": ["1", 2] } + }, + "8": { + "class_type": "SaveImage", + "inputs": { "images": ["7", 0], "filename_prefix": "evoo/image" } + } +} diff --git a/packages/core/src/media/workflows/image-to-video.json b/packages/core/src/media/workflows/image-to-video.json new file mode 100644 index 0000000..6bf53a7 --- /dev/null +++ b/packages/core/src/media/workflows/image-to-video.json @@ -0,0 +1,86 @@ +{ + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": { "ckpt_name": "ltx-video-2b-v0.9.5.safetensors" } + }, + "2": { + "class_type": "CLIPLoader", + "inputs": { "clip_name": "t5xxl_fp16.safetensors", "type": "ltxv" } + }, + "3": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{prompt}}", "clip": ["2", 0] } + }, + "4": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{negative_prompt}}", "clip": ["2", 0] } + }, + "5": { + "class_type": "LoadImage", + "inputs": { "image": "{{image}}" } + }, + "6": { + "class_type": "LTXVImgToVideo", + "inputs": { + "positive": ["3", 0], + "negative": ["4", 0], + "vae": ["1", 2], + "image": ["5", 0], + "width": "{{width}}", + "height": "{{height}}", + "length": "{{frames}}", + "batch_size": 1, + "strength": 1 + } + }, + "7": { + "class_type": "LTXVConditioning", + "inputs": { "positive": ["6", 0], "negative": ["6", 1], "frame_rate": "{{fps}}" } + }, + "8": { + "class_type": "LTXVScheduler", + "inputs": { + "steps": "{{steps}}", + "max_shift": 2.05, + "base_shift": 0.95, + "stretch": true, + "terminal": 0.1, + "latent": ["6", 2] + } + }, + "9": { + "class_type": "KSamplerSelect", + "inputs": { "sampler_name": "euler" } + }, + "10": { + "class_type": "SamplerCustom", + "inputs": { + "model": ["1", 0], + "add_noise": true, + "noise_seed": "{{seed}}", + "cfg": 3, + "positive": ["7", 0], + "negative": ["7", 1], + "sampler": ["9", 0], + "sigmas": ["8", 0], + "latent_image": ["6", 2] + } + }, + "11": { + "class_type": "VAEDecode", + "inputs": { "samples": ["10", 0], "vae": ["1", 2] } + }, + "12": { + "class_type": "CreateVideo", + "inputs": { "images": ["11", 0], "fps": "{{fps}}" } + }, + "13": { + "class_type": "SaveVideo", + "inputs": { + "video": ["12", 0], + "filename_prefix": "evoo/video", + "format": "mp4", + "codec": "h264" + } + } +} diff --git a/packages/core/src/media/workflows/text-to-image.json b/packages/core/src/media/workflows/text-to-image.json new file mode 100644 index 0000000..633cdc5 --- /dev/null +++ b/packages/core/src/media/workflows/text-to-image.json @@ -0,0 +1,41 @@ +{ + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": { "ckpt_name": "{{model}}" } + }, + "2": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{prompt}}", "clip": ["1", 1] } + }, + "3": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{negative_prompt}}", "clip": ["1", 1] } + }, + "4": { + "class_type": "EmptyLatentImage", + "inputs": { "width": "{{width}}", "height": "{{height}}", "batch_size": 1 } + }, + "5": { + "class_type": "KSampler", + "inputs": { + "model": ["1", 0], + "positive": ["2", 0], + "negative": ["3", 0], + "latent_image": ["4", 0], + "seed": "{{seed}}", + "steps": "{{steps}}", + "cfg": 7, + "sampler_name": "dpmpp_2m", + "scheduler": "karras", + "denoise": 1 + } + }, + "6": { + "class_type": "VAEDecode", + "inputs": { "samples": ["5", 0], "vae": ["1", 2] } + }, + "7": { + "class_type": "SaveImage", + "inputs": { "images": ["6", 0], "filename_prefix": "evoo/image" } + } +} diff --git a/packages/core/src/media/workflows/text-to-video.json b/packages/core/src/media/workflows/text-to-video.json new file mode 100644 index 0000000..55a3b3c --- /dev/null +++ b/packages/core/src/media/workflows/text-to-video.json @@ -0,0 +1,77 @@ +{ + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": { "ckpt_name": "ltx-video-2b-v0.9.5.safetensors" } + }, + "2": { + "class_type": "CLIPLoader", + "inputs": { "clip_name": "t5xxl_fp16.safetensors", "type": "ltxv" } + }, + "3": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{prompt}}", "clip": ["2", 0] } + }, + "4": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{negative_prompt}}", "clip": ["2", 0] } + }, + "5": { + "class_type": "LTXVConditioning", + "inputs": { "positive": ["3", 0], "negative": ["4", 0], "frame_rate": "{{fps}}" } + }, + "6": { + "class_type": "EmptyLTXVLatentVideo", + "inputs": { + "width": "{{width}}", + "height": "{{height}}", + "length": "{{frames}}", + "batch_size": 1 + } + }, + "7": { + "class_type": "LTXVScheduler", + "inputs": { + "steps": "{{steps}}", + "max_shift": 2.05, + "base_shift": 0.95, + "stretch": true, + "terminal": 0.1, + "latent": ["6", 0] + } + }, + "8": { + "class_type": "KSamplerSelect", + "inputs": { "sampler_name": "euler" } + }, + "9": { + "class_type": "SamplerCustom", + "inputs": { + "model": ["1", 0], + "add_noise": true, + "noise_seed": "{{seed}}", + "cfg": 3, + "positive": ["5", 0], + "negative": ["5", 1], + "sampler": ["8", 0], + "sigmas": ["7", 0], + "latent_image": ["6", 0] + } + }, + "10": { + "class_type": "VAEDecode", + "inputs": { "samples": ["9", 0], "vae": ["1", 2] } + }, + "11": { + "class_type": "CreateVideo", + "inputs": { "images": ["10", 0], "fps": "{{fps}}" } + }, + "12": { + "class_type": "SaveVideo", + "inputs": { + "video": ["11", 0], + "filename_prefix": "evoo/video", + "format": "mp4", + "codec": "h264" + } + } +} diff --git a/packages/testkit/src/cli/harness.ts b/packages/testkit/src/cli/harness.ts index 6e26f8b..184a79a 100644 --- a/packages/testkit/src/cli/harness.ts +++ b/packages/testkit/src/cli/harness.ts @@ -106,6 +106,7 @@ async function main() { "packages/testkit/src/connections.test.ts", "packages/testkit/src/bot-secrets.test.ts", "packages/testkit/src/delegation.postgres.test.ts", + "packages/testkit/src/media.postgres.test.ts", "packages/testkit/src/user-profile.postgres.test.ts", "packages/db/src/space-membership.postgres.test.ts", "packages/db/src/main-agent.postgres.test.ts", @@ -117,7 +118,9 @@ async function main() { "packages/adapters/src/realtime.postgres.test.ts", "packages/adapters/src/job-reconciler.postgres.test.ts", "packages/adapters/src/cloud-agent.postgres.test.ts", + "packages/adapters/src/media/image-operations.postgres.test.ts", "packages/adapters/src/media/transcription.postgres.test.ts", + "packages/adapters/src/media-runtime.postgres.test.ts", "packages/adapters/src/builtin-memory-provider.postgres.test.ts", "packages/adapters/src/profile-consolidation.postgres.test.ts", ]; diff --git a/packages/testkit/src/media.postgres.test.ts b/packages/testkit/src/media.postgres.test.ts new file mode 100644 index 0000000..f5de0c6 --- /dev/null +++ b/packages/testkit/src/media.postgres.test.ts @@ -0,0 +1,378 @@ +import { mkdtempSync, rmSync } from "node:fs"; +import { writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import path from "node:path"; +import { ComposioEmulator, type MediaOperationRegistration } from "@rakazo/adapters"; +import { + type MediaJobBlock, + MediaOutputSchema, + type MediaSettingsView, + type MessageBlock, +} from "@rakazo/contracts"; +import { createThreadMessage } from "@rakazo/db"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import type { createApp } from "../../../apps/api/src/app.ts"; +import { sessionCookieHeader } from "./index.js"; + +type App = { request: (input: string, init?: RequestInit) => Response | Promise }; +type AppHandles = Awaited>; + +process.env.WAKEUP_DRIVER = "memory"; +process.env.SANDBOX_PROVIDER = "fake"; +process.env.AGENT_RUNTIME = "scripted"; + +const hasDb = process.env.VERIFY_DATABASE === "1" && Boolean(process.env.DATABASE_URL); +const describeWithDatabase = hasDb ? describe.sequential : describe.skip; + +const PNG_BASE64 = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 1, 2, 3]).toString( + "base64", +); + +// testkit has no zod of its own; the contracts' schemas carry it. +const SourceInputSchema = MediaOutputSchema.pick({ artifactId: true }); + +/** Fake card operations: an export that copies its source, and a transcription that waits. */ +function fakeOperations(): MediaOperationRegistration[] { + const exportVideo: MediaOperationRegistration<{ artifactId: string }> = { + operation: "export_video", + inputSchema: SourceInputSchema, + title: () => "Export", + available: async () => true, + artifactIds: (input) => [input.artifactId], + handler: { + operation: "export_video", + run: async (context) => { + const { artifactId } = context.input as { artifactId: string }; + const source = await context.fetchArtifact(artifactId); + const out = path.join(context.tempDir, "export.png"); + await writeFile(out, Buffer.from(PNG_BASE64, "base64")); + await context.addOutput({ + path: out, + name: "export.png", + mimeType: "image/png", + parentArtifactId: source.artifactId, + preset: "square", + }); + return { summary: "Exported." }; + }, + }, + }; + const transcribe: MediaOperationRegistration<{ artifactId: string }> = { + operation: "transcribe", + inputSchema: SourceInputSchema, + title: () => "Transcript", + available: async () => true, + artifactIds: (input) => [input.artifactId], + handler: { + operation: "transcribe", + run: async (context) => { + await new Promise((resolve) => { + context.signal.addEventListener("abort", () => resolve(), { once: true }); + }); + context.signal.throwIfAborted(); + return { summary: "" }; + }, + }, + }; + return [exportVideo, transcribe] as MediaOperationRegistration[]; +} + +describeWithDatabase("media RPCs (PostgreSQL)", () => { + let handles: AppHandles; + let app: App; + const stamp = `${Date.now()}-${Math.random().toString(36).slice(2)}`; + const dataDir = mkdtempSync(path.join(tmpdir(), "rakazo-media-")); + + beforeAll(async () => { + const { createApp } = await import("../../../apps/api/src/app.ts"); + handles = await createApp({ + databaseUrl: process.env.DATABASE_URL!, + dataDir, + sandboxProvider: "fake", + agentRuntime: "scripted", + wakeupDriver: "memory", + signupsEnabled: "true", + composio: new ComposioEmulator(), + mediaOperations: fakeOperations(), + }); + app = handles.app; + }); + + afterAll(async () => { + await handles?.prisma.integrationProviderConfig.deleteMany({ where: { id: "media" } }); + await handles?.stop(); + rmSync(dataDir, { recursive: true, force: true }); + }); + + async function user(label: string) { + const cookie = await signup(app, `${label}-${stamp}@rakazo.test`, label); + const me = await rpc<{ userId: string; spaceId: string }>(app, cookie, "me"); + const bot = await rpc<{ id: string }>(app, cookie, "bots/create", { + name: `${label} bot`, + title: "", + description: "", + instructions: "", + notifyOnFinish: false, + }); + const thread = await handles.prisma.thread.findUniqueOrThrow({ where: { botId: bot.id } }); + const card = await createThreadMessage(handles.prisma, { + threadId: thread.id, + role: "bot", + botId: bot.id, + blocks: [{ kind: "text", text: "Here is your clip." }], + }); + const artifact = await rpc<{ id: string }>(app, cookie, "artifacts/create", { + botId: bot.id, + name: "clip.png", + mimeType: "image/png", + contentBase64: PNG_BASE64, + }); + return { cookie, me, bot, thread, card, artifact }; + } + + async function waitForStatus(jobId: string, status: string) { + const start = Date.now(); + while (Date.now() - start < 20_000) { + const job = await handles.prisma.mediaJob.findUnique({ where: { id: jobId } }); + if (job?.status === status) return job; + await new Promise((resolve) => setTimeout(resolve, 50)); + } + throw new Error(`media job ${jobId} never reached ${status}`); + } + + it("lets only the deployment owner read and change media settings, and never returns keys", async () => { + const owner = await user("media-owner"); + const other = await user("media-other"); + await handles.prisma.deploymentSettings.upsert({ + where: { id: "default" }, + create: { id: "default", ownerUserId: owner.me.userId }, + update: { ownerUserId: owner.me.userId }, + }); + + const update = { + comfyui: { + baseUrl: "http://127.0.0.1:8188/", + apiKey: "comfy-test-key", + imageModel: "sd_xl_base_1.0.safetensors", + imageWorkflow: '{"3":{"class_type":"KSampler","inputs":{}}}', + videoWorkflow: "", + }, + transcription: { baseUrl: "http://127.0.0.1:8080", model: "whisper-1", apiKey: "" }, + }; + expect((await raw(app, other.cookie, "media/settings/get", {})).status).toBe(403); + expect((await raw(app, other.cookie, "media/settings/update", update)).status).toBe(403); + expect(await handles.prisma.integrationProviderConfig.count({ where: { id: "media" } })).toBe( + 0, + ); + + const saved = await rpc(app, owner.cookie, "media/settings/update", update); + expect(saved).toEqual({ + comfyui: { + baseUrl: "http://127.0.0.1:8188", + hasApiKey: true, + imageModel: "sd_xl_base_1.0.safetensors", + imageWorkflow: update.comfyui.imageWorkflow, + }, + transcription: { baseUrl: "http://127.0.0.1:8080/v1", model: "whisper-1", hasApiKey: false }, + }); + expect(JSON.stringify(saved)).not.toContain("comfy-test-key"); + expect(await rpc(app, owner.cookie, "media/settings/get")).toEqual(saved); + const row = await handles.prisma.integrationProviderConfig.findUniqueOrThrow({ + where: { id: "media" }, + }); + expect(row.ciphertext).not.toContain("comfy-test-key"); + expect(row.ciphertext).not.toContain("127.0.0.1"); + + // An omitted key stays with the same server, never follows the settings elsewhere. + const { apiKey: _key, ...withoutKey } = update.comfyui; + const kept = await rpc(app, owner.cookie, "media/settings/update", { + ...update, + comfyui: withoutKey, + }); + expect(kept.comfyui?.hasApiKey).toBe(true); + const moved = await rpc(app, owner.cookie, "media/settings/update", { + ...update, + comfyui: { ...withoutKey, baseUrl: "http://127.0.0.1:8189" }, + }); + expect(moved.comfyui?.hasApiKey).toBe(false); + + for (const baseUrl of ["http://169.254.169.254", "http://user:pass@127.0.0.1:8188"]) { + const refused = await raw(app, owner.cookie, "media/settings/update", { + ...update, + comfyui: { ...update.comfyui, baseUrl }, + }); + expect(refused.status, baseUrl).toBe(400); + } + const badWorkflow = await raw(app, owner.cookie, "media/settings/update", { + ...update, + comfyui: { ...update.comfyui, imageWorkflow: "[1, 2" }, + }); + expect(badWorkflow.status).toBe(400); + + const previous = process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC; + try { + delete process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC; + expect( + ( + await raw(app, owner.cookie, "media/settings/update", { + comfyui: null, + transcription: { baseUrl: "http://transcribe.example.com" }, + }) + ).status, + ).toBe(400); + process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC = "1"; + // A key is never sent in cleartext to a public host. + expect( + ( + await raw(app, owner.cookie, "media/settings/update", { + comfyui: null, + transcription: { baseUrl: "http://transcribe.example.com", apiKey: "secret" }, + }) + ).status, + ).toBe(400); + const https = await rpc(app, owner.cookie, "media/settings/update", { + comfyui: null, + transcription: { baseUrl: "https://transcribe.example.com", apiKey: "secret" }, + }); + expect(https.transcription).toMatchObject({ + baseUrl: "https://transcribe.example.com/v1", + hasApiKey: true, + }); + } finally { + if (previous === undefined) delete process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC; + else process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC = previous; + } + + expect( + await rpc(app, owner.cookie, "media/settings/update", { comfyui: null, transcription: null }), + ).toEqual({ comfyui: null, transcription: null }); + expect(await handles.prisma.integrationProviderConfig.count({ where: { id: "media" } })).toBe( + 0, + ); + }); + + it("starts, reads and cancels card jobs only for the actor's own threads and files", async () => { + const alice = await user("media-alice"); + const mallory = await user("media-mallory"); + + const started = await rpc<{ jobId: string; messageId: string }>( + app, + alice.cookie, + "media/jobs/start", + { + botId: alice.bot.id, + messageId: alice.card.id, + operation: "export_video", + input: { artifactId: alice.artifact.id }, + }, + ); + const job = await waitForStatus(started.jobId, "done"); + expect(job).toMatchObject({ + spaceId: alice.me.spaceId, + userId: alice.me.userId, + botId: alice.bot.id, + threadId: alice.thread.id, + runId: null, + continueAfter: false, + }); + const snapshot = await rpc<{ status: string; outputs: MediaJobBlock["outputs"] }>( + app, + alice.cookie, + "media/jobs/get", + { jobId: started.jobId }, + ); + expect(snapshot.status).toBe("done"); + expect(snapshot.outputs).toEqual([ + expect.objectContaining({ + name: "export.png", + parentArtifactId: alice.artifact.id, + preset: "square", + }), + ]); + const card = await handles.prisma.message.findUniqueOrThrow({ + where: { id: started.messageId }, + }); + expect(card).toMatchObject({ threadId: alice.thread.id, botId: alice.bot.id, role: "bot" }); + expect((card.blocks as MessageBlock[])[0]).toMatchObject({ kind: "media", status: "done" }); + + // Another user can neither read nor cancel the job, nor start one in Alice's conversation + // or on her files. + expect( + (await raw(app, mallory.cookie, "media/jobs/get", { jobId: started.jobId })).status, + ).toBe(404); + expect( + (await raw(app, mallory.cookie, "media/jobs/cancel", { jobId: started.jobId })).status, + ).toBe(404); + for (const body of [ + { + botId: alice.bot.id, + messageId: alice.card.id, + operation: "export_video", + input: { artifactId: mallory.artifact.id }, + }, + { + botId: mallory.bot.id, + messageId: alice.card.id, + operation: "export_video", + input: { artifactId: mallory.artifact.id }, + }, + { + botId: mallory.bot.id, + messageId: mallory.card.id, + operation: "export_video", + input: { artifactId: alice.artifact.id }, + }, + { + botId: mallory.bot.id, + messageId: mallory.card.id, + operation: "generate_image", + input: { prompt: "anything" }, + }, + ]) { + const response = await raw(app, mallory.cookie, "media/jobs/start", body); + expect(response.status, JSON.stringify(body)).toBeGreaterThanOrEqual(400); + } + expect(await handles.prisma.mediaJob.count({ where: { userId: mallory.me.userId } })).toBe(0); + + const slow = await rpc<{ jobId: string }>(app, alice.cookie, "media/jobs/start", { + botId: alice.bot.id, + messageId: alice.card.id, + operation: "transcribe", + input: { artifactId: alice.artifact.id }, + }); + await waitForStatus(slow.jobId, "running"); + const cancelled = await rpc<{ status: string }>(app, alice.cookie, "media/jobs/cancel", { + jobId: slow.jobId, + }); + expect(["running", "cancelled"]).toContain(cancelled.status); + await waitForStatus(slow.jobId, "cancelled"); + }); +}); + +async function signup(app: App, email: string, name: string) { + const res = await app.request("/api/auth/sign-up/email", { + method: "POST", + headers: { "content-type": "application/json", origin: "http://127.0.0.1:5173" }, + body: JSON.stringify({ email, password: "password12", name }), + }); + if (res.status >= 400) throw new Error(`signup failed ${res.status}: ${await res.text()}`); + return sessionCookieHeader(res); +} + +function raw(app: App, cookie: string, proc: string, body: unknown = {}) { + return app.request(`/rpc/${proc}`, { + method: "POST", + headers: { "content-type": "application/json", cookie, origin: "http://127.0.0.1:5173" }, + body: JSON.stringify({ json: body }), + }); +} + +async function rpc(app: App, cookie: string, proc: string, body: unknown = {}): Promise { + const res = await raw(app, cookie, proc, body); + const text = await res.text(); + const parsed = JSON.parse(text) as { json?: T; error?: { message?: string } }; + if (res.status >= 400 || parsed.error) { + throw new Error(`${proc} ${res.status}: ${parsed.error?.message ?? text}`); + } + return parsed.json as T; +}