image-studio/operations/generation/inpaintPrep.ts

import type { Asset } from "@core/asset";
import type { ImageDocument } from "@core/document";
import type { Rect } from "@core/geometry";
import type { AssetId, InpaintRegionId, LayerId } from "@core/id";
import type { Layer } from "@core/layer";
import type { SelectionState } from "@editor/state";
import type { GenerateSettings } from "@editor/tools";
import { createDocumentReadIndex, resolveIndexedLayerBounds } from "@editor/document-indexes";
import { createNormalizedMaskSource, cropCanvas, cropMaskValuesToDataUrl, expandRectWithinBounds, loadImageCanvas, sampleDocumentCanvasInLayerSpace } from "@platform/browser/maskRaster";
import { renderArtboardCanvas } from "@platform/browser/exportArtboardPng";

export type InpaintBundle = {
  inputImage: string;
  sourceImage: string;
  contextImage: string;
  maskImage: string;
  editMaskImage: string;
  blendMaskImage: string;
  width: number;
  height: number;
  targetLayerId: LayerId;
  regionId: InpaintRegionId;
  sourceAssetId: AssetId;
  maskAssetId: AssetId;
  revision: { source: string; mask: string };
  crop: {
    assetBounds: Rect;
    documentBounds: Rect;
    padding: number;
    maskedAreaOnly: boolean;
  };
  mask: {
    polarity: GenerateSettings["inpaint"]["maskPolarity"];
    activeBounds: Rect;
  };
  placement: {
    artboardId: string;
    layerName: string;
    transform: Extract<Layer, { type: "image" | "raster" }>["transform"];
  };
  backend: {
    growMaskBy: number;
    maskedContent: GenerateSettings["inpaint"]["maskedContent"];
    maskBlur: number;
    maskFeather: number;
    maskExpand: number;
    cropPadding: number;
  };
};

type InpaintTarget = {
  artboardId: string;
  layer: Extract<Layer, { type: "image" | "raster" }>;
  asset: Asset;
  maskAsset: Asset;
  regionId: InpaintRegionId;
};

const modelMultiple = 8;
const minModelSize = 64;
const maxModelSize = 4096;

