304 lines
14 KiB
TypeScript
304 lines
14 KiB
TypeScript
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";
|
|
|
|
export type GenerateActionControlsProps = {
|
|
document: ImageDocument;
|
|
selection: SelectionState;
|
|
viewport: ViewportState;
|
|
settings: GenerateSettings;
|
|
generation: GenerationState;
|
|
dispatch: AppStore["dispatch"];
|
|
};
|
|
|
|
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 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"}
|
|
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));
|
|
}}
|
|
>
|
|
{busy === "Generating" ? `Generating ${formatElapsed(elapsedSeconds)}` : "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}
|
|
{candidate ? (
|
|
<>
|
|
<CandidatePicker generation={generation} dispatch={dispatch} />
|
|
<CandidateControls
|
|
document={document}
|
|
candidate={candidate}
|
|
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>
|
|
);
|
|
}
|
|
|
|
function CandidatePicker({ generation, dispatch }: { generation: GenerationState; dispatch: AppStore["dispatch"] }) {
|
|
if (generation.candidates.length < 2) return null;
|
|
return (
|
|
<div className="subtle-scrollbar flex max-w-[min(28rem,calc(100vw-2rem))] items-center gap-1 overflow-x-auto rounded-full bg-white/[0.04] px-1.5 py-1 ring-1 ring-white/[0.05]" role="group" aria-label="Generation candidates">
|
|
{generation.candidates.map((candidate) => {
|
|
const selected = candidate.id === (generation.selectedCandidateId ?? generation.candidates[0]?.id);
|
|
return (
|
|
<button
|
|
key={candidate.id}
|
|
type="button"
|
|
className={`h-9 w-9 overflow-hidden rounded-full ring-2 transition ${selected ? "ring-white" : "ring-white/10 hover:ring-white/40"}`}
|
|
title={`Candidate seed ${candidate.seed}`}
|
|
onClick={() => dispatch(commandIds.generationSelectCandidate, { candidateId: candidate.id })}
|
|
>
|
|
<img src={candidate.source} alt="" className="h-full w-full object-cover" />
|
|
</button>
|
|
);
|
|
})}
|
|
</div>
|
|
);
|
|
}
|
|
|
|
function CandidateControls({
|
|
document,
|
|
candidate,
|
|
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;
|
|
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));
|
|
};
|
|
const disabled = Boolean(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="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="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));
|
|
}}
|
|
/>
|
|
<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));
|
|
}}
|
|
/>
|
|
<CandidateButton
|
|
disabled={disabled || !candidate.inpaint}
|
|
label="Edit mask"
|
|
title={candidate.inpaint ? "Paint refinement mask over the source layer" : "Only inpaint candidates have a mask to edit"}
|
|
onClick={() => {
|
|
if (!candidate.inpaint) return;
|
|
dispatch(commandIds.selectionSet, { artboardId: candidate.placement.artboardId, layerIds: [candidate.inpaint.targetLayerId] });
|
|
dispatch(commandIds.toolEnterMaskEdit, { targetLayerId: candidate.inpaint.targetLayerId, maskLayerId: candidate.inpaint.maskLayerId });
|
|
}}
|
|
/>
|
|
<CandidateButton disabled={disabled} label="Dismiss" title="Remove this candidate" onClick={() => dispatch(commandIds.generationRemoveCandidate, { candidateId: candidate.id })} />
|
|
{settings.seed !== candidate.seed ? null : <span className="sr-only">Current settings reuse this seed</span>}
|
|
</div>
|
|
);
|
|
}
|
|
|
|
function CandidateCompareControls({ compareMode, disabled, dispatch }: { compareMode: GenerationCompareMode; disabled: boolean; dispatch: AppStore["dispatch"] }) {
|
|
return (
|
|
<span className="flex items-center gap-1 rounded-full bg-black/20 p-1" aria-label="Compare candidate">
|
|
{generationCompareOptions.map((option) => {
|
|
const active = compareMode === option.mode;
|
|
return (
|
|
<button
|
|
key={option.mode}
|
|
type="button"
|
|
className={`h-7 rounded-full px-2 text-[0.7rem] font-semibold transition ${
|
|
active ? "bg-white text-black" : "text-white/55 hover:bg-white/10 hover:text-white"
|
|
} disabled:pointer-events-none disabled:opacity-35`}
|
|
disabled={disabled}
|
|
title={option.title}
|
|
onClick={() => dispatch(commandIds.generationSetCompareMode, { mode: option.mode })}
|
|
>
|
|
{option.label}
|
|
</button>
|
|
);
|
|
})}
|
|
</span>
|
|
);
|
|
}
|
|
|
|
const generationCompareOptions: Array<{ mode: GenerationCompareMode; label: string; title: string }> = [
|
|
{ mode: "result", label: "After", title: "Show the generated result over the document" },
|
|
{ mode: "before", label: "Before", title: "Hide the generated result and show the source document" },
|
|
{ mode: "split", label: "Split", title: "Compare source on the left with result on the right" },
|
|
];
|
|
|
|
function CandidatePreview({ candidate }: { candidate: GenerationCandidate }) {
|
|
if (!candidate.inputImage) {
|
|
return <img src={candidate.source} alt="" className="h-10 w-10 rounded-full bg-black/25 object-cover ring-1 ring-white/10" />;
|
|
}
|
|
|
|
return (
|
|
<span className="flex items-center -space-x-2" title={candidate.maskImage ? "Input crop, normalized mask, generated candidate" : "Before and generated candidate"}>
|
|
<img src={candidate.inputImage} alt="" className="h-10 w-10 rounded-full bg-black/25 object-cover ring-1 ring-white/10" />
|
|
{candidate.maskImage ? <img src={candidate.maskImage} alt="" className="h-10 w-10 rounded-full bg-black/25 object-cover ring-1 ring-white/20" /> : null}
|
|
<img src={candidate.source} alt="" className="h-10 w-10 rounded-full bg-black/25 object-cover ring-2 ring-white/40" />
|
|
</span>
|
|
);
|
|
}
|
|
|
|
function CandidateButton({ label, title, disabled, busy, onClick }: { label: string; title: string; disabled?: boolean; busy?: boolean; onClick: () => void }) {
|
|
return (
|
|
<button
|
|
type="button"
|
|
className="h-9 rounded-full bg-white/5 px-3 text-xs font-semibold text-white/70 transition hover:bg-white/10 hover:text-white disabled:pointer-events-none disabled:opacity-35"
|
|
disabled={disabled}
|
|
title={title}
|
|
onClick={onClick}
|
|
>
|
|
{busy ? "..." : label}
|
|
</button>
|
|
);
|
|
}
|
|
|
|
function selectedCandidate(generation: GenerationState): GenerationCandidate | undefined {
|
|
return generation.candidates.find((candidate) => candidate.id === generation.selectedCandidateId) ?? generation.candidates[0];
|
|
}
|
|
|
|
function applyCandidateAsLayer(candidate: GenerationCandidate, dispatch: AppStore["dispatch"]) {
|
|
applyCandidateAsLayerWithIds(candidate, { layerId: crypto.randomUUID(), assetId: crypto.randomUUID() }, dispatch);
|
|
}
|
|
|
|
async function applyCandidateAsRefinementLayer(candidate: GenerationCandidate, dispatch: AppStore["dispatch"]) {
|
|
const layerId = crypto.randomUUID();
|
|
const maskLayerId = crypto.randomUUID();
|
|
const maskAssetId = crypto.randomUUID();
|
|
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");
|
|
|
|
dispatch(commandIds.documentAddLayerMask, {
|
|
layerId,
|
|
asset: {
|
|
id: maskAssetId,
|
|
name: `${candidate.placement.layerName} refinement mask`,
|
|
mimeType: "image/png",
|
|
source,
|
|
intrinsicSize: { w: width, h: height },
|
|
},
|
|
maskLayer: {
|
|
id: maskLayerId,
|
|
type: "raster",
|
|
name: `${candidate.placement.layerName} refinement mask`,
|
|
visible: true,
|
|
locked: false,
|
|
opacity: 1,
|
|
assetId: maskAssetId,
|
|
transform: {
|
|
position: { ...candidate.placement.transform.position },
|
|
scale: { ...candidate.placement.transform.scale },
|
|
rotation: candidate.placement.transform.rotation,
|
|
},
|
|
},
|
|
});
|
|
dispatch(commandIds.toolSetActive, { tool: "eraser" });
|
|
}
|
|
|
|
function applyCandidateAsLayerWithIds(candidate: GenerationCandidate, ids: { layerId: string; assetId: string }, dispatch: AppStore["dispatch"]) {
|
|
dispatch(commandIds.generationApplyCandidateAsLayer, {
|
|
candidateId: candidate.id,
|
|
assetId: ids.assetId,
|
|
layerId: ids.layerId,
|
|
});
|
|
}
|
|
|
|
function formatElapsed(seconds: number) {
|
|
const minutes = Math.floor(seconds / 60);
|
|
const remainder = seconds % 60;
|
|
return `${minutes}:${remainder.toString().padStart(2, "0")}`;
|
|
}
|