evOo-local/packages/adapters/src/media/image-operations.test.ts

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);
    const titled = await run(
      operation,
      "edit_image",
      { artifactId: "photo-1", title: "Titelbild: Sommer", rotate: 180 },
      { "photo-1": photo },
    );
    expect(titled.outputs[0]?.name).toBe("Titelbild Sommer.jpg");
    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("1×1 · 90° · 9:16 · PNG");
    expect(
      editImageTitle({
        artifactId: "a",
        flip: "vertical",
        extend: { top: 4, right: 0, bottom: 0, left: 0 },
      }),
    ).toBe("↕");
    expect(
      editImageTitle({ artifactId: "a", title: "Titelbild", prompt: "snow", rotate: 90 }),
    ).toBe("Titelbild");
    expect(editImageTitle({ artifactId: "a", resize: { width: 800, fit: "crop" } })).toBe("800×…");
    expect(editImageTitle({ artifactId: "a", prompt: "make it snow" })).toBe("make it snow");
  });
});