image-studio/operations/generation/workflow.ts

import { commandIds } from "@commands/ids";
import type { Layer } from "@core/layer";
import type { GenerationCandidate, GenerationJobKind } from "@editor/state";
import type { AppStore } from "@editor/store";
import type { GenerateSettings } from "@editor/tools";
import { createInpaintRegionMask, createRefinementMask } from "@operations/masks/rasterActions";
import { createMaskedPixelReplacementSource } from "./candidateActions";
import { runGenerationJob } from "./generationJob";
import { loadGenerationResources } from "./loadResources";
import { runGenerate, runGenerateFromCandidate } from "./runGenerate";
import { checkGenerationPreconditions } from "./preconditions";

export type GenerationWorkflow = ReturnType<typeof createGenerationWorkflow>;

export type GenerationWorkflowDependencies = {
  runGenerate: typeof runGenerate;
  runGenerateFromCandidate: typeof runGenerateFromCandidate;
  createMaskedPixelReplacementSource: typeof createMaskedPixelReplacementSource;
  createRefinementMask: typeof createRefinementMask;
  createInpaintRegionMask: typeof createInpaintRegionMask;
  loadGenerationResources: typeof loadGenerationResources;
  createId(): string;
};

const defaultDependencies: GenerationWorkflowDependencies = {
  runGenerate,
  runGenerateFromCandidate,
  createMaskedPixelReplacementSource,
  createRefinementMask,
  createInpaintRegionMask,
  loadGenerationResources,
  createId: () => crypto.randomUUID(),
};