export async function buildInpaintBundle(document: ImageDocument, selection: SelectionState, settings: GenerateSettings): Promise<InpaintBundle> {
  const target = resolveInpaintTarget(document, selection);
  validateInpaintTarget(target);

  const width = Math.max(1, Math.round(target.asset.intrinsicSize.w));
  const height = Math.max(1, Math.round(target.asset.intrinsicSize.h));
  const sourceLimit = target.layer.sourceRect ?? { x: 0, y: 0, w: width, h: height };
  const editMask = await createNormalizedMaskSource(target.maskAsset.source, width, height, {
    polarity: "revealed",
    despeckle: settings.inpaint.maskDespeckle,
    limit: sourceLimit,
  });
  const noiseMask = await createNormalizedMaskSource(target.maskAsset.source, width, height, {
    polarity: "revealed",
    expand: settings.inpaint.maskExpand,
    blur: settings.inpaint.maskBlur,
    despeckle: settings.inpaint.maskDespeckle,
    limit: sourceLimit,
  });
  const blendMask = await createNormalizedMaskSource(target.maskAsset.source, width, height, {
    polarity: "revealed",
    feather: settings.inpaint.maskFeather,
    despeckle: settings.inpaint.maskDespeckle,
    limit: sourceLimit,
  });

  if (!editMask.bounds || !noiseMask.bounds) throw new Error("Paint over the area you want AI to replace.");

  const crop = settings.inpaint.maskedAreaOnly
    ? expandRectWithinBounds(noiseMask.bounds, settings.inpaint.cropPadding, { w: width, h: height }, modelMultiple, minModelSize)
    : { x: 0, y: 0, w: width, h: height };
  const outputWidth = toModelSize(crop.w, "width");
  const outputHeight = toModelSize(crop.h, "height");

  const sourceCanvas = await loadImageCanvas(target.asset.source, width, height);
  const sourceImage = cropCanvas(sourceCanvas, crop, outputWidth, outputHeight);
  const preparedCanvas = prepareMaskedContentInputCanvas(await loadImageCanvas(target.asset.source, width, height), editMask.values, settings.inpaint.maskedContent);
  const preparedSourceImage = cropCanvas(preparedCanvas, crop, outputWidth, outputHeight);
  const artboard = document.artboards.find((candidate) => candidate.id === target.artboardId);
  if (!artboard) throw new Error("The target artboard no longer exists.");
  const contextImage = sampleDocumentCanvasInLayerSpace(await renderArtboardCanvas(artboard, document.assets), artboard.bounds, target.layer, target.asset.intrinsicSize, crop, outputWidth, outputHeight);
  const inputImage = await mergePreparedRegionIntoContext(contextImage, preparedSourceImage, cropMaskValuesToDataUrl(editMask.values, width, height, crop, outputWidth, outputHeight), outputWidth, outputHeight);
  const maskImage = cropMaskValuesToDataUrl(noiseMask.values, width, height, crop, outputWidth, outputHeight);
  const editMaskImage = cropMaskValuesToDataUrl(editMask.values, width, height, crop, outputWidth, outputHeight);
  const blendMaskImage = cropMaskValuesToDataUrl(blendMask.values, width, height, crop, outputWidth, outputHeight);
  const scaleX = target.layer.transform.scale.x;
  const scaleY = target.layer.transform.scale.y;
  const documentBounds = {
    x: target.layer.transform.position.x + crop.x * scaleX,
    y: target.layer.transform.position.y + crop.y * scaleY,
    w: outputWidth * scaleX,
    h: outputHeight * scaleY,
  };

  return {
    inputImage,
    sourceImage,
    contextImage,
    maskImage,
    editMaskImage,
    blendMaskImage,
    width: outputWidth,
    height: outputHeight,
    targetLayerId: target.layer.id,
    regionId: target.regionId,
    sourceAssetId: target.asset.id,
    maskAssetId: target.maskAsset.id,
    crop: {
      assetBounds: crop,
      documentBounds,
      padding: settings.inpaint.cropPadding,
      maskedAreaOnly: settings.inpaint.maskedAreaOnly,
    },
    mask: {
      polarity: "revealed",
      activeBounds: editMask.bounds,
    },
    placement: {
      artboardId: target.artboardId,
      layerName: `${target.layer.name} inpaint`,
      transform: {
        position: rotatedCropPosition(target.layer, target.asset, { ...crop, w: outputWidth, h: outputHeight }),
        scale: { x: documentBounds.w / outputWidth, y: documentBounds.h / outputHeight },
        rotation: target.layer.transform.rotation,
      },
    },
    backend: {
      growMaskBy: settings.inpaint.growMaskBy,
      maskedContent: settings.inpaint.maskedContent,
      maskBlur: settings.inpaint.maskBlur,
      maskFeather: settings.inpaint.maskFeather,
      maskExpand: settings.inpaint.maskExpand,
      cropPadding: settings.inpaint.cropPadding,
    },
    revision: {
      source: await createContentRevision(target.asset.source),
      mask: await createContentRevision(target.maskAsset.source),
    },
  };
}

function resolveInpaintTarget(document: ImageDocument, selection: SelectionState): InpaintTarget {
  if (selection.layerIds.length !== 1 || !selection.layerIds[0]) throw new Error("Select one image or raster layer to inpaint.");

  const documentIndex = createDocumentReadIndex(document);
  const layerInfo = documentIndex.layerInfoById.get(selection.layerIds[0]);
  if (!layerInfo || (layerInfo.layer.type !== "image" && layerInfo.layer.type !== "raster")) throw new Error("Select one image or raster layer to inpaint. Text must be rasterized first.");

  const asset = documentIndex.assetById.get(layerInfo.layer.assetId);
  if (!asset) throw new Error("The selected layer is missing its source image.");
  const region = document.inpaintRegions.find((candidate) => candidate.targetLayerId === layerInfo.layer.id && candidate.enabled);
  if (!region) throw new Error("Add an AI edit region before running inpaint.");
  const maskAsset = documentIndex.assetById.get(region.maskAssetId);
  if (!maskAsset) throw new Error("The AI edit region is missing its mask data.");

  if (!resolveIndexedLayerBounds(documentIndex, layerInfo.layer)) throw new Error("Unable to resolve the selected layer bounds.");

  return { artboardId: layerInfo.artboardId, layer: layerInfo.layer, asset, maskAsset, regionId: region.id };
}

