7188569672
- Implemented ComfyUI API for generating images with various modes (text-to-image, image-to-image, inpaint, outpaint). - Created GenerateSheet and associated controls for user input on generation settings. - Added subtle scrollbar styles for improved UI experience. - Enhanced canvas input handling to ignore key events when focused on editable elements. - Optimized canvas resizing logic to prevent unnecessary dispatches. - Introduced error handling for generation failures and loading models. - Added functionality to upload images and masks for inpainting.
335 lines
18 KiB
TypeScript
335 lines
18 KiB
TypeScript
import { useEffect, useRef, useState, type RefObject } from "react";
|
||
import { CaretDown, CaretUp } from "@phosphor-icons/react";
|
||
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 { GenerateMode, GenerateModel, GenerateSettings } from "@editor/tools";
|
||
import { BottomControlSelectMenu, type BottomControlSelectOption } from "./SelectMenu";
|
||
import { BottomControlSlider } from "./Slider";
|
||
|
||
const modes = [
|
||
{ value: "text-to-image", label: "Text → image" },
|
||
{ value: "image-to-image", label: "Image → image" },
|
||
{ value: "inpaint", label: "Inpaint" },
|
||
{ value: "outpaint", label: "Outpaint" },
|
||
] satisfies readonly BottomControlSelectOption<GenerateMode>[];
|
||
|
||
const sizePresets = [
|
||
{ label: "1:1", w: 1024, h: 1024 },
|
||
{ label: "4:3", w: 1152, h: 896 },
|
||
{ label: "3:4", w: 896, h: 1152 },
|
||
{ label: "16:9", w: 1344, h: 768 },
|
||
{ label: "9:16", w: 768, h: 1344 },
|
||
] as const;
|
||
|
||
export type GenerateControlsProps = {
|
||
document: ImageDocument;
|
||
selection: SelectionState;
|
||
viewport: ViewportState;
|
||
settings: GenerateSettings;
|
||
dispatch: AppStore["dispatch"];
|
||
};
|
||
|
||
export function GenerateControls({ settings, dispatch }: GenerateControlsProps) {
|
||
const [models, setModels] = useState<readonly BottomControlSelectOption<GenerateModel>[]>([{ value: "auto", label: "Auto" }]);
|
||
const [samplers, setSamplers] = useState<readonly BottomControlSelectOption<string>[]>([{ value: settings.sampler, label: settings.sampler }]);
|
||
const [schedulers, setSchedulers] = useState<readonly BottomControlSelectOption<string>[]>([{ value: settings.scheduler, label: settings.scheduler }]);
|
||
const [advancedOpen, setAdvancedOpen] = useState(false);
|
||
const [outpaintOpen, setOutpaintOpen] = useState(false);
|
||
const [sizeOpen, setSizeOpen] = useState(false);
|
||
const sizeRef = useRef<HTMLDivElement>(null);
|
||
const [error, setError] = useState<string>();
|
||
|
||
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: { models?: string[]; samplers?: string[]; schedulers?: string[] }) => {
|
||
if (cancelled) return;
|
||
setModels([{ value: "auto", label: "Auto" }, ...(body.models ?? []).map((model) => ({ value: model, label: model }))]);
|
||
if (body.samplers?.length) setSamplers(body.samplers.map((sampler) => ({ value: sampler, label: sampler })));
|
||
if (body.schedulers?.length) setSchedulers(body.schedulers.map((scheduler) => ({ value: scheduler, label: scheduler })));
|
||
})
|
||
.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) => {
|
||
if (!sizeRef.current?.contains(event.target as Node)) setSizeOpen(false);
|
||
};
|
||
window.addEventListener("pointerdown", close);
|
||
return () => window.removeEventListener("pointerdown", close);
|
||
}, [sizeOpen]);
|
||
|
||
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}
|
||
|
||
<section className={panelSectionClass()}>
|
||
<label className="grid gap-2">
|
||
<span className={panelLabelClass()}>Prompt</span>
|
||
<textarea
|
||
aria-label="Generate prompt"
|
||
value={settings.prompt}
|
||
rows={6}
|
||
placeholder="Describe the image you want to create…"
|
||
onChange={(event) => dispatch(commandIds.toolSetGenerateSettings, { prompt: event.currentTarget.value })}
|
||
className={panelTextAreaClass("min-h-36")}
|
||
/>
|
||
</label>
|
||
<label className="grid gap-2">
|
||
<span className={panelLabelClass()}>Negative prompt</span>
|
||
<textarea
|
||
aria-label="Negative prompt"
|
||
value={settings.negativePrompt}
|
||
rows={3}
|
||
placeholder="Things to avoid, e.g. bad quality, artifacts…"
|
||
onChange={(event) => dispatch(commandIds.toolSetGenerateSettings, { negativePrompt: event.currentTarget.value })}
|
||
className={panelTextAreaClass("min-h-20")}
|
||
/>
|
||
</label>
|
||
</section>
|
||
|
||
<section className={panelSectionClass()}>
|
||
<SectionTitle title="Essentials" />
|
||
<PanelSelect label="Model" value={settings.model} options={models} ariaLabel="Generate model" onValueChange={(model) => dispatch(commandIds.toolSetGenerateSettings, { model })} />
|
||
<SizeControl refRoot={sizeRef} open={sizeOpen} setOpen={setSizeOpen} settings={settings} dispatch={dispatch} />
|
||
<PanelNumber label="Seed" aria-label="Generate seed" min={-1} max={Number.MAX_SAFE_INTEGER} value={settings.seed} onValueChange={(seed) => dispatch(commandIds.toolSetGenerateSettings, { seed })} />
|
||
</section>
|
||
|
||
<section className={panelSectionClass()}>
|
||
<button type="button" className={sectionToggleClass()} aria-expanded={advancedOpen} aria-controls="generate-advanced-controls" onClick={() => setAdvancedOpen((open) => !open)}>
|
||
<span>
|
||
<span className="block text-sm font-semibold text-white/85">Advanced</span>
|
||
<span className="block text-xs text-white/40">Mode, sampler, scheduler, steps, CFG, strength</span>
|
||
</span>
|
||
{advancedOpen ? <CaretUp size={18} weight="bold" /> : <CaretDown size={18} weight="bold" />}
|
||
</button>
|
||
<div id="generate-advanced-controls" className={`grid gap-2 overflow-hidden transition-all duration-200 ${advancedOpen ? "max-h-[32rem] pt-2 opacity-100" : "max-h-0 opacity-0"}`}>
|
||
<PanelSelect label="Mode" value={settings.mode} options={modes} ariaLabel="Generate mode" onValueChange={(mode) => dispatch(commandIds.toolSetGenerateSettings, { mode })} />
|
||
<PanelSelect label="Sampler" value={settings.sampler} options={samplers} ariaLabel="Generate sampler" onValueChange={(sampler) => dispatch(commandIds.toolSetGenerateSettings, { sampler })} />
|
||
<PanelSelect label="Scheduler" value={settings.scheduler} options={schedulers} ariaLabel="Generate scheduler" onValueChange={(scheduler) => dispatch(commandIds.toolSetGenerateSettings, { scheduler })} />
|
||
<div className="grid grid-cols-2 gap-2">
|
||
<PanelNumber label="Steps" aria-label="Generate steps" value={settings.steps} onValueChange={(steps) => dispatch(commandIds.toolSetGenerateSettings, { steps })} />
|
||
<PanelNumber label="CFG" aria-label="Generate CFG" value={settings.cfg} onValueChange={(cfg) => dispatch(commandIds.toolSetGenerateSettings, { cfg })} />
|
||
</div>
|
||
<div className={panelRowClass()}>
|
||
<span className={panelLabelClass()}>Strength</span>
|
||
<BottomControlSlider min={0} max={100} value={settings.strength} className="w-32" aria-label="Generate strength" onValueChange={(strength) => dispatch(commandIds.toolSetGenerateSettings, { strength })} />
|
||
<span className="w-10 text-right text-sm text-white">{Math.round(settings.strength)}</span>
|
||
</div>
|
||
</div>
|
||
</section>
|
||
|
||
<section className={panelSectionClass()}>
|
||
<button type="button" className={sectionToggleClass()} aria-expanded={outpaintOpen} aria-controls="generate-outpaint-controls" onClick={() => setOutpaintOpen((open) => !open)}>
|
||
<span>
|
||
<span className="block text-sm font-semibold text-white/85">Outpaint</span>
|
||
<span className="block text-xs text-white/40">Padding and feathering for outpaint mode</span>
|
||
</span>
|
||
{outpaintOpen ? <CaretUp size={18} weight="bold" /> : <CaretDown size={18} weight="bold" />}
|
||
</button>
|
||
<div id="generate-outpaint-controls" className={`grid gap-2 overflow-hidden transition-all duration-200 ${outpaintOpen ? "max-h-72 pt-2 opacity-100" : "max-h-0 opacity-0"}`}>
|
||
<div className="grid grid-cols-2 gap-2">
|
||
<PanelNumber label="Left" aria-label="Outpaint left padding" value={settings.outpaint.left} onValueChange={(left) => dispatch(commandIds.toolSetGenerateSettings, { outpaint: { ...settings.outpaint, left } })} />
|
||
<PanelNumber label="Top" aria-label="Outpaint top padding" value={settings.outpaint.top} onValueChange={(top) => dispatch(commandIds.toolSetGenerateSettings, { outpaint: { ...settings.outpaint, top } })} />
|
||
<PanelNumber label="Right" aria-label="Outpaint right padding" value={settings.outpaint.right} onValueChange={(right) => dispatch(commandIds.toolSetGenerateSettings, { outpaint: { ...settings.outpaint, right } })} />
|
||
<PanelNumber label="Bottom" aria-label="Outpaint bottom padding" value={settings.outpaint.bottom} onValueChange={(bottom) => dispatch(commandIds.toolSetGenerateSettings, { outpaint: { ...settings.outpaint, bottom } })} />
|
||
</div>
|
||
<PanelNumber label="Feather" aria-label="Outpaint feathering" value={settings.outpaint.feathering} onValueChange={(feathering) => dispatch(commandIds.toolSetGenerateSettings, { outpaint: { ...settings.outpaint, feathering } })} />
|
||
</div>
|
||
</section>
|
||
</div>
|
||
);
|
||
}
|
||
|
||
function SectionTitle({ title }: { title: string }) {
|
||
return <div className="px-1 text-xs font-semibold uppercase tracking-[0.18em] text-white/35">{title}</div>;
|
||
}
|
||
|
||
function panelSectionClass() {
|
||
return "grid gap-3 rounded-[1.75rem] bg-white/[0.035] p-3 ring-1 ring-white/[0.04]";
|
||
}
|
||
|
||
function sectionToggleClass() {
|
||
return "flex w-full items-center justify-between gap-3 rounded-[1.25rem] px-2 py-1 text-left transition hover:bg-white/[0.04] focus:outline-none focus-visible:ring-2 focus-visible:ring-white/30";
|
||
}
|
||
|
||
function PanelSelect<TValue extends string>({ label, value, options, ariaLabel, onValueChange }: { label: string; value: TValue; options: readonly BottomControlSelectOption<TValue>[]; ariaLabel: string; onValueChange: (value: TValue) => void }) {
|
||
return (
|
||
<div className="rounded-[1.25rem] bg-white/[0.04] p-1">
|
||
<BottomControlSelectMenu label={label} value={value} options={options} aria-label={ariaLabel} placement="inline" onValueChange={onValueChange} />
|
||
</div>
|
||
);
|
||
}
|
||
|
||
function PanelNumber({ label, value, onValueChange, min = 0, max = 4096, ...props }: { label: string; "aria-label": string; value: number; min?: number; max?: number; onValueChange: (value: number) => void }) {
|
||
return (
|
||
<label className={panelRowClass()}>
|
||
<span className={panelLabelClass()}>{label}</span>
|
||
<NumberInput {...props} min={min} max={max} value={value} onValueChange={onValueChange} />
|
||
</label>
|
||
);
|
||
}
|
||
|
||
function SizeControl({ refRoot, open, setOpen, settings, dispatch }: { refRoot: RefObject<HTMLDivElement | null>; open: boolean; setOpen: (open: boolean) => void; settings: GenerateSettings; dispatch: AppStore["dispatch"] }) {
|
||
return (
|
||
<div ref={refRoot} className="relative">
|
||
<button type="button" className={compactRowButtonClass()} aria-label={`Generate size ${settings.width} by ${settings.height}`} aria-expanded={open} onClick={() => setOpen(!open)}>
|
||
<span className={panelLabelClass()}>Size</span>
|
||
<span className="min-w-0 flex-1 text-right text-white">{settings.width} × {settings.height}</span>
|
||
<CaretDown size={16} weight="bold" />
|
||
</button>
|
||
{open ? (
|
||
<div className="absolute right-0 top-full z-30 mt-2 grid w-full gap-3 rounded-[1.5rem] bg-slate-950/90 p-3 text-white shadow-2xl ring-1 ring-white/10 backdrop-blur-xl">
|
||
<div className="grid grid-cols-2 gap-2">
|
||
<PanelNumber label="W" aria-label="Generate width" value={settings.width} onValueChange={(width) => dispatch(commandIds.toolSetGenerateSettings, { width })} />
|
||
<PanelNumber label="H" aria-label="Generate height" value={settings.height} onValueChange={(height) => dispatch(commandIds.toolSetGenerateSettings, { height })} />
|
||
</div>
|
||
<div className="grid grid-cols-5 gap-2">
|
||
{sizePresets.map((preset) => (
|
||
<button key={preset.label} type="button" className="rounded-full bg-white/5 px-2 py-2 text-xs text-white/75 transition hover:bg-white/10 hover:text-white focus:outline-none focus-visible:ring-2 focus-visible:ring-white/30" onClick={() => dispatch(commandIds.toolSetGenerateSettings, { width: preset.w, height: preset.h })}>
|
||
{preset.label}
|
||
</button>
|
||
))}
|
||
</div>
|
||
</div>
|
||
) : null}
|
||
</div>
|
||
);
|
||
}
|
||
|
||
function NumberInput({ value, onValueChange, min = 0, max = 4096, ...props }: { "aria-label": string; value: number; min?: number; max?: number; onValueChange: (value: number) => void }) {
|
||
return <input {...props} type="number" min={min} max={max} step={1} value={Math.round(value)} className="h-9 w-16 rounded-full bg-white/5 px-2 text-right text-sm text-white outline-none transition hover:bg-white/10 focus:bg-white/10 focus:ring-2 focus:ring-white/30" onChange={(event) => onValueChange(Number(event.currentTarget.value))} />;
|
||
}
|
||
|
||
function panelRowClass() {
|
||
return "flex min-h-11 items-center justify-between gap-3 rounded-full bg-white/[0.04] px-4";
|
||
}
|
||
|
||
function compactRowButtonClass() {
|
||
return `${panelRowClass()} w-full text-left transition hover:bg-white/[0.07] focus:outline-none focus-visible:ring-2 focus-visible:ring-white/30`;
|
||
}
|
||
|
||
function panelLabelClass() {
|
||
return "shrink-0 text-sm font-medium text-white/55";
|
||
}
|
||
|
||
function panelTextAreaClass(extra = "") {
|
||
return `${extra} resize-none rounded-[1.25rem] bg-white/5 px-4 py-3 text-sm text-white outline-none transition placeholder:text-white/25 focus:bg-white/[0.07] focus:ring-2 focus:ring-white/30`;
|
||
}
|
||
|
||
async function generateImage(options: GenerateControlsProps & { setBusy: (busy: boolean) => void; setError: (error: string | undefined) => void }) {
|
||
const { document, selection, viewport, settings, dispatch, setBusy, setError } = options;
|
||
const artboard = selection.artboardId ? document.artboards.find((candidate) => candidate.id === selection.artboardId) : document.artboards[0];
|
||
if (!artboard) return;
|
||
|
||
setBusy(true);
|
||
setError(undefined);
|
||
try {
|
||
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: viewport.center.x - intrinsicSize.w / 2, y: viewport.center.y - intrinsicSize.h / 2 }, scale: { x: 1, y: 1 }, rotation: 0 },
|
||
},
|
||
});
|
||
dispatch(commandIds.selectionSet, { artboardId: artboard.id, layerIds: [layerId] });
|
||
} catch (reason) {
|
||
setError(reason instanceof Error ? reason.message : "Generation failed");
|
||
} finally {
|
||
setBusy(false);
|
||
}
|
||
}
|
||
|
||
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;
|
||
});
|
||
}
|