image-studio/operations/generation/runGenerate.ts

import { commandIds } from "@commands/ids";
import type { ImageDocument } from "@core/document";
import type { Layer } from "@core/layer";
import { getLayerMask } from "@core/layer-mask-utils";
import type { AppStore } from "@editor/store";
import type { GenerationCandidate, SelectionState, ViewportState } from "@editor/state";
import type { GenerateSettings } from "@editor/tools";
import { buildInpaintBundle, type InpaintBundle } from "./inpaintPrep";
import { imageSourceToDataUrl, loadImageSize } from "@platform/browser/imageRaster";
import { requestGeneration } from "@platform/comfy/generationClient";
import { checkGenerationPreconditions } from "./preconditions";
import { resolveGeneratedOutputPlacement } from "./outputPlacement";

export async function runGenerate(options: {
  document: ImageDocument;
  selection: SelectionState;
  viewport: ViewportState;
  settings: GenerateSettings;
  dispatch: AppStore["dispatch"];
  signal?: AbortSignal;
  onProgress?: (progress: number, detail: string) => void;
}) {
  const { document, selection, settings, dispatch } = options;
  const precondition = checkGenerationPreconditions(document, selection, settings);
  if (!precondition.ready) throw new Error(precondition.message);
  const artboard = selection.artboardId ? document.artboards.find((candidate) => candidate.id === selection.artboardId) : document.artboards[0];
  if (!artboard) throw new Error("Create an artboard before generating an image.");

  const target = resolveSelectedImage(document, selection);
  const inpaintBundle = settings.mode === "inpaint" ? await buildInpaintBundle(document, selection, settings) : undefined;
  const inputImage = inpaintBundle?.inputImage ?? (target && settings.mode !== "text-to-image" ? await imageSourceToDataUrl(target.asset.source) : undefined);
  const maskImage = inpaintBundle?.maskImage;
  const seed = resolveSeed(settings.seed);
  const requestSettings = { ...settings, seed };
  const width = inpaintBundle?.width ?? settings.width;
  const height = inpaintBundle?.height ?? settings.height;
  const generated = await requestGenerate({
    settings: requestSettings,
    width,
    height,
    inputImage,
    maskImage,
    inpaintBundle,
    signal: options.signal,
    onProgress: options.onProgress,
  });
  for (const result of generated.results) {
    const intrinsicSize = await loadImageSize(result.source);
    const placement = resolveGeneratedOutputPlacement({ document, selection, settings, intrinsicSize, inpaintBundle });
    dispatch(commandIds.generationAddCandidate, {
      candidate: createGenerationCandidate({
        source: result.source,
        mimeType: result.mimeType,
        intrinsicSize,
        settings: requestSettings,
        seed: result.seed || seed,
        width,
        height,
        inputImage,
        maskImage,
        placement,
        inpaintBundle,
      }),
    });
  }
}

export async function runGenerateFromCandidate(options: {
  candidate: GenerationCandidate;
  settings?: GenerateSettings;
  dispatch: AppStore["dispatch"];
  signal?: AbortSignal;
  onProgress?: (progress: number, detail: string) => void;
}) {
  const settings = options.settings ?? options.candidate.settings;
  const seed = resolveSeed(settings.seed);
  const requestSettings = { ...settings, seed };
  const generated = await requestGenerate({
    settings: requestSettings,
    width: options.candidate.width,
    height: options.candidate.height,
    inputImage: options.candidate.inputImage,
    maskImage: options.candidate.maskImage,
    inpaintCandidate: options.candidate,
    signal: options.signal,
    onProgress: options.onProgress,
  });
  for (const result of generated.results) {
    const intrinsicSize = await loadImageSize(result.source);
    options.dispatch(commandIds.generationAddCandidate, {
      candidate: {
        ...options.candidate,
        id: crypto.randomUUID(),
        source: result.source,
        mimeType: result.mimeType,
        intrinsicSize,
        settings: requestSettings,
        seed: result.seed || seed,
      },
    });
  }
}

function createGenerationCandidate(options: {
  source: string;
  mimeType: string;
  intrinsicSize: { w: number; h: number };
  settings: GenerateSettings;
  seed: number;
  width: number;
  height: number;
  inputImage?: string;
  maskImage?: string;
  placement: GenerationCandidate["placement"];
  inpaintBundle?: InpaintBundle;
}): GenerationCandidate {
  return {
    id: crypto.randomUUID(),
    source: options.source,
    mimeType: options.mimeType,
    intrinsicSize: options.intrinsicSize,
    mode: options.settings.mode,
    settings: options.settings,
    seed: options.seed,
    width: options.width,
    height: options.height,
    inputImage: options.inputImage,
    maskImage: options.maskImage,
    blendMaskImage: options.inpaintBundle?.blendMaskImage,
    placement: options.placement,
    inpaint: options.inpaintBundle
      ? {
          targetLayerId: options.inpaintBundle.targetLayerId,
          regionId: options.inpaintBundle.regionId,
          sourceAssetId: options.inpaintBundle.sourceAssetId,
          maskAssetId: options.inpaintBundle.maskAssetId,
          inputImage: options.inpaintBundle.inputImage,
          maskImage: options.inpaintBundle.maskImage,
          editMaskImage: options.inpaintBundle.editMaskImage,
          blendMaskImage: options.inpaintBundle.blendMaskImage,
          revision: options.inpaintBundle.revision,
          crop: options.inpaintBundle.crop,
          mask: options.inpaintBundle.mask,
          backend: options.inpaintBundle.backend,
        }
      : undefined,
  };
}

