evOo-local/packages/adapters/src/image-generation-conformance.test.ts

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: () => {
    const emulator = new ComfyUiEmulator();
    return createComfyUiProvider({
      baseUrl: "http://127.0.0.1:8188",
      fetch: emulator.fetch,
      webSocket: emulator.openWebSocket,
      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);
      // Progress never goes backwards and reaches 1 only once every output is in.
      expect(progress).toEqual([...progress].sort((a, b) => a - b));
      expect(progress.filter((fraction) => fraction >= 1)).toEqual([1]);
      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();
    });
  });
}