image-studio/commands/tool.ts

import type { ImageDocument } from "@core/document";
import type { Vec2D } from "@core/geometry";
import type { LayerId, ArtboardId, AssetId, InpaintRegionId } from "@core/id";
import type { Layer } from "@core/layer";
import { getLayerMask } from "@core/layer-mask-utils";
import type { MaskViewMode } from "@editor/state";
import { generateArchitectureDefaults, inpaintProfileDefaults } from "@editor/tools";
import type { BrushSettings, ChromaKeySettings, GenerateIntent, GenerateSettings, MagicWandSettings, ToolId } from "@editor/tools";
import type { Command } from "./command";
import { commandIds } from "./ids";

export type ToolSetActivePayload = {
  tool: ToolId;
};

export type ToolSetBrushSettingsPayload = Partial<BrushSettings>;

export type ToolSetGenerateSettingsPayload = Omit<Partial<GenerateSettings>, "inpaint" | "outpaint"> & {
  inpaint?: Partial<GenerateSettings["inpaint"]>;
  outpaint?: Partial<GenerateSettings["outpaint"]>;
};
export type ToolChooseGenerateIntentPayload = { intent: GenerateIntent };

export type ToolSetChromaKeySettingsPayload = Partial<ChromaKeySettings>;

export type ToolSetMagicWandSettingsPayload = Partial<MagicWandSettings>;

export type ToolSetBrushPreviewPayload = { position: Vec2D } | undefined;

export type ToolSetBrushStrokePreviewPayload =
  | {
      layerId: LayerId;
      assetId: AssetId;
      source: string;
    }
  | undefined;

export type ToolSetMaskViewModePayload = {
  mode: MaskViewMode;
};
export type ToolBeginMaskShapePayload = { point: Vec2D; mode: "replace" | "add" | "subtract" };
export type ToolAppendMaskShapePayload = { point: Vec2D };

export type ToolEnterMaskEditPayload = {
  targetLayerId: LayerId;
  maskLayerId: LayerId;
};

export type ToolEnterInpaintRegionEditPayload = {
  targetLayerId: LayerId;
  regionId: InpaintRegionId;
};

export const toolSetActiveCommand: Command<ToolSetActivePayload> = {
  id: commandIds.toolSetActive,
  name: "Set active tool",
  execute({ state }, payload) {
    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          activeTool: payload.tool,
          interactionMode: { type: "tool", tool: payload.tool },
        },
        brushPreview: undefined,
        brushStrokePreview: undefined,
        maskShapeSession: undefined,
        workspace: state.editor.workspace.panel === "generate" || state.editor.workspace.panel === "chromaKey"
          ? { panel: "none" }
          : state.editor.workspace,
      },
    };
  },
};