export function createGenerationWorkflow(store: AppStore, dependencies: GenerationWorkflowDependencies = defaultDependencies) {
  let activeController: AbortController | undefined;
  const job = async (kind: GenerationJobKind, label: string, task: (signal: AbortSignal, report: (progress: number, detail: string) => void) => Promise<void>) => {
    if (store.getState().editor.generation.jobs.some((candidate) => candidate.status === "running")) return;
    const controller = new AbortController();
    activeController = controller;
    try {
      await runGenerationJob({ kind, label, dispatch: store.dispatch, signal: controller.signal, task });
    } finally {
      if (activeController === controller) activeController = undefined;
    }
  };

  return {
    precondition: () => {
      const state = store.getState();
      return checkGenerationPreconditions(state.document, state.editor.selection, state.editor.tools.generate);
    },

    loadResources: () => dependencies.loadGenerationResources(store),

    prepareInpaintMask: async () => {
      const state = store.getState();
      const layerId = state.editor.selection.layerIds.length === 1 ? state.editor.selection.layerIds[0] : undefined;
      if (!layerId) return;
      const artboard = state.document.artboards.find((candidate) => candidate.id === state.editor.selection.artboardId);
      const layer = artboard ? findLayer(artboard.layers, layerId) : undefined;
      if (!layer || (layer.type !== "image" && layer.type !== "raster")) return;
      const asset = state.document.assets.find((candidate) => candidate.id === layer.assetId);
      if (!asset) return;
      const existing = state.document.inpaintRegions.find((region) => region.targetLayerId === layerId && region.enabled);
      if (existing) {
        store.dispatch(commandIds.toolEnterInpaintRegionEdit, { targetLayerId: layerId, regionId: existing.id });
        return;
      }
      const source = await dependencies.createInpaintRegionMask(asset.intrinsicSize.w, asset.intrinsicSize.h);
      const maskAssetId = dependencies.createId();
      const regionId = dependencies.createId();
      store.dispatch(commandIds.documentAddInpaintRegion, {
        region: { id: regionId, name: `${layer.name} AI edit`, targetLayerId: layer.id, maskAssetId, enabled: true },
        maskAsset: { id: maskAssetId, name: `${layer.name} AI edit mask`, mimeType: "image/png", source, intrinsicSize: { ...asset.intrinsicSize } },
      });
      store.dispatch(commandIds.toolEnterInpaintRegionEdit, { targetLayerId: layer.id, regionId });
    },

    generate: () => job("generate", "Generating", async (signal, report) => {
      const state = store.getState();
      await dependencies.runGenerate({
        document: state.document,
        selection: state.editor.selection,
        viewport: state.editor.viewport,
        settings: state.editor.tools.generate,
        dispatch: store.dispatch,
        signal,
        onProgress: report,
      });
    }),

    regenerate: (candidateId: string, settings?: GenerateSettings, label = "Regenerate") =>
      job("regenerate", label, async (signal, report) => {
        const candidate = findCandidate(store, candidateId);
        const nextSettings = settings ?? candidate.settings;
        store.dispatch(commandIds.toolSetGenerateSettings, nextSettings);
        await dependencies.runGenerateFromCandidate({ candidate, settings: nextSettings, dispatch: store.dispatch, signal, onProgress: report });
      }),

    rebuildFromCurrentRegion: (candidateId: string) =>
      job("regenerate", "Rebuilding from current edit region", async (signal, report) => {
        const candidate = findCandidate(store, candidateId);
        if (!candidate.inpaint) throw new Error("Only inpaint candidates can rebuild from an edit region.");
        store.dispatch(commandIds.selectionSet, { artboardId: candidate.placement.artboardId, layerIds: [candidate.inpaint.targetLayerId] });
        store.dispatch(commandIds.toolSetGenerateSettings, candidate.settings);
        const state = store.getState();
        await dependencies.runGenerate({
          document: state.document,
          selection: state.editor.selection,
          viewport: state.editor.viewport,
          settings: candidate.settings,
          dispatch: store.dispatch,
          signal,
          onProgress: report,
        });
      }),

    applyCandidateAsLayer: (candidateId: string) => {
      store.dispatch(commandIds.generationApplyCandidateAsLayer, {
        candidateId,
        assetId: dependencies.createId(),
        layerId: dependencies.createId(),
      });
    },

    applyCandidateAsRefinementLayer: (candidateId: string) =>
      job("refine", "Adding refinement mask", async () => {
        const candidate = findCandidate(store, candidateId);
        const layerId = dependencies.createId();
        const maskAssetId = dependencies.createId();
        const width = Math.max(1, Math.round(candidate.intrinsicSize.w));
        const height = Math.max(1, Math.round(candidate.intrinsicSize.h));
        const source = await dependencies.createRefinementMask(width, height);

        store.dispatch(commandIds.generationApplyCandidateAsLayer, {
          candidateId,
          assetId: dependencies.createId(),
          layerId,
        });
        store.dispatch(commandIds.documentAddLayerMask, {
          layerId,
          asset: {
            id: maskAssetId,
            name: `${candidate.placement.layerName} refinement mask`,
            mimeType: "image/png",
            source,
            intrinsicSize: { w: width, h: height },
          },
          maskLayer: {
            id: dependencies.createId(),
            type: "raster",
            name: `${candidate.placement.layerName} refinement mask`,
            visible: true,
            locked: false,
            opacity: 1,
            assetId: maskAssetId,
            transform: {
              position: { ...candidate.placement.transform.position },
              scale: { ...candidate.placement.transform.scale },
              rotation: candidate.placement.transform.rotation,
            },
          },
        });
        store.dispatch(commandIds.toolSetActive, { tool: "eraser" });
      }),

    replaceCandidatePixels: (candidateId: string) =>
      job("replace", "Replacing pixels", async () => {
        const candidate = findCandidate(store, candidateId);
        const source = await dependencies.createMaskedPixelReplacementSource(store.getState().document, candidate);
        store.dispatch(commandIds.generationReplaceCandidatePixels, { candidateId, source, mimeType: "image/png" });
      }),

    cancel: () => activeController?.abort(),
  };
}

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

function findCandidate(store: AppStore, candidateId: string): GenerationCandidate {
  const candidate = store.getState().editor.generation.candidates.find((item) => item.id === candidateId);
  if (!candidate) throw new Error("This generation candidate is no longer available.");
  return candidate;
}