async function requestGenerate(options: {
  settings: GenerateSettings;
  width: number;
  height: number;
  inputImage?: string;
  maskImage?: string;
  inpaintBundle?: InpaintBundle;
  inpaintCandidate?: GenerationCandidate;
  signal?: AbortSignal;
  onProgress?: (progress: number, detail: string) => void;
}) {
  return requestGeneration({
      architecture: options.settings.architecture,
      mode: options.settings.mode,
      model: options.settings.model,
      textEncoder: options.settings.textEncoder,
      vae: options.settings.vae,
      prompt: options.settings.prompt,
      negativePrompt: options.settings.negativePrompt,
      strength: options.settings.strength,
      steps: options.settings.steps,
      cfg: options.settings.cfg,
      seed: options.settings.seed,
      sampler: options.settings.sampler,
      scheduler: options.settings.scheduler,
      width: options.width,
      height: options.height,
      batchSize: options.settings.batchSize,
      refinePass: options.settings.refinePass,
      refineStrength: options.settings.refineStrength,
      outpaint: options.settings.outpaint,
      inpaint: resolveInpaintRequest(options.inpaintBundle, options.inpaintCandidate, options.settings),
      inputImage: options.inputImage,
      maskImage: options.maskImage,
  }, options.signal, options.onProgress);
}

function resolveInpaintRequest(inpaintBundle: InpaintBundle | undefined, inpaintCandidate: GenerationCandidate | undefined, settings: GenerateSettings) {
  if (inpaintBundle) {
    return {
      growMaskBy: inpaintBundle.backend.growMaskBy,
      maskedContent: inpaintBundle.backend.maskedContent,
      maskBlur: inpaintBundle.backend.maskBlur,
      maskFeather: inpaintBundle.backend.maskFeather,
      maskExpand: inpaintBundle.backend.maskExpand,
      cropPadding: inpaintBundle.backend.cropPadding,
      maskPolarity: inpaintBundle.mask.polarity,
      crop: inpaintBundle.crop,
      placement: inpaintBundle.placement,
      structureControl: settings.inpaint.structureControl,
      controlStrength: settings.inpaint.controlStrength,
      controlModel: settings.inpaint.controlModel,
    };
  }

  if (inpaintCandidate?.inpaint) {
    return {
      growMaskBy: inpaintCandidate.inpaint.backend.growMaskBy,
      maskedContent: inpaintCandidate.inpaint.backend.maskedContent,
      maskBlur: inpaintCandidate.inpaint.backend.maskBlur,
      maskFeather: inpaintCandidate.inpaint.backend.maskFeather,
      maskExpand: inpaintCandidate.inpaint.backend.maskExpand,
      cropPadding: inpaintCandidate.inpaint.backend.cropPadding,
      maskPolarity: inpaintCandidate.inpaint.mask.polarity,
      crop: inpaintCandidate.inpaint.crop,
      placement: inpaintCandidate.placement,
      structureControl: settings.inpaint.structureControl,
      controlStrength: settings.inpaint.controlStrength,
      controlModel: settings.inpaint.controlModel,
    };
  }

  return settings.inpaint;
}

function resolveSeed(seed: number): number {
  return seed < 0 ? Math.floor(Math.random() * 2 ** 32) : Math.round(seed);
}

function resolveSelectedImage(document: ImageDocument, selection: SelectionState) {
  const layerId = selection.layerIds[0];
  if (!layerId) return undefined;
  const layer = findLayer(document.artboards.find((artboard) => artboard.id === selection.artboardId)?.layers ?? [], layerId);
  if (!layer || (layer.type !== "image" && layer.type !== "raster")) return undefined;
  const asset = document.assets.find((candidate) => candidate.id === layer.assetId);
  const layerMask = getLayerMask(layer);
  const maskLayer = layerMask?.enabled ? findLayer(document.artboards.flatMap((artboard) => artboard.layers), layerMask.maskLayerId) : undefined;
  const maskAsset = maskLayer && (maskLayer.type === "image" || maskLayer.type === "raster") ? document.assets.find((candidate) => candidate.id === maskLayer.assetId) : undefined;
  return asset ? { layer, asset, maskAsset } : undefined;
}

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 found = findLayer(layer.children, layerId);
      if (found) return found;
    }
  }
  return undefined;
}