324 lines
11 KiB
TypeScript
324 lines
11 KiB
TypeScript
import type { Asset } from "@core/asset";
|
|
import type { AssetGenerationProvenance, GeneratedAssetAcceptance } from "@core/asset-provenance";
|
|
import type { ImageDocument } from "@core/document";
|
|
import type { ArtboardId, AssetId, LayerId } from "@core/id";
|
|
import type { ImageLayer } from "@core/image-layer";
|
|
import type { Layer } from "@core/layer";
|
|
import type { GenerationCandidate, GenerationCompareMode, GenerationState } from "@editor/state";
|
|
import type { Command } from "./command";
|
|
import { commandIds } from "./ids";
|
|
|
|
export type GenerationAddCandidatePayload = {
|
|
candidate: GenerationCandidate;
|
|
};
|
|
|
|
export type GenerationSelectCandidatePayload = {
|
|
candidateId?: string;
|
|
};
|
|
|
|
export type GenerationSetCompareModePayload = {
|
|
mode: GenerationCompareMode;
|
|
};
|
|
|
|
export type GenerationRemoveCandidatePayload = {
|
|
candidateId: string;
|
|
};
|
|
|
|
export type GenerationApplyCandidateAsLayerPayload = {
|
|
candidateId: string;
|
|
assetId: AssetId;
|
|
layerId: LayerId;
|
|
};
|
|
|
|
export type GenerationReplaceCandidatePixelsPayload = {
|
|
candidateId: string;
|
|
source: string;
|
|
mimeType?: string;
|
|
};
|
|
|
|
const maxCandidates = 12;
|
|
const generationCompareModes = new Set<GenerationCompareMode>(["result", "before", "split"]);
|
|
|
|
export const generationAddCandidateCommand: Command<GenerationAddCandidatePayload> = {
|
|
id: commandIds.generationAddCandidate,
|
|
name: "Add generation candidate",
|
|
history: { mode: "ignore" },
|
|
execute({ state }, payload) {
|
|
const candidates = [payload.candidate, ...state.editor.generation.candidates.filter((candidate) => candidate.id !== payload.candidate.id)].slice(0, maxCandidates);
|
|
return {
|
|
...state,
|
|
editor: {
|
|
...state.editor,
|
|
generation: {
|
|
candidates,
|
|
selectedCandidateId: payload.candidate.id,
|
|
compareMode: "result",
|
|
},
|
|
},
|
|
};
|
|
},
|
|
};
|
|
|
|
export const generationSelectCandidateCommand: Command<GenerationSelectCandidatePayload> = {
|
|
id: commandIds.generationSelectCandidate,
|
|
name: "Select generation candidate",
|
|
history: { mode: "ignore" },
|
|
execute({ state }, payload) {
|
|
const selectedCandidateId = payload.candidateId && state.editor.generation.candidates.some((candidate) => candidate.id === payload.candidateId) ? payload.candidateId : undefined;
|
|
if (state.editor.generation.selectedCandidateId === selectedCandidateId) return state;
|
|
return {
|
|
...state,
|
|
editor: {
|
|
...state.editor,
|
|
generation: {
|
|
...state.editor.generation,
|
|
selectedCandidateId,
|
|
},
|
|
},
|
|
};
|
|
},
|
|
};
|
|
|
|
export const generationSetCompareModeCommand: Command<GenerationSetCompareModePayload> = {
|
|
id: commandIds.generationSetCompareMode,
|
|
name: "Set generation compare mode",
|
|
history: { mode: "ignore" },
|
|
execute({ state }, payload) {
|
|
if (!generationCompareModes.has(payload.mode)) return state;
|
|
if (state.editor.generation.compareMode === payload.mode) return state;
|
|
return {
|
|
...state,
|
|
editor: {
|
|
...state.editor,
|
|
generation: {
|
|
...state.editor.generation,
|
|
compareMode: payload.mode,
|
|
},
|
|
},
|
|
};
|
|
},
|
|
};
|
|
|
|
export const generationRemoveCandidateCommand: Command<GenerationRemoveCandidatePayload> = {
|
|
id: commandIds.generationRemoveCandidate,
|
|
name: "Remove generation candidate",
|
|
history: { mode: "ignore" },
|
|
execute({ state }, payload) {
|
|
const generation = removeGenerationCandidate(state.editor.generation, payload.candidateId);
|
|
if (generation === state.editor.generation) return state;
|
|
return {
|
|
...state,
|
|
editor: {
|
|
...state.editor,
|
|
generation,
|
|
},
|
|
};
|
|
},
|
|
};
|
|
|
|
export const generationClearCandidatesCommand: Command = {
|
|
id: commandIds.generationClearCandidates,
|
|
name: "Clear generation candidates",
|
|
history: { mode: "ignore" },
|
|
execute({ state }) {
|
|
if (state.editor.generation.candidates.length === 0 && !state.editor.generation.selectedCandidateId && state.editor.generation.compareMode === "result") return state;
|
|
return {
|
|
...state,
|
|
editor: {
|
|
...state.editor,
|
|
generation: { candidates: [], selectedCandidateId: undefined, compareMode: "result" },
|
|
},
|
|
};
|
|
},
|
|
};
|
|
|
|
export const generationApplyCandidateAsLayerCommand: Command<GenerationApplyCandidateAsLayerPayload> = {
|
|
id: commandIds.generationApplyCandidateAsLayer,
|
|
name: "Apply generation candidate as layer",
|
|
execute({ state }, payload) {
|
|
const candidate = state.editor.generation.candidates.find((item) => item.id === payload.candidateId);
|
|
if (!candidate) return state;
|
|
if (state.document.assets.some((asset) => asset.id === payload.assetId) || findLayerLocation(state.document, payload.layerId)) return state;
|
|
if (!state.document.artboards.some((artboard) => artboard.id === candidate.placement.artboardId)) return state;
|
|
|
|
const asset: Asset = {
|
|
id: payload.assetId,
|
|
name: candidate.placement.layerName,
|
|
mimeType: candidate.mimeType,
|
|
source: candidate.source,
|
|
intrinsicSize: { ...candidate.intrinsicSize },
|
|
provenance: generationProvenance(candidate, "layer"),
|
|
};
|
|
const layer: ImageLayer = {
|
|
id: payload.layerId,
|
|
type: "image",
|
|
name: candidate.placement.layerName,
|
|
visible: true,
|
|
locked: false,
|
|
opacity: 1,
|
|
assetId: asset.id,
|
|
transform: {
|
|
position: { ...candidate.placement.transform.position },
|
|
scale: { ...candidate.placement.transform.scale },
|
|
rotation: candidate.placement.transform.rotation,
|
|
},
|
|
};
|
|
|
|
return {
|
|
...state,
|
|
document: insertLayerAtTop({ ...state.document, assets: [...state.document.assets, asset] }, candidate.placement.artboardId, layer),
|
|
editor: {
|
|
...state.editor,
|
|
generation: removeGenerationCandidate(state.editor.generation, candidate.id),
|
|
selection: { artboardId: candidate.placement.artboardId, layerIds: [layer.id] },
|
|
},
|
|
};
|
|
},
|
|
};
|
|
|
|
export const generationReplaceCandidatePixelsCommand: Command<GenerationReplaceCandidatePixelsPayload> = {
|
|
id: commandIds.generationReplaceCandidatePixels,
|
|
name: "Replace masked pixels with generation candidate",
|
|
execute({ state }, payload) {
|
|
const candidate = state.editor.generation.candidates.find((item) => item.id === payload.candidateId);
|
|
if (!candidate?.inpaint || !payload.source.trim()) return state;
|
|
const targetAsset = state.document.assets.find((asset) => asset.id === candidate.inpaint?.sourceAssetId);
|
|
const targetLayerLocation = findLayerLocation(state.document, candidate.inpaint.targetLayerId);
|
|
if (!targetAsset || !targetLayerLocation) return state;
|
|
|
|
return {
|
|
...state,
|
|
document: {
|
|
...state.document,
|
|
assets: state.document.assets.map((asset) =>
|
|
asset.id === targetAsset.id
|
|
? {
|
|
...asset,
|
|
source: payload.source,
|
|
mimeType: payload.mimeType ?? asset.mimeType,
|
|
provenance: generationProvenance(candidate, "replacement"),
|
|
}
|
|
: asset,
|
|
),
|
|
},
|
|
editor: {
|
|
...state.editor,
|
|
generation: removeGenerationCandidate(state.editor.generation, candidate.id),
|
|
selection: { artboardId: targetLayerLocation.artboardId, layerIds: [candidate.inpaint.targetLayerId] },
|
|
},
|
|
};
|
|
},
|
|
};
|
|
|
|
export const generationCommands = [
|
|
generationAddCandidateCommand,
|
|
generationSelectCandidateCommand,
|
|
generationSetCompareModeCommand,
|
|
generationRemoveCandidateCommand,
|
|
generationClearCandidatesCommand,
|
|
generationApplyCandidateAsLayerCommand,
|
|
generationReplaceCandidatePixelsCommand,
|
|
] satisfies Command<unknown>[];
|
|
|
|
type LayerLocation = {
|
|
artboardId: ArtboardId;
|
|
layer: Layer;
|
|
};
|
|
|
|
function insertLayerAtTop(document: ImageDocument, artboardId: ArtboardId, layer: Layer): ImageDocument {
|
|
return {
|
|
...document,
|
|
artboards: document.artboards.map((artboard) => artboard.id === artboardId ? { ...artboard, layers: [layer, ...artboard.layers] } : artboard),
|
|
};
|
|
}
|
|
|
|
function findLayerLocation(document: ImageDocument, layerId: LayerId): LayerLocation | undefined {
|
|
for (const artboard of document.artboards) {
|
|
const layer = findLayerInTree(artboard.layers, layerId);
|
|
if (layer) return { artboardId: artboard.id, layer };
|
|
}
|
|
return undefined;
|
|
}
|
|
|
|
function findLayerInTree(layers: readonly Layer[], layerId: LayerId): Layer | undefined {
|
|
for (const layer of layers) {
|
|
if (layer.id === layerId) return layer;
|
|
if (layer.type === "group") {
|
|
const child = findLayerInTree(layer.children, layerId);
|
|
if (child) return child;
|
|
}
|
|
}
|
|
return undefined;
|
|
}
|
|
|
|
function removeGenerationCandidate(generation: GenerationState, candidateId: string): GenerationState {
|
|
const removedIndex = generation.candidates.findIndex((candidate) => candidate.id === candidateId);
|
|
if (removedIndex < 0) return generation;
|
|
|
|
const candidates = generation.candidates.filter((candidate) => candidate.id !== candidateId);
|
|
const selectionStillExists = generation.selectedCandidateId
|
|
? candidates.some((candidate) => candidate.id === generation.selectedCandidateId)
|
|
: false;
|
|
const selectedCandidateId = selectionStillExists
|
|
? generation.selectedCandidateId
|
|
: candidates[Math.min(removedIndex, candidates.length - 1)]?.id;
|
|
|
|
return {
|
|
candidates,
|
|
selectedCandidateId,
|
|
compareMode: candidates.length > 0 ? generation.compareMode : "result",
|
|
};
|
|
}
|
|
|
|
function generationProvenance(candidate: GenerationCandidate, acceptance: GeneratedAssetAcceptance): AssetGenerationProvenance {
|
|
return {
|
|
kind: "generated",
|
|
candidateId: candidate.id,
|
|
mode: candidate.mode,
|
|
acceptance,
|
|
prompt: candidate.settings.prompt,
|
|
negativePrompt: candidate.settings.negativePrompt,
|
|
seed: candidate.seed,
|
|
outputSize: { ...candidate.intrinsicSize },
|
|
settings: {
|
|
architecture: candidate.settings.architecture,
|
|
model: candidate.settings.model,
|
|
textEncoder: candidate.settings.textEncoder,
|
|
vae: candidate.settings.vae,
|
|
strength: candidate.settings.strength,
|
|
steps: candidate.settings.steps,
|
|
cfg: candidate.settings.cfg,
|
|
sampler: candidate.settings.sampler,
|
|
scheduler: candidate.settings.scheduler,
|
|
width: candidate.settings.width,
|
|
height: candidate.settings.height,
|
|
},
|
|
inpaint: candidate.inpaint
|
|
? {
|
|
targetLayerId: candidate.inpaint.targetLayerId,
|
|
maskLayerId: candidate.inpaint.maskLayerId,
|
|
sourceAssetId: candidate.inpaint.sourceAssetId,
|
|
maskAssetId: candidate.inpaint.maskAssetId,
|
|
crop: {
|
|
assetBounds: { ...candidate.inpaint.crop.assetBounds },
|
|
documentBounds: { ...candidate.inpaint.crop.documentBounds },
|
|
padding: candidate.inpaint.crop.padding,
|
|
maskedAreaOnly: candidate.inpaint.crop.maskedAreaOnly,
|
|
},
|
|
mask: {
|
|
polarity: candidate.inpaint.mask.polarity,
|
|
activeBounds: { ...candidate.inpaint.mask.activeBounds },
|
|
},
|
|
backend: {
|
|
growMaskBy: candidate.inpaint.backend.growMaskBy,
|
|
maskedContent: candidate.inpaint.backend.maskedContent,
|
|
maskBlur: candidate.inpaint.backend.maskBlur,
|
|
maskFeather: candidate.inpaint.backend.maskFeather,
|
|
maskExpand: candidate.inpaint.backend.maskExpand,
|
|
cropPadding: candidate.inpaint.backend.cropPadding,
|
|
},
|
|
}
|
|
: undefined,
|
|
};
|
|
}
|