diff --git a/frontend/libs/components/pagedraw/src/lib/PageDraw.tsx b/frontend/libs/components/pagedraw/src/lib/PageDraw.tsx index 44b3080861..7043c3b50d 100644 --- a/frontend/libs/components/pagedraw/src/lib/PageDraw.tsx +++ b/frontend/libs/components/pagedraw/src/lib/PageDraw.tsx @@ -32,6 +32,7 @@ import { EncodeImageBitmapToBase64 } from "./utilities/EncodeImageBitmapToBase64 import { compositeInWorker, maskInWorker, + setBaseBitmapInWorker, } from "./utilities/generatePipeline"; import { RefImage, usePrompt2DStore } from "@storyteller/ui-promptbox"; import { PromptsApi } from "@storyteller/api"; @@ -598,19 +599,13 @@ const PageDraw = ({ adapter }: PageDrawProps) => { pixelRatio: 1 / stageRef.current.scaleX(), }); - // Move the marker pixels and a clone of the base bitmap off the main - // thread. The worker handles the exact-size resize, compositing, and PNG - // encoding so the renderer stays responsive during a 4K generate. - const [markerBitmap, baseBitmap] = await Promise.all([ - createImageBitmap(markerLayerCanvas), - baseImageBitmap - ? createImageBitmap(baseImageBitmap) - : Promise.resolve(undefined), - ]); + // Move the marker pixels off the main thread. The base bitmap is already + // cached in the worker (synced via setBaseBitmapInWorker on baseImageBitmap + // changes), so we only need to send the marker on each bake. + const markerBitmap = await createImageBitmap(markerLayerCanvas); const blob = await compositeInWorker({ markerBitmap, - baseBitmap, width, height, }); @@ -618,6 +613,105 @@ const PageDraw = ({ adapter }: PageDrawProps) => { return new File([blob], `${uuid}.png`, { type: "image/png" }); }, [baseImageBitmap, baseImageInfo]); + // Push the current base bitmap into the worker's cache whenever it changes. + // The worker keeps it across generates so we don't pay createImageBitmap + + // postMessage(transfer) on every Generate click. + useEffect(() => { + if (!baseImageBitmap) { + setBaseBitmapInWorker(undefined).catch((err) => + console.error("Failed to clear worker base bitmap:", err), + ); + return; + } + let cancelled = false; + createImageBitmap(baseImageBitmap) + .then((clone) => { + if (cancelled) { + clone.close(); + return; + } + return setBaseBitmapInWorker(clone); + }) + .catch((err) => + console.error("Failed to send worker base bitmap:", err), + ); + return () => { + cancelled = true; + }; + }, [baseImageBitmap]); + + // Pre-bake the composite File on idle so the Generate click path is + // near-instant in the common case (user pauses >300ms before clicking). The + // bake itself still freezes the main thread for the Konva readback, but it + // happens between user actions instead of on click. A generation counter + // makes concurrent bakes correct: only the latest bake commits its result, + // and only if no canvas change happened mid-bake. + const bakedCompositeRef = useRef(null); + const bakeDirtyRef = useRef(true); + const bakeInFlightRef = useRef | null>(null); + const bakeGenRef = useRef(0); + const bakeTimeoutRef = useRef(null); + + const runCompositeBake = useCallback((): Promise => { + const myGen = ++bakeGenRef.current; + bakeDirtyRef.current = false; + const promise = (async () => { + try { + const file = await getCompositeCanvasFile(); + if (myGen === bakeGenRef.current && !bakeDirtyRef.current) { + bakedCompositeRef.current = file; + } + return file; + } catch (err) { + if (myGen === bakeGenRef.current) bakeDirtyRef.current = true; + console.error("Composite pre-bake failed:", err); + return null; + } finally { + // Only the latest bake clears the in-flight slot. Older bakes that + // finish after a newer one has started leave the newer one's promise + // in place. + if (myGen === bakeGenRef.current) bakeInFlightRef.current = null; + } + })(); + bakeInFlightRef.current = promise; + return promise; + }, [getCompositeCanvasFile]); + + const getCompositeFile = useCallback((): Promise => { + if (!bakeDirtyRef.current) { + if (bakedCompositeRef.current) { + return Promise.resolve(bakedCompositeRef.current); + } + if (bakeInFlightRef.current) return bakeInFlightRef.current; + } + return runCompositeBake(); + }, [runCompositeBake]); + + const scheduleCompositeBake = useCallback(() => { + if (bakeTimeoutRef.current !== null) { + window.clearTimeout(bakeTimeoutRef.current); + } + bakeTimeoutRef.current = window.setTimeout(() => { + bakeTimeoutRef.current = null; + runCompositeBake().catch(() => {}); + }, 300); + }, [runCompositeBake]); + + useEffect(() => { + bakeDirtyRef.current = true; + bakedCompositeRef.current = null; + scheduleCompositeBake(); + }, [drawNodes, baseImageBitmap, baseImageInfo, scheduleCompositeBake]); + + useEffect( + () => () => { + if (bakeTimeoutRef.current !== null) { + window.clearTimeout(bakeTimeoutRef.current); + } + }, + [], + ); + const handleGenerate = useCallback( async ( prompt: string, @@ -676,7 +770,7 @@ const PageDraw = ({ adapter }: PageDrawProps) => { }); } else if (selectedImageModel?.isNanoBananaModel()) { // CASE 2 - NANO BANANA - const compositeFile = await getCompositeCanvasFile(); + const compositeFile = await getCompositeFile(); if (!compositeFile) { console.error("Failed to create composite canvas"); @@ -710,7 +804,7 @@ const PageDraw = ({ adapter }: PageDrawProps) => { }); } else { // CASE 3 - DEFAULT - const compositeFile = await getCompositeCanvasFile(); + const compositeFile = await getCompositeFile(); if (!compositeFile) { console.error("Failed to create composite canvas"); @@ -754,7 +848,7 @@ const PageDraw = ({ adapter }: PageDrawProps) => { }, [ generationCount, - getCompositeCanvasFile, + getCompositeFile, selectedImageModel, adapter, baseImageInfo, diff --git a/frontend/libs/components/pagedraw/src/lib/utilities/generatePipeline.ts b/frontend/libs/components/pagedraw/src/lib/utilities/generatePipeline.ts index d0682b8b70..c117c9acde 100644 --- a/frontend/libs/components/pagedraw/src/lib/utilities/generatePipeline.ts +++ b/frontend/libs/components/pagedraw/src/lib/utilities/generatePipeline.ts @@ -1,12 +1,13 @@ import GeneratePipelineWorker from "./generatePipeline.worker.ts?worker&inline"; +type SetBaseOk = { id: number; ok: true; type: "setBase" }; type CompositeOk = { id: number; ok: true; type: "composite"; blob: Blob }; type MaskOk = { id: number; ok: true; type: "mask"; bytes: Uint8Array }; type WorkerErr = { id: number; ok: false; error: string }; -type WorkerResponse = CompositeOk | MaskOk | WorkerErr; +type WorkerResponse = SetBaseOk | CompositeOk | MaskOk | WorkerErr; type Pending = { - resolve: (value: CompositeOk | MaskOk) => void; + resolve: (value: SetBaseOk | CompositeOk | MaskOk) => void; reject: (err: Error) => void; }; @@ -39,16 +40,34 @@ const ensureWorker = (): Worker => { return worker; }; +export const setBaseBitmapInWorker = ( + baseBitmap?: ImageBitmap, +): Promise => { + const worker = ensureWorker(); + const id = nextId++; + const transfer: Transferable[] = baseBitmap ? [baseBitmap] : []; + return new Promise((resolve, reject) => { + pendingRequests.set(id, { + resolve: (res) => { + if (res.type !== "setBase") { + reject(new Error("Unexpected response type for setBase")); + return; + } + resolve(); + }, + reject, + }); + worker.postMessage({ id, type: "setBase", baseBitmap }, transfer); + }); +}; + export const compositeInWorker = (params: { markerBitmap: ImageBitmap; - baseBitmap?: ImageBitmap; width: number; height: number; }): Promise => { const worker = ensureWorker(); const id = nextId++; - const transfer: Transferable[] = [params.markerBitmap]; - if (params.baseBitmap) transfer.push(params.baseBitmap); return new Promise((resolve, reject) => { pendingRequests.set(id, { resolve: (res) => { @@ -60,7 +79,9 @@ export const compositeInWorker = (params: { }, reject, }); - worker.postMessage({ id, type: "composite", ...params }, transfer); + worker.postMessage({ id, type: "composite", ...params }, [ + params.markerBitmap, + ]); }); }; diff --git a/frontend/libs/components/pagedraw/src/lib/utilities/generatePipeline.worker.ts b/frontend/libs/components/pagedraw/src/lib/utilities/generatePipeline.worker.ts index 9b5cd49716..c06d4eca33 100644 --- a/frontend/libs/components/pagedraw/src/lib/utilities/generatePipeline.worker.ts +++ b/frontend/libs/components/pagedraw/src/lib/utilities/generatePipeline.worker.ts @@ -1,11 +1,15 @@ /// type WorkerRequest = + | { + id: number; + type: "setBase"; + baseBitmap?: ImageBitmap; + } | { id: number; type: "composite"; markerBitmap: ImageBitmap; - baseBitmap?: ImageBitmap; width: number; height: number; } @@ -18,10 +22,16 @@ type WorkerRequest = }; type WorkerResponse = + | { id: number; ok: true; type: "setBase" } | { id: number; ok: true; type: "composite"; blob: Blob } | { id: number; ok: true; type: "mask"; bytes: Uint8Array } | { id: number; ok: false; error: string }; +// The base image rarely changes (only when the user picks a new starting image +// or lands on a generation result). Caching it here avoids cloning + transferring +// it on every Generate click. +let cachedBaseBitmap: ImageBitmap | null = null; + const ctxOf = (canvas: OffscreenCanvas): OffscreenCanvasRenderingContext2D => { const ctx = canvas.getContext("2d"); if (!ctx) throw new Error("Failed to get OffscreenCanvas 2d context"); @@ -38,12 +48,15 @@ const post = (response: WorkerResponse, transfer: Transferable[] = []) => { self.addEventListener("message", async (event: MessageEvent) => { const req = event.data; try { - if (req.type === "composite") { + if (req.type === "setBase") { + cachedBaseBitmap?.close(); + cachedBaseBitmap = req.baseBitmap ?? null; + post({ id: req.id, ok: true, type: "setBase" }); + } else if (req.type === "composite") { const canvas = new OffscreenCanvas(req.width, req.height); const ctx = ctxOf(canvas); - if (req.baseBitmap) { - ctx.drawImage(req.baseBitmap, 0, 0, req.width, req.height); - req.baseBitmap.close(); + if (cachedBaseBitmap) { + ctx.drawImage(cachedBaseBitmap, 0, 0, req.width, req.height); } else { ctx.fillStyle = "white"; ctx.fillRect(0, 0, req.width, req.height);