image-studio/commands/inpaint-region.ts

import type { Asset } from "@core/asset";
import type { InpaintRegion } from "@core/inpaint-region";
import type { InpaintRegionId } from "@core/id";
import type { Command } from "./command";
import { commandIds } from "./ids";
import type { LayerMaskOperation } from "./document";
import { findLayerLocation } from "./document-tree";

export type DocumentAddInpaintRegionPayload = {
  region: InpaintRegion;
  maskAsset: Asset;
};

export type DocumentApplyInpaintRegionMaskOperationPayload = {
  regionId: InpaintRegionId;
  source: string;
  mimeType?: string;
  operation: LayerMaskOperation;
};

export type DocumentRemoveInpaintRegionPayload = {
  regionId: InpaintRegionId;
};

export const documentAddInpaintRegionCommand: Command<DocumentAddInpaintRegionPayload> = {
  id: commandIds.documentAddInpaintRegion,
  name: "Add inpaint region",
  execute({ state }, payload) {
    const target = findLayerLocation(state.document, payload.region.targetLayerId)?.layer;
    if (!target || (target.type !== "image" && target.type !== "raster")) return state;
    const targetAsset = state.document.assets.find((asset) => asset.id === target.assetId);
    if (!targetAsset || state.document.inpaintRegions.some((region) => region.id === payload.region.id)) return state;
    if (state.document.assets.some((asset) => asset.id === payload.maskAsset.id)) return state;
    if (payload.region.maskAssetId !== payload.maskAsset.id || !payload.region.name.trim()) return state;
    if (Math.round(payload.maskAsset.intrinsicSize.w) !== Math.round(targetAsset.intrinsicSize.w) || Math.round(payload.maskAsset.intrinsicSize.h) !== Math.round(targetAsset.intrinsicSize.h)) return state;

    return {
      ...state,
      document: {
        ...state.document,
        assets: [...state.document.assets, payload.maskAsset],
        inpaintRegions: [...state.document.inpaintRegions, { ...payload.region, name: payload.region.name.trim() }],
      },
    };
  },
};

export const documentApplyInpaintRegionMaskOperationCommand: Command<DocumentApplyInpaintRegionMaskOperationPayload> = {
  id: commandIds.documentApplyInpaintRegionMaskOperation,
  name: "Apply inpaint region mask operation",
  execute({ state }, payload) {
    const region = state.document.inpaintRegions.find((candidate) => candidate.id === payload.regionId);
    if (!region || !payload.source.trim()) return state;
    return {
      ...state,
      document: {
        ...state.document,
        assets: state.document.assets.map((asset) => asset.id === region.maskAssetId ? { ...asset, source: payload.source, mimeType: payload.mimeType ?? asset.mimeType } : asset),
      },
      editor: {
        ...state.editor,
        brushStrokePreview: state.editor.brushStrokePreview?.assetId === region.maskAssetId ? undefined : state.editor.brushStrokePreview,
      },
    };
  },
};

export const documentRemoveInpaintRegionCommand: Command<DocumentRemoveInpaintRegionPayload> = {
  id: commandIds.documentRemoveInpaintRegion,
  name: "Remove inpaint region",
  execute({ state }, payload) {
    const region = state.document.inpaintRegions.find((candidate) => candidate.id === payload.regionId);
    if (!region) return state;
    const inpaintRegions = state.document.inpaintRegions.filter((candidate) => candidate.id !== payload.regionId);
    const maskStillUsed = inpaintRegions.some((candidate) => candidate.maskAssetId === region.maskAssetId);
    const exitingEdit = state.editor.maskEdit?.inpaintRegionId === region.id;
    const contextualTool = state.editor.tools.activeTool === "semanticSelect" || state.editor.tools.activeTool === "maskLasso" || state.editor.tools.activeTool === "maskRectangle";
    return {
      ...state,
      document: {
        ...state.document,
        inpaintRegions,
        assets: maskStillUsed ? state.document.assets : state.document.assets.filter((asset) => asset.id !== region.maskAssetId),
      },
      editor: {
        ...state.editor,
        maskEdit: exitingEdit ? undefined : state.editor.maskEdit,
        brushStrokePreview: state.editor.brushStrokePreview?.assetId === region.maskAssetId ? undefined : state.editor.brushStrokePreview,
        maskShapeSession: exitingEdit ? undefined : state.editor.maskShapeSession,
        tools: exitingEdit && contextualTool ? { ...state.editor.tools, activeTool: "select", interactionMode: { type: "tool", tool: "select" } } : state.editor.tools,
      },
    };
  },
};

export const inpaintRegionCommands = [
  documentAddInpaintRegionCommand,
  documentApplyInpaintRegionMaskOperationCommand,
  documentRemoveInpaintRegionCommand,
] satisfies Command<unknown>[];