artcraft: pre-bake PageDraw composite on idle and cache base bitmap in worker

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<File | null>(null);
+  const bakeDirtyRef = useRef(true);
+  const bakeInFlightRef = useRef<Promise<File | null> | null>(null);
+  const bakeGenRef = useRef(0);
+  const bakeTimeoutRef = useRef<number | null>(null);
+
+  const runCompositeBake = useCallback((): Promise<File | null> => {
+    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<File | null> => {
+    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,

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<void> => {
+  const worker = ensureWorker();
+  const id = nextId++;
+  const transfer: Transferable[] = baseBitmap ? [baseBitmap] : [];
+  return new Promise<void>((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<Blob> => {
   const worker = ensureWorker();
   const id = nextId++;
-  const transfer: Transferable[] = [params.markerBitmap];
-  if (params.baseBitmap) transfer.push(params.baseBitmap);
   return new Promise<Blob>((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,
+    ]);
   });
 };
 

frontend/libs/components/pagedraw/src/lib/utilities/generatePipeline.worker.ts

@@ -1,11 +1,15 @@
 /// <reference lib="webworker" />
 
 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<WorkerRequest>) => {
   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);