diff --git a/app/app.ts b/app/app.ts index 9b3bbc9..63576da 100644 --- a/app/app.ts +++ b/app/app.ts @@ -11,6 +11,7 @@ import { viewportCommands } from "@commands/viewport"; import { workspaceCommands } from "@commands/workspace"; import { editorCommands } from "@commands/editor"; import { projectCommands } from "@commands/project"; +import { inpaintRegionCommands } from "@commands/inpaint-region"; import { createInitialAppState } from "@editor/initial-state"; import { createAppStore } from "@editor/store"; import { createGenerationWorkflow } from "@operations/generation/workflow"; @@ -19,7 +20,7 @@ import { createDocumentActions } from "./document-actions"; export type ImageStudioApp = ReturnType; export function createImageStudioApp(options?: { documentName?: string; createDefaultArtboard?: boolean }) { - const registry = createCommandRegistry([...projectCommands, ...viewportCommands, ...selectionCommands, ...documentCommands, ...toolCommands, ...generationCommands, ...transformCommands, ...historyCommands, ...commandPaletteCommands, ...workspaceCommands, ...editorCommands]); + const registry = createCommandRegistry([...projectCommands, ...viewportCommands, ...selectionCommands, ...documentCommands, ...inpaintRegionCommands, ...toolCommands, ...generationCommands, ...transformCommands, ...historyCommands, ...commandPaletteCommands, ...workspaceCommands, ...editorCommands]); const store = createAppStore(createInitialAppState(options?.documentName), registry); const generation = createGenerationWorkflow(store); const documentActions = createDocumentActions(store, generation); diff --git a/commands/document-tree.ts b/commands/document-tree.ts index e06592b..16e0c65 100644 --- a/commands/document-tree.ts +++ b/commands/document-tree.ts @@ -3,6 +3,7 @@ import type { ArtboardId, LayerId } from "@core/id"; import type { Layer } from "@core/layer"; import type { LayerGroup } from "@core/layer-group"; import { getLayerMask } from "@core/layer-mask-utils"; +import type { MaskEditState } from "@editor/state"; export type LayerLocation = { artboardId: ArtboardId; @@ -242,13 +243,18 @@ export function removeMissingMaskReferences(document: ImageDocument): ImageDocum }); } -export function isMaskEditFor(maskEdit: { targetLayerId: LayerId; maskLayerId: LayerId } | undefined, targetLayerId: LayerId, maskLayerId: LayerId) { +export function isMaskEditFor(maskEdit: MaskEditState | undefined, targetLayerId: LayerId, maskLayerId: LayerId) { return maskEdit?.targetLayerId === targetLayerId && maskEdit.maskLayerId === maskLayerId; } -export function isMaskEditValid(maskEdit: { targetLayerId: LayerId; maskLayerId: LayerId } | undefined, document: ImageDocument) { +export function isMaskEditValid(maskEdit: MaskEditState | undefined, document: ImageDocument) { if (!maskEdit) return false; const target = findLayerLocation(document, maskEdit.targetLayerId)?.layer; + if (maskEdit.kind === "inpaintRegion") { + const region = document.inpaintRegions.find((candidate) => candidate.id === maskEdit.inpaintRegionId && candidate.targetLayerId === maskEdit.targetLayerId && candidate.maskAssetId === maskEdit.maskAssetId); + return Boolean(target && region && document.assets.some((asset) => asset.id === region.maskAssetId)); + } + if (!maskEdit.maskLayerId) return false; const mask = findLayerLocation(document, maskEdit.maskLayerId)?.layer; return Boolean(target && getLayerMask(target)?.maskLayerId === maskEdit.maskLayerId && mask && mask.type !== "group"); } @@ -300,3 +306,26 @@ export function mapAllLayersInTree(layers: Layer[], mapLayer: (layer: Layer) => return mapLayer(mapped); }); } + +export function removeInpaintRegionsForTargets(document: ImageDocument, targetLayerIds: ReadonlySet): ImageDocument { + const removed = document.inpaintRegions.filter((region) => targetLayerIds.has(region.targetLayerId)); + if (removed.length === 0) return document; + const inpaintRegions = document.inpaintRegions.filter((region) => !targetLayerIds.has(region.targetLayerId)); + const removedMaskAssetIds = new Set(removed.map((region) => region.maskAssetId)); + const retainedRegionAssetIds = new Set(inpaintRegions.map((region) => region.maskAssetId)); + const layerAssetIds = new Set(); + for (const artboard of document.artboards) { + const stack = [...artboard.layers]; + while (stack.length > 0) { + const layer = stack.pop(); + if (!layer) continue; + if (layer.type === "group") stack.push(...layer.children); + else if (layer.type === "image" || layer.type === "raster") layerAssetIds.add(layer.assetId); + } + } + return { + ...document, + inpaintRegions, + assets: document.assets.filter((asset) => !removedMaskAssetIds.has(asset.id) || retainedRegionAssetIds.has(asset.id) || layerAssetIds.has(asset.id)), + }; +} diff --git a/commands/document.test.ts b/commands/document.test.ts index cac282e..2ec20af 100644 --- a/commands/document.test.ts +++ b/commands/document.test.ts @@ -285,7 +285,7 @@ describe("document commands", () => { inverted: false, }); expect(masked.document.artboards[0]?.layers[1]?.clippingMask).toEqual({ maskLayerId: "mask" }); - expect(masked.editor.maskEdit).toEqual({ targetLayerId: "target", maskLayerId: "mask" }); + expect(masked.editor.maskEdit).toEqual({ kind: "layerMask", targetLayerId: "target", maskLayerId: "mask", maskAssetId: "mask-asset" }); expect(masked.editor.tools.activeTool).toBe("brush"); }); diff --git a/commands/document.ts b/commands/document.ts index 0efa8ea..acbeedc 100644 --- a/commands/document.ts +++ b/commands/document.ts @@ -1,4 +1,4 @@ -import { findLayerLocation, isReferencedMaskLayer, mapLayerInDocument, insertLayer, replaceLayerListInDocument, replaceSelectedLayersWithGroup, removeLayerFromDocument, ungroupLayerInDocument, findGroup, removeLayerMaskReference, withLayerMask, removeUnreferencedMaskLayer, removeMissingMaskReferences, isMaskEditFor, isMaskEditValid, collectLayerIds, collectClippingMaskIds, collectAttachedMaskIds } from "./document-tree"; +import { findLayerLocation, isReferencedMaskLayer, mapLayerInDocument, insertLayer, replaceLayerListInDocument, replaceSelectedLayersWithGroup, removeLayerFromDocument, ungroupLayerInDocument, findGroup, removeLayerMaskReference, withLayerMask, removeUnreferencedMaskLayer, removeMissingMaskReferences, isMaskEditFor, isMaskEditValid, collectLayerIds, collectClippingMaskIds, collectAttachedMaskIds, removeInpaintRegionsForTargets } from "./document-tree"; import type { Asset } from "@core/asset"; import type { ImageDocument } from "@core/document"; import type { Rect } from "@core/geometry"; @@ -235,10 +235,14 @@ export const documentRemoveArtboardCommand: Command artboard.id === payload.id); + const withoutArtboard = { ...state.document, artboards: state.document.artboards.filter((artboard) => artboard.id !== payload.id), }; + const removedLayerIds = new Set(); + for (const layer of removedArtboard?.layers ?? []) collectLayerIds(layer, removedLayerIds); + const document = removeInpaintRegionsForTargets(withoutArtboard, removedLayerIds); return { ...state, @@ -630,14 +634,15 @@ export const documentAddLayerMaskCommand: Command = const existingMaskId = getLayerMask(targetLocation.layer)?.maskLayerId; if (existingMaskId) { const existingMaskLocation = findLayerLocation(state.document, existingMaskId); - if (existingMaskLocation?.layer.type === "group") return state; - if (existingMaskLocation) { + const existingMaskLayer = existingMaskLocation?.layer; + if (existingMaskLayer && existingMaskLayer.type !== "image" && existingMaskLayer.type !== "raster") return state; + if (existingMaskLocation && existingMaskLayer) { return { ...state, editor: { ...state.editor, selection: { artboardId: targetLocation.artboardId, layerIds: [payload.layerId] }, - maskEdit: { targetLayerId: payload.layerId, maskLayerId: existingMaskId }, + maskEdit: { kind: "layerMask", targetLayerId: payload.layerId, maskLayerId: existingMaskId, maskAssetId: existingMaskLayer.assetId }, tools: { ...state.editor.tools, activeTool: "brush", @@ -670,7 +675,7 @@ export const documentAddLayerMaskCommand: Command = editor: { ...state.editor, selection: { artboardId: targetLocation.artboardId, layerIds: [payload.layerId] }, - maskEdit: { targetLayerId: payload.layerId, maskLayerId: maskLayer.id }, + maskEdit: { kind: "layerMask", targetLayerId: payload.layerId, maskLayerId: maskLayer.id, maskAssetId: payload.asset.id }, tools: { ...state.editor.tools, activeTool: "brush", @@ -749,7 +754,8 @@ export const documentRemoveLayerCommand: Command = { const removedLayerIds = collectLayerIds(removed.layer); const removedMaskLayerIds = collectClippingMaskIds([removed.layer]); const cleanedReferences = removeMissingMaskReferences(removed.document); - const document = [...removedMaskLayerIds].reduce((nextDocument, maskLayerId) => removeUnreferencedMaskLayer(nextDocument, maskLayerId), cleanedReferences); + const withoutMasks = [...removedMaskLayerIds].reduce((nextDocument, maskLayerId) => removeUnreferencedMaskLayer(nextDocument, maskLayerId), cleanedReferences); + const document = removeInpaintRegionsForTargets(withoutMasks, removedLayerIds); const selection = { ...state.editor.selection, layerIds: state.editor.selection.layerIds.filter((id) => !removedLayerIds.has(id)), diff --git a/commands/generation.test.ts b/commands/generation.test.ts index 7463287..6d23665 100644 --- a/commands/generation.test.ts +++ b/commands/generation.test.ts @@ -231,11 +231,14 @@ function generationCandidate(id: string, inpaint = false): GenerationCandidate { maskImage: "mask", inpaint: { targetLayerId: "source-layer", - maskLayerId: "mask-layer", + regionId: "region", sourceAssetId: "source-asset", maskAssetId: "mask-asset", inputImage: "input", maskImage: "mask", + editMaskImage: "edit-mask", + blendMaskImage: "blend-mask", + revision: { source: "source-revision", mask: "mask-revision" }, crop: { assetBounds: { x: 0, y: 0, w: 64, h: 64 }, documentBounds: { x: 0, y: 0, w: 64, h: 64 }, diff --git a/commands/generation.ts b/commands/generation.ts index bb73b58..9292bbf 100644 --- a/commands/generation.ts +++ b/commands/generation.ts @@ -23,6 +23,7 @@ export type GenerationSetCompareModePayload = { export type GenerationRemoveCandidatePayload = { candidateId: GenerationCandidateId; }; +export type GenerationToggleCandidateFavoritePayload = { candidateId: GenerationCandidateId }; export type GenerationReuseCandidateSettingsPayload = { candidateId: GenerationCandidateId; @@ -41,6 +42,7 @@ export type GenerationReplaceCandidatePixelsPayload = { }; export type GenerationStartJobPayload = { jobId: GenerationJobId; kind: GenerationJobKind; label: string; startedAt: number }; +export type GenerationUpdateJobPayload = { jobId: GenerationJobId; progress: number; detail: string }; export type GenerationSucceedJobPayload = { jobId: GenerationJobId; finishedAt: number }; export type GenerationFailJobPayload = { jobId: GenerationJobId; finishedAt: number; error: string }; export type GenerationCancelJobPayload = { jobId: GenerationJobId; finishedAt: number }; @@ -57,7 +59,8 @@ export const generationAddCandidateCommand: Command candidate.id !== payload.candidate.id)]); + const existing = state.editor.generation.candidates.filter((candidate) => candidate.id !== payload.candidate.id); + const candidates = retainCandidateBudget([payload.candidate, ...existing.filter((candidate) => candidate.favorite), ...existing.filter((candidate) => !candidate.favorite)]); return { ...state, editor: { @@ -130,6 +133,25 @@ export const generationRemoveCandidateCommand: Command = { + id: commandIds.generationToggleCandidateFavorite, + name: "Favorite generation candidate", + history: { mode: "ignore" }, + execute({ state }, payload) { + if (!state.editor.generation.candidates.some((candidate) => candidate.id === payload.candidateId)) return state; + return { + ...state, + editor: { + ...state.editor, + generation: { + ...state.editor.generation, + candidates: state.editor.generation.candidates.map((candidate) => candidate.id === payload.candidateId ? { ...candidate, favorite: !candidate.favorite } : candidate), + }, + }, + }; + }, +}; + export const generationClearCandidatesCommand: Command = { id: commandIds.generationClearCandidates, name: "Clear generation candidates", @@ -260,6 +282,17 @@ export const generationStartJobCommand: Command = { }, }; +export const generationUpdateJobCommand: Command = { + id: commandIds.generationUpdateJob, + name: "Update generation progress", + history: { mode: "ignore" }, + execute({ state }, payload) { + const job = state.editor.generation.jobs.find((candidate) => candidate.id === payload.jobId); + if (!job || job.status !== "running" || !Number.isFinite(payload.progress) || !payload.detail.trim()) return state; + return updateJobs(state, state.editor.generation.jobs.map((candidate) => candidate.id === payload.jobId ? { ...candidate, progress: Math.max(0, Math.min(1, payload.progress)), detail: payload.detail.trim() } : candidate)); + }, +}; + export const generationSucceedJobCommand: Command = { id: commandIds.generationSucceedJob, name: "Complete generation job", @@ -322,11 +355,13 @@ export const generationCommands = [ generationSelectCandidateCommand, generationSetCompareModeCommand, generationRemoveCandidateCommand, + generationToggleCandidateFavoriteCommand, generationClearCandidatesCommand, generationReuseCandidateSettingsCommand, generationApplyCandidateAsLayerCommand, generationReplaceCandidatePixelsCommand, generationStartJobCommand, + generationUpdateJobCommand, generationSucceedJobCommand, generationFailJobCommand, generationCancelJobCommand, @@ -364,7 +399,7 @@ function retainCandidateBudget(candidates: GenerationCandidate[]): GenerationCan } function candidateRetainedBytes(candidate: GenerationCandidate): number { - return [candidate.source, candidate.inputImage, candidate.maskImage, candidate.inpaint?.inputImage, candidate.inpaint?.maskImage] + return [candidate.source, candidate.inputImage, candidate.maskImage, candidate.blendMaskImage, candidate.inpaint?.inputImage, candidate.inpaint?.maskImage, candidate.inpaint?.editMaskImage, candidate.inpaint?.blendMaskImage] .reduce((total, source) => total + (source?.length ?? 0) * 2, 0); } @@ -441,11 +476,14 @@ function generationProvenance(candidate: GenerationCandidate, acceptance: Genera scheduler: candidate.settings.scheduler, width: candidate.settings.width, height: candidate.settings.height, + batchSize: candidate.settings.batchSize, + refinePass: candidate.settings.refinePass, + refineStrength: candidate.settings.refineStrength, }, inpaint: candidate.inpaint ? { targetLayerId: candidate.inpaint.targetLayerId, - maskLayerId: candidate.inpaint.maskLayerId, + regionId: candidate.inpaint.regionId, sourceAssetId: candidate.inpaint.sourceAssetId, maskAssetId: candidate.inpaint.maskAssetId, crop: { @@ -465,6 +503,11 @@ function generationProvenance(candidate: GenerationCandidate, acceptance: Genera maskFeather: candidate.inpaint.backend.maskFeather, maskExpand: candidate.inpaint.backend.maskExpand, cropPadding: candidate.inpaint.backend.cropPadding, + profile: candidate.settings.inpaint.profile, + structureControl: candidate.settings.inpaint.structureControl, + controlStrength: candidate.settings.inpaint.controlStrength, + controlModel: candidate.settings.inpaint.controlModel, + colorMatch: candidate.settings.inpaint.colorMatch, }, } : undefined, diff --git a/commands/ids.ts b/commands/ids.ts index 54d7ecc..fed46f4 100644 --- a/commands/ids.ts +++ b/commands/ids.ts @@ -30,6 +30,9 @@ export const commandIds = { documentAddLayerMask: "document.addLayerMask", documentApplyLayerMaskOperation: "document.applyLayerMaskOperation", documentRemoveLayerMask: "document.removeLayerMask", + documentAddInpaintRegion: "document.addInpaintRegion", + documentApplyInpaintRegionMaskOperation: "document.applyInpaintRegionMaskOperation", + documentRemoveInpaintRegion: "document.removeInpaintRegion", selectionSet: "selection.set", selectionClear: "selection.clear", selectionAddLayer: "selection.addLayer", @@ -41,8 +44,12 @@ export const commandIds = { toolSetMagicWandSettings: "tool.setMagicWandSettings", toolSetBrushPreview: "tool.setBrushPreview", toolSetBrushStrokePreview: "tool.setBrushStrokePreview", + toolBeginMaskShape: "tool.beginMaskShape", + toolAppendMaskShape: "tool.appendMaskShape", + toolClearMaskShape: "tool.clearMaskShape", toolSetMaskViewMode: "tool.setMaskViewMode", toolEnterMaskEdit: "tool.enterMaskEdit", + toolEnterInpaintRegionEdit: "tool.enterInpaintRegionEdit", toolExitMaskEdit: "tool.exitMaskEdit", toolEnterTemporaryPan: "tool.enterTemporaryPan", toolExitTemporaryPan: "tool.exitTemporaryPan", @@ -50,11 +57,13 @@ export const commandIds = { generationSelectCandidate: "generation.selectCandidate", generationSetCompareMode: "generation.setCompareMode", generationRemoveCandidate: "generation.removeCandidate", + generationToggleCandidateFavorite: "generation.toggleCandidateFavorite", generationClearCandidates: "generation.clearCandidates", generationReuseCandidateSettings: "generation.reuseCandidateSettings", generationApplyCandidateAsLayer: "generation.applyCandidateAsLayer", generationReplaceCandidatePixels: "generation.replaceCandidatePixels", generationStartJob: "generation.startJob", + generationUpdateJob: "generation.updateJob", generationSucceedJob: "generation.succeedJob", generationFailJob: "generation.failJob", generationCancelJob: "generation.cancelJob", diff --git a/commands/index.ts b/commands/index.ts index 1f698c1..4579647 100644 --- a/commands/index.ts +++ b/commands/index.ts @@ -101,10 +101,12 @@ export type { CommandRegistry } from "./registry"; export { createCommandRegistry } from "./registry"; export { selectionAddLayerCommand, selectionClearCommand, selectionCommands, selectionSetCommand } from "./selection"; export type { SelectionAddLayerPayload, SelectionSetPayload } from "./selection"; -export { toolChooseGenerateIntentCommand, toolCommands, toolEnterMaskEditCommand, toolEnterTemporaryPanCommand, toolExitMaskEditCommand, toolExitTemporaryPanCommand, toolSetActiveCommand, toolSetBrushPreviewCommand, toolSetBrushSettingsCommand, toolSetBrushStrokePreviewCommand, toolSetChromaKeySettingsCommand, toolSetGenerateSettingsCommand, toolSetMagicWandSettingsCommand, toolSetMaskViewModeCommand } from "./tool"; +export { toolChooseGenerateIntentCommand, toolCommands, toolEnterInpaintRegionEditCommand, toolEnterMaskEditCommand, toolEnterTemporaryPanCommand, toolExitMaskEditCommand, toolExitTemporaryPanCommand, toolSetActiveCommand, toolSetBrushPreviewCommand, toolSetBrushSettingsCommand, toolSetBrushStrokePreviewCommand, toolSetChromaKeySettingsCommand, toolSetGenerateSettingsCommand, toolSetMagicWandSettingsCommand, toolSetMaskViewModeCommand } from "./tool"; export { transformBeginCommand, transformCommands, transformEndCommand, transformSetBoundsCommand, transformSetRotationCommand, transformUpdateCommand } from "./transform"; export type { TransformBeginPayload, TransformSetBoundsPayload, TransformSetRotationPayload, TransformUpdatePayload } from "./transform"; -export type { ToolChooseGenerateIntentPayload, ToolEnterMaskEditPayload, ToolSetActivePayload, ToolSetBrushPreviewPayload, ToolSetBrushSettingsPayload, ToolSetBrushStrokePreviewPayload, ToolSetChromaKeySettingsPayload, ToolSetGenerateSettingsPayload, ToolSetMagicWandSettingsPayload, ToolSetMaskViewModePayload } from "./tool"; +export type { ToolChooseGenerateIntentPayload, ToolEnterInpaintRegionEditPayload, ToolEnterMaskEditPayload, ToolSetActivePayload, ToolSetBrushPreviewPayload, ToolSetBrushSettingsPayload, ToolSetBrushStrokePreviewPayload, ToolSetChromaKeySettingsPayload, ToolSetGenerateSettingsPayload, ToolSetMagicWandSettingsPayload, ToolSetMaskViewModePayload } from "./tool"; +export { inpaintRegionCommands } from "./inpaint-region"; +export type { DocumentAddInpaintRegionPayload, DocumentApplyInpaintRegionMaskOperationPayload, DocumentRemoveInpaintRegionPayload } from "./inpaint-region"; export { viewportCommands, viewportPanCommand, diff --git a/commands/inpaint-region.test.ts b/commands/inpaint-region.test.ts new file mode 100644 index 0000000..b04a2a3 --- /dev/null +++ b/commands/inpaint-region.test.ts @@ -0,0 +1,32 @@ +import { describe, expect, test } from "bun:test"; +import { createInitialAppState } from "@editor/initial-state"; +import { documentAddInpaintRegionCommand, documentApplyInpaintRegionMaskOperationCommand, documentRemoveInpaintRegionCommand } from "./inpaint-region"; + +describe("inpaint region commands", () => { + test("persists an AI edit mask independently from layer visibility", () => { + const state = stateWithTarget(); + const added = documentAddInpaintRegionCommand.execute({ state }, { + region: { id: "region", name: "Target AI edit", targetLayerId: "target", maskAssetId: "mask", enabled: true }, + maskAsset: { id: "mask", name: "Mask", mimeType: "image/png", source: "black", intrinsicSize: { w: 64, h: 32 } }, + }); + expect(added.document.inpaintRegions).toHaveLength(1); + expect(added.document.artboards[0]?.layers[0]).not.toHaveProperty("layerMask"); + expect(added.document.artboards[0]?.layers[0]?.visible).toBe(true); + + const painted = documentApplyInpaintRegionMaskOperationCommand.execute({ state: added }, { regionId: "region", source: "painted", operation: { type: "paint" } }); + expect(painted.document.assets.find((asset) => asset.id === "mask")?.source).toBe("painted"); + + const removed = documentRemoveInpaintRegionCommand.execute({ state: painted }, { regionId: "region" }); + expect(removed.document.inpaintRegions).toHaveLength(0); + expect(removed.document.assets.some((asset) => asset.id === "mask")).toBe(false); + }); +}); + +function stateWithTarget() { + const state = createInitialAppState("Test"); + state.document.assets.push({ id: "source", name: "Source", mimeType: "image/png", source: "pixels", intrinsicSize: { w: 64, h: 32 } }); + state.document.artboards.push({ id: "board", name: "Board", bounds: { x: 0, y: 0, w: 64, h: 32 }, backgroundColor: "transparent", visible: true, locked: false, layers: [ + { id: "target", type: "raster", name: "Target", visible: true, locked: false, opacity: 1, assetId: "source", transform: { position: { x: 0, y: 0 }, scale: { x: 1, y: 1 }, rotation: 0 } }, + ] }); + return state; +} diff --git a/commands/inpaint-region.ts b/commands/inpaint-region.ts new file mode 100644 index 0000000..e5c0774 --- /dev/null +++ b/commands/inpaint-region.ts @@ -0,0 +1,100 @@ +import type { Asset } from "@core/asset"; +import type { InpaintRegion } from "@core/inpaint-region"; +import type { InpaintRegionId } from "@core/id"; +import type { Command } from "./command"; +import { commandIds } from "./ids"; +import type { LayerMaskOperation } from "./document"; +import { findLayerLocation } from "./document-tree"; + +export type DocumentAddInpaintRegionPayload = { + region: InpaintRegion; + maskAsset: Asset; +}; + +export type DocumentApplyInpaintRegionMaskOperationPayload = { + regionId: InpaintRegionId; + source: string; + mimeType?: string; + operation: LayerMaskOperation; +}; + +export type DocumentRemoveInpaintRegionPayload = { + regionId: InpaintRegionId; +}; + +export const documentAddInpaintRegionCommand: Command = { + id: commandIds.documentAddInpaintRegion, + name: "Add inpaint region", + execute({ state }, payload) { + const target = findLayerLocation(state.document, payload.region.targetLayerId)?.layer; + if (!target || (target.type !== "image" && target.type !== "raster")) return state; + const targetAsset = state.document.assets.find((asset) => asset.id === target.assetId); + if (!targetAsset || state.document.inpaintRegions.some((region) => region.id === payload.region.id)) return state; + if (state.document.assets.some((asset) => asset.id === payload.maskAsset.id)) return state; + if (payload.region.maskAssetId !== payload.maskAsset.id || !payload.region.name.trim()) return state; + if (Math.round(payload.maskAsset.intrinsicSize.w) !== Math.round(targetAsset.intrinsicSize.w) || Math.round(payload.maskAsset.intrinsicSize.h) !== Math.round(targetAsset.intrinsicSize.h)) return state; + + return { + ...state, + document: { + ...state.document, + assets: [...state.document.assets, payload.maskAsset], + inpaintRegions: [...state.document.inpaintRegions, { ...payload.region, name: payload.region.name.trim() }], + }, + }; + }, +}; + +export const documentApplyInpaintRegionMaskOperationCommand: Command = { + id: commandIds.documentApplyInpaintRegionMaskOperation, + name: "Apply inpaint region mask operation", + execute({ state }, payload) { + const region = state.document.inpaintRegions.find((candidate) => candidate.id === payload.regionId); + if (!region || !payload.source.trim()) return state; + return { + ...state, + document: { + ...state.document, + assets: state.document.assets.map((asset) => asset.id === region.maskAssetId ? { ...asset, source: payload.source, mimeType: payload.mimeType ?? asset.mimeType } : asset), + }, + editor: { + ...state.editor, + brushStrokePreview: state.editor.brushStrokePreview?.assetId === region.maskAssetId ? undefined : state.editor.brushStrokePreview, + }, + }; + }, +}; + +export const documentRemoveInpaintRegionCommand: Command = { + id: commandIds.documentRemoveInpaintRegion, + name: "Remove inpaint region", + execute({ state }, payload) { + const region = state.document.inpaintRegions.find((candidate) => candidate.id === payload.regionId); + if (!region) return state; + const inpaintRegions = state.document.inpaintRegions.filter((candidate) => candidate.id !== payload.regionId); + const maskStillUsed = inpaintRegions.some((candidate) => candidate.maskAssetId === region.maskAssetId); + const exitingEdit = state.editor.maskEdit?.inpaintRegionId === region.id; + const contextualTool = state.editor.tools.activeTool === "semanticSelect" || state.editor.tools.activeTool === "maskLasso" || state.editor.tools.activeTool === "maskRectangle"; + return { + ...state, + document: { + ...state.document, + inpaintRegions, + assets: maskStillUsed ? state.document.assets : state.document.assets.filter((asset) => asset.id !== region.maskAssetId), + }, + editor: { + ...state.editor, + maskEdit: exitingEdit ? undefined : state.editor.maskEdit, + brushStrokePreview: state.editor.brushStrokePreview?.assetId === region.maskAssetId ? undefined : state.editor.brushStrokePreview, + maskShapeSession: exitingEdit ? undefined : state.editor.maskShapeSession, + tools: exitingEdit && contextualTool ? { ...state.editor.tools, activeTool: "select", interactionMode: { type: "tool", tool: "select" } } : state.editor.tools, + }, + }; + }, +}; + +export const inpaintRegionCommands = [ + documentAddInpaintRegionCommand, + documentApplyInpaintRegionMaskOperationCommand, + documentRemoveInpaintRegionCommand, +] satisfies Command[]; diff --git a/commands/payloads.ts b/commands/payloads.ts index ff6f49e..9f6043b 100644 --- a/commands/payloads.ts +++ b/commands/payloads.ts @@ -35,11 +35,13 @@ import type { GenerationAddCandidatePayload, GenerationApplyCandidateAsLayerPayload, GenerationRemoveCandidatePayload, + GenerationToggleCandidateFavoritePayload, GenerationReuseCandidateSettingsPayload, GenerationReplaceCandidatePixelsPayload, GenerationSelectCandidatePayload, GenerationSetCompareModePayload, GenerationStartJobPayload, + GenerationUpdateJobPayload, GenerationSucceedJobPayload, GenerationFailJobPayload, GenerationCancelJobPayload, @@ -52,11 +54,12 @@ import type { CommandPaletteSetSelectedIndexPayload, } from "./palette"; import type { SelectionAddLayerPayload, SelectionSetPayload } from "./selection"; -import type { ToolChooseGenerateIntentPayload, ToolEnterMaskEditPayload, ToolSetActivePayload, ToolSetBrushPreviewPayload, ToolSetBrushSettingsPayload, ToolSetBrushStrokePreviewPayload, ToolSetChromaKeySettingsPayload, ToolSetGenerateSettingsPayload, ToolSetMagicWandSettingsPayload, ToolSetMaskViewModePayload } from "./tool"; +import type { ToolAppendMaskShapePayload, ToolBeginMaskShapePayload, ToolChooseGenerateIntentPayload, ToolEnterInpaintRegionEditPayload, ToolEnterMaskEditPayload, ToolSetActivePayload, ToolSetBrushPreviewPayload, ToolSetBrushSettingsPayload, ToolSetBrushStrokePreviewPayload, ToolSetChromaKeySettingsPayload, ToolSetGenerateSettingsPayload, ToolSetMagicWandSettingsPayload, ToolSetMaskViewModePayload } from "./tool"; import type { TransformBeginPayload, TransformSetBoundsPayload, TransformSetRotationPayload, TransformUpdatePayload } from "./transform"; import type { WorkspaceSetPanelPayload } from "./workspace"; import type { EditorSetPointerSessionPayload } from "./editor"; import type { ProjectOpenPayload } from "./project"; +import type { DocumentAddInpaintRegionPayload, DocumentApplyInpaintRegionMaskOperationPayload, DocumentRemoveInpaintRegionPayload } from "./inpaint-region"; import type { ViewportFitArtboardPayload, ViewportPanPayload, @@ -86,6 +89,9 @@ export type CommandPayloads = { [commandIds.documentAddLayerMask]: DocumentAddLayerMaskPayload; [commandIds.documentApplyLayerMaskOperation]: DocumentApplyLayerMaskOperationPayload; [commandIds.documentRemoveLayerMask]: DocumentRemoveLayerMaskPayload; + [commandIds.documentAddInpaintRegion]: DocumentAddInpaintRegionPayload; + [commandIds.documentApplyInpaintRegionMaskOperation]: DocumentApplyInpaintRegionMaskOperationPayload; + [commandIds.documentRemoveInpaintRegion]: DocumentRemoveInpaintRegionPayload; [commandIds.documentMoveLayer]: DocumentMoveLayerPayload; [commandIds.documentGroupLayers]: DocumentGroupLayersPayload; [commandIds.documentUngroupLayer]: DocumentUngroupLayerPayload; @@ -108,8 +114,12 @@ export type CommandPayloads = { [commandIds.toolSetMagicWandSettings]: ToolSetMagicWandSettingsPayload; [commandIds.toolSetBrushPreview]: ToolSetBrushPreviewPayload; [commandIds.toolSetBrushStrokePreview]: ToolSetBrushStrokePreviewPayload; + [commandIds.toolBeginMaskShape]: ToolBeginMaskShapePayload; + [commandIds.toolAppendMaskShape]: ToolAppendMaskShapePayload; + [commandIds.toolClearMaskShape]: void; [commandIds.toolSetMaskViewMode]: ToolSetMaskViewModePayload; [commandIds.toolEnterMaskEdit]: ToolEnterMaskEditPayload; + [commandIds.toolEnterInpaintRegionEdit]: ToolEnterInpaintRegionEditPayload; [commandIds.toolExitMaskEdit]: void; [commandIds.toolEnterTemporaryPan]: void; [commandIds.toolExitTemporaryPan]: void; @@ -117,11 +127,13 @@ export type CommandPayloads = { [commandIds.generationSelectCandidate]: GenerationSelectCandidatePayload; [commandIds.generationSetCompareMode]: GenerationSetCompareModePayload; [commandIds.generationRemoveCandidate]: GenerationRemoveCandidatePayload; + [commandIds.generationToggleCandidateFavorite]: GenerationToggleCandidateFavoritePayload; [commandIds.generationClearCandidates]: void; [commandIds.generationReuseCandidateSettings]: GenerationReuseCandidateSettingsPayload; [commandIds.generationApplyCandidateAsLayer]: GenerationApplyCandidateAsLayerPayload; [commandIds.generationReplaceCandidatePixels]: GenerationReplaceCandidatePixelsPayload; [commandIds.generationStartJob]: GenerationStartJobPayload; + [commandIds.generationUpdateJob]: GenerationUpdateJobPayload; [commandIds.generationSucceedJob]: GenerationSucceedJobPayload; [commandIds.generationFailJob]: GenerationFailJobPayload; [commandIds.generationCancelJob]: GenerationCancelJobPayload; diff --git a/commands/tool.test.ts b/commands/tool.test.ts index 060bddf..8a437e1 100644 --- a/commands/tool.test.ts +++ b/commands/tool.test.ts @@ -1,11 +1,12 @@ import { describe, expect, test } from "bun:test"; import { createInitialAppState } from "@editor/initial-state"; -import { toolChooseGenerateIntentCommand, toolEnterMaskEditCommand, toolEnterTemporaryPanCommand, toolExitMaskEditCommand, toolExitTemporaryPanCommand, toolSetActiveCommand, toolSetBrushPreviewCommand, toolSetBrushSettingsCommand, toolSetBrushStrokePreviewCommand, toolSetChromaKeySettingsCommand, toolSetGenerateSettingsCommand, toolSetMaskViewModeCommand } from "./tool"; +import { initialToolState } from "@editor/tools"; +import { toolAppendMaskShapeCommand, toolBeginMaskShapeCommand, toolChooseGenerateIntentCommand, toolEnterMaskEditCommand, toolEnterTemporaryPanCommand, toolExitMaskEditCommand, toolExitTemporaryPanCommand, toolSetActiveCommand, toolSetBrushPreviewCommand, toolSetBrushSettingsCommand, toolSetBrushStrokePreviewCommand, toolSetChromaKeySettingsCommand, toolSetGenerateSettingsCommand, toolSetMaskViewModeCommand } from "./tool"; -const defaultBrush = { color: "#111827", size: 8, hardness: 100 }; +const defaultBrush = { color: "#111827", size: 8, hardness: 100, opacity: 100, flow: 100, smoothing: 20, pressureSize: true }; const defaultChromaKey = { color: "#00ff00", tolerance: 32, softness: 24, feather: 0, choke: 0, despeckle: 0, spill: 50 }; const defaultMagicWand = { tolerance: 32, feather: 0, choke: 0, despeckle: 0, contiguous: true, mode: "replace" as const }; -const defaultGenerate = { architecture: "sdxl" as const, mode: "text-to-image" as const, model: "auto" as const, textEncoder: "auto", vae: "auto", prompt: "", negativePrompt: "", strength: 75, steps: 30, cfg: 7, seed: -1, sampler: "euler", scheduler: "normal", width: 1024, height: 1024, outpaint: { left: 128, top: 128, right: 128, bottom: 128, feathering: 32 }, inpaint: { maskedAreaOnly: true, cropPadding: 96, maskPolarity: "hidden" as const, maskedContent: "neutral" as const, growMaskBy: 6, maskExpand: 0, maskFeather: 0, maskBlur: 0, maskDespeckle: 0 } }; +const defaultGenerate = initialToolState.generate; describe("tool commands", () => { test("sets active tool", () => { @@ -24,7 +25,7 @@ describe("tool commands", () => { test("sets brush settings", () => { const next = toolSetBrushSettingsCommand.execute({ state: createInitialAppState("Test") }, { color: "#ff0000", size: 24, hardness: 50 }); - expect(next.editor.tools.brush).toEqual({ color: "#ff0000", size: 24, hardness: 50 }); + expect(next.editor.tools.brush).toEqual({ ...defaultBrush, color: "#ff0000", size: 24, hardness: 50 }); }); test("sets chroma key settings", () => { @@ -43,6 +44,11 @@ describe("tool commands", () => { expect(next.editor.tools.generate.inpaint).toEqual({ ...defaultGenerate.inpaint, cropPadding: 2048, growMaskBy: 0, maskExpand: -256, maskBlur: 256, maskPolarity: "revealed", maskedContent: "original" }); }); + test("applies outcome-oriented inpaint profile defaults", () => { + const next = toolSetGenerateSettingsCommand.execute({ state: createInitialAppState("Test") }, { mode: "inpaint", inpaint: { profile: "material" } }); + expect(next.editor.tools.generate).toMatchObject({ strength: 65, inpaint: { profile: "material", maskedContent: "edges", structureControl: "canny", controlStrength: 0.55, colorMatch: false } }); + }); + test("applies architecture defaults and filters unsupported modes", () => { const initial = toolSetGenerateSettingsCommand.execute({ state: createInitialAppState("Test") }, { mode: "inpaint" }); const next = toolSetGenerateSettingsCommand.execute({ state: initial }, { architecture: "z-image-turbo" }); @@ -69,6 +75,11 @@ describe("tool commands", () => { expect(variations.editor.tools.generate).toMatchObject({ architecture: "sdxl", mode: "image-to-image" }); }); + test("maps object removal to the dedicated inpaint profile", () => { + const removed = toolChooseGenerateIntentCommand.execute({ state: createInitialAppState("Test") }, { intent: "remove" }); + expect(removed.editor.tools.generate).toMatchObject({ architecture: "sdxl", mode: "inpaint", strength: 100, inpaint: { profile: "remove", maskedContent: "neutral", maskExpand: 12 } }); + }); + test("sets and clears brush preview", () => { const showing = toolSetBrushPreviewCommand.execute({ state: createInitialAppState("Test") }, { position: { x: 10, y: 20 } }); const cleared = toolSetBrushPreviewCommand.execute({ state: showing }, undefined); @@ -85,11 +96,25 @@ describe("tool commands", () => { expect(cleared.editor.brushStrokePreview).toBeUndefined(); }); + test("keeps rectangle mask gestures as two command-owned corner points", () => { + const state = createInitialAppState("Test"); + state.editor = { + ...state.editor, + tools: { ...state.editor.tools, activeTool: "maskRectangle", interactionMode: { type: "tool", tool: "maskRectangle" } }, + maskEdit: { kind: "inpaintRegion", targetLayerId: "target", inpaintRegionId: "region", maskAssetId: "mask" }, + }; + const begun = toolBeginMaskShapeCommand.execute({ state }, { point: { x: 10, y: 20 }, mode: "add" }); + const moved = toolAppendMaskShapeCommand.execute({ state: begun }, { point: { x: 40, y: 60 } }); + const movedAgain = toolAppendMaskShapeCommand.execute({ state: moved }, { point: { x: 50, y: 70 } }); + + expect(movedAgain.editor.maskShapeSession).toEqual({ shape: "rectangle", mode: "add", points: [{ x: 10, y: 20 }, { x: 50, y: 70 }] }); + }); + test("enters and exits mask edit", () => { const editing = toolEnterMaskEditCommand.execute({ state: stateWithMask() }, { targetLayerId: "target", maskLayerId: "mask" }); const exited = toolExitMaskEditCommand.execute({ state: editing }, undefined); - expect(editing.editor.maskEdit).toEqual({ targetLayerId: "target", maskLayerId: "mask" }); + expect(editing.editor.maskEdit).toEqual({ kind: "layerMask", targetLayerId: "target", maskLayerId: "mask", maskAssetId: "mask-asset" }); expect(editing.editor.selection).toEqual({ artboardId: "a1", layerIds: ["target"] }); expect(editing.editor.tools.activeTool).toBe("brush"); expect(exited.editor.maskEdit).toBeUndefined(); @@ -100,7 +125,7 @@ describe("tool commands", () => { const alpha = toolSetMaskViewModeCommand.execute({ state: editing }, { mode: "alpha" }); const ignored = toolSetMaskViewModeCommand.execute({ state: createInitialAppState("Test") }, { mode: "blackWhite" }); - expect(alpha.editor.maskEdit).toEqual({ targetLayerId: "target", maskLayerId: "mask", viewMode: "alpha" }); + expect(alpha.editor.maskEdit).toEqual({ kind: "layerMask", targetLayerId: "target", maskLayerId: "mask", maskAssetId: "mask-asset", viewMode: "alpha" }); expect(ignored.editor.maskEdit).toBeUndefined(); }); diff --git a/commands/tool.ts b/commands/tool.ts index 8cc4b00..67bbce7 100644 --- a/commands/tool.ts +++ b/commands/tool.ts @@ -1,10 +1,10 @@ import type { ImageDocument } from "@core/document"; import type { Vec2D } from "@core/geometry"; -import type { LayerId, ArtboardId, AssetId } from "@core/id"; +import type { LayerId, ArtboardId, AssetId, InpaintRegionId } from "@core/id"; import type { Layer } from "@core/layer"; import { getLayerMask } from "@core/layer-mask-utils"; import type { MaskViewMode } from "@editor/state"; -import { generateArchitectureDefaults } from "@editor/tools"; +import { generateArchitectureDefaults, inpaintProfileDefaults } from "@editor/tools"; import type { BrushSettings, ChromaKeySettings, GenerateIntent, GenerateSettings, MagicWandSettings, ToolId } from "@editor/tools"; import type { Command } from "./command"; import { commandIds } from "./ids"; @@ -15,7 +15,10 @@ export type ToolSetActivePayload = { export type ToolSetBrushSettingsPayload = Partial; -export type ToolSetGenerateSettingsPayload = Partial; +export type ToolSetGenerateSettingsPayload = Omit, "inpaint" | "outpaint"> & { + inpaint?: Partial; + outpaint?: Partial; +}; export type ToolChooseGenerateIntentPayload = { intent: GenerateIntent }; export type ToolSetChromaKeySettingsPayload = Partial; @@ -35,12 +38,19 @@ export type ToolSetBrushStrokePreviewPayload = export type ToolSetMaskViewModePayload = { mode: MaskViewMode; }; +export type ToolBeginMaskShapePayload = { point: Vec2D; mode: "replace" | "add" | "subtract" }; +export type ToolAppendMaskShapePayload = { point: Vec2D }; export type ToolEnterMaskEditPayload = { targetLayerId: LayerId; maskLayerId: LayerId; }; +export type ToolEnterInpaintRegionEditPayload = { + targetLayerId: LayerId; + regionId: InpaintRegionId; +}; + export const toolSetActiveCommand: Command = { id: commandIds.toolSetActive, name: "Set active tool", @@ -56,6 +66,7 @@ export const toolSetActiveCommand: Command = { }, brushPreview: undefined, brushStrokePreview: undefined, + maskShapeSession: undefined, workspace: state.editor.workspace.panel === "generate" || state.editor.workspace.panel === "chromaKey" ? { panel: "none" } : state.editor.workspace, @@ -74,6 +85,8 @@ export const toolSetGenerateSettingsCommand: Command = { create: "text-to-image", replace: "inpaint", + remove: "inpaint", extend: "outpaint", variations: "image-to-image", }; @@ -136,7 +158,7 @@ export const toolChooseGenerateIntentCommand: Command = color: payload.color ?? state.editor.tools.brush.color, size: clampNumber(payload.size ?? state.editor.tools.brush.size, 1, 200), hardness: clampNumber(payload.hardness ?? state.editor.tools.brush.hardness, 0, 100), + opacity: clampNumber(payload.opacity ?? state.editor.tools.brush.opacity, 0, 100), + flow: clampNumber(payload.flow ?? state.editor.tools.brush.flow, 1, 100), + smoothing: clampNumber(payload.smoothing ?? state.editor.tools.brush.smoothing, 0, 100), + pressureSize: payload.pressureSize ?? state.editor.tools.brush.pressureSize, }, }, }, @@ -270,6 +296,43 @@ export const toolSetMaskViewModeCommand: Command = { }, }; +export const toolBeginMaskShapeCommand: Command = { + id: commandIds.toolBeginMaskShape, + name: "Begin mask lasso", + execute({ state }, payload) { + const tool = state.editor.tools.activeTool; + if ((tool !== "maskLasso" && tool !== "maskRectangle") || state.editor.maskEdit?.kind !== "inpaintRegion") return state; + return { ...state, editor: { ...state.editor, maskShapeSession: { shape: tool === "maskRectangle" ? "rectangle" : "lasso", points: [{ ...payload.point }], mode: payload.mode } } }; + }, +}; + +export const toolAppendMaskShapeCommand: Command = { + id: commandIds.toolAppendMaskShape, + name: "Append mask lasso point", + history: { mode: "ignore" }, + execute({ state }, payload) { + const session = state.editor.maskShapeSession; + if (!session) return state; + if (session.shape === "rectangle") { + const start = session.points[0]; + if (!start) return state; + return { ...state, editor: { ...state.editor, maskShapeSession: { ...session, points: [start, { ...payload.point }] } } }; + } + const previous = session.points[session.points.length - 1]; + if (previous && Math.hypot(payload.point.x - previous.x, payload.point.y - previous.y) < 1) return state; + return { ...state, editor: { ...state.editor, maskShapeSession: { ...session, points: [...session.points, { ...payload.point }] } } }; + }, +}; + +export const toolClearMaskShapeCommand: Command = { + id: commandIds.toolClearMaskShape, + name: "Clear mask lasso", + history: { mode: "ignore" }, + execute({ state }) { + return state.editor.maskShapeSession ? { ...state, editor: { ...state.editor, maskShapeSession: undefined } } : state; + }, +}; + export const toolEnterMaskEditCommand: Command = { id: commandIds.toolEnterMaskEdit, name: "Enter mask edit", @@ -278,14 +341,41 @@ export const toolEnterMaskEditCommand: Command = { const maskLocation = findLayerLocation(state.document, payload.maskLayerId); if (!targetLocation || !maskLocation) return state; if (getLayerMask(targetLocation.layer)?.maskLayerId !== payload.maskLayerId) return state; - if (maskLocation.layer.type === "group") return state; + if (maskLocation.layer.type !== "image" && maskLocation.layer.type !== "raster") return state; return { ...state, editor: { ...state.editor, selection: { artboardId: targetLocation.artboardId, layerIds: [payload.targetLayerId] }, - maskEdit: { targetLayerId: payload.targetLayerId, maskLayerId: payload.maskLayerId }, + maskEdit: { kind: "layerMask", targetLayerId: payload.targetLayerId, maskLayerId: payload.maskLayerId, maskAssetId: maskLocation.layer.assetId }, + brushPreview: undefined, + brushStrokePreview: undefined, + tools: { + ...state.editor.tools, + activeTool: "brush", + interactionMode: { type: "tool", tool: "brush" }, + }, + }, + }; + }, +}; + +export const toolEnterInpaintRegionEditCommand: Command = { + id: commandIds.toolEnterInpaintRegionEdit, + name: "Enter inpaint region edit", + execute({ state }, payload) { + const targetLocation = findLayerLocation(state.document, payload.targetLayerId); + const region = state.document.inpaintRegions.find((candidate) => candidate.id === payload.regionId && candidate.targetLayerId === payload.targetLayerId); + const maskAsset = region ? state.document.assets.find((asset) => asset.id === region.maskAssetId) : undefined; + if (!targetLocation || !region || !maskAsset) return state; + + return { + ...state, + editor: { + ...state.editor, + selection: { artboardId: targetLocation.artboardId, layerIds: [payload.targetLayerId] }, + maskEdit: { kind: "inpaintRegion", targetLayerId: payload.targetLayerId, inpaintRegionId: region.id, maskAssetId: region.maskAssetId, viewMode: "overlay" }, brushPreview: undefined, brushStrokePreview: undefined, tools: { @@ -303,6 +393,7 @@ export const toolExitMaskEditCommand: Command = { name: "Exit mask edit", execute({ state }) { if (!state.editor.maskEdit) return state; + const contextualTool = state.editor.tools.activeTool === "semanticSelect" || state.editor.tools.activeTool === "maskLasso" || state.editor.tools.activeTool === "maskRectangle"; return { ...state, @@ -311,6 +402,8 @@ export const toolExitMaskEditCommand: Command = { maskEdit: undefined, brushPreview: undefined, brushStrokePreview: undefined, + maskShapeSession: undefined, + tools: contextualTool ? { ...state.editor.tools, activeTool: "select", interactionMode: { type: "tool", tool: "select" } } : state.editor.tools, }, }; }, @@ -366,7 +459,11 @@ export const toolCommands = [ toolSetBrushPreviewCommand, toolSetBrushStrokePreviewCommand, toolSetMaskViewModeCommand, + toolBeginMaskShapeCommand, + toolAppendMaskShapeCommand, + toolClearMaskShapeCommand, toolEnterMaskEditCommand, + toolEnterInpaintRegionEditCommand, toolExitMaskEditCommand, toolEnterTemporaryPanCommand, toolExitTemporaryPanCommand, diff --git a/core/asset-provenance.ts b/core/asset-provenance.ts index 8c058fd..f62b90f 100644 --- a/core/asset-provenance.ts +++ b/core/asset-provenance.ts @@ -1,5 +1,5 @@ import type { Rect, Size } from "./geometry"; -import type { AssetId, GenerationCandidateId, LayerId } from "./id"; +import type { AssetId, GenerationCandidateId, InpaintRegionId, LayerId } from "./id"; export type GeneratedAssetMode = "text-to-image" | "image-to-image" | "inpaint" | "outpaint"; @@ -26,10 +26,13 @@ export type AssetGenerationProvenance = { scheduler: string; width: number; height: number; + batchSize: number; + refinePass: boolean; + refineStrength: number; }; inpaint?: { targetLayerId: LayerId; - maskLayerId: LayerId; + regionId: InpaintRegionId; sourceAssetId: AssetId; maskAssetId: AssetId; crop: { @@ -49,6 +52,11 @@ export type AssetGenerationProvenance = { maskFeather: number; maskExpand: number; cropPadding: number; + profile: string; + structureControl: string; + controlStrength: number; + controlModel: string; + colorMatch: boolean; }; }; }; diff --git a/core/document.ts b/core/document.ts index a00ab9c..693dfda 100644 --- a/core/document.ts +++ b/core/document.ts @@ -1,6 +1,7 @@ import type { Artboard } from "./artboard"; import type { Asset } from "./asset"; import type { DocumentId } from "./id"; +import type { InpaintRegion } from "./inpaint-region"; export type ImageDocument = { id: DocumentId; @@ -8,4 +9,5 @@ export type ImageDocument = { version: number; artboards: Artboard[]; assets: Asset[]; + inpaintRegions: InpaintRegion[]; }; diff --git a/core/id.ts b/core/id.ts index 38d531b..004de75 100644 --- a/core/id.ts +++ b/core/id.ts @@ -5,5 +5,6 @@ export type DocumentId = OpaqueId<"DocumentId">; export type ArtboardId = OpaqueId<"ArtboardId">; export type LayerId = OpaqueId<"LayerId">; export type AssetId = OpaqueId<"AssetId">; +export type InpaintRegionId = OpaqueId<"InpaintRegionId">; export type GenerationCandidateId = OpaqueId<"GenerationCandidateId">; export type GenerationJobId = OpaqueId<"GenerationJobId">; diff --git a/core/index.ts b/core/index.ts index 6d6312d..63f77ea 100644 --- a/core/index.ts +++ b/core/index.ts @@ -13,7 +13,8 @@ export type { Transform, Vec2D, } from "./geometry"; -export type { ArtboardId, AssetId, DocumentId, GenerationCandidateId, GenerationJobId, LayerId } from "./id"; +export type { ArtboardId, AssetId, DocumentId, GenerationCandidateId, GenerationJobId, InpaintRegionId, LayerId } from "./id"; +export type { InpaintRegion } from "./inpaint-region"; export type { ImageLayer } from "./image-layer"; export type { AdjustmentLayer, ColorAdjustment } from "./adjustment-layer"; export { neutralColorAdjustment } from "./adjustment-layer"; diff --git a/core/inpaint-region.ts b/core/inpaint-region.ts new file mode 100644 index 0000000..eac8360 --- /dev/null +++ b/core/inpaint-region.ts @@ -0,0 +1,9 @@ +import type { AssetId, InpaintRegionId, LayerId } from "./id"; + +export type InpaintRegion = { + id: InpaintRegionId; + name: string; + targetLayerId: LayerId; + maskAssetId: AssetId; + enabled: boolean; +}; diff --git a/editor/document-indexes.test.ts b/editor/document-indexes.test.ts index 44c8e10..e30b40e 100644 --- a/editor/document-indexes.test.ts +++ b/editor/document-indexes.test.ts @@ -4,6 +4,7 @@ import type { Layer } from "@core/layer"; import { createDocumentReadIndex, forEachLayerBackToFront, resolveIndexedLayerBounds } from "./document-indexes"; const document: ImageDocument = { + inpaintRegions: [], id: "d1", name: "Indexed Document", version: 1, diff --git a/editor/initial-state.ts b/editor/initial-state.ts index ed5104e..8c586ca 100644 --- a/editor/initial-state.ts +++ b/editor/initial-state.ts @@ -31,6 +31,7 @@ export const initialEditorState: EditorState = { maskEdit: undefined, brushPreview: undefined, brushStrokePreview: undefined, + maskShapeSession: undefined, pointerSession: undefined, }; @@ -42,6 +43,7 @@ export function createInitialAppState(name = "Untitled"): AppState { version: 1, artboards: [], assets: [], + inpaintRegions: [], }, editor: initialEditorState, history: { past: [], future: [] }, diff --git a/editor/state.ts b/editor/state.ts index 83c4e96..b7ef6ec 100644 --- a/editor/state.ts +++ b/editor/state.ts @@ -1,6 +1,6 @@ import type { ImageDocument } from "@core/document"; import type { Angle, Rect, Size, Transform, Vec2D } from "@core/geometry"; -import type { ArtboardId, AssetId, GenerationCandidateId, GenerationJobId, LayerId } from "@core/id"; +import type { ArtboardId, AssetId, GenerationCandidateId, GenerationJobId, InpaintRegionId, LayerId } from "@core/id"; import type { GenerateArchitecture, GenerateMode, GenerateSettings, ToolState } from "./tools"; import type { TransformSession } from "./transform"; @@ -19,8 +19,11 @@ export type SelectionState = { export type MaskViewMode = "composite" | "blackWhite" | "alpha" | "overlay"; export type MaskEditState = { + kind: "layerMask" | "inpaintRegion"; targetLayerId: LayerId; - maskLayerId: LayerId; + maskAssetId: AssetId; + maskLayerId?: LayerId; + inpaintRegionId?: InpaintRegionId; viewMode?: MaskViewMode; }; @@ -34,6 +37,12 @@ export type BrushStrokePreviewState = { source: string; }; +export type MaskShapeSession = { + shape: "lasso" | "rectangle"; + points: Vec2D[]; + mode: "replace" | "add" | "subtract"; +}; + export type GenerationCandidate = { id: GenerationCandidateId; source: string; @@ -42,10 +51,12 @@ export type GenerationCandidate = { mode: GenerateSettings["mode"]; settings: GenerateSettings; seed: number; + favorite?: boolean; width: number; height: number; inputImage?: string; maskImage?: string; + blendMaskImage?: string; placement: { artboardId: ArtboardId; layerName: string; @@ -53,11 +64,14 @@ export type GenerationCandidate = { }; inpaint?: { targetLayerId: LayerId; - maskLayerId: LayerId; + regionId: InpaintRegionId; sourceAssetId: AssetId; maskAssetId: AssetId; inputImage: string; maskImage: string; + editMaskImage: string; + blendMaskImage: string; + revision: { source: string; mask: string }; crop: { assetBounds: Rect; documentBounds: Rect; @@ -81,7 +95,7 @@ export type GenerationCandidate = { export type GenerationCompareMode = "result" | "before" | "split"; -export type GenerationJobKind = "generate" | "regenerate" | "refine" | "replace"; +export type GenerationJobKind = "generate" | "regenerate" | "refine" | "replace" | "mask"; export type GenerationJobStatus = "running" | "succeeded" | "failed" | "cancelled"; export type GenerationJob = { @@ -92,6 +106,8 @@ export type GenerationJob = { startedAt: number; finishedAt?: number; error?: string; + progress?: number; + detail?: string; }; export type GenerationState = { @@ -113,10 +129,15 @@ export type GenerationArchitectureOption = { export type GenerationOptions = { architectures?: GenerationArchitectureOption[]; models?: string[]; + inpaintModels?: string[]; textEncoders?: string[]; vaes?: string[]; samplers?: string[]; schedulers?: string[]; + controlModels?: string[]; + structureControls?: Array<"canny" | "depth" | "pose">; + semanticSelection?: boolean; + sam3Models?: string[]; }; export type GenerationResourcesState = { status: "idle" | "loading" | "ready" | "failed"; options?: GenerationOptions; error?: string }; @@ -148,6 +169,7 @@ export type EditorState = { maskEdit?: MaskEditState; brushPreview?: BrushPreviewState; brushStrokePreview?: BrushStrokePreviewState; + maskShapeSession?: MaskShapeSession; pointerSession?: { type: "pan" }; }; diff --git a/editor/tools.ts b/editor/tools.ts index 8ca0d00..7dfe8c4 100644 --- a/editor/tools.ts +++ b/editor/tools.ts @@ -1,4 +1,4 @@ -export const availableToolIds = ["select", "brush", "eraser", "magicWand", "pan"] as const; +export const availableToolIds = ["select", "brush", "eraser", "magicWand", "semanticSelect", "maskLasso", "maskRectangle", "pan"] as const; export const availableOperationIds = ["generate", "chromaKey"] as const; @@ -13,6 +13,10 @@ export type BrushSettings = { color: string; size: number; hardness: number; + opacity: number; + flow: number; + smoothing: number; + pressureSize: boolean; }; export type ChromaKeySettings = { @@ -28,13 +32,15 @@ export type ChromaKeySettings = { export type MagicWandMode = "replace" | "add" | "subtract"; export type GenerateMode = "text-to-image" | "image-to-image" | "inpaint" | "outpaint"; -export type GenerateIntent = "create" | "replace" | "extend" | "variations"; +export type GenerateIntent = "create" | "replace" | "remove" | "extend" | "variations"; export const generateArchitectures = ["sdxl", "z-image", "z-image-turbo", "anima"] as const; export type GenerateArchitecture = (typeof generateArchitectures)[number]; export type GenerateModel = string; export type InpaintMaskedContent = "neutral" | "original" | "originalColor" | "edges"; +export type InpaintProfile = "remove" | "replace" | "repair" | "material" | "reshape" | "custom"; +export type InpaintStructureControl = "none" | "canny" | "depth" | "pose"; export type GenerateSettings = { architecture: GenerateArchitecture; @@ -52,6 +58,9 @@ export type GenerateSettings = { scheduler: string; width: number; height: number; + batchSize: number; + refinePass: boolean; + refineStrength: number; outpaint: { left: number; top: number; @@ -60,6 +69,7 @@ export type GenerateSettings = { feathering: number; }; inpaint: { + profile: InpaintProfile; maskedAreaOnly: boolean; cropPadding: number; maskPolarity: "hidden" | "revealed"; @@ -69,6 +79,10 @@ export type GenerateSettings = { maskFeather: number; maskBlur: number; maskDespeckle: number; + structureControl: InpaintStructureControl; + controlStrength: number; + controlModel: string; + colorMatch: boolean; }; }; @@ -109,7 +123,7 @@ export type ToolState = { export const initialToolState: ToolState = { activeTool: "select", interactionMode: { type: "tool", tool: "select" }, - brush: { color: "#111827", size: 8, hardness: 100 }, + brush: { color: "#111827", size: 8, hardness: 100, opacity: 100, flow: 100, smoothing: 20, pressureSize: true }, chromaKey: { color: "#00ff00", tolerance: 32, softness: 24, feather: 0, choke: 0, despeckle: 0, spill: 50 }, magicWand: { tolerance: 32, feather: 0, choke: 0, despeckle: 0, contiguous: true, mode: "replace" }, generate: { @@ -128,11 +142,22 @@ export const initialToolState: ToolState = { scheduler: "normal", width: 1024, height: 1024, + batchSize: 4, + refinePass: true, + refineStrength: 20, outpaint: { left: 128, top: 128, right: 128, bottom: 128, feathering: 32 }, - inpaint: { maskedAreaOnly: true, cropPadding: 96, maskPolarity: "hidden", maskedContent: "neutral", growMaskBy: 6, maskExpand: 0, maskFeather: 0, maskBlur: 0, maskDespeckle: 0 }, + inpaint: { profile: "replace", maskedAreaOnly: true, cropPadding: 128, maskPolarity: "revealed", maskedContent: "neutral", growMaskBy: 6, maskExpand: 8, maskFeather: 3, maskBlur: 0, maskDespeckle: 0, structureControl: "none", controlStrength: 0.55, controlModel: "auto", colorMatch: true }, }, }; +export const inpaintProfileDefaults: Record, Omit, "inpaint"> & { inpaint: Partial }> = { + remove: { strength: 100, inpaint: { maskedContent: "neutral", cropPadding: 160, maskExpand: 12, maskFeather: 4, structureControl: "none" } }, + replace: { strength: 85, inpaint: { maskedContent: "neutral", cropPadding: 128, maskExpand: 8, maskFeather: 3, structureControl: "none" } }, + repair: { strength: 45, refinePass: true, refineStrength: 15, inpaint: { maskedContent: "original", cropPadding: 96, maskExpand: 4, maskFeather: 2, structureControl: "none" } }, + material: { strength: 65, refinePass: true, refineStrength: 20, inpaint: { maskedContent: "edges", cropPadding: 128, maskExpand: 6, maskFeather: 3, structureControl: "canny", controlStrength: 0.55, colorMatch: false } }, + reshape: { strength: 90, inpaint: { maskedContent: "neutral", cropPadding: 160, maskExpand: 10, maskFeather: 4, structureControl: "depth", controlStrength: 0.4 } }, +}; + export function isPanInteractionMode(interactionMode: InteractionMode): boolean { return interactionMode.type === "temporary-pan" || (interactionMode.type === "tool" && interactionMode.tool === "pan"); } diff --git a/editor/transform-targets.test.ts b/editor/transform-targets.test.ts index 116c6e5..b4c5746 100644 --- a/editor/transform-targets.test.ts +++ b/editor/transform-targets.test.ts @@ -4,6 +4,7 @@ import { resolveTransformTargetBounds, selectedTransformTarget } from "./transfo import { applyTransformTargetBounds } from "@commands/transform-document"; const document: ImageDocument = { + inpaintRegions: [], id: "d1", name: "Test", version: 1, diff --git a/index.ts b/index.ts index 52bd3b9..4f37a4f 100644 --- a/index.ts +++ b/index.ts @@ -6,6 +6,7 @@ const server = serve({ routes: { "/api/comfy/models": handleComfyApi, "/api/comfy/generate": handleComfyApi, + "/api/comfy/segment": handleComfyApi, "/*": index, }, diff --git a/input/dom.ts b/input/dom.ts index 59edc05..76e18de 100644 --- a/input/dom.ts +++ b/input/dom.ts @@ -23,6 +23,7 @@ export function pointerInputEventFromPointerEvent(event: PointerEvent): PointerI pointerType: normalizePointerType(event.pointerType), position: { x: event.offsetX, y: event.offsetY }, buttons: event.buttons, + pressure: event.pointerType === "pen" ? event.pressure : 1, altKey: event.altKey, ctrlKey: event.ctrlKey, metaKey: event.metaKey, diff --git a/input/pointer.ts b/input/pointer.ts index 8569b96..45d8b8d 100644 --- a/input/pointer.ts +++ b/input/pointer.ts @@ -5,6 +5,7 @@ export type PointerInputEvent = { pointerType: "mouse" | "pen" | "touch"; position: Vec2D; buttons: number; + pressure?: number; altKey: boolean; ctrlKey: boolean; metaKey: boolean; diff --git a/input/transform-controls.ts b/input/transform-controls.ts index f73c7bf..e51d236 100644 --- a/input/transform-controls.ts +++ b/input/transform-controls.ts @@ -16,7 +16,7 @@ import type { PointerInputEvent } from "./pointer"; type TransformHandle = "body" | "nw" | "n" | "ne" | "e" | "se" | "s" | "sw" | "w"; -type InputToolId = "select" | "brush" | "eraser" | "magicWand" | "pan"; +type InputToolId = "select" | "brush" | "eraser" | "magicWand" | "semanticSelect" | "maskLasso" | "maskRectangle" | "pan"; type InputInteractionMode = | { type: "tool"; tool: InputToolId } diff --git a/operations/generation/candidateActions.ts b/operations/generation/candidateActions.ts index 000ab62..730db7f 100644 --- a/operations/generation/candidateActions.ts +++ b/operations/generation/candidateActions.ts @@ -1,16 +1,24 @@ import type { ImageDocument } from "@core/document"; import type { GenerationCandidate } from "@editor/state"; import { loadImageCanvas, maskValueFromRgba } from "@platform/browser/maskRaster"; +import { createContentRevision } from "./inpaintPrep"; export async function createMaskedPixelReplacementSource(document: ImageDocument, candidate: GenerationCandidate): Promise { if (!candidate.inpaint) throw new Error("Only inpaint candidates can replace masked pixels."); const targetAsset = document.assets.find((asset) => asset.id === candidate.inpaint?.sourceAssetId); if (!targetAsset) throw new Error("The source layer for this candidate no longer exists."); + const region = document.inpaintRegions.find((item) => item.id === candidate.inpaint?.regionId); + const maskAsset = region ? document.assets.find((asset) => asset.id === region.maskAssetId) : undefined; + if (!region || !maskAsset) throw new Error("The AI edit region for this candidate no longer exists."); + const [sourceRevision, maskRevision] = await Promise.all([createContentRevision(targetAsset.source), createContentRevision(maskAsset.source)]); + if (sourceRevision !== candidate.inpaint.revision.source || maskRevision !== candidate.inpaint.revision.mask) { + throw new Error("The source or AI edit region changed after generation. Rebuild candidates from the current edit region before replacing pixels."); + } const targetCanvas = await loadImageCanvas(targetAsset.source, targetAsset.intrinsicSize.w, targetAsset.intrinsicSize.h); const generatedCanvas = await loadImageCanvas(candidate.source, candidate.width, candidate.height); - const maskCanvas = await loadImageCanvas(candidate.inpaint.maskImage, candidate.width, candidate.height); + const maskCanvas = await loadImageCanvas(candidate.inpaint.blendMaskImage, candidate.width, candidate.height); const targetContext = require2dContext(targetCanvas); const generatedContext = require2dContext(generatedCanvas); @@ -19,6 +27,7 @@ export async function createMaskedPixelReplacementSource(document: ImageDocument const generatedData = generatedContext.getImageData(0, 0, generatedCanvas.width, generatedCanvas.height); const maskData = maskContext.getImageData(0, 0, maskCanvas.width, maskCanvas.height); const crop = candidate.inpaint.crop.assetBounds; + const colorOffset = candidate.settings.inpaint.colorMatch ? boundaryColorOffset(targetData.data, generatedData.data, maskData.data, candidate.width, candidate.height, crop, targetCanvas.width, targetCanvas.height) : [0, 0, 0]; for (let y = 0; y < candidate.height; y += 1) { for (let x = 0; x < candidate.width; x += 1) { @@ -33,7 +42,8 @@ export async function createMaskedPixelReplacementSource(document: ImageDocument for (let channel = 0; channel < 4; channel += 1) { const previous = targetData.data[targetIndex + channel] ?? 0; - const next = generatedData.data[generatedIndex + channel] ?? previous; + const rawNext = generatedData.data[generatedIndex + channel] ?? previous; + const next = channel < 3 ? Math.max(0, Math.min(255, rawNext + (colorOffset[channel] ?? 0))) : rawNext; targetData.data[targetIndex + channel] = Math.round(previous * (1 - mask) + next * mask); } } @@ -43,6 +53,28 @@ export async function createMaskedPixelReplacementSource(document: ImageDocument return targetCanvas.toDataURL("image/png"); } +function boundaryColorOffset(target: Uint8ClampedArray, generated: Uint8ClampedArray, mask: Uint8ClampedArray, width: number, height: number, crop: { x: number; y: number }, targetWidth: number, targetHeight: number): number[] { + const targetTotal = [0, 0, 0]; + const generatedTotal = [0, 0, 0]; + let count = 0; + for (let y = 0; y < height; y += 1) for (let x = 0; x < width; x += 1) { + const index = (y * width + x) * 4; + const amount = maskValueFromRgba(mask, index) / 255; + if (amount <= 0.05 || amount >= 0.65) continue; + const targetX = Math.round(crop.x) + x; + const targetY = Math.round(crop.y) + y; + if (targetX < 0 || targetY < 0 || targetX >= targetWidth || targetY >= targetHeight) continue; + const targetIndex = (targetY * targetWidth + targetX) * 4; + for (let channel = 0; channel < 3; channel += 1) { + targetTotal[channel] = (targetTotal[channel] ?? 0) + (target[targetIndex + channel] ?? 0); + generatedTotal[channel] = (generatedTotal[channel] ?? 0) + (generated[index + channel] ?? 0); + } + count += 1; + } + if (count < 16) return [0, 0, 0]; + return targetTotal.map((total, channel) => Math.max(-32, Math.min(32, total / count - (generatedTotal[channel] ?? 0) / count))); +} + function require2dContext(canvas: HTMLCanvasElement): CanvasRenderingContext2D { const context = canvas.getContext("2d"); if (!context) throw new Error("Unable to prepare generated candidate"); diff --git a/operations/generation/generationJob.ts b/operations/generation/generationJob.ts index e01d1e5..c9aac04 100644 --- a/operations/generation/generationJob.ts +++ b/operations/generation/generationJob.ts @@ -7,7 +7,7 @@ export async function runGenerationJob(options: { label: string; dispatch: AppStore["dispatch"]; signal: AbortSignal; - task: (signal: AbortSignal) => Promise; + task: (signal: AbortSignal, report: (progress: number, detail: string) => void) => Promise; }): Promise { const jobId = crypto.randomUUID(); const startedAt = Date.now(); @@ -15,7 +15,7 @@ export async function runGenerationJob(options: { if (!nextState.editor.generation.jobs.some((job) => job.id === jobId && job.status === "running")) return; try { - await options.task(options.signal); + await options.task(options.signal, (progress, detail) => options.dispatch(commandIds.generationUpdateJob, { jobId, progress, detail })); options.dispatch(commandIds.generationSucceedJob, { jobId, finishedAt: Date.now() }); } catch (reason: unknown) { if (options.signal.aborted) { diff --git a/operations/generation/inpaintPrep.ts b/operations/generation/inpaintPrep.ts index 248890c..8f48d1b 100644 --- a/operations/generation/inpaintPrep.ts +++ b/operations/generation/inpaintPrep.ts @@ -1,22 +1,28 @@ import type { Asset } from "@core/asset"; import type { ImageDocument } from "@core/document"; import type { Rect } from "@core/geometry"; +import type { AssetId, InpaintRegionId, LayerId } from "@core/id"; import type { Layer } from "@core/layer"; -import { getLayerMask } from "@core/layer-mask-utils"; import type { SelectionState } from "@editor/state"; import type { GenerateSettings } from "@editor/tools"; import { createDocumentReadIndex, resolveIndexedLayerBounds } from "@editor/document-indexes"; -import { createNormalizedMaskSource, cropCanvas, cropMaskValuesToDataUrl, expandRectWithinBounds, loadImageCanvas } from "@platform/browser/maskRaster"; +import { createNormalizedMaskSource, cropCanvas, cropMaskValuesToDataUrl, expandRectWithinBounds, loadImageCanvas, sampleDocumentCanvasInLayerSpace } from "@platform/browser/maskRaster"; +import { renderArtboardCanvas } from "@platform/browser/exportArtboardPng"; export type InpaintBundle = { inputImage: string; + sourceImage: string; + contextImage: string; maskImage: string; + editMaskImage: string; + blendMaskImage: string; width: number; height: number; - targetLayerId: string; - maskLayerId: string; - sourceAssetId: string; - maskAssetId: string; + targetLayerId: LayerId; + regionId: InpaintRegionId; + sourceAssetId: AssetId; + maskAssetId: AssetId; + revision: { source: string; mask: string }; crop: { assetBounds: Rect; documentBounds: Rect; @@ -46,10 +52,8 @@ type InpaintTarget = { artboardId: string; layer: Extract; asset: Asset; - bounds: Rect; - maskLayer: Extract; maskAsset: Asset; - maskBounds: Rect; + regionId: InpaintRegionId; }; const modelMultiple = 8; @@ -62,41 +66,65 @@ export async function buildInpaintBundle(document: ImageDocument, selection: Sel const width = Math.max(1, Math.round(target.asset.intrinsicSize.w)); const height = Math.max(1, Math.round(target.asset.intrinsicSize.h)); - const normalizedMask = await createNormalizedMaskSource(target.maskAsset.source, width, height, { - polarity: settings.inpaint.maskPolarity, + const sourceLimit = target.layer.sourceRect ?? { x: 0, y: 0, w: width, h: height }; + const editMask = await createNormalizedMaskSource(target.maskAsset.source, width, height, { + polarity: "revealed", + despeckle: settings.inpaint.maskDespeckle, + limit: sourceLimit, + }); + const noiseMask = await createNormalizedMaskSource(target.maskAsset.source, width, height, { + polarity: "revealed", expand: settings.inpaint.maskExpand, - feather: settings.inpaint.maskFeather, blur: settings.inpaint.maskBlur, despeckle: settings.inpaint.maskDespeckle, + limit: sourceLimit, + }); + const blendMask = await createNormalizedMaskSource(target.maskAsset.source, width, height, { + polarity: "revealed", + feather: settings.inpaint.maskFeather, + despeckle: settings.inpaint.maskDespeckle, + limit: sourceLimit, }); - if (!normalizedMask.bounds) throw new Error("The selected layer mask has no inpaint pixels."); + if (!editMask.bounds || !noiseMask.bounds) throw new Error("Paint over the area you want AI to replace."); const crop = settings.inpaint.maskedAreaOnly - ? expandRectWithinBounds(normalizedMask.bounds, settings.inpaint.cropPadding, { w: width, h: height }, modelMultiple, minModelSize) + ? expandRectWithinBounds(noiseMask.bounds, settings.inpaint.cropPadding, { w: width, h: height }, modelMultiple, minModelSize) : { x: 0, y: 0, w: width, h: height }; const outputWidth = toModelSize(crop.w, "width"); const outputHeight = toModelSize(crop.h, "height"); - const inputCanvas = prepareMaskedContentInputCanvas(await loadImageCanvas(target.asset.source, width, height), normalizedMask.values, settings.inpaint.maskedContent); - const inputImage = cropCanvas(inputCanvas, crop, outputWidth, outputHeight); - const maskImage = cropMaskValuesToDataUrl(normalizedMask.values, width, height, crop, outputWidth, outputHeight); - const scaleX = target.bounds.w / width; - const scaleY = target.bounds.h / height; + const sourceCanvas = await loadImageCanvas(target.asset.source, width, height); + const sourceImage = cropCanvas(sourceCanvas, crop, outputWidth, outputHeight); + const preparedCanvas = prepareMaskedContentInputCanvas(await loadImageCanvas(target.asset.source, width, height), editMask.values, settings.inpaint.maskedContent); + const preparedSourceImage = cropCanvas(preparedCanvas, crop, outputWidth, outputHeight); + const artboard = document.artboards.find((candidate) => candidate.id === target.artboardId); + if (!artboard) throw new Error("The target artboard no longer exists."); + const contextImage = sampleDocumentCanvasInLayerSpace(await renderArtboardCanvas(artboard, document.assets), artboard.bounds, target.layer, target.asset.intrinsicSize, crop, outputWidth, outputHeight); + const inputImage = await mergePreparedRegionIntoContext(contextImage, preparedSourceImage, cropMaskValuesToDataUrl(editMask.values, width, height, crop, outputWidth, outputHeight), outputWidth, outputHeight); + const maskImage = cropMaskValuesToDataUrl(noiseMask.values, width, height, crop, outputWidth, outputHeight); + const editMaskImage = cropMaskValuesToDataUrl(editMask.values, width, height, crop, outputWidth, outputHeight); + const blendMaskImage = cropMaskValuesToDataUrl(blendMask.values, width, height, crop, outputWidth, outputHeight); + const scaleX = target.layer.transform.scale.x; + const scaleY = target.layer.transform.scale.y; const documentBounds = { - x: target.bounds.x + crop.x * scaleX, - y: target.bounds.y + crop.y * scaleY, + x: target.layer.transform.position.x + crop.x * scaleX, + y: target.layer.transform.position.y + crop.y * scaleY, w: outputWidth * scaleX, h: outputHeight * scaleY, }; return { inputImage, + sourceImage, + contextImage, maskImage, + editMaskImage, + blendMaskImage, width: outputWidth, height: outputHeight, targetLayerId: target.layer.id, - maskLayerId: target.maskLayer.id, + regionId: target.regionId, sourceAssetId: target.asset.id, maskAssetId: target.maskAsset.id, crop: { @@ -106,14 +134,14 @@ export async function buildInpaintBundle(document: ImageDocument, selection: Sel maskedAreaOnly: settings.inpaint.maskedAreaOnly, }, mask: { - polarity: settings.inpaint.maskPolarity, - activeBounds: normalizedMask.bounds, + polarity: "revealed", + activeBounds: editMask.bounds, }, placement: { artboardId: target.artboardId, layerName: `${target.layer.name} inpaint`, transform: { - position: { x: documentBounds.x, y: documentBounds.y }, + position: rotatedCropPosition(target.layer, target.asset, { ...crop, w: outputWidth, h: outputHeight }), scale: { x: documentBounds.w / outputWidth, y: documentBounds.h / outputHeight }, rotation: target.layer.transform.rotation, }, @@ -126,6 +154,10 @@ export async function buildInpaintBundle(document: ImageDocument, selection: Sel maskExpand: settings.inpaint.maskExpand, cropPadding: settings.inpaint.cropPadding, }, + revision: { + source: await createContentRevision(target.asset.source), + mask: await createContentRevision(target.maskAsset.source), + }, }; } @@ -138,20 +170,14 @@ function resolveInpaintTarget(document: ImageDocument, selection: SelectionState const asset = documentIndex.assetById.get(layerInfo.layer.assetId); if (!asset) throw new Error("The selected layer is missing its source image."); - const layerMask = getLayerMask(layerInfo.layer); - if (!layerMask?.enabled) throw new Error("Add a layer mask before running inpaint."); + const region = document.inpaintRegions.find((candidate) => candidate.targetLayerId === layerInfo.layer.id && candidate.enabled); + if (!region) throw new Error("Add an AI edit region before running inpaint."); + const maskAsset = documentIndex.assetById.get(region.maskAssetId); + if (!maskAsset) throw new Error("The AI edit region is missing its mask data."); - const maskLayer = documentIndex.layerById.get(layerMask.maskLayerId); - if (!maskLayer || (maskLayer.type !== "image" && maskLayer.type !== "raster")) throw new Error("The selected layer mask is missing."); + if (!resolveIndexedLayerBounds(documentIndex, layerInfo.layer)) throw new Error("Unable to resolve the selected layer bounds."); - const maskAsset = documentIndex.assetById.get(maskLayer.assetId); - if (!maskAsset) throw new Error("The selected layer mask is missing its image data."); - - const bounds = resolveIndexedLayerBounds(documentIndex, layerInfo.layer); - const maskBounds = resolveIndexedLayerBounds(documentIndex, maskLayer); - if (!bounds || !maskBounds) throw new Error("Unable to resolve the selected layer and mask bounds."); - - return { artboardId: layerInfo.artboardId, layer: layerInfo.layer, asset, bounds, maskLayer, maskAsset, maskBounds }; + return { artboardId: layerInfo.artboardId, layer: layerInfo.layer, asset, maskAsset, regionId: region.id }; } function validateInpaintTarget(target: InpaintTarget) { @@ -160,13 +186,6 @@ function validateInpaintTarget(target: InpaintTarget) { if (Math.round(target.asset.intrinsicSize.w) !== Math.round(target.maskAsset.intrinsicSize.w) || Math.round(target.asset.intrinsicSize.h) !== Math.round(target.maskAsset.intrinsicSize.h)) { throw new Error("The selected layer and mask image sizes do not match."); } - if (!rectsAligned(target.bounds, target.maskBounds) || Math.abs(target.layer.transform.rotation - target.maskLayer.transform.rotation) > 0.001) { - throw new Error("The selected layer and mask are not aligned."); - } -} - -function rectsAligned(a: Rect, b: Rect): boolean { - return Math.abs(a.x - b.x) <= 0.5 && Math.abs(a.y - b.y) <= 0.5 && Math.abs(a.w - b.w) <= 0.5 && Math.abs(a.h - b.h) <= 0.5; } function toModelSize(value: number, axis: "width" | "height"): number { @@ -258,3 +277,52 @@ function blendChannel(previous: number, next: number, amount: number): number { function clampNumber(value: number, min: number, max: number): number { return Math.max(min, Math.min(max, value)); } + +async function mergePreparedRegionIntoContext(contextSource: string, preparedSource: string, editMaskSource: string, width: number, height: number): Promise { + const [contextCanvas, preparedCanvas, maskCanvas] = await Promise.all([ + loadImageCanvas(contextSource, width, height), + loadImageCanvas(preparedSource, width, height), + loadImageCanvas(editMaskSource, width, height), + ]); + const context = contextCanvas.getContext("2d"); + const prepared = preparedCanvas.getContext("2d"); + const mask = maskCanvas.getContext("2d"); + if (!context || !prepared || !mask) return contextSource; + const contextData = context.getImageData(0, 0, width, height); + const preparedData = prepared.getImageData(0, 0, width, height); + const maskData = mask.getImageData(0, 0, width, height); + for (let pixel = 0; pixel < width * height; pixel += 1) { + const amount = (maskData.data[pixel * 4] ?? 0) / 255; + if (amount <= 0) continue; + for (let channel = 0; channel < 4; channel += 1) { + const index = pixel * 4 + channel; + contextData.data[index] = blendChannel(contextData.data[index] ?? 0, preparedData.data[index] ?? 0, amount); + } + } + context.putImageData(contextData, 0, 0); + return contextCanvas.toDataURL("image/png"); +} + +function rotatedCropPosition(layer: InpaintTarget["layer"], asset: Asset, crop: Rect) { + const source = layer.sourceRect ?? { x: 0, y: 0, ...asset.intrinsicSize }; + const scale = layer.transform.scale; + const originalCenter = { + x: layer.transform.position.x + (source.x + source.w / 2) * scale.x, + y: layer.transform.position.y + (source.y + source.h / 2) * scale.y, + }; + const cropCenter = { + x: layer.transform.position.x + (crop.x + crop.w / 2) * scale.x, + y: layer.transform.position.y + (crop.y + crop.h / 2) * scale.y, + }; + const dx = cropCenter.x - originalCenter.x; + const dy = cropCenter.y - originalCenter.y; + const cos = Math.cos(layer.transform.rotation); + const sin = Math.sin(layer.transform.rotation); + const rotatedCenter = { x: originalCenter.x + dx * cos - dy * sin, y: originalCenter.y + dx * sin + dy * cos }; + return { x: rotatedCenter.x - crop.w * scale.x / 2, y: rotatedCenter.y - crop.h * scale.y / 2 }; +} + +export async function createContentRevision(source: string): Promise { + const digest = await crypto.subtle.digest("SHA-256", new TextEncoder().encode(source)); + return Array.from(new Uint8Array(digest), (byte) => byte.toString(16).padStart(2, "0")).join(""); +} diff --git a/operations/generation/outputPlacement.test.ts b/operations/generation/outputPlacement.test.ts index 06506f9..bb3a3b1 100644 --- a/operations/generation/outputPlacement.test.ts +++ b/operations/generation/outputPlacement.test.ts @@ -42,8 +42,8 @@ describe("generated output placement", () => { settings: settings("inpaint"), intrinsicSize: { w: 128, h: 64 }, inpaintBundle: { - inputImage: "input", maskImage: "mask", width: 256, height: 128, - targetLayerId: "source", maskLayerId: "mask", sourceAssetId: "source-asset", maskAssetId: "mask-asset", + inputImage: "input", sourceImage: "source", contextImage: "context", maskImage: "mask", editMaskImage: "edit", blendMaskImage: "blend", width: 256, height: 128, + targetLayerId: "source", regionId: "region", sourceAssetId: "source-asset", maskAssetId: "mask-asset", revision: { source: "source-revision", mask: "mask-revision" }, crop: { assetBounds: { x: 0, y: 0, w: 256, h: 128 }, documentBounds: { x: 140, y: 150, w: 384, h: 64 }, padding: 16, maskedAreaOnly: true }, mask: { polarity: "hidden", activeBounds: { x: 20, y: 20, w: 40, h: 40 } }, placement: { artboardId: "artboard", layerName: "Source inpaint", transform: { position: { x: 140, y: 150 }, scale: { x: 1.5, y: 0.5 }, rotation: 12 } }, @@ -81,7 +81,7 @@ function selected() { function document(): ImageDocument { return { - id: "document", name: "Test", version: 1, + id: "document", name: "Test", version: 1, inpaintRegions: [], assets: [ { id: "source-asset", name: "Source", mimeType: "image/png", source: "source", intrinsicSize: { w: 100, h: 500 } }, { id: "mask-asset", name: "Mask", mimeType: "image/png", source: "mask", intrinsicSize: { w: 100, h: 500 } }, diff --git a/operations/generation/preconditions.test.ts b/operations/generation/preconditions.test.ts index 2d40ddf..a95c89d 100644 --- a/operations/generation/preconditions.test.ts +++ b/operations/generation/preconditions.test.ts @@ -15,7 +15,7 @@ describe("generation preconditions", () => { }); test("requires an enabled mask for inpaint", () => { - expect(checkGenerationPreconditions(document(), selected(), settings("inpaint"))).toEqual({ ready: false, message: "Paint a mask over the area you want AI to replace.", repair: "add-mask" }); + expect(checkGenerationPreconditions(document(), selected(), settings("inpaint"))).toEqual({ ready: false, message: "Paint an AI edit region over the area you want to replace.", repair: "add-mask" }); }); test("allows inpaint when the selected source has an aligned enabled mask", () => { @@ -46,6 +46,7 @@ function document(masked = false): ImageDocument { { id: "source-asset", name: "Source", mimeType: "image/png", source: "source", intrinsicSize: { w: 100, h: 100 } }, { id: "mask-asset", name: "Mask", mimeType: "image/png", source: "mask", intrinsicSize: { w: 100, h: 100 } }, ], + inpaintRegions: masked ? [{ id: "region", name: "AI edit", targetLayerId: "source", maskAssetId: "mask-asset", enabled: true }] : [], artboards: [{ id: "artboard", name: "Artboard", @@ -64,7 +65,6 @@ function document(masked = false): ImageDocument { opacity: 1, assetId: "source-asset", transform, - layerMask: masked ? { kind: "raster", maskLayerId: "mask", enabled: true, inverted: false } : undefined, }, ], }], diff --git a/operations/generation/preconditions.ts b/operations/generation/preconditions.ts index a0e0563..9b882aa 100644 --- a/operations/generation/preconditions.ts +++ b/operations/generation/preconditions.ts @@ -1,5 +1,4 @@ import type { ImageDocument } from "@core/document"; -import { getLayerMask } from "@core/layer-mask-utils"; import { createDocumentReadIndex, resolveIndexedLayerBounds } from "@editor/document-indexes"; import type { SelectionState } from "@editor/state"; import type { GenerateSettings } from "@editor/tools"; @@ -44,21 +43,15 @@ export function checkGenerationPreconditions( if (settings.mode !== "inpaint") return { ready: true }; - const layerMask = getLayerMask(layerInfo.layer); - if (!layerMask?.enabled) return missing("Paint a mask over the area you want AI to replace.", "add-mask"); - const maskLayer = index.layerById.get(layerMask.maskLayerId); - if (!maskLayer || (maskLayer.type !== "image" && maskLayer.type !== "raster")) return missing("The selected layer mask is missing."); - const maskAsset = index.assetById.get(maskLayer.assetId); - if (!maskAsset) return missing("The selected layer mask is missing its image data."); + const region = document.inpaintRegions.find((candidate) => candidate.targetLayerId === layerInfo.layer.id && candidate.enabled); + if (!region) return missing("Paint an AI edit region over the area you want to replace.", "add-mask"); + const maskAsset = index.assetById.get(region.maskAssetId); + if (!maskAsset) return missing("The AI edit region is missing its mask data."); if (Math.round(asset.intrinsicSize.w) !== Math.round(maskAsset.intrinsicSize.w) || Math.round(asset.intrinsicSize.h) !== Math.round(maskAsset.intrinsicSize.h)) { return missing("The selected layer and mask image sizes must match."); } - const layerBounds = resolveIndexedLayerBounds(index, layerInfo.layer); - const maskBounds = resolveIndexedLayerBounds(index, maskLayer); - if (!layerBounds || !maskBounds || !rectsAligned(layerBounds, maskBounds) || Math.abs(layerInfo.layer.transform.rotation - maskLayer.transform.rotation) > 0.001) { - return missing("Align the selected layer and its mask before inpainting."); - } + if (!resolveIndexedLayerBounds(index, layerInfo.layer)) return missing("The selected layer has invalid geometry."); return { ready: true }; } @@ -72,7 +65,3 @@ function modeSelectionMessage(mode: GenerateSettings["mode"]): string { function missing(message: string, repair?: Extract["repair"]): GenerationPrecondition { return { ready: false, message, ...(repair ? { repair } : {}) }; } - -function rectsAligned(a: { x: number; y: number; w: number; h: number }, b: { x: number; y: number; w: number; h: number }): boolean { - return Math.abs(a.x - b.x) <= 0.5 && Math.abs(a.y - b.y) <= 0.5 && Math.abs(a.w - b.w) <= 0.5 && Math.abs(a.h - b.h) <= 0.5; -} diff --git a/operations/generation/runGenerate.ts b/operations/generation/runGenerate.ts index 294f8a6..6f318e3 100644 --- a/operations/generation/runGenerate.ts +++ b/operations/generation/runGenerate.ts @@ -18,6 +18,7 @@ export async function runGenerate(options: { settings: GenerateSettings; dispatch: AppStore["dispatch"]; signal?: AbortSignal; + onProgress?: (progress: number, detail: string) => void; }) { const { document, selection, settings, dispatch } = options; const precondition = checkGenerationPreconditions(document, selection, settings); @@ -41,25 +42,27 @@ export async function runGenerate(options: { maskImage, inpaintBundle, signal: options.signal, + onProgress: options.onProgress, }); - const intrinsicSize = await loadImageSize(generated.source); - const placement = resolveGeneratedOutputPlacement({ document, selection, settings, intrinsicSize, inpaintBundle }); - - dispatch(commandIds.generationAddCandidate, { - candidate: createGenerationCandidate({ - source: generated.source, - mimeType: generated.mimeType, - intrinsicSize, - settings: requestSettings, - seed, - width, - height, - inputImage, - maskImage, - placement, - inpaintBundle, - }), - }); + for (const result of generated.results) { + const intrinsicSize = await loadImageSize(result.source); + const placement = resolveGeneratedOutputPlacement({ document, selection, settings, intrinsicSize, inpaintBundle }); + dispatch(commandIds.generationAddCandidate, { + candidate: createGenerationCandidate({ + source: result.source, + mimeType: result.mimeType, + intrinsicSize, + settings: requestSettings, + seed: result.seed || seed, + width, + height, + inputImage, + maskImage, + placement, + inpaintBundle, + }), + }); + } } export async function runGenerateFromCandidate(options: { @@ -67,6 +70,7 @@ export async function runGenerateFromCandidate(options: { settings?: GenerateSettings; dispatch: AppStore["dispatch"]; signal?: AbortSignal; + onProgress?: (progress: number, detail: string) => void; }) { const settings = options.settings ?? options.candidate.settings; const seed = resolveSeed(settings.seed); @@ -79,20 +83,22 @@ export async function runGenerateFromCandidate(options: { maskImage: options.candidate.maskImage, inpaintCandidate: options.candidate, signal: options.signal, + onProgress: options.onProgress, }); - 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, - }, - }); + for (const result of generated.results) { + const intrinsicSize = await loadImageSize(result.source); + options.dispatch(commandIds.generationAddCandidate, { + candidate: { + ...options.candidate, + id: crypto.randomUUID(), + source: result.source, + mimeType: result.mimeType, + intrinsicSize, + settings: requestSettings, + seed: result.seed || seed, + }, + }); + } } function createGenerationCandidate(options: { @@ -120,15 +126,19 @@ function createGenerationCandidate(options: { height: options.height, inputImage: options.inputImage, maskImage: options.maskImage, + blendMaskImage: options.inpaintBundle?.blendMaskImage, placement: options.placement, inpaint: options.inpaintBundle ? { targetLayerId: options.inpaintBundle.targetLayerId, - maskLayerId: options.inpaintBundle.maskLayerId, + regionId: options.inpaintBundle.regionId, sourceAssetId: options.inpaintBundle.sourceAssetId, maskAssetId: options.inpaintBundle.maskAssetId, inputImage: options.inpaintBundle.inputImage, maskImage: options.inpaintBundle.maskImage, + editMaskImage: options.inpaintBundle.editMaskImage, + blendMaskImage: options.inpaintBundle.blendMaskImage, + revision: options.inpaintBundle.revision, crop: options.inpaintBundle.crop, mask: options.inpaintBundle.mask, backend: options.inpaintBundle.backend, @@ -146,6 +156,7 @@ async function requestGenerate(options: { inpaintBundle?: InpaintBundle; inpaintCandidate?: GenerationCandidate; signal?: AbortSignal; + onProgress?: (progress: number, detail: string) => void; }) { return requestGeneration({ architecture: options.settings.architecture, @@ -163,11 +174,14 @@ async function requestGenerate(options: { scheduler: options.settings.scheduler, width: options.width, height: options.height, + batchSize: options.settings.batchSize, + refinePass: options.settings.refinePass, + refineStrength: options.settings.refineStrength, outpaint: options.settings.outpaint, inpaint: resolveInpaintRequest(options.inpaintBundle, options.inpaintCandidate, options.settings), inputImage: options.inputImage, maskImage: options.maskImage, - }, options.signal); + }, options.signal, options.onProgress); } function resolveInpaintRequest(inpaintBundle: InpaintBundle | undefined, inpaintCandidate: GenerationCandidate | undefined, settings: GenerateSettings) { @@ -182,6 +196,9 @@ function resolveInpaintRequest(inpaintBundle: InpaintBundle | undefined, inpaint maskPolarity: inpaintBundle.mask.polarity, crop: inpaintBundle.crop, placement: inpaintBundle.placement, + structureControl: settings.inpaint.structureControl, + controlStrength: settings.inpaint.controlStrength, + controlModel: settings.inpaint.controlModel, }; } @@ -196,6 +213,9 @@ function resolveInpaintRequest(inpaintBundle: InpaintBundle | undefined, inpaint maskPolarity: inpaintCandidate.inpaint.mask.polarity, crop: inpaintCandidate.inpaint.crop, placement: inpaintCandidate.placement, + structureControl: settings.inpaint.structureControl, + controlStrength: settings.inpaint.controlStrength, + controlModel: settings.inpaint.controlModel, }; } diff --git a/operations/generation/workflow.test.ts b/operations/generation/workflow.test.ts index f877dcb..131b966 100644 --- a/operations/generation/workflow.test.ts +++ b/operations/generation/workflow.test.ts @@ -1,5 +1,6 @@ import { describe, expect, test } from "bun:test"; import { documentCommands } from "@commands/document"; +import { inpaintRegionCommands } from "@commands/inpaint-region"; import { generationCommands } from "@commands/generation"; import { commandIds } from "@commands/ids"; import { createCommandRegistry } from "@commands/registry"; @@ -46,14 +47,14 @@ describe("generation workflow", () => { state.document.assets.push({ id: "source-asset", name: "Source", mimeType: "image/png", source: "source", intrinsicSize: { w: 80, h: 60 } }); state.document.artboards[0]!.layers.push({ id: "source", type: "raster", name: "Source", visible: true, locked: false, opacity: 1, assetId: "source-asset", transform: { position: { x: 4, y: 8 }, scale: { x: 2, y: 2 }, rotation: 0 } }); state.editor.selection = { artboardId: "artboard", layerIds: ["source"] }; - const ids = ["mask-asset", "mask-layer"]; + const ids = ["mask-asset", "region"]; const workflow = createGenerationWorkflow(app.store, dependencies({ createId: () => ids.shift() ?? "unused" })); await workflow.prepareInpaintMask(); const next = app.store.getState(); - expect(next.document.artboards[0]?.layers[0]).toMatchObject({ id: "mask-layer", transform: { position: { x: 4, y: 8 }, scale: { x: 2, y: 2 }, rotation: 0 } }); - expect(next.editor.maskEdit).toEqual({ targetLayerId: "source", maskLayerId: "mask-layer" }); + expect(next.document.inpaintRegions).toContainEqual({ id: "region", name: "Source AI edit", targetLayerId: "source", maskAssetId: "mask-asset", enabled: true }); + expect(next.editor.maskEdit).toEqual({ kind: "inpaintRegion", targetLayerId: "source", inpaintRegionId: "region", maskAssetId: "mask-asset", viewMode: "overlay" }); }); test("cancels the active operation and records a cancelled job", async () => { @@ -82,6 +83,7 @@ function dependencies(overrides: Partial): Gener runGenerateFromCandidate: async () => undefined, createMaskedPixelReplacementSource: async () => "replacement", createRefinementMask: async () => "mask", + createInpaintRegionMask: async () => "mask", loadGenerationResources: async () => undefined, createId: () => crypto.randomUUID(), ...overrides, @@ -118,6 +120,6 @@ function createTestApp() { locked: false, layers: [], }); - const registry = createCommandRegistry([...documentCommands, ...toolCommands, ...generationCommands]); + const registry = createCommandRegistry([...documentCommands, ...inpaintRegionCommands, ...toolCommands, ...generationCommands]); return { store: createAppStore(state, registry) }; } diff --git a/operations/generation/workflow.ts b/operations/generation/workflow.ts index 11e2890..a42f1ba 100644 --- a/operations/generation/workflow.ts +++ b/operations/generation/workflow.ts @@ -3,7 +3,7 @@ import type { Layer } from "@core/layer"; import type { GenerationCandidate, GenerationJobKind } from "@editor/state"; import type { AppStore } from "@editor/store"; import type { GenerateSettings } from "@editor/tools"; -import { createRefinementMask } from "@operations/masks/rasterActions"; +import { createInpaintRegionMask, createRefinementMask } from "@operations/masks/rasterActions"; import { createMaskedPixelReplacementSource } from "./candidateActions"; import { runGenerationJob } from "./generationJob"; import { loadGenerationResources } from "./loadResources"; @@ -17,6 +17,7 @@ export type GenerationWorkflowDependencies = { runGenerateFromCandidate: typeof runGenerateFromCandidate; createMaskedPixelReplacementSource: typeof createMaskedPixelReplacementSource; createRefinementMask: typeof createRefinementMask; + createInpaintRegionMask: typeof createInpaintRegionMask; loadGenerationResources: typeof loadGenerationResources; createId(): string; }; @@ -26,13 +27,14 @@ const defaultDependencies: GenerationWorkflowDependencies = { runGenerateFromCandidate, createMaskedPixelReplacementSource, createRefinementMask, + createInpaintRegionMask, loadGenerationResources, createId: () => crypto.randomUUID(), }; export function createGenerationWorkflow(store: AppStore, dependencies: GenerationWorkflowDependencies = defaultDependencies) { let activeController: AbortController | undefined; - const job = async (kind: GenerationJobKind, label: string, task: (signal: AbortSignal) => Promise) => { + const job = async (kind: GenerationJobKind, label: string, task: (signal: AbortSignal, report: (progress: number, detail: string) => void) => Promise) => { if (store.getState().editor.generation.jobs.some((candidate) => candidate.status === "running")) return; const controller = new AbortController(); activeController = controller; @@ -60,26 +62,22 @@ export function createGenerationWorkflow(store: AppStore, dependencies: Generati if (!layer || (layer.type !== "image" && layer.type !== "raster")) return; const asset = state.document.assets.find((candidate) => candidate.id === layer.assetId); if (!asset) return; - const source = await dependencies.createRefinementMask(asset.intrinsicSize.w, asset.intrinsicSize.h); + const existing = state.document.inpaintRegions.find((region) => region.targetLayerId === layerId && region.enabled); + if (existing) { + store.dispatch(commandIds.toolEnterInpaintRegionEdit, { targetLayerId: layerId, regionId: existing.id }); + return; + } + const source = await dependencies.createInpaintRegionMask(asset.intrinsicSize.w, asset.intrinsicSize.h); const maskAssetId = dependencies.createId(); - const maskLayerId = dependencies.createId(); - store.dispatch(commandIds.documentAddLayerMask, { - layerId, - asset: { id: maskAssetId, name: `${layer.name} AI edit mask`, mimeType: "image/png", source, intrinsicSize: { ...asset.intrinsicSize } }, - maskLayer: { - id: maskLayerId, - type: "raster", - name: `${layer.name} AI edit mask`, - visible: true, - locked: false, - opacity: 1, - assetId: maskAssetId, - transform: { position: { ...layer.transform.position }, scale: { ...layer.transform.scale }, rotation: layer.transform.rotation }, - }, + const regionId = dependencies.createId(); + store.dispatch(commandIds.documentAddInpaintRegion, { + region: { id: regionId, name: `${layer.name} AI edit`, targetLayerId: layer.id, maskAssetId, enabled: true }, + maskAsset: { id: maskAssetId, name: `${layer.name} AI edit mask`, mimeType: "image/png", source, intrinsicSize: { ...asset.intrinsicSize } }, }); + store.dispatch(commandIds.toolEnterInpaintRegionEdit, { targetLayerId: layer.id, regionId }); }, - generate: () => job("generate", "Generating", async (signal) => { + generate: () => job("generate", "Generating", async (signal, report) => { const state = store.getState(); await dependencies.runGenerate({ document: state.document, @@ -88,15 +86,34 @@ export function createGenerationWorkflow(store: AppStore, dependencies: Generati settings: state.editor.tools.generate, dispatch: store.dispatch, signal, + onProgress: report, }); }), regenerate: (candidateId: string, settings?: GenerateSettings, label = "Regenerate") => - job("regenerate", label, async (signal) => { + job("regenerate", label, async (signal, report) => { const candidate = findCandidate(store, candidateId); const nextSettings = settings ?? candidate.settings; store.dispatch(commandIds.toolSetGenerateSettings, nextSettings); - await dependencies.runGenerateFromCandidate({ candidate, settings: nextSettings, dispatch: store.dispatch, signal }); + await dependencies.runGenerateFromCandidate({ candidate, settings: nextSettings, dispatch: store.dispatch, signal, onProgress: report }); + }), + + rebuildFromCurrentRegion: (candidateId: string) => + job("regenerate", "Rebuilding from current edit region", async (signal, report) => { + const candidate = findCandidate(store, candidateId); + if (!candidate.inpaint) throw new Error("Only inpaint candidates can rebuild from an edit region."); + store.dispatch(commandIds.selectionSet, { artboardId: candidate.placement.artboardId, layerIds: [candidate.inpaint.targetLayerId] }); + store.dispatch(commandIds.toolSetGenerateSettings, candidate.settings); + const state = store.getState(); + await dependencies.runGenerate({ + document: state.document, + selection: state.editor.selection, + viewport: state.editor.viewport, + settings: candidate.settings, + dispatch: store.dispatch, + signal, + onProgress: report, + }); }), applyCandidateAsLayer: (candidateId: string) => { diff --git a/operations/masks/lasso.ts b/operations/masks/lasso.ts new file mode 100644 index 0000000..0912c17 --- /dev/null +++ b/operations/masks/lasso.ts @@ -0,0 +1,52 @@ +import { commandIds } from "@commands/ids"; +import type { Vec2D } from "@core/geometry"; +import type { Layer } from "@core/layer"; +import type { AppStore } from "@editor/store"; +import { applyPolygonMask } from "@platform/browser/maskRaster"; + +export async function commitInpaintLasso(store: AppStore) { + const state = store.getState(); + const edit = state.editor.maskEdit; + const session = state.editor.maskShapeSession; + const minimumPoints = session?.shape === "rectangle" ? 2 : 3; + if (edit?.kind !== "inpaintRegion" || !edit.inpaintRegionId || !session || session.points.length < minimumPoints) { + store.dispatch(commandIds.toolClearMaskShape, undefined); + return; + } + const region = state.document.inpaintRegions.find((candidate) => candidate.id === edit.inpaintRegionId); + const target = findLayer(state.document.artboards.flatMap((artboard) => artboard.layers), edit.targetLayerId); + const asset = region ? state.document.assets.find((candidate) => candidate.id === region.maskAssetId) : undefined; + if (!region || !target || (target.type !== "image" && target.type !== "raster") || !asset) { + store.dispatch(commandIds.toolClearMaskShape, undefined); + return; + } + const documentPoints = session.shape === "rectangle" ? rectanglePoints(session.points[0]!, session.points[1]!) : session.points; + const points = documentPoints.map((point) => documentPointToAssetPoint(point, target, asset.intrinsicSize)); + const source = await applyPolygonMask(asset.source, asset.intrinsicSize.w, asset.intrinsicSize.h, points, session.mode); + store.dispatch(commandIds.documentApplyInpaintRegionMaskOperation, { regionId: region.id, source, mimeType: "image/png", operation: { type: "paint" } }); + store.dispatch(commandIds.toolClearMaskShape, undefined); +} + +function rectanglePoints(start: Vec2D, end: Vec2D): Vec2D[] { + return [start, { x: end.x, y: start.y }, end, { x: start.x, y: end.y }]; +} + +function documentPointToAssetPoint(point: Vec2D, layer: Extract, intrinsicSize: { w: number; h: number }): Vec2D { + const source = layer.sourceRect ?? { x: 0, y: 0, ...intrinsicSize }; + const destination = { x: layer.transform.position.x + source.x * layer.transform.scale.x, y: layer.transform.position.y + source.y * layer.transform.scale.y, w: source.w * layer.transform.scale.x, h: source.h * layer.transform.scale.y }; + const center = { x: destination.x + destination.w / 2, y: destination.y + destination.h / 2 }; + const dx = point.x - center.x; + const dy = point.y - center.y; + const cos = Math.cos(-layer.transform.rotation); + const sin = Math.sin(-layer.transform.rotation); + const x = center.x + dx * cos - dy * sin; + const y = center.y + dx * sin + dy * cos; + return { x: source.x + (x - destination.x) / Math.max(0.0001, layer.transform.scale.x), y: source.y + (y - destination.y) / Math.max(0.0001, layer.transform.scale.y) }; +} + +function findLayer(layers: readonly Layer[], id: string): Layer | undefined { + for (const layer of layers) { + if (layer.id === id) return layer; + if (layer.type === "group") { const child = findLayer(layer.children, id); if (child) return child; } + } +} diff --git a/operations/masks/magic-wand.ts b/operations/masks/magic-wand.ts index 4f57484..aa0c47d 100644 --- a/operations/masks/magic-wand.ts +++ b/operations/masks/magic-wand.ts @@ -11,12 +11,17 @@ import { createWandMask } from "@platform/browser/magicWandRaster"; export async function applyMagicWandAt(store: AppStore, point: Vec2D, modeOverride?: EditorState["tools"]["magicWand"]["mode"]) { const state = store.getState(); if (state.editor.tools.activeTool !== "magicWand") return false; - const target = resolveTarget(state.document, state.editor); + const target = resolveMaskSelectionTarget(state.document, state.editor); if (!target) return true; - const x = Math.floor((point.x - target.layer.transform.position.x) / Math.max(0.0001, target.layer.transform.scale.x)); - const y = Math.floor((point.y - target.layer.transform.position.y) / Math.max(0.0001, target.layer.transform.scale.y)); + const assetPoint = documentPointToAssetPoint(point, target.layer, target.asset.intrinsicSize); + const x = Math.floor(assetPoint.x); + const y = Math.floor(assetPoint.y); if (x < 0 || y < 0 || x >= target.asset.intrinsicSize.w || y >= target.asset.intrinsicSize.h) return true; - const source = await createWandMask(target.asset.source, target.maskAsset?.source, Math.round(target.asset.intrinsicSize.w), Math.round(target.asset.intrinsicSize.h), x, y, { ...state.editor.tools.magicWand, mode: modeOverride ?? state.editor.tools.magicWand.mode }); + const source = await createWandMask(target.asset.source, target.maskAsset?.source, Math.round(target.asset.intrinsicSize.w), Math.round(target.asset.intrinsicSize.h), x, y, { ...state.editor.tools.magicWand, mode: modeOverride ?? state.editor.tools.magicWand.mode, target: target.inpaintRegion ? "inpaint" : "visibility" }); + if (target.inpaintRegion && target.maskAsset) { + store.dispatch(commandIds.documentApplyInpaintRegionMaskOperation, { regionId: target.inpaintRegion.id, source, mimeType: "image/png", operation: { type: "magicWand" } }); + return true; + } if (target.maskAsset && target.maskLayer && (target.maskLayer.type === "image" || target.maskLayer.type === "raster")) { store.dispatch(commandIds.documentApplyLayerMaskOperation, { maskLayerId: target.maskLayer.id, source, mimeType: "image/png", operation: { type: "magicWand" } }); return true; @@ -35,7 +40,7 @@ export async function applyMagicWandAt(store: AppStore, point: Vec2D, modeOverri return true; } -function resolveTarget(document: ImageDocument, editor: EditorState) { +export function resolveMaskSelectionTarget(document: ImageDocument, editor: EditorState) { const layerId = editor.selection.layerIds[0]; if (!layerId || editor.selection.layerIds.length !== 1) return undefined; const layer = findLayer(document.artboards.flatMap((artboard) => artboard.layers), layerId); @@ -43,9 +48,23 @@ function resolveTarget(document: ImageDocument, editor: EditorState) { const asset = document.assets.find((candidate) => candidate.id === layer.assetId); const bounds = resolveTransformTargetBounds(document, { type: "layer", id: layer.id }); const layerMask = getLayerMask(layer); + const inpaintRegion = editor.maskEdit?.kind === "inpaintRegion" ? document.inpaintRegions.find((candidate) => candidate.id === editor.maskEdit?.inpaintRegionId && candidate.targetLayerId === layer.id) : undefined; const maskLayer = layerMask?.enabled ? findLayer(document.artboards.flatMap((artboard) => artboard.layers), layerMask.maskLayerId) : undefined; - const maskAsset = maskLayer && (maskLayer.type === "image" || maskLayer.type === "raster") ? document.assets.find((candidate) => candidate.id === maskLayer.assetId) : undefined; - return asset && bounds ? { layer, asset, bounds, maskLayer, maskAsset } : undefined; + const maskAsset = inpaintRegion ? document.assets.find((candidate) => candidate.id === inpaintRegion.maskAssetId) : maskLayer && (maskLayer.type === "image" || maskLayer.type === "raster") ? document.assets.find((candidate) => candidate.id === maskLayer.assetId) : undefined; + return asset && bounds ? { layer, asset, bounds, maskLayer, maskAsset, inpaintRegion } : undefined; +} + +export function documentPointToAssetPoint(point: Vec2D, layer: Extract, intrinsicSize: { w: number; h: number }): Vec2D { + const source = layer.sourceRect ?? { x: 0, y: 0, ...intrinsicSize }; + const destination = { x: layer.transform.position.x + source.x * layer.transform.scale.x, y: layer.transform.position.y + source.y * layer.transform.scale.y, w: source.w * layer.transform.scale.x, h: source.h * layer.transform.scale.y }; + const center = { x: destination.x + destination.w / 2, y: destination.y + destination.h / 2 }; + const dx = point.x - center.x; + const dy = point.y - center.y; + const cos = Math.cos(-layer.transform.rotation); + const sin = Math.sin(-layer.transform.rotation); + const x = center.x + dx * cos - dy * sin; + const y = center.y + dx * sin + dy * cos; + return { x: source.x + (x - destination.x) / Math.max(0.0001, layer.transform.scale.x), y: source.y + (y - destination.y) / Math.max(0.0001, layer.transform.scale.y) }; } function findLayer(layers: readonly Layer[], layerId: string): Layer | undefined { diff --git a/operations/masks/rasterActions.ts b/operations/masks/rasterActions.ts index feed84e..58b4b21 100644 --- a/operations/masks/rasterActions.ts +++ b/operations/masks/rasterActions.ts @@ -1,10 +1,13 @@ import { commandIds } from "@commands/ids"; import type { Asset } from "@core/asset"; import type { LayerId } from "@core/id"; +import type { InpaintRegionId } from "@core/id"; import type { AppStore } from "@editor/store"; import { analyzeMaskSource, applyMaskRasterOperation, createSolidMaskSource, type MaskAnalysis, type MaskRasterOperation } from "@platform/browser/maskRaster"; export type { MaskAnalysis, MaskRasterOperation }; export function analyzeMask(asset: Asset): Promise { return analyzeMaskSource(asset.source, asset.intrinsicSize.w, asset.intrinsicSize.h); } export async function runMaskOperation(maskLayerId: LayerId, asset: Asset, operation: MaskRasterOperation, dispatch: AppStore["dispatch"]) { const source = await applyMaskRasterOperation(asset.source, asset.intrinsicSize.w, asset.intrinsicSize.h, operation); dispatch(commandIds.documentApplyLayerMaskOperation, { maskLayerId, source, mimeType: "image/png", operation }); } +export async function runInpaintRegionOperation(regionId: InpaintRegionId, asset: Asset, operation: MaskRasterOperation, dispatch: AppStore["dispatch"]) { const source = await applyMaskRasterOperation(asset.source, asset.intrinsicSize.w, asset.intrinsicSize.h, operation); dispatch(commandIds.documentApplyInpaintRegionMaskOperation, { regionId, source, mimeType: "image/png", operation }); } export function createRefinementMask(width: number, height: number) { return createSolidMaskSource(width, height, "white"); } +export function createInpaintRegionMask(width: number, height: number) { return createSolidMaskSource(width, height, "black"); } diff --git a/operations/masks/semantic-select.ts b/operations/masks/semantic-select.ts new file mode 100644 index 0000000..13ca923 --- /dev/null +++ b/operations/masks/semantic-select.ts @@ -0,0 +1,34 @@ +import { commandIds } from "@commands/ids"; +import type { Vec2D } from "@core/geometry"; +import type { AppStore } from "@editor/store"; +import { mergeMaskSources } from "@platform/browser/maskRaster"; +import { requestSemanticSelection } from "@platform/comfy/generationClient"; +import { runGenerationJob } from "@operations/generation/generationJob"; +import { documentPointToAssetPoint, resolveMaskSelectionTarget } from "./magic-wand"; + +export async function applySemanticSelectionAt(store: AppStore, point: Vec2D, mode: "replace" | "add" | "subtract") { + const state = store.getState(); + if (state.editor.tools.activeTool !== "semanticSelect") return false; + const target = resolveMaskSelectionTarget(state.document, state.editor); + if (!target?.inpaintRegion || !target.maskAsset) return true; + const region = target.inpaintRegion; + const maskAsset = target.maskAsset; + const assetPoint = documentPointToAssetPoint(point, target.layer, target.asset.intrinsicSize); + if (assetPoint.x < 0 || assetPoint.y < 0 || assetPoint.x >= target.asset.intrinsicSize.w || assetPoint.y >= target.asset.intrinsicSize.h) return true; + const controller = new AbortController(); + await runGenerationJob({ + kind: "mask", + label: "Selecting object", + dispatch: store.dispatch, + signal: controller.signal, + task: async (signal, report) => { + report(0.1, "Sending point to SAM3"); + const result = await requestSemanticSelection({ inputImage: target.asset.source, x: assetPoint.x, y: assetPoint.y }, signal); + report(0.85, "Merging object mask"); + const source = await mergeMaskSources(maskAsset.source, result.source, Math.round(target.asset.intrinsicSize.w), Math.round(target.asset.intrinsicSize.h), mode); + store.dispatch(commandIds.documentApplyInpaintRegionMaskOperation, { regionId: region.id, source, mimeType: "image/png", operation: { type: "magicWand" } }); + report(1, "Object selected"); + }, + }); + return true; +} diff --git a/operations/paint/brush.ts b/operations/paint/brush.ts index 5f5fbca..d1afc31 100644 --- a/operations/paint/brush.ts +++ b/operations/paint/brush.ts @@ -25,6 +25,7 @@ export type BrushSession = { previewInFlight?: boolean; previewFrame?: number; previewSource?: string; + targetLayer: RasterLayer; }; export type BrushTargetEditorState = { @@ -34,11 +35,9 @@ export type BrushTargetEditorState = { }; export function beginBrushSession(document: ImageDocument, editor: BrushTargetEditorState, point: Vec2D): BrushSession | undefined { - const layer = resolveBrushTargetLayer(document, editor); - if (!layer || (editor.tools.activeTool !== "brush" && editor.tools.activeTool !== "eraser")) return undefined; - - const asset = document.assets.find((candidate) => candidate.id === layer.assetId); - if (!asset) return undefined; + const target = resolveBrushTarget(document, editor); + if (!target || (editor.tools.activeTool !== "brush" && editor.tools.activeTool !== "eraser")) return undefined; + const { layer, asset } = target; const surface = createBrushSurface(asset.intrinsicSize.w, asset.intrinsicSize.h, asset.source); if (!surface) return undefined; @@ -52,17 +51,18 @@ export function beginBrushSession(document: ImageDocument, editor: BrushTargetEd ready: surface.ready, previousPoint: point, mode: editor.tools.activeTool, + targetLayer: layer, }; return session; } export function canPreviewBrush(document: ImageDocument, editor: BrushTargetEditorState): boolean { - return Boolean(resolveBrushTargetLayer(document, editor)); + return Boolean(resolveBrushTarget(document, editor)); } export function brushUnavailableHint(document: ImageDocument, editor: BrushTargetEditorState): string | undefined { if (isPanInteractionMode(editor.tools.interactionMode) || (editor.tools.activeTool !== "brush" && editor.tools.activeTool !== "eraser")) return undefined; - if (resolveBrushTargetLayer(document, editor)) return undefined; + if (resolveBrushTarget(document, editor)) return undefined; const layerId = editor.maskEdit?.maskLayerId ?? editor.selection.layerIds[0]; if (!layerId) { @@ -79,14 +79,24 @@ export function brushUnavailableHint(document: ImageDocument, editor: BrushTarge return "Select a raster layer or layer mask to paint."; } -function resolveBrushTargetLayer(document: ImageDocument, editor: BrushTargetEditorState): RasterLayer | undefined { +function resolveBrushTarget(document: ImageDocument, editor: BrushTargetEditorState): { layer: RasterLayer; asset: ImageDocument["assets"][number] } | undefined { if (isPanInteractionMode(editor.tools.interactionMode) || (editor.tools.activeTool !== "brush" && editor.tools.activeTool !== "eraser")) return undefined; + if (editor.maskEdit?.kind === "inpaintRegion") { + const target = findLayer(document.artboards.flatMap((artboard) => artboard.layers), editor.maskEdit.targetLayerId); + const asset = document.assets.find((candidate) => candidate.id === editor.maskEdit?.maskAssetId); + if (!target || (target.type !== "image" && target.type !== "raster") || target.locked || !asset) return undefined; + return { + layer: { ...target, type: "raster", assetId: asset.id }, + asset, + }; + } const editingMask = Boolean(editor.maskEdit); const layerId = editor.maskEdit?.maskLayerId ?? editor.selection.layerIds[0]; if (!layerId) return undefined; const layer = findRasterLayer(document.artboards.flatMap((artboard) => artboard.layers), layerId); if (!layer || layer.locked || (!editingMask && !layer.visible)) return undefined; - return layer; + const asset = document.assets.find((candidate) => candidate.id === layer.assetId); + return asset ? { layer, asset } : undefined; } export function updateBrushSession(options: { @@ -96,13 +106,21 @@ export function updateBrushSession(options: { color: string; size: number; hardness: number; + opacity: number; + flow: number; + smoothing: number; + pressure: number; + pressureSize: boolean; }): BrushSession { const state = options.store.getState(); - const layer = findRasterLayer(state.document.artboards.flatMap((artboard) => artboard.layers), options.session.layerId); - if (!layer || layer.assetId !== options.session.assetId) return options.session; + const asset = state.document.assets.find((candidate) => candidate.id === options.session.assetId); + if (!asset) return options.session; + const layer = options.session.targetLayer; const from = options.session.previousPoint; - const to = options.point; + const smoothing = Math.max(0, Math.min(100, options.smoothing)) / 100; + const follow = 1 - smoothing * 0.85; + const to = { x: from.x + (options.point.x - from.x) * follow, y: from.y + (options.point.y - from.y) * follow }; options.session.previousPoint = to; options.session.pending = (options.session.pending ?? Promise.resolve()) .then(async () => { @@ -113,8 +131,10 @@ export function updateBrushSession(options: { from: documentPointToAssetPoint(from, layer, options.session.width, options.session.height), to: documentPointToAssetPoint(to, layer, options.session.width, options.session.height), color: state.editor.maskEdit ? "#ffffff" : options.color, - size: options.size, + size: options.size * (options.pressureSize ? Math.max(0.1, options.pressure) : 1), hardness: options.hardness, + opacity: options.opacity, + flow: options.flow, mode: options.session.mode, }); @@ -135,7 +155,9 @@ export async function commitBrushSession(options: { store: AppStore; session: Br if (source) { const state = options.store.getState(); const maskEdit = state.editor.maskEdit; - if (maskEdit?.maskLayerId === options.session.layerId) { + if (maskEdit?.kind === "inpaintRegion" && maskEdit.inpaintRegionId && maskEdit.maskAssetId === options.session.assetId) { + options.store.dispatch(commandIds.documentApplyInpaintRegionMaskOperation, { regionId: maskEdit.inpaintRegionId, source, mimeType: "image/png", operation: { type: "paint" } }); + } else if (maskEdit?.maskLayerId === options.session.layerId) { options.store.dispatch(commandIds.documentApplyLayerMaskOperation, { maskLayerId: maskEdit.maskLayerId, source, mimeType: "image/png", operation: { type: "paint" } }); } else { options.store.dispatch(commandIds.documentUpdateAssetSource, { assetId: options.session.assetId, source }); @@ -152,9 +174,22 @@ export function cancelBrushSession(options: { store: AppStore; session: BrushSes } function documentPointToAssetPoint(point: Vec2D, layer: RasterLayer, width: number, height: number): Vec2D { + const source = layer.sourceRect ?? { x: 0, y: 0, w: width, h: height }; + const destination = { + x: layer.transform.position.x + source.x * layer.transform.scale.x, + y: layer.transform.position.y + source.y * layer.transform.scale.y, + w: source.w * layer.transform.scale.x, + h: source.h * layer.transform.scale.y, + }; + const center = { x: destination.x + destination.w / 2, y: destination.y + destination.h / 2 }; + const cos = Math.cos(-layer.transform.rotation); + const sin = Math.sin(-layer.transform.rotation); + const dx = point.x - center.x; + const dy = point.y - center.y; + const unrotated = { x: center.x + dx * cos - dy * sin, y: center.y + dx * sin + dy * cos }; return { - x: ((point.x - layer.transform.position.x) / Math.max(0.0001, layer.transform.scale.x) / width) * width, - y: ((point.y - layer.transform.position.y) / Math.max(0.0001, layer.transform.scale.y) / height) * height, + x: source.x + (unrotated.x - destination.x) / Math.max(0.0001, layer.transform.scale.x), + y: source.y + (unrotated.y - destination.y) / Math.max(0.0001, layer.transform.scale.y), }; } diff --git a/operations/project/format.test.ts b/operations/project/format.test.ts index 4826d62..e66a552 100644 --- a/operations/project/format.test.ts +++ b/operations/project/format.test.ts @@ -14,7 +14,7 @@ describe("project format", () => { test("migrates a legacy bare document", () => { const result = parseProject(JSON.stringify(projectDocument())); - expect(result.version).toBe(1); + expect(result.version).toBe(CURRENT_PROJECT_VERSION); expect(result.savedAt).toBe("1970-01-01T00:00:00.000Z"); }); @@ -76,6 +76,7 @@ function projectDocument(): ImageDocument { name: "Test Project", version: 1, assets: [{ id: "asset-1", name: "pixels.png", mimeType: "image/png", source: "data:image/png;base64,AA==", intrinsicSize: { w: 10, h: 20 } }], + inpaintRegions: [], artboards: [{ id: "artboard-1", name: "Board", diff --git a/operations/project/format.ts b/operations/project/format.ts index 756d74b..7b4e5e4 100644 --- a/operations/project/format.ts +++ b/operations/project/format.ts @@ -4,7 +4,7 @@ import type { Rect } from "@core/geometry"; import { isValidTextStyle, type TextStyle } from "@core/text-layer"; export const PROJECT_FORMAT = "image-studio-project"; -export const CURRENT_PROJECT_VERSION = 1; +export const CURRENT_PROJECT_VERSION = 2; export type ProjectFile = { format: typeof PROJECT_FORMAT; @@ -41,6 +41,14 @@ function migrateProject(value: unknown): ProjectFile { if (!isRecord(value)) throw new Error("The project file must contain an object."); if (value.format === PROJECT_FORMAT) { + if (value.version === 1 && isRecord(value.document)) { + return { + format: PROJECT_FORMAT, + version: CURRENT_PROJECT_VERSION, + savedAt: typeof value.savedAt === "string" ? value.savedAt : new Date(0).toISOString(), + document: { ...value.document, inpaintRegions: [] } as unknown as ImageDocument, + }; + } if (value.version !== CURRENT_PROJECT_VERSION) { throw new Error(`Unsupported project version: ${String(value.version)}.`); } @@ -48,13 +56,13 @@ function migrateProject(value: unknown): ProjectFile { return value as ProjectFile; } - // Legacy exports stored ImageDocument directly. Loading upgrades them to v1. + // Legacy exports stored ImageDocument directly. Loading upgrades them to the current format. if (looksLikeDocument(value)) { return { format: PROJECT_FORMAT, version: CURRENT_PROJECT_VERSION, savedAt: new Date(0).toISOString(), - document: value as ImageDocument, + document: { ...value, inpaintRegions: Array.isArray(value.inpaintRegions) ? value.inpaintRegions : [] } as ImageDocument, }; } @@ -65,7 +73,7 @@ function assertImageDocument(value: unknown): asserts value is ImageDocument { if (!isRecord(value) || typeof value.id !== "string" || typeof value.name !== "string" || typeof value.version !== "number") { throw new Error("The project contains an invalid document."); } - if (!Array.isArray(value.artboards) || !Array.isArray(value.assets)) throw new Error("The project document is incomplete."); + if (!Array.isArray(value.artboards) || !Array.isArray(value.assets) || !Array.isArray(value.inpaintRegions)) throw new Error("The project document is incomplete."); for (const asset of value.assets) { if (!isRecord(asset) || typeof asset.id !== "string" || typeof asset.name !== "string" || typeof asset.mimeType !== "string" || typeof asset.source !== "string" || !isSize(asset.intrinsicSize)) { @@ -78,6 +86,11 @@ function assertImageDocument(value: unknown): asserts value is ImageDocument { } assertLayers(artboard.layers, false); } + for (const region of value.inpaintRegions) { + if (!isRecord(region) || typeof region.id !== "string" || typeof region.name !== "string" || typeof region.targetLayerId !== "string" || typeof region.maskAssetId !== "string" || typeof region.enabled !== "boolean") { + throw new Error("The project contains an invalid inpaint region."); + } + } } export function assertDocumentAssetOwnership(document: ImageDocument): void { @@ -98,6 +111,10 @@ export function assertDocumentAssetOwnership(document: ImageDocument): void { } }); } + for (const region of document.inpaintRegions) { + if (!assetIds.has(region.maskAssetId)) throw new Error(`Inpaint region ${region.name} references missing mask asset ${region.maskAssetId}.`); + if (!layerIds.has(region.targetLayerId)) throw new Error(`Inpaint region ${region.name} references missing target layer ${region.targetLayerId}.`); + } } function isOwnedAssetSource(source: string): boolean { diff --git a/platform/browser/brushRaster.ts b/platform/browser/brushRaster.ts index 1932021..60e7086 100644 --- a/platform/browser/brushRaster.ts +++ b/platform/browser/brushRaster.ts @@ -15,9 +15,9 @@ export function createBrushSurface(width: number, height: number, source: string return surface; } -export function drawBrushSegment(surface: BrushSurface, options: { from: { x: number; y: number }; to: { x: number; y: number }; color: string; size: number; hardness: number; mode: "brush" | "eraser" }) { +export function drawBrushSegment(surface: BrushSurface, options: { from: { x: number; y: number }; to: { x: number; y: number }; color: string; size: number; hardness: number; opacity: number; flow: number; mode: "brush" | "eraser" }) { const context = internal(surface).context; const hardness = Math.max(0, Math.min(100, options.hardness)) / 100; - context.save(); context.globalCompositeOperation = options.mode === "eraser" ? "destination-out" : "source-over"; context.strokeStyle = options.color; context.shadowColor = options.mode === "eraser" ? "rgba(0,0,0,1)" : options.color; context.shadowBlur = (1 - hardness) * options.size; context.lineWidth = options.size; context.lineCap = "round"; context.lineJoin = "round"; context.beginPath(); context.moveTo(options.from.x, options.from.y); context.lineTo(options.to.x, options.to.y); context.stroke(); context.restore(); + context.save(); context.globalAlpha = Math.max(0, Math.min(1, options.opacity / 100)) * Math.max(0.01, Math.min(1, options.flow / 100)); context.globalCompositeOperation = options.mode === "eraser" ? "destination-out" : "source-over"; context.strokeStyle = options.color; context.shadowColor = options.mode === "eraser" ? "rgba(0,0,0,1)" : options.color; context.shadowBlur = (1 - hardness) * options.size; context.lineWidth = options.size; context.lineCap = "round"; context.lineJoin = "round"; context.beginPath(); context.moveTo(options.from.x, options.from.y); context.lineTo(options.to.x, options.to.y); context.stroke(); context.restore(); } export function brushSurfaceDataUrl(surface: BrushSurface) { try { return internal(surface).canvas.toDataURL("image/png"); } catch { return undefined; } } diff --git a/platform/browser/exportArtboardPng.ts b/platform/browser/exportArtboardPng.ts index 7b95dd2..762cba0 100644 --- a/platform/browser/exportArtboardPng.ts +++ b/platform/browser/exportArtboardPng.ts @@ -7,6 +7,15 @@ import { getLayerMask } from "@core/layer-mask-utils"; import { measureTextLayer } from "@core/text-layer"; export async function downloadArtboardPng(artboard: Artboard, assets: readonly Asset[]) { + const canvas = await renderArtboardCanvas(artboard, assets); + const url = canvas.toDataURL("image/png"); + const link = document.createElement("a"); + link.href = url; + link.download = `${safeFilename(artboard.name)}.png`; + link.click(); +} + +export async function renderArtboardCanvas(artboard: Artboard, assets: readonly Asset[]): Promise { const width = Math.max(1, Math.round(artboard.bounds.w)); const height = Math.max(1, Math.round(artboard.bounds.h)); const canvas = document.createElement("canvas"); @@ -27,11 +36,7 @@ export async function downloadArtboardPng(artboard: Artboard, assets: readonly A for (const layer of renderStack(artboard.layers)) await drawLayer(context, layer, artboard.layers, assets, artboard.bounds, { maskLayerIds }); context.restore(); - const url = canvas.toDataURL("image/png"); - const link = document.createElement("a"); - link.href = url; - link.download = `${safeFilename(artboard.name)}.png`; - link.click(); + return canvas; } async function drawLayer( diff --git a/platform/browser/magicWandRaster.ts b/platform/browser/magicWandRaster.ts index 1a39b6e..f552c54 100644 --- a/platform/browser/magicWandRaster.ts +++ b/platform/browser/magicWandRaster.ts @@ -1,6 +1,6 @@ import { blurMaskValues, despeckleMaskValues, dilateMaskValues, erodeMaskValues, maskValueFromRgba } from "./maskRaster"; -export type MagicWandRasterSettings = { tolerance: number; feather: number; choke: number; despeckle: number; contiguous: boolean; mode: "replace" | "add" | "subtract" }; +export type MagicWandRasterSettings = { tolerance: number; feather: number; choke: number; despeckle: number; contiguous: boolean; mode: "replace" | "add" | "subtract"; target?: "visibility" | "inpaint" }; export async function createWandMask(source: string, existingMaskSource: string | undefined, width: number, height: number, startX: number, startY: number, settings: MagicWandRasterSettings) { const canvas = document.createElement("canvas"); @@ -10,8 +10,9 @@ export async function createWandMask(source: string, existingMaskSource: string context.drawImage(await loadImage(source), 0, 0, canvas.width, canvas.height); const data = context.getImageData(0, 0, canvas.width, canvas.height); const start = (startY * canvas.width + startX) * 4; - const key = [data.data[start] ?? 0, data.data[start + 1] ?? 0, data.data[start + 2] ?? 0]; - const selected = settings.contiguous ? floodSelect(data, width, height, startX, startY, key, settings.tolerance) : globalSelect(data, key, settings.tolerance); + const key = rgbToLab(data.data[start] ?? 0, data.data[start + 1] ?? 0, data.data[start + 2] ?? 0); + const keyAlpha = data.data[start + 3] ?? 255; + const selected = settings.contiguous ? floodSelect(data, width, height, startX, startY, key, keyAlpha, settings.tolerance) : globalSelect(data, key, keyAlpha, settings.tolerance); let values: Uint8ClampedArray = toValues(selected); if (settings.despeckle > 0) values = despeckleMaskValues(values, width, height, Math.round(settings.despeckle)); if (settings.choke > 0) values = erodeMaskValues(values, width, height, Math.round(settings.choke)); @@ -19,20 +20,31 @@ export async function createWandMask(source: string, existingMaskSource: string if (settings.feather > 0) values = blurMaskValues(values, width, height, Math.round(settings.feather)); const existing = existingMaskSource ? await loadMask(existingMaskSource, width, height) : undefined; for (let pixel = 0; pixel < values.length; pixel++) { - const current = existing?.[pixel] ?? 255; const selectedValue = values[pixel] ?? 0; - const alpha = settings.mode === "add" ? Math.min(current, 255 - selectedValue) : settings.mode === "subtract" ? Math.max(current, selectedValue) : 255 - selectedValue; + const current = existing?.[pixel] ?? (settings.target === "inpaint" ? 0 : 255); const selectedValue = values[pixel] ?? 0; + const alpha = settings.target === "inpaint" + ? settings.mode === "add" ? Math.max(current, selectedValue) : settings.mode === "subtract" ? Math.min(current, 255 - selectedValue) : selectedValue + : settings.mode === "add" ? Math.min(current, 255 - selectedValue) : settings.mode === "subtract" ? Math.max(current, selectedValue) : 255 - selectedValue; const index = pixel * 4; data.data[index] = data.data[index + 1] = data.data[index + 2] = 255; data.data[index + 3] = alpha; } context.putImageData(data, 0, 0); return canvas.toDataURL("image/png"); } -function floodSelect(data: ImageData, width: number, height: number, x: number, y: number, key: number[], tolerance: number) { +function floodSelect(data: ImageData, width: number, height: number, x: number, y: number, key: readonly number[], keyAlpha: number, tolerance: number) { const result = new Uint8Array(width * height); const queue: Array<[number, number]> = [[x, y]]; - while (queue.length) { const [px, py] = queue.pop()!; if (px < 0 || py < 0 || px >= width || py >= height) continue; const i = py * width + px; if (result[i] || !matches(data, i, key, tolerance)) continue; result[i] = 1; queue.push([px + 1, py], [px - 1, py], [px, py + 1], [px, py - 1]); } + while (queue.length) { const [px, py] = queue.pop()!; if (px < 0 || py < 0 || px >= width || py >= height) continue; const i = py * width + px; if (result[i] || !matches(data, i, key, keyAlpha, tolerance)) continue; result[i] = 1; queue.push([px + 1, py], [px - 1, py], [px, py + 1], [px, py - 1]); } return result; } -function globalSelect(data: ImageData, key: number[], tolerance: number) { const result = new Uint8Array(data.width * data.height); for (let i = 0; i < result.length; i++) if (matches(data, i, key, tolerance)) result[i] = 1; return result; } -function matches(data: ImageData, pixel: number, key: number[], tolerance: number) { const i = pixel * 4; return Math.hypot((data.data[i] ?? 0) - (key[0] ?? 0), (data.data[i + 1] ?? 0) - (key[1] ?? 0), (data.data[i + 2] ?? 0) - (key[2] ?? 0)) <= tolerance; } +function globalSelect(data: ImageData, key: readonly number[], keyAlpha: number, tolerance: number) { const result = new Uint8Array(data.width * data.height); for (let i = 0; i < result.length; i++) if (matches(data, i, key, keyAlpha, tolerance)) result[i] = 1; return result; } +function matches(data: ImageData, pixel: number, key: readonly number[], keyAlpha: number, tolerance: number) { const i = pixel * 4; const sample = rgbToLab(data.data[i] ?? 0, data.data[i + 1] ?? 0, data.data[i + 2] ?? 0); const colorDistance = Math.hypot((sample[0] ?? 0) - (key[0] ?? 0), (sample[1] ?? 0) - (key[1] ?? 0), (sample[2] ?? 0) - (key[2] ?? 0)); const alphaDistance = Math.abs((data.data[i + 3] ?? 255) - keyAlpha) / 2.55; return Math.hypot(colorDistance, alphaDistance) <= tolerance * 0.45; } + +function rgbToLab(red: number, green: number, blue: number): [number, number, number] { + const linear = [red, green, blue].map((value) => { const channel = value / 255; return channel <= 0.04045 ? channel / 12.92 : ((channel + 0.055) / 1.055) ** 2.4; }); + const x = ((linear[0] ?? 0) * 0.4124 + (linear[1] ?? 0) * 0.3576 + (linear[2] ?? 0) * 0.1805) / 0.95047; + const y = ((linear[0] ?? 0) * 0.2126 + (linear[1] ?? 0) * 0.7152 + (linear[2] ?? 0) * 0.0722); + const z = ((linear[0] ?? 0) * 0.0193 + (linear[1] ?? 0) * 0.1192 + (linear[2] ?? 0) * 0.9505) / 1.08883; + const f = (value: number) => value > 0.008856 ? Math.cbrt(value) : 7.787 * value + 16 / 116; + return [116 * f(y) - 16, 500 * (f(x) - f(y)), 200 * (f(y) - f(z))]; +} function toValues(selected: Uint8Array) { const values = new Uint8ClampedArray(selected.length); for (let i = 0; i < selected.length; i++) values[i] = selected[i] ? 255 : 0; return values; } async function loadMask(source: string, width: number, height: number) { const canvas = document.createElement("canvas"); canvas.width = width; canvas.height = height; const context = canvas.getContext("2d"); if (!context) return undefined; context.drawImage(await loadImage(source), 0, 0, width, height); const data = context.getImageData(0, 0, width, height); const values = new Uint8ClampedArray(width * height); for (let i = 0; i < values.length; i++) values[i] = maskValueFromRgba(data.data, i * 4); return values; } function loadImage(source: string) { 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; }); } diff --git a/platform/browser/maskRaster.test.ts b/platform/browser/maskRaster.test.ts index f805795..ed69515 100644 --- a/platform/browser/maskRaster.test.ts +++ b/platform/browser/maskRaster.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from "bun:test"; -import { blurMaskValues, cropMaskValuesToRgba, dilateMaskValues, erodeMaskValues, expandRectWithinBounds, invertMaskValues } from "./maskRaster"; +import { blurMaskValues, cropMaskValuesToRgba, dilateMaskValues, erodeMaskValues, expandRectWithinBounds, invertMaskValues, limitMaskValues } from "./maskRaster"; describe("mask raster utilities", () => { test("exports the normalized drawn mask without filling the whole crop", () => { @@ -35,6 +35,15 @@ describe("mask raster utilities", () => { expect([...erodeMaskValues(new Uint8ClampedArray(new Array(9).fill(255)), 3, 3, 1)]).toEqual(new Array(9).fill(255)); expect([...blurMaskValues(values, 3, 3, 1)]).toEqual(new Array(9).fill(28)); }); + + test("clips AI edit regions to a retained source crop", () => { + expect([...limitMaskValues(new Uint8ClampedArray(16).fill(255), 4, 4, { x: 1, y: 1, w: 2, h: 2 })]).toEqual([ + 0, 0, 0, 0, + 0, 255, 255, 0, + 0, 255, 255, 0, + 0, 0, 0, 0, + ]); + }); }); function activeRedPixels(rgba: Uint8ClampedArray) { diff --git a/platform/browser/maskRaster.ts b/platform/browser/maskRaster.ts index e488a28..9788ad2 100644 --- a/platform/browser/maskRaster.ts +++ b/platform/browser/maskRaster.ts @@ -30,6 +30,7 @@ export type NormalizedMaskOptions = { feather?: number; blur?: number; despeckle?: number; + limit?: Rect; }; export async function createSolidMaskSource(width: number, height: number, fill: MaskFill): Promise { @@ -58,6 +59,36 @@ export async function applyMaskRasterOperation(source: string, width: number, he return maskValuesToDataUrl(next, mask.width, mask.height); } +export async function applyPolygonMask(source: string, width: number, height: number, points: readonly { x: number; y: number }[], mode: "replace" | "add" | "subtract"): Promise { + assertProcessingRasterSize(width, height, "Mask"); + if (points.length < 3) return source; + const canvas = await loadImageCanvas(source, width, height); + const context = require2dContext(canvas); + if (mode === "replace") context.clearRect(0, 0, canvas.width, canvas.height); + context.save(); + context.globalCompositeOperation = mode === "subtract" ? "destination-out" : "source-over"; + context.fillStyle = "#ffffff"; + context.beginPath(); + context.moveTo(points[0]?.x ?? 0, points[0]?.y ?? 0); + for (let index = 1; index < points.length; index += 1) context.lineTo(points[index]?.x ?? 0, points[index]?.y ?? 0); + context.closePath(); + context.fill(); + context.restore(); + return canvas.toDataURL("image/png"); +} + +export async function mergeMaskSources(existingSource: string, selectedSource: string, width: number, height: number, mode: "replace" | "add" | "subtract"): Promise { + assertProcessingRasterSize(width, height, "Mask"); + const [existing, selected] = await Promise.all([loadMaskValues(existingSource, width, height), loadMaskValues(selectedSource, width, height)]); + const values = new Uint8ClampedArray(width * height); + for (let pixel = 0; pixel < values.length; pixel += 1) { + const current = existing.values[pixel] ?? 0; + const choice = selected.values[pixel] ?? 0; + values[pixel] = mode === "add" ? Math.max(current, choice) : mode === "subtract" ? Math.min(current, 255 - choice) : choice; + } + return maskValuesToDataUrl(values, width, height); +} + export async function analyzeMaskSource(source: string, width: number, height: number): Promise { assertProcessingRasterSize(width, height, "Mask"); const mask = await loadMaskValues(source, width, height); @@ -81,6 +112,7 @@ export async function createNormalizedMaskSource(source: string, width: number, assertProcessingRasterSize(width, height, "Mask"); const mask = await loadMaskValues(source, width, height); let values = options.polarity === "hidden" ? invertMaskValues(mask.values) : new Uint8ClampedArray(mask.values); + if (options.limit) values = limitMaskValues(values, mask.width, mask.height, options.limit); const despeckle = Math.round(clampNumber(options.despeckle ?? 0, 0, 64)); const expand = Math.round(clampNumber(options.expand ?? 0, -256, 256)); @@ -97,6 +129,16 @@ export async function createNormalizedMaskSource(source: string, width: number, return { source: maskValuesToDataUrl(values, mask.width, mask.height), values, bounds: analysis.bounds }; } +export function limitMaskValues(values: Uint8ClampedArray, width: number, height: number, limit: Rect): Uint8ClampedArray { + const next = new Uint8ClampedArray(values.length); + const x1 = Math.max(0, Math.floor(limit.x)); + const y1 = Math.max(0, Math.floor(limit.y)); + const x2 = Math.min(width, Math.ceil(limit.x + limit.w)); + const y2 = Math.min(height, Math.ceil(limit.y + limit.h)); + for (let y = y1; y < y2; y += 1) for (let x = x1; x < x2; x += 1) next[y * width + x] = values[y * width + x] ?? 0; + return next; +} + export async function loadImageCanvas(source: string, width?: number, height?: number): Promise { const image = await loadImage(source); assertCanvasRasterSize(width ?? image.naturalWidth, height ?? image.naturalHeight); @@ -122,6 +164,33 @@ export function cropCanvas(sourceCanvas: HTMLCanvasElement, crop: Rect, outputWi return canvas.toDataURL("image/png"); } +export function sampleDocumentCanvasInLayerSpace( + documentCanvas: HTMLCanvasElement, + artboardBounds: Rect, + layer: { transform: { position: { x: number; y: number }; scale: { x: number; y: number }; rotation: number }; sourceRect?: Rect }, + intrinsicSize: { w: number; h: number }, + crop: Rect, + outputWidth = crop.w, + outputHeight = crop.h, +): string { + assertProcessingRasterSize(outputWidth, outputHeight, "Context crop"); + const canvas = createCanvas(outputWidth, outputHeight); + const context = require2dContext(canvas); + const source = layer.sourceRect ?? { x: 0, y: 0, ...intrinsicSize }; + const center = { + x: layer.transform.position.x + (source.x + source.w / 2) * layer.transform.scale.x, + y: layer.transform.position.y + (source.y + source.h / 2) * layer.transform.scale.y, + }; + context.translate(-crop.x, -crop.y); + context.scale(1 / Math.max(0.0001, layer.transform.scale.x), 1 / Math.max(0.0001, layer.transform.scale.y)); + context.translate(-layer.transform.position.x, -layer.transform.position.y); + context.translate(center.x, center.y); + context.rotate(-layer.transform.rotation); + context.translate(-center.x, -center.y); + context.drawImage(documentCanvas, artboardBounds.x, artboardBounds.y); + return canvas.toDataURL("image/png"); +} + export function cropMaskValuesToDataUrl(values: Uint8ClampedArray, width: number, height: number, crop: Rect, outputWidth = crop.w, outputHeight = crop.h): string { assertProcessingRasterSize(outputWidth, outputHeight, "Mask crop"); const canvas = createCanvas(outputWidth, outputHeight); diff --git a/platform/comfy/generationClient.ts b/platform/comfy/generationClient.ts index 1c18738..9ecc3c0 100644 --- a/platform/comfy/generationClient.ts +++ b/platform/comfy/generationClient.ts @@ -4,8 +4,40 @@ export async function fetchGenerationOptions(): Promise { return response.json() as Promise; } -export async function requestGeneration(body: unknown, signal?: AbortSignal): Promise<{ source: string; mimeType: string }> { - const response = await fetch("/api/comfy/generate", { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify(body), signal }); +export type GenerationResult = { source: string; mimeType: string; seed: number }; + +export async function requestSemanticSelection(body: { inputImage: string; x: number; y: number; model?: string }, signal?: AbortSignal): Promise<{ source: string; mimeType: string }> { + const response = await fetch("/api/comfy/segment", { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify(body), signal }); if (!response.ok) throw new Error(await response.text()); return response.json() as Promise<{ source: string; mimeType: string }>; } + +export async function requestGeneration(body: unknown, signal?: AbortSignal, onProgress?: (progress: number, detail: string) => void): Promise<{ results: GenerationResult[] }> { + const response = await fetch("/api/comfy/generate", { method: "POST", headers: { "content-type": "application/json", accept: "application/x-ndjson" }, body: JSON.stringify(body), signal }); + if (!response.ok) throw new Error(await response.text()); + const contentType = response.headers.get("content-type") ?? ""; + if (contentType.includes("application/x-ndjson") && response.body) { + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + let pending = ""; + while (true) { + const { value, done } = await reader.read(); + pending += decoder.decode(value, { stream: !done }); + const lines = pending.split("\n"); + pending = lines.pop() ?? ""; + for (const line of lines) { + if (!line.trim()) continue; + const event = JSON.parse(line) as { type: "progress"; progress: number; detail: string } | { type: "result"; results: GenerationResult[] } | { type: "error"; message: string }; + if (event.type === "progress") onProgress?.(event.progress, event.detail); + if (event.type === "error") throw new Error(event.message); + if (event.type === "result") return { results: event.results }; + } + if (done) break; + } + throw new Error("Generation stream ended without results"); + } + const payload = await response.json() as { results?: GenerationResult[]; source?: string; mimeType?: string }; + if (payload.results?.length) return { results: payload.results }; + if (payload.source) return { results: [{ source: payload.source, mimeType: payload.mimeType ?? "image/png", seed: 0 }] }; + throw new Error("Generation returned no results"); +} diff --git a/renderer/brush-preview.ts b/renderer/brush-preview.ts index 6906c91..91c446b 100644 --- a/renderer/brush-preview.ts +++ b/renderer/brush-preview.ts @@ -1,7 +1,6 @@ import type { ImageDocument } from "@core/document"; import type { Vec2D } from "@core/geometry"; import type { Layer } from "@core/layer"; -import type { RasterLayer } from "@core/raster-layer"; import type { EditorState } from "@editor/state"; import type { RgbaColor, WebGlRendererContext } from "./types"; @@ -97,12 +96,12 @@ function resolveBrushPreview(document: ImageDocument, editor: EditorState, canva }; } -function resolveBrushTargetLayer(document: ImageDocument, editor: EditorState): RasterLayer | undefined { +function resolveBrushTargetLayer(document: ImageDocument, editor: EditorState): Extract | undefined { const editingMask = Boolean(editor.maskEdit); - const layerId = editor.maskEdit?.maskLayerId ?? editor.selection.layerIds[0]; + const layerId = editor.maskEdit?.kind === "inpaintRegion" ? editor.maskEdit.targetLayerId : editor.maskEdit?.maskLayerId ?? editor.selection.layerIds[0]; if (!layerId) return undefined; - const layer = findRasterLayer(document.artboards.flatMap((artboard) => artboard.layers), layerId); + const layer = findPaintableLayer(document.artboards.flatMap((artboard) => artboard.layers), layerId); if (!layer || layer.locked || (!editingMask && !layer.visible)) return undefined; return layer; } @@ -124,11 +123,11 @@ function previewVertices(center: Vec2D, radius: Vec2D) { return new Float32Array([x1, y1, x2, y1, x1, y2, x1, y2, x2, y1, x2, y2]); } -function findRasterLayer(layers: readonly Layer[], layerId: string): RasterLayer | undefined { +function findPaintableLayer(layers: readonly Layer[], layerId: string): Extract | undefined { for (const layer of layers) { - if (layer.id === layerId && layer.type === "raster") return layer; + if (layer.id === layerId && (layer.type === "image" || layer.type === "raster")) return layer; if (layer.type === "group") { - const child = findRasterLayer(layer.children, layerId); + const child = findPaintableLayer(layer.children, layerId); if (child) return child; } } diff --git a/renderer/image-texture-programs.ts b/renderer/image-texture-programs.ts index c4b3186..e26e24f 100644 --- a/renderer/image-texture-programs.ts +++ b/renderer/image-texture-programs.ts @@ -80,6 +80,8 @@ export function maskVisualizationModeValue(mode: MaskVisualizationMode) { return 1; case "hiddenOverlay": return 2; + case "activeOverlay": + return 3; } } @@ -119,7 +121,7 @@ function createMaskVisualizationProgram(gl: WebGL2RenderingContext): MaskVisuali return; } - float alpha = (1.0 - maskAlpha) * u_color.a; + float alpha = (u_mode == 3 ? maskAlpha : (1.0 - maskAlpha)) * u_color.a; outColor = vec4(u_color.rgb * alpha, alpha); }`, ); diff --git a/renderer/image-textures.ts b/renderer/image-textures.ts index 9876797..554f189 100644 --- a/renderer/image-textures.ts +++ b/renderer/image-textures.ts @@ -5,7 +5,7 @@ import type { Rect } from "@core/geometry"; import { rotatedRectBounds, rotatedRectCorners } from "./rotated-rect"; import { textureCoordinatesForCrop, textureCoordinatesForRect } from "./texture-coordinates"; -export type MaskVisualizationMode = "blackWhite" | "alpha" | "hiddenOverlay"; +export type MaskVisualizationMode = "blackWhite" | "alpha" | "hiddenOverlay" | "activeOverlay"; export type ImageTextureRenderer = { syncAssets(assets: readonly Pick[]): void; diff --git a/renderer/layers.ts b/renderer/layers.ts index 8e13175..9dfab12 100644 --- a/renderer/layers.ts +++ b/renderer/layers.ts @@ -15,6 +15,7 @@ import { rasterizedTextLayerRenderAsset } from "./text-asset"; const imageLayerColor: RgbaColor = [0.38, 0.42, 0.5, 1]; const imageLayerInsetColor: RgbaColor = [0.48, 0.54, 0.64, 1]; const hiddenMaskOverlayColor: RgbaColor = [1, 0.08, 0.08, 0.45]; +const inpaintRegionOverlayColor: RgbaColor = [1, 0.12, 0.18, 0.48]; const comparisonDividerColor: RgbaColor = [1, 1, 1, 0.9]; const maskRevealPreviewOpacity = 0.28; @@ -99,6 +100,24 @@ function renderLeafLayer( const rect = documentRectToScreenRect(context.canvas, bounds, editor.viewport); const asset = assetWithBrushStrokePreview(documentIndex.assetById.get(layer.assetId), editor); + const activeInpaintRegion = editor.maskEdit?.kind === "inpaintRegion" && editor.maskEdit.targetLayerId === layer.id + ? documentIndex.assetById.get(editor.maskEdit.maskAssetId) + : undefined; + if (asset && activeInpaintRegion) { + const inpaintMaskAsset = assetWithBrushStrokePreview(activeInpaintRegion, editor); + const inpaintMaskRect = documentRectToScreenRect(context.canvas, { + x: layer.transform.position.x, + y: layer.transform.position.y, + w: inpaintMaskAsset.intrinsicSize.w * layer.transform.scale.x, + h: inpaintMaskAsset.intrinsicSize.h * layer.transform.scale.y, + }, editor.viewport); + if (maskViewMode === "blackWhite" && imageTextureRenderer.renderMaskVisualization(inpaintMaskAsset, inpaintMaskRect, "blackWhite", undefined, effectiveClipRect)) return; + if (maskViewMode === "alpha" && imageTextureRenderer.renderMaskVisualization(inpaintMaskAsset, inpaintMaskRect, "alpha", undefined, effectiveClipRect)) return; + if (imageTextureRenderer.render(asset, rect, effectiveClipRect, effectiveOpacity, layer.transform.rotation, layer.sourceRect)) { + imageTextureRenderer.renderMaskVisualization(inpaintMaskAsset, inpaintMaskRect, "activeOverlay", inpaintRegionOverlayColor, effectiveClipRect); + return; + } + } const layerMask = getLayerMask(layer); const maskLayer = !editingMaskLayer && layerMask?.enabled ? documentIndex.layerById.get(layerMask.maskLayerId) : undefined; const maskAsset = assetWithBrushStrokePreview(maskLayer && (maskLayer.type === "image" || maskLayer.type === "raster") ? documentIndex.assetById.get(maskLayer.assetId) : undefined, editor); diff --git a/renderer/mask-shape-preview.ts b/renderer/mask-shape-preview.ts new file mode 100644 index 0000000..3b85d24 --- /dev/null +++ b/renderer/mask-shape-preview.ts @@ -0,0 +1,32 @@ +import type { EditorState } from "@editor/state"; +import { clearScreenRect } from "./clear-rect"; +import type { WebGlRendererContext } from "./types"; + +const previewColor = [1, 0.18, 0.24, 0.95] as const; + +export function renderMaskShapePreview(context: WebGlRendererContext, editor: EditorState) { + const session = editor.maskShapeSession; + const points = session?.shape === "rectangle" && session.points.length > 1 + ? rectanglePoints(session.points[0]!, session.points[1]!) + : session?.points; + if (!points || points.length === 0) return; + const screenPoints = points.map((point) => ({ + x: context.canvas.width / 2 + (point.x - editor.viewport.center.x) * editor.viewport.zoom, + y: context.canvas.height / 2 + (point.y - editor.viewport.center.y) * editor.viewport.zoom, + })); + const closed = screenPoints.length > 2 ? [...screenPoints, screenPoints[0]!] : screenPoints; + for (let index = 1; index < closed.length; index += 1) { + const from = closed[index - 1]!; + const to = closed[index]!; + const distance = Math.max(1, Math.hypot(to.x - from.x, to.y - from.y)); + const steps = Math.max(1, Math.ceil(distance / 3)); + for (let step = 0; step <= steps; step += 1) { + const amount = step / steps; + clearScreenRect(context, { x: from.x + (to.x - from.x) * amount - 1.5, y: from.y + (to.y - from.y) * amount - 1.5, w: 3, h: 3 }, previewColor); + } + } +} + +function rectanglePoints(start: { x: number; y: number }, end: { x: number; y: number }) { + return [start, { x: end.x, y: start.y }, end, { x: start.x, y: end.y }]; +} diff --git a/renderer/renderer.ts b/renderer/renderer.ts index c49c5a2..46204cd 100644 --- a/renderer/renderer.ts +++ b/renderer/renderer.ts @@ -10,6 +10,7 @@ import { renderSelectionOverlay } from "./selection"; import { renderTransformControls } from "./transform-controls"; import type { WebGlRendererContext } from "./types"; import { createAdjustmentPass } from "./adjustment-pass"; +import { renderMaskShapePreview } from "./mask-shape-preview"; export type RenderFrame = { document: ImageDocument; @@ -77,6 +78,7 @@ export function createRenderer(canvas: HTMLCanvasElement, backend: RendererBacke renderTransformControls(rendererContext, frame.document, frame.editor); } brushPreviewRenderer?.render(frame.document, frame.editor); + renderMaskShapePreview(rendererContext, frame.editor); context.disable(context.SCISSOR_TEST); }, diff --git a/server/comfy-routes.ts b/server/comfy-routes.ts index c8b0201..a50556b 100644 --- a/server/comfy-routes.ts +++ b/server/comfy-routes.ts @@ -1,16 +1,34 @@ -import { generate, listGenerationOptions, type ComfyGenerateRequest } from "./comfy"; +import { generate, listGenerationOptions, segment, type ComfyGenerateRequest, type ComfySegmentRequest } from "./comfy"; export async function handleComfyApi(request: Request) { try { const url = new URL(request.url); if (url.pathname === "/api/comfy/models" && request.method === "GET") return json(await listGenerationOptions()); - if (url.pathname === "/api/comfy/generate" && request.method === "POST") return json(await generate(await request.json() as ComfyGenerateRequest, request.signal)); + if (url.pathname === "/api/comfy/generate" && request.method === "POST") { + const body = await request.json() as ComfyGenerateRequest; + if (request.headers.get("accept")?.includes("application/x-ndjson")) return generationStream(body, request.signal); + return json(await generate(body, request.signal)); + } + if (url.pathname === "/api/comfy/segment" && request.method === "POST") return json(await segment(await request.json() as ComfySegmentRequest, request.signal)); return new Response("Not found", { status: 404 }); } catch (error) { return new Response(error instanceof Error ? error.message : "ComfyUI request failed", { status: 500 }); } } +function generationStream(body: ComfyGenerateRequest, signal: AbortSignal) { + const encoder = new TextEncoder(); + return new Response(new ReadableStream({ + start(controller) { + const send = (value: unknown) => controller.enqueue(encoder.encode(`${JSON.stringify(value)}\n`)); + void generate(body, signal, (event) => send({ type: "progress", ...event })) + .then((result) => { send({ type: "result", ...result }); controller.close(); }) + .catch((error) => { send({ type: "error", message: error instanceof Error ? error.message : "Generation failed" }); controller.close(); }); + }, + cancel() {}, + }), { headers: { "content-type": "application/x-ndjson; charset=utf-8", "cache-control": "no-store" } }); +} + function json(value: unknown) { return new Response(JSON.stringify(value), { headers: { "content-type": "application/json" } }); } diff --git a/server/comfy.test.ts b/server/comfy.test.ts index 50550d9..4fb652f 100644 --- a/server/comfy.test.ts +++ b/server/comfy.test.ts @@ -1,5 +1,5 @@ import { describe, expect, test } from "bun:test"; -import { buildAnimaWorkflow, buildSdxlWorkflow, buildZImageTurboWorkflow, buildZImageWorkflow, selectGeneratedOutputImage } from "./comfy"; +import { buildAnimaWorkflow, buildSdxlWorkflow, buildSemanticSelectionWorkflow, buildZImageTurboWorkflow, buildZImageWorkflow, selectGeneratedOutputImage, selectGeneratedOutputImages } from "./comfy"; import { handleComfyApi } from "./comfy-routes"; describe("Comfy adapter", () => { @@ -26,6 +26,13 @@ describe("Comfy adapter", () => { expect(image?.filename).toBe("image-studio-inpaint_00002_.png"); }); + test("retains every generated image from a batch", () => { + expect(selectGeneratedOutputImages({ outputs: { "8": { images: [ + { filename: "image-studio-inpaint_1.png", type: "output" }, + { filename: "image-studio-inpaint_2.png", type: "output" }, + ] } } })).toHaveLength(2); + }); + test("builds neutral inpaint with VAEEncodeForInpaint", () => { const workflow = buildSdxlWorkflow(inpaintRequest({ maskedContent: "neutral" })); @@ -45,6 +52,32 @@ describe("Comfy adapter", () => { expect(workflow["6"]?.inputs.latent_image).toEqual(["13", 0]); }); + test("repeats inpaint latents for batches and applies Canny ControlNet", () => { + const workflow = buildSdxlWorkflow({ + ...inpaintRequest({ maskedContent: "neutral" }), + batchSize: 4, + inpaint: { maskedContent: "neutral", structureControl: "canny", controlModel: "controlnet-canny.safetensors", controlStrength: 0.6 }, + }); + expect(workflow["19"]).toMatchObject({ class_type: "RepeatLatentBatch", inputs: { amount: 4 } }); + expect(workflow["20"]).toMatchObject({ class_type: "ControlNetLoader", inputs: { control_net_name: "controlnet-canny.safetensors" } }); + expect(workflow["21"]?.class_type).toBe("Canny"); + expect(workflow["22"]).toMatchObject({ class_type: "ControlNetApplyAdvanced", inputs: { strength: 0.6 } }); + expect(workflow["6"]?.inputs).toMatchObject({ latent_image: ["19", 0], positive: ["22", 0], negative: ["22", 1] }); + }); + + test("adds an optional low-denoise detail pass", () => { + const workflow = buildSdxlWorkflow({ ...inpaintRequest({ maskedContent: "neutral" }), refinePass: true, refineStrength: 18, seed: 40 }); + expect(workflow["30"]).toMatchObject({ class_type: "KSampler", inputs: { seed: 41, denoise: 0.18, latent_image: ["6", 0] } }); + expect(workflow["7"]?.inputs.samples).toEqual(["30", 0]); + }); + + test("builds native SAM3 point selection as a mask output", () => { + const workflow = buildSemanticSelectionWorkflow({ inputImage: "input.png", model: "sam3.safetensors", x: 24.4, y: 18.6 }); + expect(workflow["1"]).toMatchObject({ class_type: "UNETLoader", inputs: { unet_name: "sam3.safetensors" } }); + expect(workflow["3"]).toMatchObject({ class_type: "SAM3_Detect", inputs: { positive_coords: '[{"x":24,"y":19}]', refine_iterations: 2 } }); + expect(workflow["4"]).toMatchObject({ class_type: "MaskToImage", inputs: { mask: ["3", 0] } }); + }); + test("builds Z-Image text-to-image with separated model loaders", () => { const workflow = buildZImageWorkflow(textRequest({ architecture: "z-image", model: "z_image_bf16.safetensors", steps: 30, cfg: 4 })); diff --git a/server/comfy.ts b/server/comfy.ts index 4ff54b8..10ab83c 100644 --- a/server/comfy.ts +++ b/server/comfy.ts @@ -18,6 +18,9 @@ export type ComfyGenerateRequest = { scheduler?: string; width?: number; height?: number; + batchSize?: number; + refinePass?: boolean; + refineStrength?: number; outpaint?: { left?: number; top?: number; @@ -35,17 +38,30 @@ export type ComfyGenerateRequest = { maskedContent?: "neutral" | "original" | "originalColor" | "edges"; crop?: unknown; placement?: unknown; + structureControl?: "none" | "canny" | "depth" | "pose"; + controlStrength?: number; + controlModel?: string; }; inputImage?: string; maskImage?: string; }; +export type ComfyProgress = { progress: number; detail: string }; +export type ComfySegmentRequest = { inputImage: string; x: number; y: number; model?: string }; + type ComfyObjectInfo = { CheckpointLoaderSimple?: { input?: { required?: { ckpt_name?: [string[]] } } }; KSampler?: { input?: { required?: { sampler_name?: [string[]]; scheduler?: [string[]] } } }; UNETLoader?: { input?: { required?: { unet_name?: [string[]] } } }; CLIPLoader?: { input?: { required?: { clip_name?: [string[]] } } }; VAELoader?: { input?: { required?: { vae_name?: [string[]] } } }; + ControlNetLoader?: { input?: { required?: { control_net_name?: [string[]] } } }; + Canny?: unknown; + "MiDaS-DepthMapPreprocessor"?: unknown; + OpenposePreprocessor?: unknown; + ControlNetApplyAdvanced?: unknown; + SAM3_Detect?: unknown; + MaskToImage?: unknown; }; const comfyBaseUrl = process.env.COMFYUI_URL ?? "http://127.0.0.1:8188"; @@ -59,9 +75,7 @@ const defaultModels: Record = { }; export async function listGenerationOptions() { - const response = await fetch(`${comfyBaseUrl}/object_info`); - if (!response.ok) throw new Error(`ComfyUI option lookup failed: ${response.status}`); - const info = await response.json() as ComfyObjectInfo; + const info = await fetchObjectInfo(); const checkpointModels = info.CheckpointLoaderSimple?.input?.required?.ckpt_name?.[0] ?? []; const diffusionModels = info.UNETLoader?.input?.required?.unet_name?.[0] ?? []; const textEncoders = info.CLIPLoader?.input?.required?.clip_name?.[0] ?? []; @@ -69,11 +83,20 @@ export async function listGenerationOptions() { return { models: checkpointModels, + inpaintModels: checkpointModels.filter((model) => /inpaint|fill/i.test(model)), samplers: info.KSampler?.input?.required?.sampler_name?.[0] ?? [], schedulers: info.KSampler?.input?.required?.scheduler?.[0] ?? [], diffusionModels, textEncoders, vaes, + controlModels: info.ControlNetLoader?.input?.required?.control_net_name?.[0] ?? [], + structureControls: [ + ...(info.Canny && info.ControlNetApplyAdvanced && info.ControlNetLoader ? ["canny"] : []), + ...(info["MiDaS-DepthMapPreprocessor"] && info.ControlNetApplyAdvanced && info.ControlNetLoader ? ["depth"] : []), + ...(info.OpenposePreprocessor && info.ControlNetApplyAdvanced && info.ControlNetLoader ? ["pose"] : []), + ], + semanticSelection: Boolean(info.SAM3_Detect && info.MaskToImage && diffusionModels.some((model) => /sam.?3/i.test(model))), + sam3Models: diffusionModels.filter((model) => /sam.?3/i.test(model)), architectures: [ { value: "sdxl", @@ -107,21 +130,60 @@ export async function listGenerationOptions() { }; } +export async function segment(request: ComfySegmentRequest, signal?: AbortSignal) { + if (!request.inputImage) throw new Error("Semantic selection requires an image"); + if (!Number.isFinite(request.x) || !Number.isFinite(request.y)) throw new Error("Semantic selection requires a valid point"); + const info = await fetchObjectInfo(); + const models = info.UNETLoader?.input?.required?.unet_name?.[0] ?? []; + const model = request.model && request.model !== "auto" ? request.model : models.find((candidate) => /sam.?3/i.test(candidate)); + if (!info.SAM3_Detect || !info.MaskToImage || !model || !models.includes(model)) throw new Error("SAM3 semantic selection is not installed in ComfyUI. Install a SAM3 model and enable the native SAM3 nodes."); + const uploaded = await uploadDataUrl(request.inputImage, `image-studio-segment-${crypto.randomUUID()}.png`, signal); + const prompt = buildSemanticSelectionWorkflow({ ...request, model, inputImage: uploaded }); + const queued = await fetch(`${comfyBaseUrl}/prompt`, { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ client_id: crypto.randomUUID(), prompt }), signal }); + if (!queued.ok) throw new Error(`ComfyUI semantic selection failed: ${queued.status} ${await queued.text()}`); + const queuedBody = await queued.json() as { prompt_id?: string; node_errors?: unknown }; + const nodeError = nodeErrorsMessage(queuedBody.node_errors); + if (nodeError) throw new Error(`ComfyUI rejected semantic selection: ${nodeError}`); + if (!queuedBody.prompt_id) throw new Error("ComfyUI did not return a semantic selection prompt id"); + const history = await waitForHistory(queuedBody.prompt_id, signal); + const historyError = historyErrorMessage(history); + if (historyError) throw new Error(`ComfyUI semantic selection failed: ${historyError}`); + const image = selectGeneratedOutputImage(history); + if (!image) throw new Error("ComfyUI did not return a semantic selection mask"); + const response = await fetch(`${comfyBaseUrl}/view?${new URLSearchParams({ filename: image.filename, subfolder: image.subfolder ?? "", type: image.type ?? "output" })}`, { signal }); + if (!response.ok) throw new Error(`ComfyUI mask fetch failed: ${response.status}`); + const bytes = Buffer.from(await response.arrayBuffer()); + return { source: `data:image/png;base64,${bytes.toString("base64")}`, mimeType: "image/png" }; +} + +export function buildSemanticSelectionWorkflow(request: ComfySegmentRequest & { model: string }): Workflow { + return { + "1": { class_type: "UNETLoader", inputs: { unet_name: request.model, weight_dtype: "default" } }, + "2": { class_type: "LoadImage", inputs: { image: request.inputImage } }, + "3": { class_type: "SAM3_Detect", inputs: { model: ["1", 0], image: ["2", 0], positive_coords: JSON.stringify([{ x: Math.round(request.x), y: Math.round(request.y) }]), threshold: 0.5, refine_iterations: 2, individual_masks: false } }, + "4": { class_type: "MaskToImage", inputs: { mask: ["3", 0] } }, + "8": { class_type: "SaveImage", inputs: { filename_prefix: "image-studio-segment", images: ["4", 0] } }, + }; +} + async function listCheckpointModels() { return (await listGenerationOptions()).models; } -export async function generate(request: ComfyGenerateRequest, signal?: AbortSignal) { +export async function generate(request: ComfyGenerateRequest, signal?: AbortSignal, onProgress?: (event: ComfyProgress) => void) { if (!request.prompt?.trim()) throw new Error("Prompt is required"); + onProgress?.({ progress: 0.03, detail: "Preparing workflow" }); const architecture = normalizeArchitecture(request.architecture); if (request.mode !== "text-to-image" && architecture !== "sdxl") throw new Error(`${architectureLabel(architecture)} currently supports text-to-image only`); - if (!request.model || request.model === "auto") request.model = await defaultModelForArchitecture(architecture); + if (!request.model || request.model === "auto") request.model = await defaultModelForArchitecture(architecture, request.mode); const clientId = crypto.randomUUID(); const uploaded = request.inputImage ? await uploadDataUrl(request.inputImage, `image-studio-${crypto.randomUUID()}.png`, signal) : undefined; const mask = request.maskImage ? await uploadDataUrl(request.maskImage, `image-studio-mask-${crypto.randomUUID()}.png`, signal) : undefined; + onProgress?.({ progress: 0.16, detail: "Inputs uploaded" }); if (request.mode === "inpaint" && (!uploaded || !mask)) throw new Error("Inpaint requires normalized input and mask images"); - const prompt = buildComfyWorkflow({ ...request, architecture, inputImage: uploaded, maskImage: mask }); + const resolvedRequest = await resolveStructureControl({ ...request, architecture, inputImage: uploaded, maskImage: mask }); + const prompt = buildComfyWorkflow(resolvedRequest); const queued = await fetch(`${comfyBaseUrl}/prompt`, { method: "POST", @@ -134,23 +196,28 @@ export async function generate(request: ComfyGenerateRequest, signal?: AbortSign const nodeError = nodeErrorsMessage(queuedBody.node_errors); if (nodeError) throw new Error(`ComfyUI rejected the workflow: ${nodeError}`); if (!queuedBody.prompt_id) throw new Error("ComfyUI did not return a prompt id"); + onProgress?.({ progress: 0.24, detail: "Queued in ComfyUI" }); const prompt_id = queuedBody.prompt_id; let history: unknown; try { - history = await waitForHistory(prompt_id, signal); + history = await waitForHistory(prompt_id, signal, onProgress); } catch (error) { if (signal?.aborted) await cancelComfyPrompt(prompt_id); throw error; } const historyError = historyErrorMessage(history); if (historyError) throw new Error(`ComfyUI generation failed: ${historyError}`); - const image = selectGeneratedOutputImage(history); - if (!image) throw new Error("ComfyUI did not return an image"); - - const imageResponse = await fetch(`${comfyBaseUrl}/view?${new URLSearchParams({ filename: image.filename, subfolder: image.subfolder ?? "", type: image.type ?? "output" })}`, { signal }); - if (!imageResponse.ok) throw new Error(`ComfyUI image fetch failed: ${imageResponse.status}`); - const bytes = Buffer.from(await imageResponse.arrayBuffer()); - return { source: `data:image/png;base64,${bytes.toString("base64")}`, mimeType: "image/png" }; + const images = selectGeneratedOutputImages(history); + if (images.length === 0) throw new Error("ComfyUI did not return an image"); + onProgress?.({ progress: 0.9, detail: "Downloading results" }); + const results = await Promise.all(images.map(async (image, index) => { + const imageResponse = await fetch(`${comfyBaseUrl}/view?${new URLSearchParams({ filename: image.filename, subfolder: image.subfolder ?? "", type: image.type ?? "output" })}`, { signal }); + if (!imageResponse.ok) throw new Error(`ComfyUI image fetch failed: ${imageResponse.status}`); + const bytes = Buffer.from(await imageResponse.arrayBuffer()); + return { source: `data:image/png;base64,${bytes.toString("base64")}`, mimeType: "image/png", seed: Math.round((resolvedRequest.seed ?? 0) + index) }; + })); + onProgress?.({ progress: 1, detail: "Results ready" }); + return { results }; } async function uploadDataUrl(dataUrl: string, filename: string, signal?: AbortSignal) { @@ -167,7 +234,7 @@ async function uploadDataUrl(dataUrl: string, filename: string, signal?: AbortSi return uploaded.name; } -async function waitForHistory(promptId: string, signal?: AbortSignal) { +async function waitForHistory(promptId: string, signal?: AbortSignal, onProgress?: (event: ComfyProgress) => void) { const startedAt = Date.now(); let attempts = 0; while (Date.now() - startedAt < comfyHistoryTimeoutMs) { @@ -177,6 +244,8 @@ async function waitForHistory(promptId: string, signal?: AbortSignal) { } const response = await fetch(`${comfyBaseUrl}/history/${promptId}`, { signal }); attempts += 1; + const elapsedRatio = Math.min(1, (Date.now() - startedAt) / comfyHistoryTimeoutMs); + onProgress?.({ progress: 0.25 + elapsedRatio * 0.6, detail: attempts <= 1 ? "Generating" : `Generating · check ${attempts}` }); if (response.ok) { const history = await response.json() as Record; if (history[promptId]) return history[promptId]; @@ -209,14 +278,18 @@ function abortableSleep(ms: number, signal?: AbortSignal): Promise { } export function selectGeneratedOutputImage(history: unknown): { filename: string; subfolder?: string; type?: string } | undefined { + return selectGeneratedOutputImages(history)[0]; +} + +export function selectGeneratedOutputImages(history: unknown): Array<{ filename: string; subfolder?: string; type?: string }> { const outputs = (history as { outputs?: Record }).outputs ?? {}; - const saveImageOutput = outputs["8"]?.images?.find(isGeneratedImage); - if (saveImageOutput) return saveImageOutput; + const saveImageOutput = outputs["8"]?.images?.filter(isGeneratedImage) ?? []; + if (saveImageOutput.length > 0) return saveImageOutput; - const prefixedOutput = Object.values(outputs).flatMap((output) => output.images ?? []).find((image) => image.filename.startsWith("image-studio-") && image.type !== "input"); - if (prefixedOutput) return prefixedOutput; + const prefixedOutput = Object.values(outputs).flatMap((output) => output.images ?? []).filter((image) => image.filename.startsWith("image-studio-") && image.type !== "input"); + if (prefixedOutput.length > 0) return prefixedOutput; - return Object.values(outputs).flatMap((output) => output.images ?? []).find(isGeneratedImage); + return Object.values(outputs).flatMap((output) => output.images ?? []).filter(isGeneratedImage); } function isGeneratedImage(image: { filename: string; subfolder?: string; type?: string }) { @@ -260,8 +333,8 @@ export function buildSdxlWorkflow(request: ComfyGenerateRequest): Workflow { }; if (request.mode === "text-to-image" || !request.inputImage) { - workflow["5"] = { class_type: "EmptyLatentImage", inputs: { width, height, batch_size: 1 } }; - return workflow; + workflow["5"] = { class_type: "EmptyLatentImage", inputs: { width, height, batch_size: resolveBatchSize(request) } }; + return finalizeSdxlWorkflow(workflow, samplerInputs, request); } workflow["4"] = { class_type: "LoadImage", inputs: { image: request.inputImage } }; @@ -277,19 +350,96 @@ export function buildSdxlWorkflow(request: ComfyGenerateRequest): Workflow { } else { workflow["5"] = { class_type: "VAEEncodeForInpaint", inputs: { pixels: ["4", 0], vae: ["1", 2], mask: ["11", 0], grow_mask_by: resolveGrowMaskBy(request) } }; } - return workflow; + return finalizeSdxlWorkflow(workflow, samplerInputs, request); } if (request.mode === "outpaint") { workflow["10"] = { class_type: "ImagePadForOutpaint", inputs: { image: ["4", 0], left: Math.round(request.outpaint?.left ?? 0), top: Math.round(request.outpaint?.top ?? 0), right: Math.round(request.outpaint?.right ?? 0), bottom: Math.round(request.outpaint?.bottom ?? 0), feathering: Math.round(request.outpaint?.feathering ?? 0) } }; workflow["5"] = { class_type: "VAEEncodeForInpaint", inputs: { pixels: ["10", 0], vae: ["1", 2], mask: ["10", 1], grow_mask_by: resolveGrowMaskBy(request) } }; - return workflow; + return finalizeSdxlWorkflow(workflow, samplerInputs, request); } workflow["5"] = { class_type: "VAEEncode", inputs: { pixels: ["4", 0], vae: ["1", 2] } }; + return finalizeSdxlWorkflow(workflow, samplerInputs, request); +} + +function finalizeSdxlWorkflow(workflow: Workflow, samplerInputs: Record, request: ComfyGenerateRequest): Workflow { + const batchSize = resolveBatchSize(request); + if (batchSize > 1 && workflow["5"]?.class_type !== "EmptyLatentImage") { + const latent = samplerInputs.latent_image; + workflow["19"] = { class_type: "RepeatLatentBatch", inputs: { samples: latent, amount: batchSize } }; + samplerInputs.latent_image = ["19", 0]; + } + + const control = request.inpaint?.structureControl ?? "none"; + if (request.mode !== "inpaint" || control === "none" || !request.inpaint?.controlModel || request.inpaint.controlModel === "auto") return addRefinementPass(workflow, samplerInputs, request); + workflow["20"] = { class_type: "ControlNetLoader", inputs: { control_net_name: request.inpaint.controlModel } }; + if (control === "canny") { + workflow["21"] = { class_type: "Canny", inputs: { image: ["4", 0], low_threshold: 100, high_threshold: 200 } }; + } else if (control === "depth") { + workflow["21"] = { class_type: "MiDaS-DepthMapPreprocessor", inputs: { image: ["4", 0], a: 6.283, bg_threshold: 0.1, resolution: Math.max(request.width ?? 512, request.height ?? 512) } }; + } else { + workflow["21"] = { class_type: "OpenposePreprocessor", inputs: { image: ["4", 0], detect_hand: "enable", detect_body: "enable", detect_face: "enable", resolution: Math.max(request.width ?? 512, request.height ?? 512) } }; + } + workflow["22"] = { + class_type: "ControlNetApplyAdvanced", + inputs: { + positive: ["2", 0], + negative: ["3", 0], + control_net: ["20", 0], + image: ["21", 0], + strength: Math.max(0, Math.min(1, request.inpaint.controlStrength ?? 0.55)), + start_percent: 0, + end_percent: 0.85, + vae: ["1", 2], + }, + }; + samplerInputs.positive = ["22", 0]; + samplerInputs.negative = ["22", 1]; + return addRefinementPass(workflow, samplerInputs, request); +} + +function addRefinementPass(workflow: Workflow, samplerInputs: Record, request: ComfyGenerateRequest): Workflow { + if (!request.refinePass) return workflow; + workflow["30"] = { + class_type: "KSampler", + inputs: { + ...samplerInputs, + seed: Math.round((request.seed ?? 0) + 1), + steps: Math.max(6, Math.round((request.steps ?? 30) / 3)), + denoise: Math.max(0, Math.min(1, (request.refineStrength ?? 20) / 100)), + latent_image: ["6", 0], + }, + }; + if (workflow["7"]) workflow["7"].inputs.samples = ["30", 0]; return workflow; } +function resolveBatchSize(request: ComfyGenerateRequest) { + return Math.round(Math.max(1, Math.min(8, request.batchSize ?? 1))); +} + +async function fetchObjectInfo(): Promise { + const response = await fetch(`${comfyBaseUrl}/object_info`); + if (!response.ok) throw new Error(`ComfyUI option lookup failed: ${response.status}`); + return response.json() as Promise; +} + +async function resolveStructureControl(request: ComfyGenerateRequest): Promise { + const control = request.inpaint?.structureControl ?? "none"; + if (request.mode !== "inpaint" || control === "none") return request; + const info = await fetchObjectInfo(); + const requiredPreprocessor = control === "canny" ? info.Canny : control === "depth" ? info["MiDaS-DepthMapPreprocessor"] : info.OpenposePreprocessor; + if (!requiredPreprocessor || !info.ControlNetLoader || !info.ControlNetApplyAdvanced) { + throw new Error(`${control === "canny" ? "Canny" : control === "depth" ? "Depth" : "Pose"} structural control is not installed in ComfyUI.`); + } + const models = info.ControlNetLoader.input?.required?.control_net_name?.[0] ?? []; + const requested = request.inpaint?.controlModel; + const model = requested && requested !== "auto" ? requested : models.find((candidate) => candidate.toLowerCase().includes(control)); + if (!model || !models.includes(model)) throw new Error(`Install or select a ${control} ControlNet model before using structural control.`); + return { ...request, inpaint: { ...request.inpaint, controlModel: model } }; +} + export function buildZImageWorkflow(request: ComfyGenerateRequest): Workflow { return buildSeparatedTextToImageWorkflow(request, { architecture: "z-image", @@ -408,9 +558,10 @@ function usesOriginalLatentContent(request: ComfyGenerateRequest): boolean { return request.inpaint?.maskedContent === "original" || request.inpaint?.maskedContent === "originalColor" || request.inpaint?.maskedContent === "edges"; } -async function defaultModelForArchitecture(architecture: GenerateArchitecture) { +async function defaultModelForArchitecture(architecture: GenerateArchitecture, mode: GenerateMode) { if (architecture !== "sdxl") return defaultModels[architecture]; const models = await listCheckpointModels(); + if (mode === "inpaint") return models.find((model) => /inpaint|fill/i.test(model)) ?? models[0] ?? defaultModels.sdxl; return models[0] ?? defaultModels.sdxl; } diff --git a/view/App.tsx b/view/App.tsx index 96337f9..2de965a 100644 --- a/view/App.tsx +++ b/view/App.tsx @@ -194,12 +194,16 @@ export function App({ app }: AppProps) { activeTool={tools.activeTool} interactionMode={tools.interactionMode} panel={workspace.panel} + inpaintMaskEditing={maskEdit?.kind === "inpaintRegion"} dispatch={app.store.dispatch} /> @@ -216,7 +220,7 @@ export function App({ app }: AppProps) { document={document} selection={selection} viewport={viewport} - visible={generateOpen || chromaKeyOpen || tools.activeTool === "brush" || tools.activeTool === "eraser" || tools.activeTool === "magicWand" || Boolean(transformBounds) || viewportActivityIsland.visible} + visible={generateOpen || chromaKeyOpen || tools.activeTool === "brush" || tools.activeTool === "eraser" || tools.activeTool === "magicWand" || tools.activeTool === "semanticSelect" || tools.activeTool === "maskLasso" || tools.activeTool === "maskRectangle" || Boolean(transformBounds) || viewportActivityIsland.visible} action={viewportActivityIsland.action} activeTool={tools.activeTool} operation={generateOpen ? "generate" : chromaKeyOpen ? "chromaKey" : undefined} @@ -226,6 +230,7 @@ export function App({ app }: AppProps) { chromaKeySettings={tools.chromaKey} magicWandSettings={tools.magicWand} editingMask={Boolean(maskEdit)} + maskKind={maskEdit?.kind} maskViewMode={maskEdit?.viewMode ?? "composite"} brushHint={brushHint} transformBounds={viewportActivityIsland.visible ? undefined : transformBounds} diff --git a/view/BottomControlsIsland.tsx b/view/BottomControlsIsland.tsx index 6dc3257..08de1a1 100644 --- a/view/BottomControlsIsland.tsx +++ b/view/BottomControlsIsland.tsx @@ -31,6 +31,7 @@ export type BottomControlsIslandProps = { chromaKeySettings: ChromaKeySettings; magicWandSettings: MagicWandSettings; editingMask?: boolean; + maskKind?: "layerMask" | "inpaintRegion"; maskViewMode?: MaskViewMode; transformBounds?: Rect; transformTarget?: TransformTarget; @@ -40,7 +41,7 @@ export type BottomControlsIslandProps = { documentActions: DocumentActions; }; -export function BottomControlsIsland({ document, selection, viewport, visible, action, activeTool, operation, brushSettings, generateSettings, generation, chromaKeySettings, magicWandSettings, editingMask = false, maskViewMode = "composite", transformBounds, transformTarget, brushHint, dispatch, generationWorkflow, documentActions }: BottomControlsIslandProps) { +export function BottomControlsIsland({ document, selection, viewport, visible, action, activeTool, operation, brushSettings, generateSettings, generation, chromaKeySettings, magicWandSettings, editingMask = false, maskKind, maskViewMode = "composite", transformBounds, transformTarget, brushHint, dispatch, generationWorkflow, documentActions }: BottomControlsIslandProps) { const documentIndex = useMemo(() => createDocumentReadIndex(document), [document]); const selectedLayerInfo = selection.layerIds.length === 1 && selection.layerIds[0] ? documentIndex.layerInfoById.get(selection.layerIds[0]) : undefined; const zoomPercent = Math.round(viewport.zoom * 100); @@ -61,9 +62,13 @@ export function BottomControlsIsland({ document, selection, viewport, visible, a ) : (activeTool === "brush" || activeTool === "eraser") && brushHint ? ( ) : activeTool === "brush" || activeTool === "eraser" ? ( - + ) : activeTool === "magicWand" ? ( + ) : activeTool === "semanticSelect" ? ( +
AI object selectClick an object · Shift-click adds · Option-click protects
+ ) : activeTool === "maskLasso" || activeTool === "maskRectangle" ? ( +
AI region {activeTool === "maskRectangle" ? "rectangle" : "lasso"}Drag to replace · Shift-drag adds · Option-drag protects
) : transformBounds && transformTarget ? ( ) : action === "pan" ? ( diff --git a/view/CanvasViewport.tsx b/view/CanvasViewport.tsx index 45e487d..770c9f5 100644 --- a/view/CanvasViewport.tsx +++ b/view/CanvasViewport.tsx @@ -105,5 +105,5 @@ function interactionModesEqual(a: InteractionMode, b: InteractionMode): boolean function maskEditStatesEqual(a: MaskEditState | undefined, b: MaskEditState | undefined): boolean { if (a === b) return true; if (!a || !b) return false; - return a.targetLayerId === b.targetLayerId && a.maskLayerId === b.maskLayerId && a.viewMode === b.viewMode; + return a.kind === b.kind && a.targetLayerId === b.targetLayerId && a.maskAssetId === b.maskAssetId && a.maskLayerId === b.maskLayerId && a.inpaintRegionId === b.inpaintRegionId && a.viewMode === b.viewMode; } diff --git a/view/GenerateSheet.tsx b/view/GenerateSheet.tsx index 0a69e9f..728bcdf 100644 --- a/view/GenerateSheet.tsx +++ b/view/GenerateSheet.tsx @@ -1,16 +1,22 @@ import type { GenerateSettings } from "@editor/tools"; import type { AppStore } from "@editor/store"; -import type { GenerationResourcesState } from "@editor/state"; +import type { GenerationResourcesState, GenerationState } from "@editor/state"; import { GenerateControls } from "./bottom-controls/GenerateControls"; +import type { ImageDocument } from "@core/document"; +import type { SelectionState } from "@editor/state"; +import { CandidateReviewPanel } from "./inpaint/CandidateReviewPanel"; export type GenerateSheetProps = { settings: GenerateSettings; + document: ImageDocument; + selection: SelectionState; resources: GenerationResourcesState; + generation: GenerationState; open: boolean; dispatch: AppStore["dispatch"]; }; -export function GenerateSheet({ settings, resources, open, dispatch }: GenerateSheetProps) { +export function GenerateSheet({ settings, document, selection, resources, generation, open, dispatch }: GenerateSheetProps) { return ( diff --git a/view/GenerationJobStatus.tsx b/view/GenerationJobStatus.tsx index 5fc1476..30ee988 100644 --- a/view/GenerationJobStatus.tsx +++ b/view/GenerationJobStatus.tsx @@ -17,7 +17,7 @@ export function GenerationJobStatus({ generation, compact = false }: { generatio if (!job) return null; const elapsed = Math.max(0, Math.floor(((job.finishedAt ?? Date.now()) - job.startedAt) / 1000)); - const label = job.status === "running" ? `${job.label} ${formatElapsed(elapsed)}` : job.status === "failed" ? job.error ?? `${job.label} failed` : job.status === "cancelled" ? `${job.label} cancelled` : `${job.label} complete`; + const label = job.status === "running" ? `${job.detail ?? job.label}${job.progress !== undefined ? ` · ${Math.round(job.progress * 100)}%` : ""} ${formatElapsed(elapsed)}` : job.status === "failed" ? job.error ?? `${job.label} failed` : job.status === "cancelled" ? `${job.label} cancelled` : `${job.label} complete`; const tone = job.status === "failed" ? "bg-red-500/15 text-red-100" : job.status === "running" ? "bg-white/10 text-white/70" : job.status === "cancelled" ? "bg-amber-500/15 text-amber-100" : "bg-emerald-500/15 text-emerald-100"; return {label}; diff --git a/view/LayersSheet.tsx b/view/LayersSheet.tsx index c9fbaa1..a4764eb 100644 --- a/view/LayersSheet.tsx +++ b/view/LayersSheet.tsx @@ -346,8 +346,11 @@ function LayerRow({ const layerMask = getLayerMask(layer); const maskLayer = layerMask ? documentIndex.layerById.get(layerMask.maskLayerId) : undefined; const maskAsset = maskLayer && (maskLayer.type === "image" || maskLayer.type === "raster") ? documentIndex.assetById.get(maskLayer.assetId) : undefined; + const inpaintRegion = document.inpaintRegions.find((region) => region.targetLayerId === layer.id && region.enabled); + const inpaintMaskAsset = inpaintRegion ? documentIndex.assetById.get(inpaintRegion.maskAssetId) : undefined; const canAddMask = Boolean(layerInfo && (layer.type === "image" || layer.type === "raster") && !layerMask); const editingMask = Boolean(maskEdit && layerMask && maskEdit.targetLayerId === layer.id && maskEdit.maskLayerId === layerMask.maskLayerId); + const editingInpaintRegion = Boolean(maskEdit?.kind === "inpaintRegion" && inpaintRegion && maskEdit.inpaintRegionId === inpaintRegion.id); const thumbnail = thumbnailByLayerId.get(layer.id) ?? { kind: "empty" }; const maskThumbnail = maskLayer ? thumbnailByLayerId.get(maskLayer.id) : undefined; const rowPadding = 12 + depth * 16; @@ -479,6 +482,16 @@ function LayerRow({ ) : null} + {inpaintRegion ? ( +
+
+ ) : null} {layer.type === "group" ? layer.children.map((child) => ( - {availableToolIds.map((tool) => { + {availableToolIds.filter((tool) => inpaintMaskEditing || (tool !== "semanticSelect" && tool !== "maskLasso" && tool !== "maskRectangle")).map((tool) => { const active = !isOperationWorkspacePanel(panel) && isToolHighlighted(tool, activeTool, interactionMode); const Icon = iconForTool(tool); @@ -67,6 +68,12 @@ function iconForTool(tool: ToolId) { return Eraser; case "magicWand": return MagicWand; + case "semanticSelect": + return Selection; + case "maskLasso": + return Polygon; + case "maskRectangle": + return Rectangle; case "pan": return Hand; case "select": diff --git a/view/bottom-controls/BrushControls.tsx b/view/bottom-controls/BrushControls.tsx index 69b8605..bbf4477 100644 --- a/view/bottom-controls/BrushControls.tsx +++ b/view/bottom-controls/BrushControls.tsx @@ -13,6 +13,7 @@ export type BrushControlsProps = { tool: Extract; settings: BrushSettings; editingMask?: boolean; + maskKind?: "layerMask" | "inpaintRegion"; maskViewMode?: MaskViewMode; dispatch: AppStore["dispatch"]; }; @@ -24,7 +25,8 @@ const maskViewModeOptions = [ { value: "overlay", label: "Overlay" }, ] satisfies readonly BottomControlSelectOption[]; -export function BrushControls({ tool, settings, editingMask = false, maskViewMode = "composite", dispatch }: BrushControlsProps) { +export function BrushControls({ tool, settings, editingMask = false, maskKind = "layerMask", maskViewMode = "composite", dispatch }: BrushControlsProps) { + const inpaintRegion = editingMask && maskKind === "inpaintRegion"; return (
@@ -33,8 +35,8 @@ export function BrushControls({ tool, settings, editingMask = false, maskViewMod {editingMask ? ( <> - - + + ) : null} @@ -61,6 +63,13 @@ export function BrushControls({ tool, settings, editingMask = false, maskViewMod {Math.round(settings.size)} + + + + + + + rerun("Regenerate", candidate.settings)} /> + void workflow.rebuildFromCurrentRegion(candidate.id)} /> { 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 }); + dispatch(commandIds.toolEnterInpaintRegionEdit, { targetLayerId: candidate.inpaint.targetLayerId, regionId: candidate.inpaint.regionId }); }} /> dispatch(commandIds.generationRemoveCandidate, { candidateId: candidate.id })} /> diff --git a/view/bottom-controls/GenerateControls.tsx b/view/bottom-controls/GenerateControls.tsx index d99dacd..16e7c39 100644 --- a/view/bottom-controls/GenerateControls.tsx +++ b/view/bottom-controls/GenerateControls.tsx @@ -7,6 +7,9 @@ import type { GenerateArchitecture, GenerateIntent, GenerateMode, GenerateSettin import { resolveGenerationModeOptions, resolveGenerationModelOptions, resolveGenerationStringOptions, resolveGenerationSupportOptions } from "@operations/generation/options"; import { BottomControlSelectMenu, type BottomControlSelectOption } from "./SelectMenu"; import { BottomControlSlider } from "./Slider"; +import type { ImageDocument } from "@core/document"; +import type { SelectionState } from "@editor/state"; +import { InpaintRegionPanel } from "../inpaint/InpaintRegionPanel"; const architectures = [ { value: "sdxl", label: "SDXL" }, @@ -30,11 +33,6 @@ const sizePresets = [ { label: "9:16", w: 768, h: 1344 }, ] as const; -const inpaintPolarityOptions = [ - { value: "hidden", label: "Hidden / erased" }, - { value: "revealed", label: "Revealed / painted" }, -] satisfies readonly BottomControlSelectOption[]; - const inpaintMaskedContentOptions = [ { value: "neutral", label: "Neutral fill" }, { value: "original", label: "Original gray" }, @@ -42,13 +40,31 @@ const inpaintMaskedContentOptions = [ { value: "edges", label: "Edge map" }, ] satisfies readonly BottomControlSelectOption[]; +const inpaintProfileOptions = [ + { value: "remove", label: "Remove object" }, + { value: "replace", label: "Replace object" }, + { value: "repair", label: "Repair detail" }, + { value: "material", label: "Change material" }, + { value: "reshape", label: "Change shape" }, + { value: "custom", label: "Custom" }, +] satisfies readonly BottomControlSelectOption[]; + +const structureControlOptions = [ + { value: "none", label: "None" }, + { value: "canny", label: "Canny edges" }, + { value: "depth", label: "Depth" }, + { value: "pose", label: "Pose" }, +] satisfies readonly BottomControlSelectOption[]; + export type GenerateControlsProps = { settings: GenerateSettings; + document: ImageDocument; + selection: SelectionState; resources: GenerationResourcesState; dispatch: AppStore["dispatch"]; }; -export function GenerateControls({ settings, resources, dispatch }: GenerateControlsProps) { +export function GenerateControls({ settings, document, selection, resources, dispatch }: GenerateControlsProps) { const comfyOptions = resources.options; const [advancedOpen, setAdvancedOpen] = useState(false); const [outpaintOpen, setOutpaintOpen] = useState(false); @@ -60,6 +76,7 @@ export function GenerateControls({ settings, resources, dispatch }: GenerateCont const samplerOptions = resolveGenerationStringOptions(comfyOptions?.samplers, settings.sampler); const schedulerOptions = resolveGenerationStringOptions(comfyOptions?.schedulers, settings.scheduler); const modeOptions = resolveGenerationModeOptions(settings, comfyOptions, modes); + const availableStructureControls = structureControlOptions.filter((option) => option.value === "none" || option.value === settings.inpaint.structureControl || comfyOptions?.structureControls?.includes(option.value)); useEffect(() => { if (!sizeOpen) return; @@ -84,18 +101,14 @@ export function GenerateControls({ settings, resources, dispatch }: GenerateCont ))} -
@@ -134,6 +147,9 @@ export function GenerateControls({ settings, resources, dispatch }: GenerateCont {Math.round(settings.strength)} ) : null} + dispatch(commandIds.toolSetGenerateSettings, { batchSize })} /> + + {settings.refinePass ? dispatch(commandIds.toolSetGenerateSettings, { refineStrength })} /> : null}
@@ -184,6 +200,7 @@ export function GenerateControls({ settings, resources, dispatch }: GenerateCont
: null} {settings.mode === "inpaint" ?
+
- dispatch(commandIds.toolSetGenerateSettings, { inpaint: { ...settings.inpaint, maskPolarity } })} - /> + dispatch(commandIds.toolSetGenerateSettings, { inpaint: { profile } })} /> dispatch(commandIds.toolSetGenerateSettings, { inpaint: { ...settings.inpaint, maskedContent } })} /> + dispatch(commandIds.toolSetGenerateSettings, { inpaint: { structureControl, profile: "custom" } })} /> + {settings.inpaint.structureControl !== "none" ? ( +
+ dispatch(commandIds.toolSetGenerateSettings, { inpaint: { controlStrength: controlStrength / 100, profile: "custom" } })} /> + dispatch(commandIds.toolSetGenerateSettings, { inpaint: { controlModel } })} /> +
+ ) : null} +
dispatch(commandIds.toolSetGenerateSettings, { inpaint: { ...settings.inpaint, cropPadding } })} /> dispatch(commandIds.toolSetGenerateSettings, { inpaint: { ...settings.inpaint, growMaskBy } })} /> @@ -227,10 +248,18 @@ export function GenerateControls({ settings, resources, dispatch }: GenerateCont const intentOptions: ReadonlyArray<{ value: GenerateIntent; mode: GenerateMode; label: string; description: string }> = [ { value: "create", mode: "text-to-image", label: "Create", description: "Make a new image from your prompt." }, { value: "replace", mode: "inpaint", label: "Replace", description: "Regenerate the masked part of one layer." }, + { value: "remove", mode: "inpaint", label: "Remove", description: "Erase an object and reconstruct its background." }, { value: "extend", mode: "outpaint", label: "Extend", description: "Grow one selected image beyond its edges." }, { value: "variations", mode: "image-to-image", label: "Variations", description: "Explore alternatives based on one image." }, ]; +function isIntentActive(intent: GenerateIntent, mode: GenerateMode, settings: GenerateSettings) { + if (settings.mode !== mode) return false; + if (intent === "remove") return settings.inpaint.profile === "remove"; + if (intent === "replace") return settings.inpaint.profile !== "remove"; + return true; +} + function SectionTitle({ title }: { title: string }) { return
{title}
; } diff --git a/view/canvas/cursor.ts b/view/canvas/cursor.ts index c6cf5f1..1071515 100644 --- a/view/canvas/cursor.ts +++ b/view/canvas/cursor.ts @@ -8,6 +8,6 @@ export function canvasCursorClass(interactionMode: InteractionMode, isPanning: b if (!canBrush) return "cursor-not-allowed"; return hasBrushPreview ? "cursor-none" : "cursor-crosshair"; } - if (interactionMode.type === "tool" && interactionMode.tool === "magicWand") return "cursor-crosshair"; + if (interactionMode.type === "tool" && (interactionMode.tool === "magicWand" || interactionMode.tool === "semanticSelect" || interactionMode.tool === "maskLasso" || interactionMode.tool === "maskRectangle")) return "cursor-crosshair"; return "cursor-default"; } diff --git a/view/canvas/renderFrame.test.ts b/view/canvas/renderFrame.test.ts index 012e48f..e0479d9 100644 --- a/view/canvas/renderFrame.test.ts +++ b/view/canvas/renderFrame.test.ts @@ -137,7 +137,7 @@ const visualEditorChanges: Array<[string, (state: AppState) => AppState]> = [ ...state, editor: { ...state.editor, - maskEdit: { targetLayerId: "target", maskLayerId: "mask", viewMode: "overlay" }, + maskEdit: { kind: "layerMask", targetLayerId: "target", maskLayerId: "mask", maskAssetId: "mask-asset", viewMode: "overlay" }, }, }), ], diff --git a/view/canvas/renderFrame.ts b/view/canvas/renderFrame.ts index 22a8743..4970207 100644 --- a/view/canvas/renderFrame.ts +++ b/view/canvas/renderFrame.ts @@ -23,6 +23,7 @@ function visualEditorStatesEqual(a: EditorState, b: EditorState): boolean { maskEditStatesEqual(a.maskEdit, b.maskEdit) && brushPreviewStatesEqual(a.brushPreview, b.brushPreview) && brushStrokePreviewStatesEqual(a.brushStrokePreview, b.brushStrokePreview) && + a.maskShapeSession === b.maskShapeSession && generationStatesEqual(a.generation, b.generation) && visualToolStatesEqual(a.tools, b.tools) ); @@ -50,7 +51,7 @@ function transformTargetsEqual(a: TransformTarget, b: TransformTarget): boolean function maskEditStatesEqual(a: MaskEditState | undefined, b: MaskEditState | undefined): boolean { if (a === b) return true; if (!a || !b) return false; - return a.targetLayerId === b.targetLayerId && a.maskLayerId === b.maskLayerId && a.viewMode === b.viewMode; + return a.kind === b.kind && a.targetLayerId === b.targetLayerId && a.maskAssetId === b.maskAssetId && a.maskLayerId === b.maskLayerId && a.inpaintRegionId === b.inpaintRegionId && a.viewMode === b.viewMode; } function brushPreviewStatesEqual(a: BrushPreviewState | undefined, b: BrushPreviewState | undefined): boolean { @@ -80,7 +81,7 @@ function interactionModesEqual(a: InteractionMode, b: InteractionMode): boolean } function brushSettingsEqual(a: BrushSettings, b: BrushSettings): boolean { - return a.color === b.color && a.size === b.size && a.hardness === b.hardness; + return a.color === b.color && a.size === b.size && a.hardness === b.hardness && a.opacity === b.opacity && a.flow === b.flow && a.smoothing === b.smoothing && a.pressureSize === b.pressureSize; } function vec2Equal(a: Vec2D, b: Vec2D): boolean { diff --git a/view/canvas/useCanvasInput.ts b/view/canvas/useCanvasInput.ts index d9cebb9..70d8ae1 100644 --- a/view/canvas/useCanvasInput.ts +++ b/view/canvas/useCanvasInput.ts @@ -15,6 +15,8 @@ import { } from "@input/index"; import { beginBrushSession, canPreviewBrush, cancelBrushSession, commitBrushSession, updateBrushSession, type BrushSession } from "@operations/paint/brush"; import { applyMagicWandAt } from "@operations/masks/magic-wand"; +import { commitInpaintLasso } from "@operations/masks/lasso"; +import { applySemanticSelectionAt } from "@operations/masks/semantic-select"; export type CanvasInputOptions = { globalKeybindConsumer: GlobalKeybindConsumer; @@ -122,6 +124,12 @@ export function useCanvasInput( const state = store.getState(); const documentPoint = viewportPointToDocumentPoint(inputEvent.position, state.editor.viewport); + if (!isOperationWorkspacePanel(state.editor.workspace.panel) && (state.editor.tools.activeTool === "maskLasso" || state.editor.tools.activeTool === "maskRectangle") && state.editor.maskEdit?.kind === "inpaintRegion" && (inputEvent.buttons & 1) === 1) { + store.dispatch(commandIds.toolBeginMaskShape, { point: documentPoint, mode: inputEvent.shiftKey ? "add" : inputEvent.altKey ? "subtract" : "replace" }); + canvas.setPointerCapture(event.pointerId); + event.preventDefault(); + return; + } const brush = !isOperationWorkspacePanel(state.editor.workspace.panel) && (inputEvent.buttons & 1) === 1 && !isPanInteractionMode(state.editor.tools.interactionMode) ? beginBrushSession(state.document, state.editor, documentPoint) : undefined; @@ -139,6 +147,11 @@ export function useCanvasInput( event.preventDefault(); return; } + if (!isOperationWorkspacePanel(state.editor.workspace.panel) && state.editor.tools.activeTool === "semanticSelect") { + void applySemanticSelectionAt(store, documentPoint, inputEvent.shiftKey ? "add" : inputEvent.altKey ? "subtract" : "replace"); + event.preventDefault(); + return; + } const currentState = store.getState(); const selectionToolActive = currentState.editor.tools.activeTool === "select"; @@ -153,6 +166,12 @@ export function useCanvasInput( const handlePointerMove = (event: PointerEvent) => { const inputEvent = pointerInputEventFromPointerEvent(event); + if (store.getState().editor.maskShapeSession) { + const point = viewportPointToDocumentPoint(inputEvent.position, store.getState().editor.viewport); + store.dispatch(commandIds.toolAppendMaskShape, { point }); + event.preventDefault(); + return; + } if (brushSession.current) { if ((inputEvent.buttons & 1) !== 1 || isPanInteractionMode(store.getState().editor.tools.interactionMode)) { commitActiveBrushSession(inputEvent.position); @@ -163,7 +182,7 @@ export function useCanvasInput( const point = viewportPointToDocumentPoint(inputEvent.position, store.getState().editor.viewport); store.dispatch(commandIds.toolSetBrushPreview, { position: point }); const settings = store.getState().editor.tools.brush; - brushSession.current = updateBrushSession({ store, session: brushSession.current, point, color: settings.color, size: settings.size, hardness: settings.hardness }); + brushSession.current = updateBrushSession({ store, session: brushSession.current, point, color: settings.color, size: settings.size, hardness: settings.hardness, opacity: settings.opacity, flow: settings.flow, smoothing: settings.smoothing, pressure: inputEvent.pressure ?? 1, pressureSize: settings.pressureSize }); event.preventDefault(); return; } @@ -187,6 +206,13 @@ export function useCanvasInput( const handlePointerUp = (event: PointerEvent) => { const inputEvent = pointerInputEventFromPointerEvent(event); + if (store.getState().editor.maskShapeSession) { + const point = viewportPointToDocumentPoint(inputEvent.position, store.getState().editor.viewport); + store.dispatch(commandIds.toolAppendMaskShape, { point }); + void commitInpaintLasso(store); + event.preventDefault(); + return; + } if (brushSession.current) { commitActiveBrushSession(inputEvent.position); event.preventDefault(); diff --git a/view/inpaint/CandidateReviewPanel.tsx b/view/inpaint/CandidateReviewPanel.tsx new file mode 100644 index 0000000..f9a57e9 --- /dev/null +++ b/view/inpaint/CandidateReviewPanel.tsx @@ -0,0 +1,25 @@ +import { Star } from "@phosphor-icons/react"; +import { commandIds } from "@commands/ids"; +import type { GenerationState } from "@editor/state"; +import type { AppStore } from "@editor/store"; + +export function CandidateReviewPanel({ generation, dispatch }: { generation: GenerationState; dispatch: AppStore["dispatch"] }) { + if (generation.candidates.length === 0) return null; + const selectedId = generation.selectedCandidateId ?? generation.candidates[0]?.id; + return ( +
+
Results{generation.candidates.length} candidates
+
+ {generation.candidates.map((candidate) => ( +
+ + +
+ ))} +
+
+ ); +} diff --git a/view/inpaint/InpaintRegionPanel.tsx b/view/inpaint/InpaintRegionPanel.tsx new file mode 100644 index 0000000..3982bd3 --- /dev/null +++ b/view/inpaint/InpaintRegionPanel.tsx @@ -0,0 +1,86 @@ +import { useEffect, useMemo, useState } from "react"; +import { commandIds } from "@commands/ids"; +import type { ImageDocument } from "@core/document"; +import type { SelectionState } from "@editor/state"; +import type { AppStore } from "@editor/store"; +import type { GenerateSettings } from "@editor/tools"; +import { runInpaintRegionOperation } from "@operations/masks/rasterActions"; +import { createNormalizedMaskSource } from "@platform/browser/maskRaster"; + +type ProcessedMasks = { edit: string; noise: string; blend: string; coverage: number; bounds?: { x: number; y: number; w: number; h: number } }; + +export function InpaintRegionPanel({ document, selection, settings, dispatch }: { document: ImageDocument; selection: SelectionState; settings: GenerateSettings; dispatch: AppStore["dispatch"] }) { + const targetLayerId = selection.layerIds.length === 1 ? selection.layerIds[0] : undefined; + const region = document.inpaintRegions.find((candidate) => candidate.targetLayerId === targetLayerId && candidate.enabled); + const asset = region ? document.assets.find((candidate) => candidate.id === region.maskAssetId) : undefined; + const [processed, setProcessed] = useState(); + const [busy, setBusy] = useState(false); + const processingKey = useMemo(() => asset ? [asset.source, settings.inpaint.maskExpand, settings.inpaint.maskFeather, settings.inpaint.maskBlur, settings.inpaint.maskDespeckle].join(":") : "", [asset, settings.inpaint]); + + useEffect(() => { + if (!asset) { + setProcessed(undefined); + return; + } + let cancelled = false; + const width = Math.max(1, Math.round(asset.intrinsicSize.w)); + const height = Math.max(1, Math.round(asset.intrinsicSize.h)); + void Promise.all([ + createNormalizedMaskSource(asset.source, width, height, { polarity: "revealed", despeckle: settings.inpaint.maskDespeckle }), + createNormalizedMaskSource(asset.source, width, height, { polarity: "revealed", expand: settings.inpaint.maskExpand, blur: settings.inpaint.maskBlur, despeckle: settings.inpaint.maskDespeckle }), + createNormalizedMaskSource(asset.source, width, height, { polarity: "revealed", feather: settings.inpaint.maskFeather, despeckle: settings.inpaint.maskDespeckle }), + ]).then(([edit, noise, blend]) => { + if (cancelled) return; + const active = edit.values.reduce((sum, value) => sum + value / 255, 0); + setProcessed({ edit: edit.source, noise: noise.source, blend: blend.source, coverage: active / edit.values.length, bounds: edit.bounds }); + }).catch(() => !cancelled && setProcessed(undefined)); + return () => { cancelled = true; }; + }, [processingKey, asset, settings.inpaint.maskDespeckle, settings.inpaint.maskExpand, settings.inpaint.maskBlur, settings.inpaint.maskFeather]); + + if (!targetLayerId) return

Select one image or raster layer to create an AI edit region.

; + if (!region || !asset) return

No AI edit region yet. Use Add region beside Generate, then paint where pixels should be replaced.

; + + const operation = async (type: "invert" | "clear" | "fill") => { + setBusy(true); + try { + await runInpaintRegionOperation(region.id, asset, type === "invert" ? { type: "invert" } : { type: "fill", fill: type === "fill" ? "white" : "black" }, dispatch); + } finally { + setBusy(false); + } + }; + + return ( +
+
+
AI edit regionPainted pixels are replaced; unpainted pixels are protected.
+ +
+ {processed ? ( +
+ + + +
+ ) : Preparing mask previews…} +
+ {processed ? Replace {Math.round(processed.coverage * 100)}%{processed.bounds ? ` · ${Math.round(processed.bounds.w)}×${Math.round(processed.bounds.h)} px` : " · empty"} : } + void operation("invert")} /> + { dispatch(commandIds.toolEnterInpaintRegionEdit, { targetLayerId, regionId: region.id }); dispatch(commandIds.toolSetActive, { tool: "magicWand" }); }} /> + { dispatch(commandIds.toolEnterInpaintRegionEdit, { targetLayerId, regionId: region.id }); dispatch(commandIds.toolSetActive, { tool: "semanticSelect" }); }} /> + { dispatch(commandIds.toolEnterInpaintRegionEdit, { targetLayerId, regionId: region.id }); dispatch(commandIds.toolSetActive, { tool: "maskLasso" }); }} /> + { dispatch(commandIds.toolEnterInpaintRegionEdit, { targetLayerId, regionId: region.id }); dispatch(commandIds.toolSetActive, { tool: "maskRectangle" }); }} /> + void operation("clear")} /> + void operation("fill")} /> + dispatch(commandIds.documentRemoveInpaintRegion, { regionId: region.id })} /> +
+
+ ); +} + +function MaskPreview({ label, source }: { label: string; source: string }) { + return
{`${label}{label}
; +} + +function RegionButton({ label, disabled, danger, onClick }: { label: string; disabled?: boolean; danger?: boolean; onClick: () => void }) { + return ; +} diff --git a/view/layers/MaskControls.tsx b/view/layers/MaskControls.tsx index 242c86c..da1fd0f 100644 --- a/view/layers/MaskControls.tsx +++ b/view/layers/MaskControls.tsx @@ -3,7 +3,7 @@ import type { Asset } from "@core/asset"; import type { AppStore } from "@editor/store"; import { analyzeMask, runMaskOperation, type MaskAnalysis, type MaskRasterOperation } from "@operations/masks/rasterActions"; -export function MaskStatus({ asset }: { asset: Asset }) { +export function MaskStatus({ asset, purpose = "visibility" }: { asset: Asset; purpose?: "visibility" | "inpaint" }) { const [analysis, setAnalysis] = useState(); useEffect(() => { @@ -24,8 +24,7 @@ export function MaskStatus({ asset }: { asset: Asset }) { return ( - Reveal {formatPercent(analysis.coverage)} - Inpaint {formatPercent(analysis.hiddenCoverage)} + {purpose === "inpaint" ? Replace {formatPercent(analysis.coverage)} : <>Reveal {formatPercent(analysis.coverage)}Hidden {formatPercent(analysis.hiddenCoverage)}} ); } diff --git a/view/layers/thumbnailModel.test.ts b/view/layers/thumbnailModel.test.ts index 6da1f1e..b8510d2 100644 --- a/view/layers/thumbnailModel.test.ts +++ b/view/layers/thumbnailModel.test.ts @@ -36,7 +36,7 @@ describe("layer thumbnail read model", () => { test("indexes nested group previews once for row lookup", () => { const child = { ...base("child-group", "Child"), type: "group" as const, children: [raster("nested", "asset-b")] }; const group = { ...base("root", "Root"), type: "group" as const, children: [child] }; - const document = { id: "doc", version: 1, name: "Doc", assets: [...assets.values()], artboards: [{ id: "artboard", name: "Artboard", bounds: { x: 0, y: 0, w: 100, h: 100 }, backgroundColor: "#000000", visible: true, locked: false, layers: [group] }] }; + const document = { id: "doc", version: 1, name: "Doc", assets: [...assets.values()], inpaintRegions: [], artboards: [{ id: "artboard", name: "Artboard", bounds: { x: 0, y: 0, w: 100, h: 100 }, backgroundColor: "#000000", visible: true, locked: false, layers: [group] }] }; const index = createLayerThumbnailIndex(document, assets); expect(index.get("nested")?.kind).toBe("raster"); expect(index.get("child-group")).toMatchObject({ kind: "group", previews: [{ source: "blob:b" }] }); diff --git a/view/paletteItems.tsx b/view/paletteItems.tsx index 7aee3a1..93b23ef 100644 --- a/view/paletteItems.tsx +++ b/view/paletteItems.tsx @@ -345,6 +345,12 @@ function toolIcon(tool: ToolId) { return ; case "magicWand": return ; + case "semanticSelect": + return ; + case "maskLasso": + return ; + case "maskRectangle": + return ; case "pan": return ; } diff --git a/view/toolLabels.ts b/view/toolLabels.ts index 5b53653..6993cba 100644 --- a/view/toolLabels.ts +++ b/view/toolLabels.ts @@ -6,6 +6,12 @@ export function labelForTool(tool: ToolId): string { return "Brush"; case "magicWand": return "Magic wand"; + case "semanticSelect": + return "AI object select"; + case "maskLasso": + return "AI region lasso"; + case "maskRectangle": + return "AI region rectangle"; case "eraser": return "Eraser"; case "pan":