267 lines
9.7 KiB
TypeScript
267 lines
9.7 KiB
TypeScript
import { commandIds } from "@commands/ids";
|
|
import type { ImageDocument } from "@core/document";
|
|
import type { Transform } from "@core/geometry";
|
|
import type { Layer } from "@core/layer";
|
|
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";
|
|
|
|
export async function runGenerate(options: {
|
|
document: ImageDocument;
|
|
selection: SelectionState;
|
|
viewport: ViewportState;
|
|
settings: GenerateSettings;
|
|
dispatch: AppStore["dispatch"];
|
|
}) {
|
|
const { document, selection, 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 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,
|
|
});
|
|
const intrinsicSize = await loadImageSize(generated.source);
|
|
const targetArtboardId = inpaintBundle?.placement.artboardId ?? artboard.id;
|
|
const placement = {
|
|
artboardId: targetArtboardId,
|
|
layerName: inpaintBundle?.placement.layerName ?? "Generated image",
|
|
transform: generatedLayerTransform(inpaintBundle, intrinsicSize),
|
|
};
|
|
|
|
dispatch(commandIds.generationAddCandidate, {
|
|
candidate: createGenerationCandidate({
|
|
source: generated.source,
|
|
mimeType: generated.mimeType,
|
|
intrinsicSize,
|
|
settings: requestSettings,
|
|
seed,
|
|
width,
|
|
height,
|
|
inputImage,
|
|
maskImage,
|
|
placement,
|
|
inpaintBundle,
|
|
}),
|
|
});
|
|
}
|
|
|
|
export async function runGenerateFromCandidate(options: {
|
|
candidate: GenerationCandidate;
|
|
settings?: GenerateSettings;
|
|
dispatch: AppStore["dispatch"];
|
|
}) {
|
|
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,
|
|
});
|
|
const intrinsicSize = await loadImageSize(generated.source);
|
|
|
|
options.dispatch(commandIds.generationAddCandidate, {
|
|
candidate: {
|
|
...options.candidate,
|
|
id: crypto.randomUUID(),
|
|
source: generated.source,
|
|
mimeType: generated.mimeType,
|
|
intrinsicSize,
|
|
settings: requestSettings,
|
|
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,
|
|
placement: options.placement,
|
|
inpaint: options.inpaintBundle
|
|
? {
|
|
targetLayerId: options.inpaintBundle.targetLayerId,
|
|
maskLayerId: options.inpaintBundle.maskLayerId,
|
|
sourceAssetId: options.inpaintBundle.sourceAssetId,
|
|
maskAssetId: options.inpaintBundle.maskAssetId,
|
|
inputImage: options.inpaintBundle.inputImage,
|
|
maskImage: options.inpaintBundle.maskImage,
|
|
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;
|
|
}) {
|
|
const response = await fetch("/api/comfy/generate", {
|
|
method: "POST",
|
|
headers: { "content-type": "application/json" },
|
|
body: JSON.stringify({
|
|
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,
|
|
outpaint: options.settings.outpaint,
|
|
inpaint: resolveInpaintRequest(options.inpaintBundle, options.inpaintCandidate, options.settings),
|
|
inputImage: options.inputImage,
|
|
maskImage: options.maskImage,
|
|
}),
|
|
});
|
|
if (!response.ok) throw new Error(await response.text());
|
|
return await response.json() as { source: string; mimeType: string };
|
|
}
|
|
|
|
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,
|
|
};
|
|
}
|
|
|
|
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,
|
|
};
|
|
}
|
|
|
|
return settings.inpaint;
|
|
}
|
|
|
|
function generatedLayerTransform(inpaintBundle: InpaintBundle | undefined, intrinsicSize: { w: number; h: number }): Transform {
|
|
if (!inpaintBundle) return { position: { x: 0, y: 0 }, scale: { x: 1, y: 1 }, rotation: 0 };
|
|
return {
|
|
position: { ...inpaintBundle.placement.transform.position },
|
|
scale: {
|
|
x: (inpaintBundle.placement.transform.scale.x * inpaintBundle.width) / Math.max(1, intrinsicSize.w),
|
|
y: (inpaintBundle.placement.transform.scale.y * inpaintBundle.height) / Math.max(1, intrinsicSize.h),
|
|
},
|
|
rotation: inpaintBundle.placement.transform.rotation,
|
|
};
|
|
}
|
|
|
|
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 === "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;
|
|
});
|
|
}
|