export const toolSetGenerateSettingsCommand: Command<ToolSetGenerateSettingsPayload> = {
  id: commandIds.toolSetGenerateSettings,
  name: "Set generate settings",
  execute({ state }, payload) {
    const current = state.editor.tools.generate;
    const architecture = payload.architecture ?? current.architecture;
    const architectureChanged = architecture !== current.architecture;
    const defaults = generateArchitectureDefaults[architecture];
    const requestedMode = payload.mode ?? current.mode;
    const mode = defaults.supportedModes.includes(requestedMode) ? requestedMode : defaults.supportedModes[0] ?? "text-to-image";
    const requestedProfile = payload.inpaint?.profile ?? current.inpaint.profile;
    const profileDefaults = requestedProfile !== current.inpaint.profile && requestedProfile !== "custom" ? inpaintProfileDefaults[requestedProfile] : undefined;
    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          generate: {
            architecture,
            mode,
            model: payload.model ?? (architectureChanged ? defaults.model : current.model),
            textEncoder: payload.textEncoder ?? (architectureChanged ? defaults.textEncoder : current.textEncoder),
            vae: payload.vae ?? (architectureChanged ? defaults.vae : current.vae),
            prompt: payload.prompt ?? current.prompt,
            negativePrompt: payload.negativePrompt ?? current.negativePrompt,
            strength: clampNumber(payload.strength ?? profileDefaults?.strength ?? current.strength, 0, 100),
            steps: Math.round(clampNumber(payload.steps ?? (architectureChanged ? defaults.steps : current.steps), 1, 150)),
            cfg: clampNumber(payload.cfg ?? (architectureChanged ? defaults.cfg : current.cfg), 0, 30),
            seed: Math.round(clampNumber(payload.seed ?? current.seed, -1, Number.MAX_SAFE_INTEGER)),
            sampler: payload.sampler ?? (architectureChanged ? defaults.sampler : current.sampler),
            scheduler: payload.scheduler ?? (architectureChanged ? defaults.scheduler : current.scheduler),
            width: Math.round(clampNumber(payload.width ?? current.width, 64, 4096)),
            height: Math.round(clampNumber(payload.height ?? current.height, 64, 4096)),
            batchSize: Math.round(clampNumber(payload.batchSize ?? current.batchSize, 1, 8)),
            refinePass: payload.refinePass ?? profileDefaults?.refinePass ?? current.refinePass,
            refineStrength: clampNumber(payload.refineStrength ?? profileDefaults?.refineStrength ?? current.refineStrength, 0, 100),
            outpaint: {
              left: Math.round(clampNumber(payload.outpaint?.left ?? current.outpaint.left, 0, 2048)),
              top: Math.round(clampNumber(payload.outpaint?.top ?? current.outpaint.top, 0, 2048)),
              right: Math.round(clampNumber(payload.outpaint?.right ?? current.outpaint.right, 0, 2048)),
              bottom: Math.round(clampNumber(payload.outpaint?.bottom ?? current.outpaint.bottom, 0, 2048)),
              feathering: Math.round(clampNumber(payload.outpaint?.feathering ?? current.outpaint.feathering, 0, 512)),
            },
            inpaint: {
              profile: requestedProfile,
              maskedAreaOnly: payload.inpaint?.maskedAreaOnly ?? current.inpaint.maskedAreaOnly,
              cropPadding: Math.round(clampNumber(payload.inpaint?.cropPadding ?? profileDefaults?.inpaint.cropPadding ?? current.inpaint.cropPadding, 0, 2048)),
              maskPolarity: payload.inpaint?.maskPolarity ?? current.inpaint.maskPolarity,
              maskedContent: payload.inpaint?.maskedContent ?? profileDefaults?.inpaint.maskedContent ?? current.inpaint.maskedContent,
              growMaskBy: Math.round(clampNumber(payload.inpaint?.growMaskBy ?? current.inpaint.growMaskBy, 0, 256)),
              maskExpand: Math.round(clampNumber(payload.inpaint?.maskExpand ?? profileDefaults?.inpaint.maskExpand ?? current.inpaint.maskExpand, -256, 256)),
              maskFeather: Math.round(clampNumber(payload.inpaint?.maskFeather ?? profileDefaults?.inpaint.maskFeather ?? current.inpaint.maskFeather, 0, 256)),
              maskBlur: Math.round(clampNumber(payload.inpaint?.maskBlur ?? current.inpaint.maskBlur, 0, 256)),
              maskDespeckle: Math.round(clampNumber(payload.inpaint?.maskDespeckle ?? current.inpaint.maskDespeckle, 0, 64)),
              structureControl: payload.inpaint?.structureControl ?? profileDefaults?.inpaint.structureControl ?? current.inpaint.structureControl,
              controlStrength: clampNumber(payload.inpaint?.controlStrength ?? profileDefaults?.inpaint.controlStrength ?? current.inpaint.controlStrength, 0, 1),
              controlModel: payload.inpaint?.controlModel ?? current.inpaint.controlModel,
              colorMatch: payload.inpaint?.colorMatch ?? profileDefaults?.inpaint.colorMatch ?? current.inpaint.colorMatch,
            },
          },
        },
      },
    };
  },
};

const modeByGenerateIntent: Record<GenerateIntent, GenerateSettings["mode"]> = {
  create: "text-to-image",
  replace: "inpaint",
  remove: "inpaint",
  extend: "outpaint",
  variations: "image-to-image",
};

