diff --git a/commands/generation.test.ts b/commands/generation.test.ts index 029d5e2..7463287 100644 --- a/commands/generation.test.ts +++ b/commands/generation.test.ts @@ -12,6 +12,7 @@ import { generationFailJobCommand, generationStartJobCommand, generationSucceedJobCommand, + generationCancelJobCommand, } from "./generation"; describe("generation commands", () => { @@ -40,6 +41,16 @@ describe("generation commands", () => { expect(next.editor.generation.compareMode).toBe("split"); }); + test("bounds retained candidate source memory as well as candidate count", () => { + const largeSource = "x".repeat(34 * 1024 * 1024); + const first = { ...generationCandidate("candidate-1"), source: largeSource }; + const second = { ...generationCandidate("candidate-2"), source: largeSource }; + const withFirst = generationAddCandidateCommand.execute({ state: createInitialAppState("Test") }, { candidate: first }); + const withSecond = generationAddCandidateCommand.execute({ state: withFirst }, { candidate: second }); + + expect(withSecond.editor.generation.candidates.map((candidate) => candidate.id)).toEqual(["candidate-2"]); + }); + test("clears the candidate session and resets comparison", () => { const withCandidate = generationAddCandidateCommand.execute( { state: createInitialAppState("Test") }, @@ -159,6 +170,13 @@ describe("generation commands", () => { expect(failed.editor.generation.jobs[0]).toMatchObject({ id: "job-1", status: "failed", error: "Backend unavailable", finishedAt: 125 }); }); + + test("records explicit cancellation separately from failure", () => { + const running = generationStartJobCommand.execute({ state: createInitialAppState("Test") }, { jobId: "job-1", kind: "generate", label: "Generating", startedAt: 100 }); + const cancelled = generationCancelJobCommand.execute({ state: running }, { jobId: "job-1", finishedAt: 110 }); + + expect(cancelled.editor.generation.jobs[0]).toMatchObject({ id: "job-1", status: "cancelled", finishedAt: 110 }); + }); }); function documentWithSourceLayer() { diff --git a/commands/generation.ts b/commands/generation.ts index 47f3db9..bb73b58 100644 --- a/commands/generation.ts +++ b/commands/generation.ts @@ -43,10 +43,12 @@ export type GenerationReplaceCandidatePixelsPayload = { export type GenerationStartJobPayload = { jobId: GenerationJobId; kind: GenerationJobKind; label: string; startedAt: number }; export type GenerationSucceedJobPayload = { jobId: GenerationJobId; finishedAt: number }; export type GenerationFailJobPayload = { jobId: GenerationJobId; finishedAt: number; error: string }; +export type GenerationCancelJobPayload = { jobId: GenerationJobId; finishedAt: number }; export type GenerationSetResourcesPayload = { options: GenerationOptions }; export type GenerationFailResourcesPayload = { error: string }; const maxCandidates = 12; +const maxCandidateSourceBytes = 96 * 1024 * 1024; const maxJobs = 20; const generationCompareModes = new Set(["result", "before", "split"]); @@ -55,7 +57,7 @@ export const generationAddCandidateCommand: Command candidate.id !== payload.candidate.id)].slice(0, maxCandidates); + const candidates = retainCandidateBudget([payload.candidate, ...state.editor.generation.candidates.filter((candidate) => candidate.id !== payload.candidate.id)]); return { ...state, editor: { @@ -277,6 +279,15 @@ export const generationFailJobCommand: Command = { }, }; +export const generationCancelJobCommand: Command = { + id: commandIds.generationCancelJob, + name: "Cancel generation job", + history: { mode: "ignore" }, + execute({ state }, payload) { + return settleJob(state, payload.jobId, payload.finishedAt, "cancelled"); + }, +}; + export const generationLoadResourcesCommand: Command = { id: commandIds.generationLoadResources, name: "Load generation resources", @@ -318,6 +329,7 @@ export const generationCommands = [ generationStartJobCommand, generationSucceedJobCommand, generationFailJobCommand, + generationCancelJobCommand, generationLoadResourcesCommand, generationSetResourcesCommand, generationFailResourcesCommand, @@ -331,13 +343,31 @@ function updateResources(state: AppState, resources: GenerationState["resources" return { ...state, editor: { ...state.editor, generation: { ...state.editor.generation, resources } } }; } -function settleJob(state: AppState, jobId: GenerationJobId, finishedAt: number, status: "succeeded" | "failed", error?: string): AppState { +function settleJob(state: AppState, jobId: GenerationJobId, finishedAt: number, status: "succeeded" | "failed" | "cancelled", error?: string): AppState { if (!Number.isFinite(finishedAt)) return state; const job = state.editor.generation.jobs.find((candidate) => candidate.id === jobId); if (!job || job.status !== "running" || finishedAt < job.startedAt) return state; return updateJobs(state, state.editor.generation.jobs.map((candidate) => candidate.id === jobId ? { ...candidate, status, finishedAt, error } : candidate)); } +function retainCandidateBudget(candidates: GenerationCandidate[]): GenerationCandidate[] { + const retained: GenerationCandidate[] = []; + let bytes = 0; + for (const candidate of candidates) { + const candidateBytes = candidateRetainedBytes(candidate); + if (retained.length > 0 && bytes + candidateBytes > maxCandidateSourceBytes) continue; + retained.push(candidate); + bytes += candidateBytes; + if (retained.length >= maxCandidates) break; + } + return retained; +} + +function candidateRetainedBytes(candidate: GenerationCandidate): number { + return [candidate.source, candidate.inputImage, candidate.maskImage, candidate.inpaint?.inputImage, candidate.inpaint?.maskImage] + .reduce((total, source) => total + (source?.length ?? 0) * 2, 0); +} + type LayerLocation = { artboardId: ArtboardId; layer: Layer; diff --git a/commands/ids.ts b/commands/ids.ts index 5e1f8ba..5cf0a22 100644 --- a/commands/ids.ts +++ b/commands/ids.ts @@ -48,6 +48,7 @@ export const commandIds = { generationStartJob: "generation.startJob", generationSucceedJob: "generation.succeedJob", generationFailJob: "generation.failJob", + generationCancelJob: "generation.cancelJob", generationLoadResources: "generation.loadResources", generationSetResources: "generation.setResources", generationFailResources: "generation.failResources", diff --git a/commands/payloads.ts b/commands/payloads.ts index 0a8e692..2f91aca 100644 --- a/commands/payloads.ts +++ b/commands/payloads.ts @@ -34,6 +34,7 @@ import type { GenerationStartJobPayload, GenerationSucceedJobPayload, GenerationFailJobPayload, + GenerationCancelJobPayload, GenerationSetResourcesPayload, GenerationFailResourcesPayload, } from "./generation"; @@ -106,6 +107,7 @@ export type CommandPayloads = { [commandIds.generationStartJob]: GenerationStartJobPayload; [commandIds.generationSucceedJob]: GenerationSucceedJobPayload; [commandIds.generationFailJob]: GenerationFailJobPayload; + [commandIds.generationCancelJob]: GenerationCancelJobPayload; [commandIds.generationLoadResources]: void; [commandIds.generationSetResources]: GenerationSetResourcesPayload; [commandIds.generationFailResources]: GenerationFailResourcesPayload; diff --git a/editor/document-indexes.test.ts b/editor/document-indexes.test.ts index 2c69c57..928674e 100644 --- a/editor/document-indexes.test.ts +++ b/editor/document-indexes.test.ts @@ -67,6 +67,17 @@ describe("document read indexes", () => { expect(visited).toEqual(["group", "target", "mask"]); expect(layers.map((layer) => layer.id)).toEqual(["mask", "target", "group"]); }); + + test("indexes pathological nesting without recursive stack overflow", () => { + let layers: Layer[] = [raster("leaf", "Leaf", "asset-nested", { x: 80, y: 10 }, { x: 2, y: 3 })]; + for (let depth = 0; depth < 10_000; depth += 1) layers = [group(`group-${depth}`, `Group ${depth}`, layers)]; + const deepDocument = { ...document, artboards: [{ ...document.artboards[0]!, layers }] }; + + const index = createDocumentReadIndex(deepDocument); + + expect(index.layerById.size).toBe(10_001); + expect(resolveIndexedLayerBounds(index, layers[0]!)).toEqual({ x: 80, y: 10, w: 40, h: 30 }); + }); }); function raster( diff --git a/editor/document-indexes.ts b/editor/document-indexes.ts index f9674df..cdd8427 100644 --- a/editor/document-indexes.ts +++ b/editor/document-indexes.ts @@ -99,55 +99,58 @@ function indexLayerTree(options: { documentMaskLayerIds: Set; maskLayerIdsByLayerList: Map>; }): Set { - const layerListMaskLayerIds = new Set(); - - for (let index = 0; index < options.layers.length; index += 1) { - const layer = options.layers[index]; - if (!layer) continue; - - options.layerById.set(layer.id, layer); - options.layerInfoById.set(layer.id, { - artboardId: options.artboardId, - parentGroupId: options.parentGroupId, - layer, - siblings: options.layers, - index, - }); - - const layerMask = getLayerMask(layer); - if (layerMask) { - options.documentMaskLayerIds.add(layerMask.maskLayerId); - layerListMaskLayerIds.add(layerMask.maskLayerId); + type Frame = { layers: readonly Layer[]; parentGroupId?: LayerId; visited: boolean }; + const stack: Frame[] = [{ layers: options.layers, parentGroupId: options.parentGroupId, visited: false }]; + while (stack.length > 0) { + const frame = stack.pop(); + if (!frame) continue; + if (frame.visited) { + const ids = new Set(); + for (const layer of frame.layers) { + const mask = getLayerMask(layer); + if (mask) ids.add(mask.maskLayerId); + if (layer.type === "group") for (const id of options.maskLayerIdsByLayerList.get(layer.children) ?? []) ids.add(id); + } + options.maskLayerIdsByLayerList.set(frame.layers, ids); + continue; } - - if (layer.type === "group") { - const childMaskLayerIds = indexLayerTree({ - ...options, - layers: layer.children, - parentGroupId: layer.id, - }); - for (const maskLayerId of childMaskLayerIds) layerListMaskLayerIds.add(maskLayerId); + stack.push({ ...frame, visited: true }); + for (let index = frame.layers.length - 1; index >= 0; index -= 1) { + const layer = frame.layers[index]; + if (!layer) continue; + options.layerById.set(layer.id, layer); + options.layerInfoById.set(layer.id, { artboardId: options.artboardId, parentGroupId: frame.parentGroupId, layer, siblings: frame.layers, index }); + const mask = getLayerMask(layer); + if (mask) options.documentMaskLayerIds.add(mask.maskLayerId); + if (layer.type === "group") stack.push({ layers: layer.children, parentGroupId: layer.id, visited: false }); } } - - options.maskLayerIdsByLayerList.set(options.layers, layerListMaskLayerIds); - return layerListMaskLayerIds; + return new Set(options.maskLayerIdsByLayerList.get(options.layers) ?? []); } function countDisplayLayers(layers: readonly Layer[], maskLayerIds: ReadonlySet): number { let count = 0; - for (const layer of layers) { + const stack = [...layers]; + while (stack.length > 0) { + const layer = stack.pop(); + if (!layer) continue; if (maskLayerIds.has(layer.id)) continue; count += 1; - if (layer.type === "group") count += countDisplayLayers(layer.children, maskLayerIds); + if (layer.type === "group") stack.push(...layer.children); } return count; } function unionLayerBounds(index: DocumentReadIndex, layers: readonly Layer[]): Rect | undefined { let bounds: Rect | undefined; - - for (const layer of layers) { + const stack = [...layers]; + while (stack.length > 0) { + const layer = stack.pop(); + if (!layer) continue; + if (layer.type === "group") { + stack.push(...layer.children); + continue; + } const layerBounds = resolveIndexedLayerBounds(index, layer); if (!layerBounds) continue; bounds = bounds ? unionRects(bounds, layerBounds) : layerBounds; diff --git a/editor/state.ts b/editor/state.ts index 3134b6d..a4bf685 100644 --- a/editor/state.ts +++ b/editor/state.ts @@ -82,7 +82,7 @@ export type GenerationCandidate = { export type GenerationCompareMode = "result" | "before" | "split"; export type GenerationJobKind = "generate" | "regenerate" | "refine" | "replace"; -export type GenerationJobStatus = "running" | "succeeded" | "failed"; +export type GenerationJobStatus = "running" | "succeeded" | "failed" | "cancelled"; export type GenerationJob = { id: GenerationJobId; diff --git a/operations/generation/generationJob.ts b/operations/generation/generationJob.ts index 9ac0c27..e01d1e5 100644 --- a/operations/generation/generationJob.ts +++ b/operations/generation/generationJob.ts @@ -6,7 +6,8 @@ export async function runGenerationJob(options: { kind: GenerationJobKind; label: string; dispatch: AppStore["dispatch"]; - task: () => Promise; + signal: AbortSignal; + task: (signal: AbortSignal) => Promise; }): Promise { const jobId = crypto.randomUUID(); const startedAt = Date.now(); @@ -14,9 +15,13 @@ export async function runGenerationJob(options: { if (!nextState.editor.generation.jobs.some((job) => job.id === jobId && job.status === "running")) return; try { - await options.task(); + await options.task(options.signal); options.dispatch(commandIds.generationSucceedJob, { jobId, finishedAt: Date.now() }); } catch (reason: unknown) { + if (options.signal.aborted) { + options.dispatch(commandIds.generationCancelJob, { jobId, finishedAt: Date.now() }); + return; + } options.dispatch(commandIds.generationFailJob, { jobId, finishedAt: Date.now(), diff --git a/operations/generation/runGenerate.ts b/operations/generation/runGenerate.ts index f074ab4..b95792f 100644 --- a/operations/generation/runGenerate.ts +++ b/operations/generation/runGenerate.ts @@ -17,6 +17,7 @@ export async function runGenerate(options: { viewport: ViewportState; settings: GenerateSettings; dispatch: AppStore["dispatch"]; + signal?: AbortSignal; }) { const { document, selection, settings, dispatch } = options; const precondition = checkGenerationPreconditions(document, selection, settings); @@ -39,6 +40,7 @@ export async function runGenerate(options: { inputImage, maskImage, inpaintBundle, + signal: options.signal, }); const intrinsicSize = await loadImageSize(generated.source); const placement = resolveGeneratedOutputPlacement({ document, selection, settings, intrinsicSize, inpaintBundle }); @@ -64,6 +66,7 @@ export async function runGenerateFromCandidate(options: { candidate: GenerationCandidate; settings?: GenerateSettings; dispatch: AppStore["dispatch"]; + signal?: AbortSignal; }) { const settings = options.settings ?? options.candidate.settings; const seed = resolveSeed(settings.seed); @@ -75,6 +78,7 @@ export async function runGenerateFromCandidate(options: { inputImage: options.candidate.inputImage, maskImage: options.candidate.maskImage, inpaintCandidate: options.candidate, + signal: options.signal, }); const intrinsicSize = await loadImageSize(generated.source); @@ -141,6 +145,7 @@ async function requestGenerate(options: { maskImage?: string; inpaintBundle?: InpaintBundle; inpaintCandidate?: GenerationCandidate; + signal?: AbortSignal; }) { return requestGeneration({ architecture: options.settings.architecture, @@ -162,7 +167,7 @@ async function requestGenerate(options: { inpaint: resolveInpaintRequest(options.inpaintBundle, options.inpaintCandidate, options.settings), inputImage: options.inputImage, maskImage: options.maskImage, - }); + }, options.signal); } function resolveInpaintRequest(inpaintBundle: InpaintBundle | undefined, inpaintCandidate: GenerationCandidate | undefined, settings: GenerateSettings) { diff --git a/operations/generation/workflow.test.ts b/operations/generation/workflow.test.ts index 74ac004..fc27d73 100644 --- a/operations/generation/workflow.test.ts +++ b/operations/generation/workflow.test.ts @@ -39,6 +39,25 @@ describe("generation workflow", () => { expect(app.store.getState().document.artboards[0]?.layers.some((layer) => layer.id === "layer-new")).toBe(true); expect(app.store.getState().editor.generation.candidates).toHaveLength(0); }); + + test("cancels the active operation and records a cancelled job", async () => { + const app = createTestApp(); + let receivedSignal: AbortSignal | undefined; + const workflow = createGenerationWorkflow(app.store, dependencies({ + runGenerate: async (options) => { + receivedSignal = options.signal; + await new Promise((_resolve, reject) => options.signal?.addEventListener("abort", () => reject(new DOMException("Cancelled", "AbortError")), { once: true })); + }, + })); + + const running = workflow.generate(); + await Promise.resolve(); + workflow.cancel(); + await running; + + expect(receivedSignal?.aborted).toBe(true); + expect(app.store.getState().editor.generation.jobs[0]?.status).toBe("cancelled"); + }); }); function dependencies(overrides: Partial): GenerationWorkflowDependencies { diff --git a/operations/generation/workflow.ts b/operations/generation/workflow.ts index f2ff48a..a3a01bb 100644 --- a/operations/generation/workflow.ts +++ b/operations/generation/workflow.ts @@ -30,8 +30,17 @@ const defaultDependencies: GenerationWorkflowDependencies = { }; export function createGenerationWorkflow(store: AppStore, dependencies: GenerationWorkflowDependencies = defaultDependencies) { - const job = (kind: GenerationJobKind, label: string, task: () => Promise) => - runGenerationJob({ kind, label, dispatch: store.dispatch, task }); + let activeController: AbortController | undefined; + const job = async (kind: GenerationJobKind, label: string, task: (signal: AbortSignal) => Promise) => { + if (store.getState().editor.generation.jobs.some((candidate) => candidate.status === "running")) return; + const controller = new AbortController(); + activeController = controller; + try { + await runGenerationJob({ kind, label, dispatch: store.dispatch, signal: controller.signal, task }); + } finally { + if (activeController === controller) activeController = undefined; + } + }; return { precondition: () => { @@ -41,7 +50,7 @@ export function createGenerationWorkflow(store: AppStore, dependencies: Generati loadResources: () => dependencies.loadGenerationResources(store), - generate: () => job("generate", "Generating", async () => { + generate: () => job("generate", "Generating", async (signal) => { const state = store.getState(); await dependencies.runGenerate({ document: state.document, @@ -49,15 +58,16 @@ export function createGenerationWorkflow(store: AppStore, dependencies: Generati viewport: state.editor.viewport, settings: state.editor.tools.generate, dispatch: store.dispatch, + signal, }); }), regenerate: (candidateId: string, settings?: GenerateSettings, label = "Regenerate") => - job("regenerate", label, async () => { + job("regenerate", label, async (signal) => { const candidate = findCandidate(store, candidateId); const nextSettings = settings ?? candidate.settings; store.dispatch(commandIds.toolSetGenerateSettings, nextSettings); - await dependencies.runGenerateFromCandidate({ candidate, settings: nextSettings, dispatch: store.dispatch }); + await dependencies.runGenerateFromCandidate({ candidate, settings: nextSettings, dispatch: store.dispatch, signal }); }), applyCandidateAsLayer: (candidateId: string) => { @@ -115,6 +125,8 @@ export function createGenerationWorkflow(store: AppStore, dependencies: Generati const source = await dependencies.createMaskedPixelReplacementSource(store.getState().document, candidate); store.dispatch(commandIds.generationReplaceCandidatePixels, { candidateId, source, mimeType: "image/png" }); }), + + cancel: () => activeController?.abort(), }; } diff --git a/platform/browser/imageFiles.ts b/platform/browser/imageFiles.ts index b3cbb90..d1e4faf 100644 --- a/platform/browser/imageFiles.ts +++ b/platform/browser/imageFiles.ts @@ -8,8 +8,16 @@ export type BrowserImageFile = { export async function decodeBrowserImageFile(file: File): Promise { if (!file.type.startsWith("image/")) return undefined; + if (file.size > 128 * 1024 * 1024) throw new Error("Image file is too large. The import limit is 128 MB."); + const objectUrl = URL.createObjectURL(file); + let intrinsicSize: { w: number; h: number }; + try { + intrinsicSize = await loadImageSize(objectUrl); + assertCanvasRasterSize(intrinsicSize.w, intrinsicSize.h, "Imported image"); + } finally { + URL.revokeObjectURL(objectUrl); + } const source = await fileToDataUrl(file); - const intrinsicSize = await loadImageSize(source); return { name: file.name, mimeType: file.type, source, intrinsicSize, release: () => undefined }; } @@ -30,3 +38,4 @@ function loadImageSize(source: string): Promise<{ w: number; h: number }> { image.src = source; }); } +import { assertCanvasRasterSize } from "./rasterLimits"; diff --git a/platform/browser/maskRaster.test.ts b/platform/browser/maskRaster.test.ts index e8ed2a4..f805795 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 { cropMaskValuesToRgba, expandRectWithinBounds, invertMaskValues } from "./maskRaster"; +import { blurMaskValues, cropMaskValuesToRgba, dilateMaskValues, erodeMaskValues, expandRectWithinBounds, invertMaskValues } from "./maskRaster"; describe("mask raster utilities", () => { test("exports the normalized drawn mask without filling the whole crop", () => { @@ -28,6 +28,13 @@ describe("mask raster utilities", () => { h: 64, }); }); + + test("uses separable mask kernels without changing square-kernel semantics", () => { + const values = new Uint8ClampedArray([0, 0, 0, 0, 255, 0, 0, 0, 0]); + expect([...dilateMaskValues(values, 3, 3, 1)]).toEqual(new Array(9).fill(255)); + 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)); + }); }); function activeRedPixels(rgba: Uint8ClampedArray) { diff --git a/platform/browser/maskRaster.ts b/platform/browser/maskRaster.ts index 7a7f55b..e488a28 100644 --- a/platform/browser/maskRaster.ts +++ b/platform/browser/maskRaster.ts @@ -1,4 +1,5 @@ import type { Rect } from "@core/geometry"; +import { assertCanvasRasterSize, assertProcessingRasterSize } from "./rasterLimits"; export type MaskFill = "white" | "black" | "clear"; @@ -32,6 +33,7 @@ export type NormalizedMaskOptions = { }; export async function createSolidMaskSource(width: number, height: number, fill: MaskFill): Promise { + assertProcessingRasterSize(width, height, "Mask"); const canvas = createCanvas(width, height); const context = require2dContext(canvas); context.clearRect(0, 0, canvas.width, canvas.height); @@ -48,6 +50,7 @@ export async function createSolidMaskSource(width: number, height: number, fill: } export async function applyMaskRasterOperation(source: string, width: number, height: number, operation: MaskRasterOperation): Promise { + assertProcessingRasterSize(width, height, "Mask"); if (operation.type === "fill") return createSolidMaskSource(width, height, operation.fill); const mask = await loadMaskValues(source, width, height); @@ -56,6 +59,7 @@ export async function applyMaskRasterOperation(source: string, width: number, he } export async function analyzeMaskSource(source: string, width: number, height: number): Promise { + assertProcessingRasterSize(width, height, "Mask"); const mask = await loadMaskValues(source, width, height); const revealed = analyzeValues(mask.values, mask.width, mask.height, false); const hiddenValues = invertMaskValues(mask.values); @@ -74,6 +78,7 @@ export async function analyzeMaskSource(source: string, width: number, height: n } export async function createNormalizedMaskSource(source: string, width: number, height: number, options: NormalizedMaskOptions): Promise<{ source: string; values: Uint8ClampedArray; bounds?: Rect }> { + assertProcessingRasterSize(width, height, "Mask"); const mask = await loadMaskValues(source, width, height); let values = options.polarity === "hidden" ? invertMaskValues(mask.values) : new Uint8ClampedArray(mask.values); @@ -94,6 +99,7 @@ export async function createNormalizedMaskSource(source: string, width: number, export async function loadImageCanvas(source: string, width?: number, height?: number): Promise { const image = await loadImage(source); + assertCanvasRasterSize(width ?? image.naturalWidth, height ?? image.naturalHeight); const canvas = createCanvas(width ?? image.naturalWidth, height ?? image.naturalHeight); const context = require2dContext(canvas); context.clearRect(0, 0, canvas.width, canvas.height); @@ -108,6 +114,7 @@ export async function imageSourceToPngDataUrl(source: string): Promise { } export function cropCanvas(sourceCanvas: HTMLCanvasElement, crop: Rect, outputWidth = crop.w, outputHeight = crop.h): string { + assertProcessingRasterSize(outputWidth, outputHeight, "Crop"); const canvas = createCanvas(outputWidth, outputHeight); const context = require2dContext(canvas); context.clearRect(0, 0, outputWidth, outputHeight); @@ -116,6 +123,7 @@ export function cropCanvas(sourceCanvas: HTMLCanvasElement, crop: Rect, outputWi } 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); const context = require2dContext(canvas); const imageData = context.createImageData(outputWidth, outputHeight); @@ -203,52 +211,53 @@ export function invertMaskValues(values: Uint8ClampedArray): Uint8ClampedArray { export function erodeMaskValues(values: Uint8ClampedArray, width: number, height: number, radius: number): Uint8ClampedArray { const safeRadius = Math.round(clampNumber(radius, 0, 256)); if (safeRadius <= 0) return new Uint8ClampedArray(values); - const next = new Uint8ClampedArray(values.length); - for (let y = 0; y < height; y += 1) { - for (let x = 0; x < width; x += 1) { - let value = 255; - for (let oy = -safeRadius; oy <= safeRadius; oy += 1) { - for (let ox = -safeRadius; ox <= safeRadius; ox += 1) value = Math.min(value, values[clampInt(y + oy, 0, height - 1) * width + clampInt(x + ox, 0, width - 1)] ?? 0); - } - next[y * width + x] = value; - } - } - return next; + return separableExtrema(values, width, height, safeRadius, Math.min, 255); } export function dilateMaskValues(values: Uint8ClampedArray, width: number, height: number, radius: number): Uint8ClampedArray { const safeRadius = Math.round(clampNumber(radius, 0, 256)); if (safeRadius <= 0) return new Uint8ClampedArray(values); - const next = new Uint8ClampedArray(values.length); - for (let y = 0; y < height; y += 1) { - for (let x = 0; x < width; x += 1) { - let value = 0; - for (let oy = -safeRadius; oy <= safeRadius; oy += 1) { - for (let ox = -safeRadius; ox <= safeRadius; ox += 1) value = Math.max(value, values[clampInt(y + oy, 0, height - 1) * width + clampInt(x + ox, 0, width - 1)] ?? 0); - } - next[y * width + x] = value; - } - } - return next; + return separableExtrema(values, width, height, safeRadius, Math.max, 0); } export function blurMaskValues(values: Uint8ClampedArray, width: number, height: number, radius: number): Uint8ClampedArray { const safeRadius = Math.round(clampNumber(radius, 0, 256)); if (safeRadius <= 0) return new Uint8ClampedArray(values); + const horizontal = new Float64Array(values.length); const next = new Uint8ClampedArray(values.length); + const span = safeRadius * 2 + 1; for (let y = 0; y < height; y += 1) { + let total = 0; + for (let ox = -safeRadius; ox <= safeRadius; ox += 1) total += values[y * width + clampInt(ox, 0, width - 1)] ?? 0; for (let x = 0; x < width; x += 1) { - let total = 0; - let count = 0; - for (let oy = -safeRadius; oy <= safeRadius; oy += 1) { - for (let ox = -safeRadius; ox <= safeRadius; ox += 1) { - total += values[clampInt(y + oy, 0, height - 1) * width + clampInt(x + ox, 0, width - 1)] ?? 0; - count += 1; - } - } - next[y * width + x] = Math.round(total / count); + horizontal[y * width + x] = total / span; + total += (values[y * width + clampInt(x + safeRadius + 1, 0, width - 1)] ?? 0) - (values[y * width + clampInt(x - safeRadius, 0, width - 1)] ?? 0); } } + for (let x = 0; x < width; x += 1) { + let total = 0; + for (let oy = -safeRadius; oy <= safeRadius; oy += 1) total += horizontal[clampInt(oy, 0, height - 1) * width + x] ?? 0; + for (let y = 0; y < height; y += 1) { + next[y * width + x] = Math.round(total / span); + total += (horizontal[clampInt(y + safeRadius + 1, 0, height - 1) * width + x] ?? 0) - (horizontal[clampInt(y - safeRadius, 0, height - 1) * width + x] ?? 0); + } + } + return next; +} + +function separableExtrema(values: Uint8ClampedArray, width: number, height: number, radius: number, combine: (a: number, b: number) => number, initial: number) { + const horizontal = new Uint8ClampedArray(values.length); + const next = new Uint8ClampedArray(values.length); + for (let y = 0; y < height; y += 1) for (let x = 0; x < width; x += 1) { + let value = initial; + for (let offset = -radius; offset <= radius; offset += 1) value = combine(value, values[y * width + clampInt(x + offset, 0, width - 1)] ?? 0); + horizontal[y * width + x] = value; + } + for (let y = 0; y < height; y += 1) for (let x = 0; x < width; x += 1) { + let value = initial; + for (let offset = -radius; offset <= radius; offset += 1) value = combine(value, horizontal[clampInt(y + offset, 0, height - 1) * width + x] ?? 0); + next[y * width + x] = value; + } return next; } diff --git a/platform/browser/rasterLimits.ts b/platform/browser/rasterLimits.ts new file mode 100644 index 0000000..840eabc --- /dev/null +++ b/platform/browser/rasterLimits.ts @@ -0,0 +1,20 @@ +const maxCanvasDimension = 16_384; +const maxCanvasPixels = 64 * 1024 * 1024; +const maxProcessingPixels = 24 * 1024 * 1024; + +export function assertCanvasRasterSize(width: number, height: number, label = "Image") { + assertRasterSize(width, height, maxCanvasPixels, label); +} + +export function assertProcessingRasterSize(width: number, height: number, label = "Image") { + assertRasterSize(width, height, maxProcessingPixels, label); +} + +function assertRasterSize(width: number, height: number, maxPixels: number, label: string) { + const w = Math.round(width); + const h = Math.round(height); + if (!Number.isSafeInteger(w) || !Number.isSafeInteger(h) || w <= 0 || h <= 0) throw new Error(`${label} dimensions are invalid.`); + if (w > maxCanvasDimension || h > maxCanvasDimension || w * h > maxPixels) { + throw new Error(`${label} is too large (${w} x ${h}). Maximum processing size is ${maxCanvasDimension} px per side and ${Math.floor(maxPixels / 1_000_000)} megapixels.`); + } +} diff --git a/platform/comfy/generationClient.ts b/platform/comfy/generationClient.ts index e0bb41b..1c18738 100644 --- a/platform/comfy/generationClient.ts +++ b/platform/comfy/generationClient.ts @@ -4,8 +4,8 @@ export async function fetchGenerationOptions(): Promise { return response.json() as Promise; } -export async function requestGeneration(body: unknown): Promise<{ source: string; mimeType: string }> { - const response = await fetch("/api/comfy/generate", { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify(body) }); +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 }); if (!response.ok) throw new Error(await response.text()); return response.json() as Promise<{ source: string; mimeType: string }>; } diff --git a/renderer/layers.ts b/renderer/layers.ts index a8f53cc..16c6a92 100644 --- a/renderer/layers.ts +++ b/renderer/layers.ts @@ -3,7 +3,7 @@ import type { ImageDocument } from "@core/document"; import type { Rect } from "@core/geometry"; import type { Layer } from "@core/layer"; import { getLayerMask } from "@core/layer-mask-utils"; -import { createDocumentReadIndex, forEachLayerBackToFront, resolveIndexedLayerBounds, type DocumentReadIndex } from "@editor/document-indexes"; +import { createDocumentReadIndex, resolveIndexedLayerBounds, type DocumentReadIndex } from "@editor/document-indexes"; import type { EditorState, GenerationCandidate, MaskViewMode, ViewportState } from "@editor/state"; import { clearScreenRect } from "./clear-rect"; import type { ImageTextureRenderer } from "./image-textures"; @@ -24,7 +24,7 @@ export function renderLayers(context: WebGlRendererContext, document: ImageDocum if (!artboard.visible) continue; const clipRect = documentRectToScreenRect(context.canvas, artboard.bounds, editor.viewport); const maskLayerIds = documentIndex.maskLayerIdsByArtboardId.get(artboard.id) ?? emptyLayerIds; - forEachLayerBackToFront(artboard.layers, (layer) => renderLayer(context, documentIndex, editor, layer, imageTextureRenderer, clipRect, maskLayerIds, 1)); + renderLayerTree(context, documentIndex, editor, artboard.layers, imageTextureRenderer, clipRect, maskLayerIds); if (generationCandidate?.placement.artboardId === artboard.id) renderGenerationCandidatePreview(context, editor, generationCandidate, imageTextureRenderer, clipRect); } } @@ -33,31 +33,46 @@ export function generationCandidatePreviewAssets(editor: EditorState): Asset[] { return editor.generation.candidates.map((candidate) => generationCandidateAsset(candidate)); } -function renderLayer( +function renderLayerTree( context: WebGlRendererContext, documentIndex: DocumentReadIndex, editor: EditorState, - layer: Layer, + layers: readonly Layer[], imageTextureRenderer: ImageTextureRenderer, clipRect: ScreenRect, maskLayerIds: ReadonlySet, - inheritedOpacity: number, +) { + const stack: Array<{ layer: Layer; clipRect: ScreenRect; inheritedOpacity: number }> = []; + for (const layer of layers) stack.push({ layer, clipRect, inheritedOpacity: 1 }); + while (stack.length > 0) { + const frame = stack.pop(); + if (!frame) continue; + const { layer } = frame; + if (!layer.visible || maskLayerIds.has(layer.id)) continue; + const effectiveOpacity = frame.inheritedOpacity * layer.opacity; + if (effectiveOpacity <= 0) continue; + const effectiveClipRect = resolveLayerClipRect(context, documentIndex, editor.viewport, layer, frame.clipRect); + if (!effectiveClipRect) continue; + if (layer.type === "group") { + for (const child of layer.children) stack.push({ layer: child, clipRect: effectiveClipRect, inheritedOpacity: effectiveOpacity }); + continue; + } + renderLeafLayer(context, documentIndex, editor, layer, imageTextureRenderer, effectiveClipRect, effectiveOpacity); + } +} + +function renderLeafLayer( + context: WebGlRendererContext, + documentIndex: DocumentReadIndex, + editor: EditorState, + layer: Exclude, + imageTextureRenderer: ImageTextureRenderer, + effectiveClipRect: ScreenRect, + effectiveOpacity: number, ) { const editingMaskLayer = editor.maskEdit?.maskLayerId === layer.id; const maskViewMode = editor.maskEdit?.viewMode ?? "composite"; const isolatedMaskView = isIsolatedMaskView(maskViewMode); - if (!layer.visible || maskLayerIds.has(layer.id)) return; - const effectiveOpacity = inheritedOpacity * layer.opacity; - if (effectiveOpacity <= 0) return; - - const effectiveClipRect = resolveLayerClipRect(context, documentIndex, editor.viewport, layer, clipRect); - if (!effectiveClipRect) return; - - if (layer.type === "group") { - forEachLayerBackToFront(layer.children, (child) => renderLayer(context, documentIndex, editor, child, imageTextureRenderer, effectiveClipRect, maskLayerIds, effectiveOpacity)); - return; - } - if (isolatedMaskView && layer.id !== editor.maskEdit?.targetLayerId) return; const bounds = resolveIndexedLayerBounds(documentIndex, layer); diff --git a/server/comfy-routes.ts b/server/comfy-routes.ts index f49112b..c8b0201 100644 --- a/server/comfy-routes.ts +++ b/server/comfy-routes.ts @@ -4,7 +4,7 @@ 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)); + if (url.pathname === "/api/comfy/generate" && request.method === "POST") return json(await generate(await request.json() as ComfyGenerateRequest, request.signal)); return new Response("Not found", { status: 404 }); } catch (error) { return new Response(error instanceof Error ? error.message : "ComfyUI request failed", { status: 500 }); diff --git a/server/comfy.ts b/server/comfy.ts index bae4ba6..4ff54b8 100644 --- a/server/comfy.ts +++ b/server/comfy.ts @@ -111,15 +111,15 @@ async function listCheckpointModels() { return (await listGenerationOptions()).models; } -export async function generate(request: ComfyGenerateRequest) { +export async function generate(request: ComfyGenerateRequest, signal?: AbortSignal) { if (!request.prompt?.trim()) throw new Error("Prompt is required"); 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); const clientId = crypto.randomUUID(); - const uploaded = request.inputImage ? await uploadDataUrl(request.inputImage, `image-studio-${crypto.randomUUID()}.png`) : undefined; - const mask = request.maskImage ? await uploadDataUrl(request.maskImage, `image-studio-mask-${crypto.randomUUID()}.png`) : undefined; + 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; 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 }); @@ -127,6 +127,7 @@ export async function generate(request: ComfyGenerateRequest) { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ client_id: clientId, prompt }), + signal, }); if (!queued.ok) throw new Error(`ComfyUI prompt failed: ${queued.status} ${await queued.text()}`); const queuedBody = await queued.json() as { prompt_id?: string; node_errors?: unknown }; @@ -134,19 +135,25 @@ export async function generate(request: ComfyGenerateRequest) { if (nodeError) throw new Error(`ComfyUI rejected the workflow: ${nodeError}`); if (!queuedBody.prompt_id) throw new Error("ComfyUI did not return a prompt id"); const prompt_id = queuedBody.prompt_id; - const history = await waitForHistory(prompt_id); + let history: unknown; + try { + history = await waitForHistory(prompt_id, signal); + } 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" })}`); + 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" }; } -async function uploadDataUrl(dataUrl: string, filename: string) { +async function uploadDataUrl(dataUrl: string, filename: string, signal?: AbortSignal) { const match = /^data:([^;]+);base64,(.+)$/.exec(dataUrl); if (!match) throw new Error("Expected a base64 data URL image"); const mimeType = match[1] ?? "image/png"; @@ -154,27 +161,53 @@ async function uploadDataUrl(dataUrl: string, filename: string) { const form = new FormData(); form.append("image", new File([new Uint8Array(Buffer.from(base64, "base64"))], filename, { type: mimeType })); form.append("overwrite", "true"); - const response = await fetch(`${comfyBaseUrl}/upload/image`, { method: "POST", body: form }); + const response = await fetch(`${comfyBaseUrl}/upload/image`, { method: "POST", body: form, signal }); if (!response.ok) throw new Error(`ComfyUI upload failed: ${response.status}`); const uploaded = await response.json() as { name: string }; return uploaded.name; } -async function waitForHistory(promptId: string) { +async function waitForHistory(promptId: string, signal?: AbortSignal) { const startedAt = Date.now(); let attempts = 0; while (Date.now() - startedAt < comfyHistoryTimeoutMs) { - const response = await fetch(`${comfyBaseUrl}/history/${promptId}`); + if (signal?.aborted) { + await cancelComfyPrompt(promptId); + throw signal.reason ?? new DOMException("Generation cancelled", "AbortError"); + } + const response = await fetch(`${comfyBaseUrl}/history/${promptId}`, { signal }); attempts += 1; if (response.ok) { const history = await response.json() as Record; if (history[promptId]) return history[promptId]; } - await Bun.sleep(comfyHistoryPollIntervalMs); + await abortableSleep(comfyHistoryPollIntervalMs, signal); } throw new Error(`Timed out waiting for ComfyUI prompt ${promptId} after ${Math.round(comfyHistoryTimeoutMs / 1000)} seconds and ${attempts} checks`); } +async function cancelComfyPrompt(promptId: string) { + const queueResponse = await fetch(`${comfyBaseUrl}/queue`).catch(() => undefined); + const queue = queueResponse?.ok ? await queueResponse.json() as { queue_running?: unknown[][]; queue_pending?: unknown[][] } : undefined; + const isRunning = queue?.queue_running?.some((entry) => entry.includes(promptId)) ?? false; + const isPending = queue?.queue_pending?.some((entry) => entry.includes(promptId)) ?? true; + const requests: Promise[] = []; + if (isPending) requests.push(fetch(`${comfyBaseUrl}/queue`, { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ delete: [promptId] }) })); + if (isRunning) requests.push(fetch(`${comfyBaseUrl}/interrupt`, { method: "POST" })); + await Promise.allSettled(requests); +} + +function abortableSleep(ms: number, signal?: AbortSignal): Promise { + if (!signal) return Bun.sleep(ms); + return new Promise((resolve, reject) => { + const timeout = setTimeout(resolve, ms); + signal.addEventListener("abort", () => { + clearTimeout(timeout); + reject(signal.reason ?? new DOMException("Generation cancelled", "AbortError")); + }, { once: true }); + }); +} + export function selectGeneratedOutputImage(history: unknown): { filename: string; subfolder?: string; type?: string } | undefined { const outputs = (history as { outputs?: Record }).outputs ?? {}; const saveImageOutput = outputs["8"]?.images?.find(isGeneratedImage); diff --git a/view/GenerationJobStatus.tsx b/view/GenerationJobStatus.tsx index f21000b..26cc2b4 100644 --- a/view/GenerationJobStatus.tsx +++ b/view/GenerationJobStatus.tsx @@ -17,8 +17,8 @@ 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.label} complete`; - const tone = job.status === "failed" ? "bg-red-500/15 text-red-100" : job.status === "running" ? "bg-white/10 text-white/70" : "bg-emerald-500/15 text-emerald-100"; + 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 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/bottom-controls/GenerateActionControls.tsx b/view/bottom-controls/GenerateActionControls.tsx index 2b2439c..28481f5 100644 --- a/view/bottom-controls/GenerateActionControls.tsx +++ b/view/bottom-controls/GenerateActionControls.tsx @@ -36,6 +36,7 @@ export function GenerateActionControls({ settings, generation, dispatch, workflo {preconditionMessage ? {preconditionMessage} : null} + {busy ? : null} {candidate ? ( <>