image-studio/commands/generation.test.ts

import { describe, expect, test } from "bun:test";
import type { GenerationCandidate } from "@editor/state";
import { createInitialAppState } from "@editor/initial-state";
import {
  generationAddCandidateCommand,
  generationApplyCandidateAsLayerCommand,
  generationRemoveCandidateCommand,
  generationReplaceCandidatePixelsCommand,
} from "./generation";

describe("generation commands", () => {
  test("adds, selects, and removes candidates", () => {
    const state = createInitialAppState("Test");
    const candidate = generationCandidate("candidate-1");

    const added = generationAddCandidateCommand.execute({ state }, { candidate });
    const removed = generationRemoveCandidateCommand.execute({ state: added }, { candidateId: candidate.id });

    expect(added.editor.generation.candidates).toEqual([candidate]);
    expect(added.editor.generation.selectedCandidateId).toBe(candidate.id);
    expect(removed.editor.generation.candidates).toEqual([]);
    expect(removed.editor.generation.selectedCandidateId).toBeUndefined();
  });

  test("applies candidates as top-level layers", () => {
    const state = generationAddCandidateCommand.execute({ state: documentWithSourceLayer() }, { candidate: generationCandidate("candidate-1") });

    const next = generationApplyCandidateAsLayerCommand.execute(
      { state },
      { candidateId: "candidate-1", assetId: "generated-asset", layerId: "generated-layer" },
    );

    expect(next.document.assets.find((asset) => asset.id === "generated-asset")?.source).toBe("generated-source");
    expect(next.document.artboards[0]?.layers[0]?.id).toBe("generated-layer");
    expect(next.editor.selection).toEqual({ artboardId: "a1", layerIds: ["generated-layer"] });
    expect(next.editor.generation).toEqual({ candidates: [], selectedCandidateId: undefined });
  });

  test("replaces source asset pixels for inpaint candidates", () => {
    const state = generationAddCandidateCommand.execute({ state: documentWithSourceLayer() }, { candidate: generationCandidate("candidate-1", true) });

    const next = generationReplaceCandidatePixelsCommand.execute(
      { state },
      { candidateId: "candidate-1", source: "composited-source", mimeType: "image/png" },
    );

    expect(next.document.assets.find((asset) => asset.id === "source-asset")?.source).toBe("composited-source");
    expect(next.document.assets.find((asset) => asset.id === "source-asset")?.mimeType).toBe("image/png");
    expect(next.editor.selection).toEqual({ artboardId: "a1", layerIds: ["source-layer"] });
    expect(next.editor.generation).toEqual({ candidates: [], selectedCandidateId: undefined });
  });
});

function documentWithSourceLayer() {
  return {
    ...createInitialAppState("Test"),
    document: {
      ...createInitialAppState("Test").document,
      assets: [
        { id: "source-asset", name: "Source", mimeType: "image/png", source: "source", intrinsicSize: { w: 100, h: 100 } },
        { id: "mask-asset", name: "Mask", mimeType: "image/png", source: "mask", intrinsicSize: { w: 100, h: 100 } },
      ],
      artboards: [
        {
          id: "a1",
          name: "Artboard 1",
          bounds: { x: 0, y: 0, w: 100, h: 100 },
          backgroundColor: "transparent",
          visible: true,
          locked: false,
          layers: [
            raster("mask-layer", "Mask", "mask-asset"),
            { ...raster("source-layer", "Source", "source-asset"), clippingMask: { maskLayerId: "mask-layer" } },
          ],
        },
      ],
    },
  };
}

function generationCandidate(id: string, inpaint = false): GenerationCandidate {
  const candidate: GenerationCandidate = {
    id,
    source: "generated-source",
    mimeType: "image/png",
    intrinsicSize: { w: 64, h: 64 },
    mode: inpaint ? "inpaint" : "text-to-image",
    settings: createInitialAppState("Test").editor.tools.generate,
    seed: 123,
    width: 64,
    height: 64,
    placement: {
      artboardId: "a1",
      layerName: "Generated",
      transform: { position: { x: 5, y: 6 }, scale: { x: 1, y: 1 }, rotation: 0 },
    },
  };

  return inpaint
    ? {
        ...candidate,
        inputImage: "input",
        maskImage: "mask",
        inpaint: {
          targetLayerId: "source-layer",
          maskLayerId: "mask-layer",
          sourceAssetId: "source-asset",
          maskAssetId: "mask-asset",
          inputImage: "input",
          maskImage: "mask",
          crop: {
            assetBounds: { x: 0, y: 0, w: 64, h: 64 },
            documentBounds: { x: 0, y: 0, w: 64, h: 64 },
            padding: 12,
            maskedAreaOnly: true,
          },
          mask: {
            polarity: "hidden",
            activeBounds: { x: 10, y: 10, w: 20, h: 20 },
          },
          backend: {
            growMaskBy: 6,
            maskedContent: "neutral",
            maskBlur: 0,
            maskFeather: 0,
            maskExpand: 0,
            cropPadding: 12,
          },
        },
      }
    : candidate;
}

function raster(id: string, name: string, assetId: string) {
  return {
    id,
    type: "raster" as const,
    name,
    visible: true,
    locked: false,
    opacity: 1,
    assetId,
    transform: { position: { x: 0, y: 0 }, scale: { x: 1, y: 1 }, rotation: 0 },
  };
}