diff --git a/packages/adapters/src/builtin-tools.ts b/packages/adapters/src/builtin-tools.ts index 0ff52fa..704e9b6 100644 --- a/packages/adapters/src/builtin-tools.ts +++ b/packages/adapters/src/builtin-tools.ts @@ -2,12 +2,25 @@ import type { ConnectorTool } from "@rakazo/adapter-kit"; import { BotSecretDestination, BotSecretName, + EditImageInputSchema, + GenerateImageInputSchema, + GenerateVideoInputSchema, MODEL_TIERS, SecretAskPurpose, SecretHttpRequest, USER_FACT_CATEGORIES, } from "@rakazo/contracts"; import { z } from "zod"; +import { mediaToolInputSchema } from "./media-tools.js"; + +/** 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."; @@ -326,6 +339,24 @@ export const builtinAgentTools: ConnectorTool[] = [ }, }, }, + { + name: "generate_image", + description: + "Generate up to 4 images from a prompt on the user's 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 user's 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. Edits apply in that order.", + inputSchema: mediaParameters(EditImageInputSchema), + }, { name: "add_mcp_server", description: diff --git a/packages/adapters/src/media-operations.ts b/packages/adapters/src/media-operations.ts index 98dbf7e..8a6d556 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 type { MediaOperationRegistration } from "./media-runtime.js"; export interface MediaOperationDeps { @@ -10,6 +11,6 @@ export interface MediaOperationDeps { * Every media operation this deployment can run. The API and the worker build their runtime * from this one list, so both processes run the same operations. */ -export function createMediaOperations(_deps: MediaOperationDeps): MediaOperationRegistration[] { - return []; +export function createMediaOperations(deps: MediaOperationDeps): MediaOperationRegistration[] { + return [...createImageOperations(deps)]; } diff --git a/packages/adapters/src/media-settings.ts b/packages/adapters/src/media-settings.ts index d237a13..b82134d 100644 --- a/packages/adapters/src/media-settings.ts +++ b/packages/adapters/src/media-settings.ts @@ -5,6 +5,7 @@ import { MediaSettingsUpdateSchema, type MediaSettingsView, } from "@rakazo/contracts"; +import { parseWorkflowTemplate } from "@rakazo/core/media/workflows"; import type { PrismaClient } from "@rakazo/db"; import { assertAllowedOpenAiCompatibleRequestUrl, @@ -82,6 +83,15 @@ function keptKey( return new URL(previous.baseUrl).origin === new URL(baseUrl).origin ? previous.apiKey : undefined; } +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)}`); + } +} + /** * 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. @@ -91,6 +101,8 @@ export async function saveMediaSettings( update: MediaSettingsUpdate, ): Promise { const parsed = MediaSettingsUpdateSchema.parse(update); + assertValidWorkflow("Image workflow", parsed.comfyui?.imageWorkflow); + assertValidWorkflow("Video workflow", parsed.comfyui?.videoWorkflow); const previous = await loadMediaSettings(deps); let comfyui: MediaSettings["comfyui"] = null; if (parsed.comfyui) { 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/contracts/src/media-generation.test.ts b/packages/contracts/src/media-generation.test.ts index 4ab1a74..32f9b78 100644 --- a/packages/contracts/src/media-generation.test.ts +++ b/packages/contracts/src/media-generation.test.ts @@ -13,7 +13,6 @@ describe("media generation inputs", () => { }); expect(GenerateVideoInputSchema.parse({ prompt: "waves" })).toEqual({ prompt: "waves", - aspect: "landscape", durationSeconds: 5, }); }); diff --git a/packages/contracts/src/media-generation.ts b/packages/contracts/src/media-generation.ts index 6d33216..2e61c42 100644 --- a/packages/contracts/src/media-generation.ts +++ b/packages/contracts/src/media-generation.ts @@ -15,9 +15,8 @@ const Seed = z .optional() .describe("Reuse a seed to reproduce a result."); const Strength = z.number().min(0).max(1).optional(); -const Aspect = MediaPresetIdSchema.describe( - "vertical 9:16, landscape 16:9, square 1:1, portrait 4:5, classic 4:3, cinema 21:9.", -); +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); @@ -25,7 +24,7 @@ export const GenerateImageInputSchema = z .object({ prompt: Prompt.describe("What the image shows."), negativePrompt: NegativePrompt, - aspect: Aspect.optional(), + 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), @@ -42,7 +41,7 @@ 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"), + 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."), @@ -64,7 +63,7 @@ export const EditImageInputSchema = z flip: z.enum(["horizontal", "vertical"]).optional(), resize: z .object({ - aspect: Aspect.optional().describe("Exact size of this format."), + 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( diff --git a/packages/testkit/src/cli/harness.ts b/packages/testkit/src/cli/harness.ts index 0ee852c..5d565e9 100644 --- a/packages/testkit/src/cli/harness.ts +++ b/packages/testkit/src/cli/harness.ts @@ -116,6 +116,7 @@ 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/builtin-memory-provider.postgres.test.ts", "packages/adapters/src/profile-consolidation.postgres.test.ts", ];