image-studio/commands/tool.ts

import type { ImageDocument } from "@core/document";
import type { Vec2D } from "@core/geometry";
import type { LayerId, ArtboardId, AssetId } from "@core/id";
import type { Layer } from "@core/layer";
import { getLayerMask } from "@core/layer-mask-utils";
import type { MaskViewMode } from "@editor/state";
import { generateArchitectureDefaults } from "@editor/tools";
import type { BrushSettings, ChromaKeySettings, 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 = Partial<GenerateSettings>;

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 ToolEnterMaskEditPayload = {
  targetLayerId: LayerId;
  maskLayerId: LayerId;
};

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

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";
    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 ?? 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)),
            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: {
              maskedAreaOnly: payload.inpaint?.maskedAreaOnly ?? current.inpaint.maskedAreaOnly,
              cropPadding: Math.round(clampNumber(payload.inpaint?.cropPadding ?? current.inpaint.cropPadding, 0, 2048)),
              maskPolarity: payload.inpaint?.maskPolarity ?? current.inpaint.maskPolarity,
              maskedContent: payload.inpaint?.maskedContent ?? current.inpaint.maskedContent,
              growMaskBy: Math.round(clampNumber(payload.inpaint?.growMaskBy ?? current.inpaint.growMaskBy, 0, 256)),
              maskExpand: Math.round(clampNumber(payload.inpaint?.maskExpand ?? current.inpaint.maskExpand, -256, 256)),
              maskFeather: Math.round(clampNumber(payload.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)),
            },
          },
        },
      },
    };
  },
};

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

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 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 === "group") return state;

    return {
      ...state,
      editor: {
        ...state.editor,
        selection: { artboardId: targetLocation.artboardId, layerIds: [payload.targetLayerId] },
        maskEdit: { targetLayerId: payload.targetLayerId, maskLayerId: payload.maskLayerId },
        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;

    return {
      ...state,
      editor: {
        ...state.editor,
        maskEdit: undefined,
        brushPreview: undefined,
        brushStrokePreview: undefined,
      },
    };
  },
};

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,
  toolSetBrushSettingsCommand,
  toolSetChromaKeySettingsCommand,
  toolSetMagicWandSettingsCommand,
  toolSetBrushPreviewCommand,
  toolSetBrushStrokePreviewCommand,
  toolSetMaskViewModeCommand,
  toolEnterMaskEditCommand,
  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;
}