image-studio/view/generate/runGenerate.ts

import { commandIds } from "@commands/ids";
import type { ImageDocument } from "@core/document";
import type { Layer } from "@core/layer";
import type { AppStore } from "@editor/store";
import type { SelectionState, ViewportState } from "@editor/state";
import type { GenerateSettings } from "@editor/tools";

export async function runGenerate(options: {
  document: ImageDocument;
  selection: SelectionState;
  viewport: ViewportState;
  settings: GenerateSettings;
  dispatch: AppStore["dispatch"];
}) {
  const { document, selection, viewport, settings, dispatch } = options;
  const artboard = selection.artboardId ? document.artboards.find((candidate) => candidate.id === selection.artboardId) : document.artboards[0];
  if (!artboard) return;

  const target = resolveSelectedImage(document, selection);
  const inputImage = target && settings.mode !== "text-to-image" ? await imageSourceToDataUrl(target.asset.source) : undefined;
  const maskImage = target?.maskAsset && settings.mode === "inpaint" ? await imageSourceToDataUrl(target.maskAsset.source) : undefined;
  const response = await fetch("/api/comfy/generate", {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({
      mode: settings.mode,
      model: settings.model,
      prompt: settings.prompt,
      negativePrompt: settings.negativePrompt,
      strength: settings.strength,
      steps: settings.steps,
      cfg: settings.cfg,
      seed: settings.seed,
      sampler: settings.sampler,
      scheduler: settings.scheduler,
      width: settings.width,
      height: settings.height,
      outpaint: settings.outpaint,
      inputImage,
      maskImage,
    }),
  });
  if (!response.ok) throw new Error(await response.text());
  const generated = await response.json() as { source: string; mimeType: string };
  const intrinsicSize = await loadImageSize(generated.source);
  const assetId = crypto.randomUUID();
  const layerId = crypto.randomUUID();
  dispatch(commandIds.documentAddAsset, { asset: { id: assetId, name: "Generated image", mimeType: generated.mimeType, source: generated.source, intrinsicSize } });
  dispatch(commandIds.documentAddImageLayer, {
    artboardId: artboard.id,
    layer: {
      id: layerId,
      type: "image",
      name: "Generated image",
      visible: true,
      locked: false,
      opacity: 1,
      assetId,
      transform: { position: { x: 0, y: 0 }, scale: { x: 1, y: 1 }, rotation: 0 },
    },
  });
  dispatch(commandIds.documentMoveLayer, { layerId, toArtboardId: artboard.id, toIndex: 0 });
  dispatch(commandIds.selectionSet, { artboardId: artboard.id, layerIds: [layerId] });
}

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 === "group") return undefined;
  const asset = document.assets.find((candidate) => candidate.id === layer.assetId);
  const maskLayer = layer.clippingMask ? findLayer(document.artboards.flatMap((artboard) => artboard.layers), layer.clippingMask.maskLayerId) : undefined;
  const maskAsset = maskLayer && maskLayer.type !== "group" ? 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;
}

async function imageSourceToDataUrl(source: string) {
  if (source.startsWith("data:")) return source;
  const image = await loadImage(source);
  const canvas = document.createElement("canvas");
  canvas.width = image.naturalWidth;
  canvas.height = image.naturalHeight;
  const context = canvas.getContext("2d");
  if (!context) throw new Error("Unable to read selected image");
  context.drawImage(image, 0, 0);
  return canvas.toDataURL("image/png");
}

function loadImageSize(source: string): Promise<{ w: number; h: number }> {
  return loadImage(source).then((image) => ({ w: image.naturalWidth, h: image.naturalHeight }));
}

function loadImage(source: string): Promise<HTMLImageElement> {
  return new Promise((resolve, reject) => {
    const image = new Image();
    image.onload = () => resolve(image);
    image.onerror = () => reject(new Error("Failed to load image"));
    image.src = source;
  });
}