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
2026-07-10 23:15:02 +02:00
parent 41bbd1e16b
commit 5e4b548ad4
80 changed files with 2114 additions and 1954 deletions

View File

@@ -2,13 +2,10 @@ import { useEffect } from "react";
import { DropHalf } from "@phosphor-icons/react";
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 { resolveTransformTargetBounds } from "@editor/transform-targets";
import type { AppStore } from "@editor/store";
import type { ChromaKeySettings } from "@editor/tools";
import type { SelectionState } from "@editor/state";
import { blurMaskValues, despeckleMaskValues, dilateMaskValues, erodeMaskValues } from "../mask/maskRaster";
import { applyChromaKeyMask, previewChromaKey, resolveChromaKeyTarget } from "@operations/masks/chromaKey";
import { BottomControlColorPicker } from "./ColorPicker";
import { BottomControlDivider } from "./Divider";
import { BottomControlSlider } from "./Slider";
@@ -38,7 +35,7 @@ export function ChromaKeyControls({ document, selection, settings, dispatch }: C
return;
}
void chromaKeySource(target.asset.source, target.asset.intrinsicSize.w, target.asset.intrinsicSize.h, settings).then((source) => {
void previewChromaKey(target.asset.source, target.asset.intrinsicSize.w, target.asset.intrinsicSize.h, settings).then((source) => {
if (cancelled) return;
dispatch(commandIds.toolSetBrushStrokePreview, { layerId: target.layer.id, assetId: target.asset.id, source });
});
@@ -153,161 +150,3 @@ export function ChromaKeyControls({ document, selection, settings, dispatch }: C
</div>
);
}
function resolveChromaKeyTarget(document: ImageDocument, selection: SelectionState) {
const layerId = selection.layerIds[0];
if (selection.layerIds.length !== 1 || !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 bounds = resolveTransformTargetBounds(document, { type: "layer", id: layer.id });
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 && bounds ? { layer, asset, bounds, maskLayer, 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 applyChromaKeyMask(target: NonNullable<ReturnType<typeof resolveChromaKeyTarget>>, settings: ChromaKeySettings, dispatch: AppStore["dispatch"]) {
const source = await chromaKeyMaskSource(target.asset.source, target.asset.intrinsicSize.w, target.asset.intrinsicSize.h, settings);
dispatch(commandIds.toolSetBrushStrokePreview, undefined);
if (target.maskAsset && target.maskLayer && target.maskLayer.type !== "group") {
dispatch(commandIds.documentApplyLayerMaskOperation, { maskLayerId: target.maskLayer.id, source, mimeType: "image/png", operation: { type: "chromaKey" } });
return;
}
const assetId = crypto.randomUUID();
const maskLayerId = crypto.randomUUID();
const width = Math.max(1, Math.round(target.asset.intrinsicSize.w));
const height = Math.max(1, Math.round(target.asset.intrinsicSize.h));
dispatch(commandIds.documentAddLayerMask, {
layerId: target.layer.id,
asset: {
id: assetId,
name: `${target.layer.name} Chroma Mask`,
mimeType: "image/png",
source,
intrinsicSize: { w: width, h: height },
},
maskLayer: {
id: maskLayerId,
type: "raster",
name: `${target.layer.name} Chroma Mask`,
visible: true,
locked: false,
opacity: 1,
assetId,
transform: {
position: { x: target.bounds.x, y: target.bounds.y },
scale: { x: target.bounds.w / width, y: target.bounds.h / height },
rotation: target.layer.transform.rotation,
},
},
});
dispatch(commandIds.toolExitMaskEdit, undefined);
dispatch(commandIds.toolSetActive, { tool: "chromaKey" });
}
async function chromaKeySource(source: string, width: number, height: number, settings: ChromaKeySettings) {
const canvas = document.createElement("canvas");
canvas.width = Math.max(1, Math.round(width));
canvas.height = Math.max(1, Math.round(height));
const context = canvas.getContext("2d");
if (!context) return source;
const image = await loadImage(source);
context.drawImage(image, 0, 0, canvas.width, canvas.height);
const data = context.getImageData(0, 0, canvas.width, canvas.height);
const alpha = chromaKeyAlpha(data, canvas.width, canvas.height, settings);
for (let pixel = 0; pixel < alpha.length; pixel++) data.data[pixel * 4 + 3] = alpha[pixel] ?? 255;
context.putImageData(data, 0, 0);
return canvas.toDataURL("image/png");
}
async function chromaKeyMaskSource(source: string, width: number, height: number, settings: ChromaKeySettings) {
const canvas = document.createElement("canvas");
canvas.width = Math.max(1, Math.round(width));
canvas.height = Math.max(1, Math.round(height));
const context = canvas.getContext("2d");
if (!context) return source;
const image = await loadImage(source);
context.drawImage(image, 0, 0, canvas.width, canvas.height);
const data = context.getImageData(0, 0, canvas.width, canvas.height);
const alpha = chromaKeyAlpha(data, canvas.width, canvas.height, settings);
for (let pixel = 0; pixel < alpha.length; pixel++) {
const index = pixel * 4;
data.data[index] = 255;
data.data[index + 1] = 255;
data.data[index + 2] = 255;
data.data[index + 3] = alpha[pixel] ?? 255;
}
context.putImageData(data, 0, 0);
return canvas.toDataURL("image/png");
}
function hexToRgb(color: string) {
const hex = color.replace("#", "");
return { r: Number.parseInt(hex.slice(0, 2), 16), g: Number.parseInt(hex.slice(2, 4), 16), b: Number.parseInt(hex.slice(4, 6), 16) };
}
function chromaKeyAlpha(data: ImageData, width: number, height: number, settings: ChromaKeySettings) {
const key = hexToRgb(settings.color);
const alpha = new Uint8ClampedArray(width * height);
for (let pixel = 0; pixel < alpha.length; pixel++) {
const index = pixel * 4;
const red = data.data[index] ?? 0;
const green = data.data[index + 1] ?? 0;
const blue = data.data[index + 2] ?? 0;
const sourceAlpha = data.data[index + 3] ?? 255;
alpha[pixel] = Math.round(sourceAlpha * chromaKeyKeepFactor(red, green, blue, key, settings));
}
return postProcessAlpha(alpha, width, height, settings);
}
function chromaKeyKeepFactor(red: number, green: number, blue: number, key: { r: number; g: number; b: number }, settings: ChromaKeySettings) {
const tolerance = Math.max(0, Math.min(255, settings.tolerance));
const softness = Math.max(0, Math.min(255, settings.softness));
const spill = Math.max(0, Math.min(100, settings.spill)) / 100;
const distance = Math.hypot(red - key.r, green - key.g, blue - key.b);
const edgeKeep = distance <= tolerance ? 0 : softness > 0 && distance < tolerance + softness ? (distance - tolerance) / softness : 1;
if (spill <= 0) return edgeKeep;
const dominant = key.g >= key.r && key.g >= key.b ? green : key.r >= key.b ? red : blue;
const neutral = key.g >= key.r && key.g >= key.b ? Math.max(red, blue) : key.r >= key.b ? Math.max(green, blue) : Math.max(red, green);
const spillAmount = Math.max(0, dominant - neutral) / 255;
return Math.max(0, Math.min(edgeKeep, 1 - spillAmount * spill));
}
function postProcessAlpha(alpha: Uint8ClampedArray, width: number, height: number, settings: ChromaKeySettings) {
let next = alpha;
const despeckle = Math.round(Math.max(0, Math.min(20, settings.despeckle)));
const choke = Math.round(Math.max(-20, Math.min(20, settings.choke)));
const feather = Math.round(Math.max(0, Math.min(20, settings.feather)));
if (despeckle > 0) next = despeckleMaskValues(next, width, height, despeckle);
if (choke > 0) next = erodeMaskValues(next, width, height, choke);
if (choke < 0) next = dilateMaskValues(next, width, height, -choke);
if (feather > 0) next = blurMaskValues(next, width, height, feather);
return next;
}
function loadImage(source: string) {
return new Promise<HTMLImageElement>((resolve, reject) => {
const image = new Image();
image.onload = () => resolve(image);
image.onerror = () => reject(new Error("Failed to load image"));
image.src = source;
});
}

View File

@@ -1,12 +1,13 @@
import { useEffect, useState } from "react";
import { commandIds } from "@commands/ids";
import type { ImageDocument } from "@core/document";
import type { GenerationCandidate, GenerationCompareMode, GenerationState, SelectionState, ViewportState } from "@editor/state";
import type { GenerateSettings } from "@editor/tools";
import type { AppStore } from "@editor/store";
import { createMaskedPixelReplacementSource } from "../generate/candidateActions";
import { runGenerate, runGenerateFromCandidate } from "../generate/runGenerate";
import { createSolidMaskSource } from "../mask/maskRaster";
import { createMaskedPixelReplacementSource } from "@operations/generation/candidateActions";
import { runGenerate, runGenerateFromCandidate } from "@operations/generation/runGenerate";
import { runGenerationJob } from "@operations/generation/generationJob";
import { currentGenerationJob, GenerationJobStatus } from "../GenerationJobStatus";
import { createRefinementMask } from "@operations/masks/rasterActions";
export type GenerateActionControlsProps = {
document: ImageDocument;
@@ -18,44 +19,25 @@ export type GenerateActionControlsProps = {
};
export function GenerateActionControls({ document, selection, viewport, settings, generation, dispatch }: GenerateActionControlsProps) {
const [busy, setBusy] = useState<string>();
const [error, setError] = useState<string>();
const [elapsedSeconds, setElapsedSeconds] = useState(0);
const job = currentGenerationJob(generation);
const busy = job?.status === "running";
const candidate = selectedCandidate(generation);
const canGenerate = Boolean(settings.prompt.trim()) && !busy;
useEffect(() => {
if (!busy) {
setElapsedSeconds(0);
return;
}
setElapsedSeconds(0);
const startedAt = Date.now();
const interval = window.setInterval(() => {
setElapsedSeconds(Math.floor((Date.now() - startedAt) / 1000));
}, 1000);
return () => window.clearInterval(interval);
}, [busy]);
return (
<div className="flex max-w-[calc(100vw-2rem)] flex-wrap items-center justify-center gap-2 px-2">
<button
type="button"
disabled={!canGenerate}
className="h-12 rounded-full bg-white px-7 text-base font-semibold !text-black transition hover:bg-white/90 focus:outline-none focus-visible:ring-2 focus-visible:ring-white/40 disabled:pointer-events-none disabled:opacity-35"
title={error ?? "Generate with ComfyUI"}
title={job?.status === "failed" ? job.error : "Generate with ComfyUI"}
onClick={() => {
setBusy("Generating");
setError(undefined);
void runGenerate({ document, selection, viewport, settings, dispatch })
.catch((reason: unknown) => setError(reason instanceof Error ? reason.message : "Generation failed"))
.finally(() => setBusy(undefined));
void runGenerationJob({ kind: "generate", label: "Generating", dispatch, task: () => runGenerate({ document, selection, viewport, settings, dispatch }) });
}}
>
{busy === "Generating" ? `Generating ${formatElapsed(elapsedSeconds)}` : "Generate"}
{busy && job?.kind === "generate" ? "Generating..." : "Generate"}
</button>
{busy ? <span className="rounded-full bg-white/10 px-3 py-2 text-xs font-medium text-white/60">{busy} {formatElapsed(elapsedSeconds)}</span> : null}
<GenerationJobStatus generation={generation} />
{candidate ? (
<>
<CandidatePicker generation={generation} dispatch={dispatch} />
@@ -65,13 +47,10 @@ export function GenerateActionControls({ document, selection, viewport, settings
compareMode={generation.compareMode ?? "result"}
settings={settings}
busy={busy}
setBusy={setBusy}
setError={setError}
dispatch={dispatch}
/>
</>
) : null}
{error ? <span className="max-w-80 truncate rounded-full bg-red-500/15 px-3 py-2 text-xs font-medium text-red-100" title={error}>{error}</span> : null}
</div>
);
}
@@ -104,70 +83,53 @@ function CandidateControls({
compareMode,
settings,
busy,
setBusy,
setError,
dispatch,
}: {
document: ImageDocument;
candidate: GenerationCandidate;
compareMode: GenerationCompareMode;
settings: GenerateSettings;
busy?: string;
setBusy: (busy: string | undefined) => void;
setError: (error: string | undefined) => void;
busy: boolean;
dispatch: AppStore["dispatch"];
}) {
const rerun = (label: string, nextSettings: GenerateSettings) => {
setBusy(label);
setError(undefined);
dispatch(commandIds.toolSetGenerateSettings, nextSettings);
void runGenerateFromCandidate({ candidate, settings: nextSettings, dispatch })
.catch((reason: unknown) => setError(reason instanceof Error ? reason.message : `${label} failed`))
.finally(() => setBusy(undefined));
void runGenerationJob({ kind: "regenerate", label, dispatch, task: () => runGenerateFromCandidate({ candidate, settings: nextSettings, dispatch }) });
};
const disabled = Boolean(busy);
const disabled = busy;
return (
<div className="flex flex-wrap items-center justify-center gap-1 rounded-full bg-white/[0.04] px-2 py-1 ring-1 ring-white/[0.05]">
<CandidatePreview candidate={candidate} />
<span className="px-2 text-xs font-medium text-white/55">Seed {candidate.seed}</span>
<CandidateCompareControls compareMode={compareMode} disabled={disabled} dispatch={dispatch} />
<CandidateButton disabled={disabled} label="Regenerate" title="Regenerate same mask and crop" busy={busy === "Regenerate"} onClick={() => rerun("Regenerate", candidate.settings)} />
<CandidateButton disabled={disabled} label="Regenerate" title="Regenerate same mask and crop" onClick={() => rerun("Regenerate", candidate.settings)} />
<CandidateButton
disabled={disabled}
label="Lower"
title="Lower strength and regenerate same mask"
busy={busy === "Lower"}
onClick={() => rerun("Lower", { ...candidate.settings, strength: Math.max(0, candidate.settings.strength - 10), seed: candidate.seed })}
/>
<CandidateButton disabled={disabled} label="Reuse seed" title="Regenerate with the same seed" busy={busy === "Reuse seed"} onClick={() => rerun("Reuse seed", { ...candidate.settings, seed: candidate.seed })} />
<CandidateButton disabled={disabled} label="New seed" title="Regenerate with a new seed" busy={busy === "New seed"} onClick={() => rerun("New seed", { ...candidate.settings, seed: -1 })} />
<CandidateButton disabled={disabled} label="Reuse seed" title="Regenerate with the same seed" onClick={() => rerun("Reuse seed", { ...candidate.settings, seed: candidate.seed })} />
<CandidateButton disabled={disabled} label="New seed" title="Regenerate with a new seed" onClick={() => rerun("New seed", { ...candidate.settings, seed: -1 })} />
<CandidateButton disabled={disabled} label="Add as layer" title="Add candidate to the document as a layer" onClick={() => applyCandidateAsLayer(candidate, dispatch)} />
<CandidateButton
disabled={disabled}
label="Add + mask"
title="Add candidate as a layer with a fresh refinement mask"
busy={busy === "Refine"}
onClick={() => {
setBusy("Refine");
setError(undefined);
void applyCandidateAsRefinementLayer(candidate, dispatch)
.catch((reason: unknown) => setError(reason instanceof Error ? reason.message : "Refine setup failed"))
.finally(() => setBusy(undefined));
void runGenerationJob({ kind: "refine", label: "Adding refinement mask", dispatch, task: () => applyCandidateAsRefinementLayer(candidate, dispatch) });
}}
/>
<CandidateButton
disabled={disabled || !candidate.inpaint}
label="Replace pixels"
title={candidate.inpaint ? "Replace masked pixels and preserve unmasked pixels" : "Only inpaint candidates can replace masked pixels"}
busy={busy === "Replace"}
onClick={() => {
setBusy("Replace");
setError(undefined);
void createMaskedPixelReplacementSource(document, candidate)
.then((source) => dispatch(commandIds.generationReplaceCandidatePixels, { candidateId: candidate.id, source, mimeType: "image/png" }))
.catch((reason: unknown) => setError(reason instanceof Error ? reason.message : "Replace failed"))
.finally(() => setBusy(undefined));
void runGenerationJob({ kind: "replace", label: "Replacing pixels", dispatch, task: async () => {
const source = await createMaskedPixelReplacementSource(document, candidate);
dispatch(commandIds.generationReplaceCandidatePixels, { candidateId: candidate.id, source, mimeType: "image/png" });
} });
}}
/>
<CandidateButton
@@ -259,7 +221,7 @@ async function applyCandidateAsRefinementLayer(candidate: GenerationCandidate, d
applyCandidateAsLayerWithIds(candidate, { layerId, assetId: crypto.randomUUID() }, dispatch);
const width = Math.max(1, Math.round(candidate.intrinsicSize.w));
const height = Math.max(1, Math.round(candidate.intrinsicSize.h));
const source = await createSolidMaskSource(width, height, "white");
const source = await createRefinementMask(width, height);
dispatch(commandIds.documentAddLayerMask, {
layerId,
@@ -295,9 +257,3 @@ function applyCandidateAsLayerWithIds(candidate: GenerationCandidate, ids: { lay
layerId: ids.layerId,
});
}
function formatElapsed(seconds: number) {
const minutes = Math.floor(seconds / 60);
const remainder = seconds % 60;
return `${minutes}:${remainder.toString().padStart(2, "0")}`;
}

View File

@@ -2,6 +2,7 @@ import { useEffect, useRef, useState, type RefObject } from "react";
import { CaretDown, CaretUp } from "@phosphor-icons/react";
import { commandIds } from "@commands/ids";
import type { AppStore } from "@editor/store";
import type { GenerationOptions, GenerationResourcesState } from "@editor/state";
import { generateArchitectureDefaults } from "@editor/tools";
import type { GenerateArchitecture, GenerateMode, GenerateModel, GenerateSettings } from "@editor/tools";
import { BottomControlSelectMenu, type BottomControlSelectOption } from "./SelectMenu";
@@ -43,56 +44,23 @@ const inpaintMaskedContentOptions = [
export type GenerateControlsProps = {
settings: GenerateSettings;
resources: GenerationResourcesState;
dispatch: AppStore["dispatch"];
};
type ComfyArchitectureOption = {
value: GenerateArchitecture;
label: string;
defaultModel: string;
models: string[];
supportedModes: GenerateMode[];
};
type ComfyOptionsResponse = {
architectures?: ComfyArchitectureOption[];
models?: string[];
textEncoders?: string[];
vaes?: string[];
samplers?: string[];
schedulers?: string[];
};
export function GenerateControls({ settings, dispatch }: GenerateControlsProps) {
const [comfyOptions, setComfyOptions] = useState<ComfyOptionsResponse>();
export function GenerateControls({ settings, resources, dispatch }: GenerateControlsProps) {
const comfyOptions = resources.options;
const [advancedOpen, setAdvancedOpen] = useState(false);
const [outpaintOpen, setOutpaintOpen] = useState(false);
const [inpaintOpen, setInpaintOpen] = useState(false);
const [sizeOpen, setSizeOpen] = useState(false);
const sizeRef = useRef<HTMLDivElement>(null);
const [error, setError] = useState<string>();
const modelOptions = resolveModelOptions(settings, comfyOptions);
const supportOptions = resolveSupportOptions(settings, comfyOptions);
const samplerOptions = resolveStringOptions(comfyOptions?.samplers, settings.sampler);
const schedulerOptions = resolveStringOptions(comfyOptions?.schedulers, settings.scheduler);
const modeOptions = resolveModeOptions(settings, comfyOptions);
useEffect(() => {
let cancelled = false;
void fetch("/api/comfy/models")
.then((response) => response.ok ? response.json() : Promise.reject(new Error("Unable to load ComfyUI models")))
.then((body: ComfyOptionsResponse) => {
if (cancelled) return;
setComfyOptions(body);
})
.catch((reason: unknown) => {
if (!cancelled) setError(reason instanceof Error ? reason.message : "Unable to load ComfyUI models");
});
return () => {
cancelled = true;
};
}, []);
useEffect(() => {
if (!sizeOpen) return;
const close = (event: PointerEvent) => {
@@ -104,7 +72,7 @@ export function GenerateControls({ settings, dispatch }: GenerateControlsProps)
return (
<div className="grid gap-5 pb-2 tabular-nums">
{error ? <p className="rounded-full bg-red-500/10 px-3 py-2 text-xs text-red-200">{error}</p> : null}
{resources.error ? <p className="rounded-full bg-red-500/10 px-3 py-2 text-xs text-red-200">{resources.error}</p> : null}
<section className={panelSectionClass()}>
<label className="grid gap-2">
@@ -229,7 +197,7 @@ export function GenerateControls({ settings, dispatch }: GenerateControlsProps)
);
}
function resolveModelOptions(settings: GenerateSettings, comfyOptions: ComfyOptionsResponse | undefined): readonly BottomControlSelectOption<GenerateModel>[] {
function resolveModelOptions(settings: GenerateSettings, comfyOptions: GenerationOptions | undefined): readonly BottomControlSelectOption<GenerateModel>[] {
const architecture = comfyOptions?.architectures?.find((option) => option.value === settings.architecture);
const models = architecture?.models ?? (settings.architecture === "sdxl" ? comfyOptions?.models : undefined) ?? [];
const fallbackModel = architecture?.defaultModel ?? generateArchitectureDefaults[settings.architecture].model;
@@ -237,7 +205,7 @@ function resolveModelOptions(settings: GenerateSettings, comfyOptions: ComfyOpti
return values.map((model) => ({ value: model, label: model === "auto" ? "Auto" : model }));
}
function resolveSupportOptions(settings: GenerateSettings, comfyOptions: ComfyOptionsResponse | undefined) {
function resolveSupportOptions(settings: GenerateSettings, comfyOptions: GenerationOptions | undefined) {
const defaults = generateArchitectureDefaults[settings.architecture];
return {
textEncoders: resolveStringOptions([...(comfyOptions?.textEncoders ?? []), defaults.textEncoder].filter((value) => value !== "auto"), settings.textEncoder),
@@ -249,7 +217,7 @@ function resolveStringOptions(values: string[] | undefined, current: string): re
return unique([...(values ?? []), current]).map((value) => ({ value, label: value }));
}
function resolveModeOptions(settings: GenerateSettings, comfyOptions: ComfyOptionsResponse | undefined): readonly BottomControlSelectOption<GenerateMode>[] {
function resolveModeOptions(settings: GenerateSettings, comfyOptions: GenerationOptions | undefined): readonly BottomControlSelectOption<GenerateMode>[] {
const architecture = comfyOptions?.architectures?.find((option) => option.value === settings.architecture);
const supportedModes = architecture?.supportedModes?.length ? architecture.supportedModes : generateArchitectureDefaults[settings.architecture].supportedModes;
const availableModes = modes.filter((mode) => supportedModes.includes(mode.value));