export const toolChooseGenerateIntentCommand: Command<ToolChooseGenerateIntentPayload> = {
  id: commandIds.toolChooseGenerateIntent,
  name: "Choose AI edit intent",
  execute({ state }, payload) {
    const mode = modeByGenerateIntent[payload.intent];
    const architecture = generateArchitectureDefaults[state.editor.tools.generate.architecture].supportedModes.includes(mode)
      ? state.editor.tools.generate.architecture
      : "sdxl";
    return toolSetGenerateSettingsCommand.execute({ state }, { architecture, mode, inpaint: payload.intent === "remove" ? { profile: "remove" } : payload.intent === "replace" && state.editor.tools.generate.inpaint.profile === "remove" ? { profile: "replace" } : undefined });
  },
};

export const toolSetBrushSettingsCommand: Command<ToolSetBrushSettingsPayload> = {
  id: commandIds.toolSetBrushSettings,
  name: "Set brush settings",
  execute({ state }, payload) {
    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          brush: {
            color: payload.color ?? state.editor.tools.brush.color,
            size: clampNumber(payload.size ?? state.editor.tools.brush.size, 1, 200),
            hardness: clampNumber(payload.hardness ?? state.editor.tools.brush.hardness, 0, 100),
            opacity: clampNumber(payload.opacity ?? state.editor.tools.brush.opacity, 0, 100),
            flow: clampNumber(payload.flow ?? state.editor.tools.brush.flow, 1, 100),
            smoothing: clampNumber(payload.smoothing ?? state.editor.tools.brush.smoothing, 0, 100),
            pressureSize: payload.pressureSize ?? state.editor.tools.brush.pressureSize,
          },
        },
      },
    };
  },
};

export const toolSetChromaKeySettingsCommand: Command<ToolSetChromaKeySettingsPayload> = {
  id: commandIds.toolSetChromaKeySettings,
  name: "Set chroma key settings",
  execute({ state }, payload) {
    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          chromaKey: {
            color: payload.color ?? state.editor.tools.chromaKey.color,
            tolerance: clampNumber(payload.tolerance ?? state.editor.tools.chromaKey.tolerance, 0, 255),
            softness: clampNumber(payload.softness ?? state.editor.tools.chromaKey.softness, 0, 255),
            feather: clampNumber(payload.feather ?? state.editor.tools.chromaKey.feather, 0, 20),
            choke: clampNumber(payload.choke ?? state.editor.tools.chromaKey.choke, -20, 20),
            despeckle: clampNumber(payload.despeckle ?? state.editor.tools.chromaKey.despeckle, 0, 20),
            spill: clampNumber(payload.spill ?? state.editor.tools.chromaKey.spill, 0, 100),
          },
        },
      },
    };
  },
};

export const toolSetMagicWandSettingsCommand: Command<ToolSetMagicWandSettingsPayload> = {
  id: commandIds.toolSetMagicWandSettings,
  name: "Set magic wand settings",
  execute({ state }, payload) {
    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          magicWand: {
            tolerance: clampNumber(payload.tolerance ?? state.editor.tools.magicWand.tolerance, 0, 255),
            feather: clampNumber(payload.feather ?? state.editor.tools.magicWand.feather, 0, 20),
            choke: clampNumber(payload.choke ?? state.editor.tools.magicWand.choke, -20, 20),
            despeckle: clampNumber(payload.despeckle ?? state.editor.tools.magicWand.despeckle, 0, 20),
            contiguous: payload.contiguous ?? state.editor.tools.magicWand.contiguous,
            mode: payload.mode ?? state.editor.tools.magicWand.mode,
          },
        },
      },
    };
  },
};

export const toolSetBrushPreviewCommand: Command<ToolSetBrushPreviewPayload> = {
  id: commandIds.toolSetBrushPreview,
  name: "Set brush preview",
  execute({ state }, payload) {
    if (!payload) {
      if (!state.editor.brushPreview) return state;
      return { ...state, editor: { ...state.editor, brushPreview: undefined } };
    }

    const currentPosition = state.editor.brushPreview?.position;
    if (currentPosition && currentPosition.x === payload.position.x && currentPosition.y === payload.position.y) return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        brushPreview: { position: { ...payload.position } },
      },
    };
  },
};

