image-studio/commands/generation.ts

import type { Asset } from "@core/asset";
import type { ImageDocument } from "@core/document";
import type { ArtboardId, AssetId, LayerId } from "@core/id";
import type { ImageLayer } from "@core/image-layer";
import type { Layer } from "@core/layer";
import type { GenerationCandidate, GenerationState } from "@editor/state";
import type { Command } from "./command";
import { commandIds } from "./ids";

export type GenerationAddCandidatePayload = {
  candidate: GenerationCandidate;
};

export type GenerationSelectCandidatePayload = {
  candidateId?: string;
};

export type GenerationRemoveCandidatePayload = {
  candidateId: string;
};

export type GenerationApplyCandidateAsLayerPayload = {
  candidateId: string;
  assetId: AssetId;
  layerId: LayerId;
  variant?: boolean;
};

export type GenerationReplaceCandidatePixelsPayload = {
  candidateId: string;
  source: string;
  mimeType?: string;
};

const maxCandidates = 12;

export const generationAddCandidateCommand: Command<GenerationAddCandidatePayload> = {
  id: commandIds.generationAddCandidate,
  name: "Add generation candidate",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    const candidates = [payload.candidate, ...state.editor.generation.candidates.filter((candidate) => candidate.id !== payload.candidate.id)].slice(0, maxCandidates);
    return {
      ...state,
      editor: {
        ...state.editor,
        generation: {
          candidates,
          selectedCandidateId: payload.candidate.id,
        },
      },
    };
  },
};

export const generationSelectCandidateCommand: Command<GenerationSelectCandidatePayload> = {
  id: commandIds.generationSelectCandidate,
  name: "Select generation candidate",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    const selectedCandidateId = payload.candidateId && state.editor.generation.candidates.some((candidate) => candidate.id === payload.candidateId) ? payload.candidateId : undefined;
    if (state.editor.generation.selectedCandidateId === selectedCandidateId) return state;
    return {
      ...state,
      editor: {
        ...state.editor,
        generation: {
          ...state.editor.generation,
          selectedCandidateId,
        },
      },
    };
  },
};

export const generationRemoveCandidateCommand: Command<GenerationRemoveCandidatePayload> = {
  id: commandIds.generationRemoveCandidate,
  name: "Remove generation candidate",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    const candidates = state.editor.generation.candidates.filter((candidate) => candidate.id !== payload.candidateId);
    if (candidates.length === state.editor.generation.candidates.length) return state;
    const selectedCandidateId = state.editor.generation.selectedCandidateId === payload.candidateId ? candidates[0]?.id : state.editor.generation.selectedCandidateId;
    return {
      ...state,
      editor: {
        ...state.editor,
        generation: {
          candidates,
          selectedCandidateId,
        },
      },
    };
  },
};

export const generationClearCandidatesCommand: Command = {
  id: commandIds.generationClearCandidates,
  name: "Clear generation candidates",
  history: { mode: "ignore" },
  execute({ state }) {
    if (state.editor.generation.candidates.length === 0 && !state.editor.generation.selectedCandidateId) return state;
    return {
      ...state,
      editor: {
        ...state.editor,
        generation: { candidates: [], selectedCandidateId: undefined },
      },
    };
  },
};

export const generationApplyCandidateAsLayerCommand: Command<GenerationApplyCandidateAsLayerPayload> = {
  id: commandIds.generationApplyCandidateAsLayer,
  name: "Apply generation candidate as layer",
  execute({ state }, payload) {
    const candidate = state.editor.generation.candidates.find((item) => item.id === payload.candidateId);
    if (!candidate) return state;
    if (state.document.assets.some((asset) => asset.id === payload.assetId) || findLayerLocation(state.document, payload.layerId)) return state;
    if (!state.document.artboards.some((artboard) => artboard.id === candidate.placement.artboardId)) return state;

    const asset: Asset = {
      id: payload.assetId,
      name: payload.variant ? `${candidate.placement.layerName} variant` : candidate.placement.layerName,
      mimeType: candidate.mimeType,
      source: candidate.source,
      intrinsicSize: { ...candidate.intrinsicSize },
    };
    const layer: ImageLayer = {
      id: payload.layerId,
      type: "image",
      name: payload.variant ? `${candidate.placement.layerName} variant` : candidate.placement.layerName,
      visible: true,
      locked: false,
      opacity: 1,
      assetId: asset.id,
      transform: {
        position: { ...candidate.placement.transform.position },
        scale: { ...candidate.placement.transform.scale },
        rotation: candidate.placement.transform.rotation,
      },
    };

    return {
      ...state,
      document: insertLayerAtTop({ ...state.document, assets: [...state.document.assets, asset] }, candidate.placement.artboardId, layer),
      editor: {
        ...state.editor,
        generation: clearCommittedGenerationPreview(state.editor.generation),
        selection: { artboardId: candidate.placement.artboardId, layerIds: [layer.id] },
      },
    };
  },
};

export const generationReplaceCandidatePixelsCommand: Command<GenerationReplaceCandidatePixelsPayload> = {
  id: commandIds.generationReplaceCandidatePixels,
  name: "Replace masked pixels with generation candidate",
  execute({ state }, payload) {
    const candidate = state.editor.generation.candidates.find((item) => item.id === payload.candidateId);
    if (!candidate?.inpaint || !payload.source.trim()) return state;
    const targetAsset = state.document.assets.find((asset) => asset.id === candidate.inpaint?.sourceAssetId);
    const targetLayerLocation = findLayerLocation(state.document, candidate.inpaint.targetLayerId);
    if (!targetAsset || !targetLayerLocation) return state;

    return {
      ...state,
      document: {
        ...state.document,
        assets: state.document.assets.map((asset) => asset.id === targetAsset.id ? { ...asset, source: payload.source, mimeType: payload.mimeType ?? asset.mimeType } : asset),
      },
      editor: {
        ...state.editor,
        generation: clearCommittedGenerationPreview(state.editor.generation),
        selection: { artboardId: targetLayerLocation.artboardId, layerIds: [candidate.inpaint.targetLayerId] },
      },
    };
  },
};

export const generationCommands = [
  generationAddCandidateCommand,
  generationSelectCandidateCommand,
  generationRemoveCandidateCommand,
  generationClearCandidatesCommand,
  generationApplyCandidateAsLayerCommand,
  generationReplaceCandidatePixelsCommand,
] satisfies Command<unknown>[];

type LayerLocation = {
  artboardId: ArtboardId;
  layer: Layer;
};

function insertLayerAtTop(document: ImageDocument, artboardId: ArtboardId, layer: Layer): ImageDocument {
  return {
    ...document,
    artboards: document.artboards.map((artboard) => artboard.id === artboardId ? { ...artboard, layers: [layer, ...artboard.layers] } : artboard),
  };
}

function findLayerLocation(document: ImageDocument, layerId: LayerId): LayerLocation | undefined {
  for (const artboard of document.artboards) {
    const layer = findLayerInTree(artboard.layers, layerId);
    if (layer) return { artboardId: artboard.id, layer };
  }
  return undefined;
}

function findLayerInTree(layers: readonly Layer[], layerId: LayerId): Layer | undefined {
  for (const layer of layers) {
    if (layer.id === layerId) return layer;
    if (layer.type === "group") {
      const child = findLayerInTree(layer.children, layerId);
      if (child) return child;
    }
  }
  return undefined;
}

function clearCommittedGenerationPreview(generation: GenerationState): GenerationState {
  if (generation.candidates.length === 0 && !generation.selectedCandidateId) return generation;
  return { candidates: [], selectedCandidateId: undefined };
}