image-studio/commands/generation.ts

import type { Asset } from "@core/asset";
import type { AssetGenerationProvenance, GeneratedAssetAcceptance } from "@core/asset-provenance";
import type { ImageDocument } from "@core/document";
import type { ArtboardId, AssetId, GenerationCandidateId, GenerationJobId, LayerId } from "@core/id";
import type { ImageLayer } from "@core/image-layer";
import type { Layer } from "@core/layer";
import type { AppState, GenerationCandidate, GenerationCompareMode, GenerationJobKind, GenerationOptions, GenerationState } from "@editor/state";
import type { Command } from "./command";
import { commandIds } from "./ids";

export type GenerationAddCandidatePayload = {
  candidate: GenerationCandidate;
};

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

export type GenerationSetCompareModePayload = {
  mode: GenerationCompareMode;
};

export type GenerationRemoveCandidatePayload = {
  candidateId: GenerationCandidateId;
};

export type GenerationReuseCandidateSettingsPayload = {
  candidateId: GenerationCandidateId;
};

export type GenerationApplyCandidateAsLayerPayload = {
  candidateId: GenerationCandidateId;
  assetId: AssetId;
  layerId: LayerId;
};

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

export type GenerationStartJobPayload = { jobId: GenerationJobId; kind: GenerationJobKind; label: string; startedAt: number };
export type GenerationSucceedJobPayload = { jobId: GenerationJobId; finishedAt: number };
export type GenerationFailJobPayload = { jobId: GenerationJobId; finishedAt: number; error: string };
export type GenerationSetResourcesPayload = { options: GenerationOptions };
export type GenerationFailResourcesPayload = { error: string };

const maxCandidates = 12;
const maxJobs = 20;
const generationCompareModes = new Set<GenerationCompareMode>(["result", "before", "split"]);

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: {
          ...state.editor.generation,
          candidates,
          selectedCandidateId: payload.candidate.id,
          compareMode: "result",
        },
      },
    };
  },
};

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 generationSetCompareModeCommand: Command<GenerationSetCompareModePayload> = {
  id: commandIds.generationSetCompareMode,
  name: "Set generation compare mode",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    if (!generationCompareModes.has(payload.mode)) return state;
    if (state.editor.generation.compareMode === payload.mode) return state;
    return {
      ...state,
      editor: {
        ...state.editor,
        generation: {
          ...state.editor.generation,
          compareMode: payload.mode,
        },
      },
    };
  },
};

export const generationRemoveCandidateCommand: Command<GenerationRemoveCandidatePayload> = {
  id: commandIds.generationRemoveCandidate,
  name: "Remove generation candidate",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    const generation = removeGenerationCandidate(state.editor.generation, payload.candidateId);
    if (generation === state.editor.generation) return state;
    return {
      ...state,
      editor: {
        ...state.editor,
        generation,
      },
    };
  },
};

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 && state.editor.generation.compareMode === "result") return state;
    return {
      ...state,
      editor: {
        ...state.editor,
        generation: { ...state.editor.generation, candidates: [], selectedCandidateId: undefined, compareMode: "result" },
      },
    };
  },
};

export const generationReuseCandidateSettingsCommand: Command<GenerationReuseCandidateSettingsPayload> = {
  id: commandIds.generationReuseCandidateSettings,
  name: "Reuse generation candidate settings",
  execute({ state }, payload) {
    const candidate = state.editor.generation.candidates.find((item) => item.id === payload.candidateId);
    if (!candidate) return state;
    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          generate: {
            ...candidate.settings,
            seed: candidate.seed,
            outpaint: { ...candidate.settings.outpaint },
            inpaint: { ...candidate.settings.inpaint },
          },
        },
      },
    };
  },
};

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: candidate.placement.layerName,
      mimeType: candidate.mimeType,
      source: candidate.source,
      intrinsicSize: { ...candidate.intrinsicSize },
      provenance: generationProvenance(candidate, "layer"),
    };
    const layer: ImageLayer = {
      id: payload.layerId,
      type: "image",
      name: 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: removeGenerationCandidate(state.editor.generation, candidate.id),
        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,
                provenance: generationProvenance(candidate, "replacement"),
              }
            : asset,
        ),
      },
      editor: {
        ...state.editor,
        generation: removeGenerationCandidate(state.editor.generation, candidate.id),
        selection: { artboardId: targetLayerLocation.artboardId, layerIds: [candidate.inpaint.targetLayerId] },
      },
    };
  },
};

export const generationStartJobCommand: Command<GenerationStartJobPayload> = {
  id: commandIds.generationStartJob,
  name: "Start generation job",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    if (!payload.jobId || !payload.label.trim() || !Number.isFinite(payload.startedAt)) return state;
    if (state.editor.generation.jobs.some((job) => job.status === "running" || job.id === payload.jobId)) return state;
    const job: GenerationState["jobs"][number] = { id: payload.jobId, kind: payload.kind, label: payload.label.trim(), status: "running", startedAt: payload.startedAt };
    return updateJobs(state, [job, ...state.editor.generation.jobs].slice(0, maxJobs));
  },
};

export const generationSucceedJobCommand: Command<GenerationSucceedJobPayload> = {
  id: commandIds.generationSucceedJob,
  name: "Complete generation job",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    return settleJob(state, payload.jobId, payload.finishedAt, "succeeded");
  },
};