export const toolSetBrushStrokePreviewCommand: Command<ToolSetBrushStrokePreviewPayload> = {
  id: commandIds.toolSetBrushStrokePreview,
  name: "Set brush stroke preview",
  execute({ state }, payload) {
    if (!payload) {
      if (!state.editor.brushStrokePreview) return state;
      return { ...state, editor: { ...state.editor, brushStrokePreview: undefined } };
    }

    const currentPreview = state.editor.brushStrokePreview;
    if (currentPreview?.layerId === payload.layerId && currentPreview.assetId === payload.assetId && currentPreview.source === payload.source) return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        brushStrokePreview: { layerId: payload.layerId, assetId: payload.assetId, source: payload.source },
      },
    };
  },
};

export const toolSetMaskViewModeCommand: Command<ToolSetMaskViewModePayload> = {
  id: commandIds.toolSetMaskViewMode,
  name: "Set mask view mode",
  execute({ state }, payload) {
    if (!state.editor.maskEdit || state.editor.maskEdit.viewMode === payload.mode) return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        maskEdit: { ...state.editor.maskEdit, viewMode: payload.mode },
      },
    };
  },
};

export const toolBeginMaskShapeCommand: Command<ToolBeginMaskShapePayload> = {
  id: commandIds.toolBeginMaskShape,
  name: "Begin mask lasso",
  execute({ state }, payload) {
    const tool = state.editor.tools.activeTool;
    if ((tool !== "maskLasso" && tool !== "maskRectangle") || state.editor.maskEdit?.kind !== "inpaintRegion") return state;
    return { ...state, editor: { ...state.editor, maskShapeSession: { shape: tool === "maskRectangle" ? "rectangle" : "lasso", points: [{ ...payload.point }], mode: payload.mode } } };
  },
};

export const toolAppendMaskShapeCommand: Command<ToolAppendMaskShapePayload> = {
  id: commandIds.toolAppendMaskShape,
  name: "Append mask lasso point",
  history: { mode: "ignore" },
  execute({ state }, payload) {
    const session = state.editor.maskShapeSession;
    if (!session) return state;
    if (session.shape === "rectangle") {
      const start = session.points[0];
      if (!start) return state;
      return { ...state, editor: { ...state.editor, maskShapeSession: { ...session, points: [start, { ...payload.point }] } } };
    }
    const previous = session.points[session.points.length - 1];
    if (previous && Math.hypot(payload.point.x - previous.x, payload.point.y - previous.y) < 1) return state;
    return { ...state, editor: { ...state.editor, maskShapeSession: { ...session, points: [...session.points, { ...payload.point }] } } };
  },
};

export const toolClearMaskShapeCommand: Command = {
  id: commandIds.toolClearMaskShape,
  name: "Clear mask lasso",
  history: { mode: "ignore" },
  execute({ state }) {
    return state.editor.maskShapeSession ? { ...state, editor: { ...state.editor, maskShapeSession: undefined } } : state;
  },
};

export const toolEnterMaskEditCommand: Command<ToolEnterMaskEditPayload> = {
  id: commandIds.toolEnterMaskEdit,
  name: "Enter mask edit",
  execute({ state }, payload) {
    const targetLocation = findLayerLocation(state.document, payload.targetLayerId);
    const maskLocation = findLayerLocation(state.document, payload.maskLayerId);
    if (!targetLocation || !maskLocation) return state;
    if (getLayerMask(targetLocation.layer)?.maskLayerId !== payload.maskLayerId) return state;
    if (maskLocation.layer.type !== "image" && maskLocation.layer.type !== "raster") return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        selection: { artboardId: targetLocation.artboardId, layerIds: [payload.targetLayerId] },
        maskEdit: { kind: "layerMask", targetLayerId: payload.targetLayerId, maskLayerId: payload.maskLayerId, maskAssetId: maskLocation.layer.assetId },
        brushPreview: undefined,
        brushStrokePreview: undefined,
        tools: {
          ...state.editor.tools,
          activeTool: "brush",
          interactionMode: { type: "tool", tool: "brush" },
        },
      },
    };
  },
};

