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