function validateInpaintTarget(target: InpaintTarget) {
  if (target.asset.intrinsicSize.w <= 0 || target.asset.intrinsicSize.h <= 0) throw new Error("The selected image has an invalid size.");
  if (target.maskAsset.intrinsicSize.w <= 0 || target.maskAsset.intrinsicSize.h <= 0) throw new Error("The selected mask has an invalid size.");
  if (Math.round(target.asset.intrinsicSize.w) !== Math.round(target.maskAsset.intrinsicSize.w) || Math.round(target.asset.intrinsicSize.h) !== Math.round(target.maskAsset.intrinsicSize.h)) {
    throw new Error("The selected layer and mask image sizes do not match.");
  }
}

function toModelSize(value: number, axis: "width" | "height"): number {
  const rounded = Math.max(minModelSize, Math.ceil(value / modelMultiple) * modelMultiple);
  if (rounded > maxModelSize) throw new Error(`The inpaint ${axis} is too large for the model. Use masked-area inpaint or a smaller source.`);
  return rounded;
}

function prepareMaskedContentInputCanvas(sourceCanvas: HTMLCanvasElement, maskValues: Uint8ClampedArray, mode: GenerateSettings["inpaint"]["maskedContent"]): HTMLCanvasElement {
  if (mode === "neutral" || mode === "originalColor") return sourceCanvas;
  const context = sourceCanvas.getContext("2d");
  if (!context) return sourceCanvas;
  const imageData = context.getImageData(0, 0, sourceCanvas.width, sourceCanvas.height);
  imageData.data.set(applyMaskedContentModeToRgba(imageData.data, sourceCanvas.width, sourceCanvas.height, maskValues, mode));
  context.putImageData(imageData, 0, 0);
  return sourceCanvas;
}

export function applyMaskedContentModeToRgba(data: Uint8ClampedArray, width: number, height: number, maskValues: Uint8ClampedArray, mode: GenerateSettings["inpaint"]["maskedContent"]): Uint8ClampedArray {
  const next = new Uint8ClampedArray(data);
  if (mode === "neutral" || mode === "originalColor") return next;

  const edgeValues = mode === "edges" ? createEdgeMapValues(data, width, height) : undefined;

  for (let pixel = 0; pixel < width * height; pixel += 1) {
    const amount = (maskValues[pixel] ?? 0) / 255;
    if (amount <= 0) continue;

    const index = pixel * 4;
    const red = data[index] ?? 0;
    const green = data[index + 1] ?? 0;
    const blue = data[index + 2] ?? 0;
    const target = edgeValues ? edgeValues[pixel] ?? 128 : luminance(red, green, blue);
    next[index] = blendChannel(red, target, amount);
    next[index + 1] = blendChannel(green, target, amount);
    next[index + 2] = blendChannel(blue, target, amount);
  }

  return next;
}

function createEdgeMapValues(data: Uint8ClampedArray, width: number, height: number): Uint8ClampedArray {
  const gray = new Float32Array(width * height);
  const edges = new Uint8ClampedArray(width * height);

  for (let pixel = 0; pixel < width * height; pixel += 1) {
    const index = pixel * 4;
    gray[pixel] = luminance(data[index] ?? 0, data[index + 1] ?? 0, data[index + 2] ?? 0);
  }

  for (let y = 0; y < height; y += 1) {
    for (let x = 0; x < width; x += 1) {
      const gx =
        -sampleGray(gray, width, height, x - 1, y - 1) +
        sampleGray(gray, width, height, x + 1, y - 1) -
        2 * sampleGray(gray, width, height, x - 1, y) +
        2 * sampleGray(gray, width, height, x + 1, y) -
        sampleGray(gray, width, height, x - 1, y + 1) +
        sampleGray(gray, width, height, x + 1, y + 1);
      const gy =
        -sampleGray(gray, width, height, x - 1, y - 1) -
        2 * sampleGray(gray, width, height, x, y - 1) -
        sampleGray(gray, width, height, x + 1, y - 1) +
        sampleGray(gray, width, height, x - 1, y + 1) +
        2 * sampleGray(gray, width, height, x, y + 1) +
        sampleGray(gray, width, height, x + 1, y + 1);
      const magnitude = Math.hypot(gx, gy);
      edges[y * width + x] = Math.round(clampNumber(128 + Math.max(0, magnitude - 24) * 0.75, 128, 255));
    }
  }

  return edges;
}