export const toolEnterInpaintRegionEditCommand: Command<ToolEnterInpaintRegionEditPayload> = {
  id: commandIds.toolEnterInpaintRegionEdit,
  name: "Enter inpaint region edit",
  execute({ state }, payload) {
    const targetLocation = findLayerLocation(state.document, payload.targetLayerId);
    const region = state.document.inpaintRegions.find((candidate) => candidate.id === payload.regionId && candidate.targetLayerId === payload.targetLayerId);
    const maskAsset = region ? state.document.assets.find((asset) => asset.id === region.maskAssetId) : undefined;
    if (!targetLocation || !region || !maskAsset) return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        selection: { artboardId: targetLocation.artboardId, layerIds: [payload.targetLayerId] },
        maskEdit: { kind: "inpaintRegion", targetLayerId: payload.targetLayerId, inpaintRegionId: region.id, maskAssetId: region.maskAssetId, viewMode: "overlay" },
        brushPreview: undefined,
        brushStrokePreview: undefined,
        tools: {
          ...state.editor.tools,
          activeTool: "brush",
          interactionMode: { type: "tool", tool: "brush" },
        },
      },
    };
  },
};

export const toolExitMaskEditCommand: Command = {
  id: commandIds.toolExitMaskEdit,
  name: "Exit mask edit",
  execute({ state }) {
    if (!state.editor.maskEdit) return state;
    const contextualTool = state.editor.tools.activeTool === "semanticSelect" || state.editor.tools.activeTool === "maskLasso" || state.editor.tools.activeTool === "maskRectangle";

    return {
      ...state,
      editor: {
        ...state.editor,
        maskEdit: undefined,
        brushPreview: undefined,
        brushStrokePreview: undefined,
        maskShapeSession: undefined,
        tools: contextualTool ? { ...state.editor.tools, activeTool: "select", interactionMode: { type: "tool", tool: "select" } } : state.editor.tools,
      },
    };
  },
};

export const toolEnterTemporaryPanCommand: Command = {
  id: commandIds.toolEnterTemporaryPan,
  name: "Enter temporary pan",
  execute({ state }) {
    if (state.editor.tools.interactionMode.type === "temporary-pan") return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          interactionMode: { type: "temporary-pan", previousTool: state.editor.tools.activeTool },
        },
      },
    };
  },
};

export const toolExitTemporaryPanCommand: Command = {
  id: commandIds.toolExitTemporaryPan,
  name: "Exit temporary pan",
  execute({ state }) {
    const mode = state.editor.tools.interactionMode;
    if (mode.type !== "temporary-pan") return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        tools: {
          ...state.editor.tools,
          activeTool: mode.previousTool,
          interactionMode: { type: "tool", tool: mode.previousTool },
        },
      },
    };
  },
};

export const toolCommands = [
  toolSetActiveCommand,
  toolSetGenerateSettingsCommand,
  toolChooseGenerateIntentCommand,
  toolSetBrushSettingsCommand,
  toolSetChromaKeySettingsCommand,
  toolSetMagicWandSettingsCommand,
  toolSetBrushPreviewCommand,
  toolSetBrushStrokePreviewCommand,
  toolSetMaskViewModeCommand,
  toolBeginMaskShapeCommand,
  toolAppendMaskShapeCommand,
  toolClearMaskShapeCommand,
  toolEnterMaskEditCommand,
  toolEnterInpaintRegionEditCommand,
  toolExitMaskEditCommand,
  toolEnterTemporaryPanCommand,
  toolExitTemporaryPanCommand,
] satisfies Command<unknown>[];

function clampNumber(value: number, min: number, max: number) {
  if (!Number.isFinite(value)) return min;
  return Math.max(min, Math.min(max, value));
}

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

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;
}