feat: add ComfyUI integration for image generation and management
- Implemented ComfyGenerateRequest type and associated functions for generating images using various architectures and modes. - Added functions for listing generation options and handling image uploads. - Created workflows for different generation modes including SDXL, Z-Image, Z-Image Turbo, and Anima. - Introduced GenerationJobStatus component to display the status of ongoing generation jobs. - Developed MaskControls for managing mask operations and displaying mask analysis. - Created palette items for tool selection, layer management, and generation settings.
This commit is contained in:
syntaxbullet
239
operations/generation/runGenerate.ts
Normal file
239
operations/generation/runGenerate.ts
Normal file
@@ -0,0 +1,239 @@
|
||||
import { commandIds } from "@commands/ids";
|
||||
import type { ImageDocument } from "@core/document";
|
||||
import type { Transform } from "@core/geometry";
|
||||
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";
|
||||
|
||||
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;
|
||||
}) {
|
||||
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,
|
||||
outpaint: options.settings.outpaint,
|
||||
inpaint: resolveInpaintRequest(options.inpaintBundle, options.inpaintCandidate, options.settings),
|
||||
inputImage: options.inputImage,
|
||||
maskImage: options.maskImage,
|
||||
});
|
||||
}
|
||||
|
||||
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 layerMask = getLayerMask(layer);
|
||||
const maskLayer = layerMask?.enabled ? findLayer(document.artboards.flatMap((artboard) => artboard.layers), layerMask.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;
|
||||
}
|
||||
Reference in New Issue
Block a user