feat(media): generate images and video on ComfyUI and edit images
Bots can now make images and short clips on the user's own GPU server and rework images locally, through the media job runtime. - generate_image: prompt, aspect preset or exact size, up to 4 images, optional source image for image-to-image. Sizes target what the default models are tuned for, and sources are fitted to that size before upload, so large photos cannot exhaust GPU memory. - generate_video: prompt, aspect, duration up to 10 s, optional start frame. - edit_image: crop, rotate, flip, resize to a size or format (crop, pad or blurred background), extend, convert, in that order, with sharp; an optional prompt repaints the result on ComfyUI. Edits keep the source as parent. - Registrations declare the input artifacts so jobs refuse ones the user cannot read, and generation declares a fixed GPU memory floor (6 GB images, 12 GB video) that a custom workflow cannot lower. - Tool definitions reuse the operation input schemas plus continue_after. - Saving media settings rejects custom workflows the provider could not fill. - Tests: sharp edits on generated fixtures, handlers against the ComfyUI emulator, and a PostgreSQL suite running the operations through the media runtime (card outputs, artifact versions, foreign artifacts, GPU wait). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u
This commit is contained in:
Claude
parent
fad25678ac
commit
4217e4b4af
@ -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<string, unknown>,
|
||||
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:
|
||||
|
||||
@ -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)];
|
||||
}
|
||||
|
||||
@ -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<MediaSettingsView> {
|
||||
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) {
|
||||
|
||||
28
packages/adapters/src/media/comfyui-settings.test.ts
Normal file
28
packages/adapters/src/media/comfyui-settings.test.ts
Normal file
@ -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<typeof saveMediaSettings>[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}}");
|
||||
});
|
||||
});
|
||||
110
packages/adapters/src/media/image-edit.test.ts
Normal file
110
packages/adapters/src/media/image-edit.test.ts
Normal file
@ -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<Buffer> {
|
||||
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);
|
||||
});
|
||||
});
|
||||
156
packages/adapters/src/media/image-edit.ts
Normal file
156
packages/adapters/src/media/image-edit.ts
Normal file
@ -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<SharpFactory> | undefined;
|
||||
|
||||
/** sharp is a native module; load it on first use so a platform without it only loses edits. */
|
||||
export function loadSharp(): Promise<SharpFactory> {
|
||||
sharpModule ??= import("sharp").then((module) => module.default);
|
||||
return sharpModule;
|
||||
}
|
||||
|
||||
export async function sharpAvailable(): Promise<boolean> {
|
||||
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<Uint8Array> {
|
||||
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<Sharp> {
|
||||
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<Buffer> {
|
||||
const sharp = await loadSharp();
|
||||
let current = await materialize(sharp, sharp(input, { autoOrient: true }));
|
||||
const step = async (apply: (image: Sharp) => Sharp | Promise<Sharp>) => {
|
||||
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<ImageFormat, string> = {
|
||||
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 };
|
||||
}
|
||||
172
packages/adapters/src/media/image-operations.postgres.test.ts
Normal file
172
packages/adapters/src/media/image-operations.postgres.test.ts
Normal file
@ -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<typeof createDb>;
|
||||
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<MediaJobBlock | undefined> {
|
||||
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);
|
||||
});
|
||||
});
|
||||
274
packages/adapters/src/media/image-operations.test.ts
Normal file
274
packages/adapters/src/media/image-operations.test.ts
Normal file
@ -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<string> {
|
||||
const dir = await mkdtemp(join(tmpdir(), "image-operations-"));
|
||||
dirs.push(dir);
|
||||
return dir;
|
||||
}
|
||||
|
||||
async function fixture(name: string, width: number, height: number): Promise<string> {
|
||||
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<typeof setup>["operation"],
|
||||
name: string,
|
||||
input: unknown,
|
||||
artifacts: Record<string, string> = {},
|
||||
) {
|
||||
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");
|
||||
});
|
||||
});
|
||||
315
packages/adapters/src/media/image-operations.ts
Normal file
315
packages/adapters/src/media/image-operations.ts
Normal file
@ -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<ComfyUiProvider | null>;
|
||||
|
||||
/** 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<ComfyUiProvider> {
|
||||
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<void> {
|
||||
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<MediaJobInputFile & { width: number; height: number }> {
|
||||
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<void> {
|
||||
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<MediaOperationResult> {
|
||||
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<MediaOperationResult> {
|
||||
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<MediaOperationResult> {
|
||||
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<GenerateImageInput> = {
|
||||
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<GenerateVideoInput> = {
|
||||
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<EditImageInput> = {
|
||||
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[];
|
||||
}
|
||||
@ -13,7 +13,6 @@ describe("media generation inputs", () => {
|
||||
});
|
||||
expect(GenerateVideoInputSchema.parse({ prompt: "waves" })).toEqual({
|
||||
prompt: "waves",
|
||||
aspect: "landscape",
|
||||
durationSeconds: 5,
|
||||
});
|
||||
});
|
||||
|
||||
@ -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<typeof GenerateImageInputSchema>;
|
||||
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(
|
||||
|
||||
@ -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",
|
||||
];
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user