export const generationFailJobCommand: Command<GenerationFailJobPayload> = {
  id: commandIds.generationFailJob,
  name: "Fail generation job",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    if (!payload.error.trim()) return state;
    return settleJob(state, payload.jobId, payload.finishedAt, "failed", payload.error.trim());
  },
};

export const generationLoadResourcesCommand: Command = {
  id: commandIds.generationLoadResources,
  name: "Load generation resources",
  history: { mode: "ignore" },
  execute({ state }) {
    if (state.editor.generation.resources.status === "loading") return state;
    return updateResources(state, { status: "loading" });
  },
};

export const generationSetResourcesCommand: Command<GenerationSetResourcesPayload> = {
  id: commandIds.generationSetResources,
  name: "Set generation resources",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    return updateResources(state, { status: "ready", options: payload.options });
  },
};

export const generationFailResourcesCommand: Command<GenerationFailResourcesPayload> = {
  id: commandIds.generationFailResources,
  name: "Fail generation resources",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    if (!payload.error.trim()) return state;
    return updateResources(state, { status: "failed", error: payload.error.trim() });
  },
};

export const generationCommands = [
  generationAddCandidateCommand,
  generationSelectCandidateCommand,
  generationSetCompareModeCommand,
  generationRemoveCandidateCommand,
  generationClearCandidatesCommand,
  generationReuseCandidateSettingsCommand,
  generationApplyCandidateAsLayerCommand,
  generationReplaceCandidatePixelsCommand,
  generationStartJobCommand,
  generationSucceedJobCommand,
  generationFailJobCommand,
  generationLoadResourcesCommand,
  generationSetResourcesCommand,
  generationFailResourcesCommand,
] satisfies Command<unknown>[];

function updateJobs(state: AppState, jobs: GenerationState["jobs"]): AppState {
  return { ...state, editor: { ...state.editor, generation: { ...state.editor.generation, jobs } } };
}

function updateResources(state: AppState, resources: GenerationState["resources"]): AppState {
  return { ...state, editor: { ...state.editor, generation: { ...state.editor.generation, resources } } };
}

function settleJob(state: AppState, jobId: GenerationJobId, finishedAt: number, status: "succeeded" | "failed", error?: string): AppState {
  if (!Number.isFinite(finishedAt)) return state;
  const job = state.editor.generation.jobs.find((candidate) => candidate.id === jobId);
  if (!job || job.status !== "running" || finishedAt < job.startedAt) return state;
  return updateJobs(state, state.editor.generation.jobs.map((candidate) => candidate.id === jobId ? { ...candidate, status, finishedAt, error } : candidate));
}

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 removeGenerationCandidate(generation: GenerationState, candidateId: GenerationCandidateId): GenerationState {
  const removedIndex = generation.candidates.findIndex((candidate) => candidate.id === candidateId);
  if (removedIndex < 0) return generation;

  const candidates = generation.candidates.filter((candidate) => candidate.id !== candidateId);
  const selectionStillExists = generation.selectedCandidateId
    ? candidates.some((candidate) => candidate.id === generation.selectedCandidateId)
    : false;
  const selectedCandidateId = selectionStillExists
    ? generation.selectedCandidateId
    : candidates[Math.min(removedIndex, candidates.length - 1)]?.id;

  return {
    ...generation,
    candidates,
    selectedCandidateId,
    compareMode: candidates.length > 0 ? generation.compareMode : "result",
  };
}

function generationProvenance(candidate: GenerationCandidate, acceptance: GeneratedAssetAcceptance): AssetGenerationProvenance {
  return {
    kind: "generated",
    candidateId: candidate.id,
    mode: candidate.mode,
    acceptance,
    prompt: candidate.settings.prompt,
    negativePrompt: candidate.settings.negativePrompt,
    seed: candidate.seed,
    outputSize: { ...candidate.intrinsicSize },
    settings: {
      architecture: candidate.settings.architecture,
      model: candidate.settings.model,
      textEncoder: candidate.settings.textEncoder,
      vae: candidate.settings.vae,
      strength: candidate.settings.strength,
      steps: candidate.settings.steps,
      cfg: candidate.settings.cfg,
      sampler: candidate.settings.sampler,
      scheduler: candidate.settings.scheduler,
      width: candidate.settings.width,
      height: candidate.settings.height,
    },
    inpaint: candidate.inpaint
      ? {
          targetLayerId: candidate.inpaint.targetLayerId,
          maskLayerId: candidate.inpaint.maskLayerId,
          sourceAssetId: candidate.inpaint.sourceAssetId,
          maskAssetId: candidate.inpaint.maskAssetId,
          crop: {
            assetBounds: { ...candidate.inpaint.crop.assetBounds },
            documentBounds: { ...candidate.inpaint.crop.documentBounds },
            padding: candidate.inpaint.crop.padding,
            maskedAreaOnly: candidate.inpaint.crop.maskedAreaOnly,
          },
          mask: {
            polarity: candidate.inpaint.mask.polarity,
            activeBounds: { ...candidate.inpaint.mask.activeBounds },
          },
          backend: {
            growMaskBy: candidate.inpaint.backend.growMaskBy,
            maskedContent: candidate.inpaint.backend.maskedContent,
            maskBlur: candidate.inpaint.backend.maskBlur,
            maskFeather: candidate.inpaint.backend.maskFeather,
            maskExpand: candidate.inpaint.backend.maskExpand,
            cropPadding: candidate.inpaint.backend.cropPadding,
          },
        }
      : undefined,
  };
}