function sampleGray(values: Float32Array, width: number, height: number, x: number, y: number): number {
  const clampedX = Math.max(0, Math.min(width - 1, x));
  const clampedY = Math.max(0, Math.min(height - 1, y));
  return values[clampedY * width + clampedX] ?? 0;
}

function luminance(red: number, green: number, blue: number): number {
  return Math.round(0.2126 * red + 0.7152 * green + 0.0722 * blue);
}

function blendChannel(previous: number, next: number, amount: number): number {
  return Math.round(previous * (1 - amount) + next * amount);
}

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

async function mergePreparedRegionIntoContext(contextSource: string, preparedSource: string, editMaskSource: string, width: number, height: number): Promise<string> {
  const [contextCanvas, preparedCanvas, maskCanvas] = await Promise.all([
    loadImageCanvas(contextSource, width, height),
    loadImageCanvas(preparedSource, width, height),
    loadImageCanvas(editMaskSource, width, height),
  ]);
  const context = contextCanvas.getContext("2d");
  const prepared = preparedCanvas.getContext("2d");
  const mask = maskCanvas.getContext("2d");
  if (!context || !prepared || !mask) return contextSource;
  const contextData = context.getImageData(0, 0, width, height);
  const preparedData = prepared.getImageData(0, 0, width, height);
  const maskData = mask.getImageData(0, 0, width, height);
  for (let pixel = 0; pixel < width * height; pixel += 1) {
    const amount = (maskData.data[pixel * 4] ?? 0) / 255;
    if (amount <= 0) continue;
    for (let channel = 0; channel < 4; channel += 1) {
      const index = pixel * 4 + channel;
      contextData.data[index] = blendChannel(contextData.data[index] ?? 0, preparedData.data[index] ?? 0, amount);
    }
  }
  context.putImageData(contextData, 0, 0);
  return contextCanvas.toDataURL("image/png");
}

function rotatedCropPosition(layer: InpaintTarget["layer"], asset: Asset, crop: Rect) {
  const source = layer.sourceRect ?? { x: 0, y: 0, ...asset.intrinsicSize };
  const scale = layer.transform.scale;
  const originalCenter = {
    x: layer.transform.position.x + (source.x + source.w / 2) * scale.x,
    y: layer.transform.position.y + (source.y + source.h / 2) * scale.y,
  };
  const cropCenter = {
    x: layer.transform.position.x + (crop.x + crop.w / 2) * scale.x,
    y: layer.transform.position.y + (crop.y + crop.h / 2) * scale.y,
  };
  const dx = cropCenter.x - originalCenter.x;
  const dy = cropCenter.y - originalCenter.y;
  const cos = Math.cos(layer.transform.rotation);
  const sin = Math.sin(layer.transform.rotation);
  const rotatedCenter = { x: originalCenter.x + dx * cos - dy * sin, y: originalCenter.y + dx * sin + dy * cos };
  return { x: rotatedCenter.x - crop.w * scale.x / 2, y: rotatedCenter.y - crop.h * scale.y / 2 };
}

export async function createContentRevision(source: string): Promise<string> {
  const digest = await crypto.subtle.digest("SHA-256", new TextEncoder().encode(source));
  return Array.from(new Uint8Array(digest), (byte) => byte.toString(16).padStart(2, "0")).join("");
}