From a6d6301b54e64ff72cf8bb4c2d23e21fb9ba3401 Mon Sep 17 00:00:00 2001
From: Tanish Manku <tanish2k09.dev@gmail.com>
Date: Wed, 7 Jan 2026 18:35:59 -0800
Subject: [PATCH] Support HistoryStack in 2D

---
 .../app/src/pages/PageDraw/PageDraw.tsx       | 58 +++++++++++++-
 .../src/pages/PageDraw/stores/SceneState.ts   | 77 +++++++++++++++++++
 2 files changed, 131 insertions(+), 4 deletions(-)

diff --git a/frontend/apps/artcraft/app/src/pages/PageDraw/PageDraw.tsx b/frontend/apps/artcraft/app/src/pages/PageDraw/PageDraw.tsx
index f69fc341dc..ea0bb01916 100644
--- a/frontend/apps/artcraft/app/src/pages/PageDraw/PageDraw.tsx
+++ b/frontend/apps/artcraft/app/src/pages/PageDraw/PageDraw.tsx
@@ -24,6 +24,7 @@ import {
 import { useCanvasBgRemovedEvent } from "@storyteller/tauri-api";
 import { HelpMenuButton } from "@storyteller/ui-help-menu";
 import { GenerationProvider } from "@storyteller/api-enums";
+import { HistoryStack } from "../PageEdit/HistoryStack";
 
 const PAGE_ID: ModelPage = ModelPage.Canvas2D;
 
@@ -59,13 +60,27 @@ const PageDraw = () => {
   const stageRef = useRef<Konva.Stage>({} as Konva.Stage);
   const transformerRefs = useRef<{ [key: string]: Konva.Transformer }>({});
   const store = useSceneStore();
+  const baseImageUrl = store.baseImageInfo?.url;
+  const [pendingGenerations, setPendingGenerations] = useState<
+    { id: string; count: number }[]
+  >([]);
+  const addHistoryImageBundle = useSceneStore(
+    (state) => state.addHistoryImageBundle,
+  );
+  const removeHistoryImage = useSceneStore((state) => state.removeHistoryImage);
+  const historyImageBundles = useSceneStore(
+    (state) => state.historyImageBundles,
+  );
 
   const selectedImageModel: ImageModel | undefined =
     useSelectedImageModel(PAGE_ID);
 
-  const selectedProvider : GenerationProvider | undefined = 
+  const selectedProvider: GenerationProvider | undefined =
     useSelectedProviderForModel(PAGE_ID, selectedImageModel?.id);
 
+  const supportsMaskedInpainting =
+    selectedImageModel?.usesInpaintingMask ?? false;
+
   useDeleteHotkeys({ onDelete: store.deleteSelectedItems });
   useUndoRedoHotkeys({ undo: store.undo, redo: store.redo });
   useCopyPasteHotkeys({
@@ -270,9 +285,44 @@ const PageDraw = () => {
     <>
       <div className="fixed inset-0 -z-10 bg-ui-background" />
       <div
-        className={`preserve-aspect-ratio fixed bottom-0 left-1/2 z-10 -translate-x-1/2 transform ${
-          isSelecting ? "pointer-events-none" : "pointer-events-auto"
-        }`}
+        className={`preserve-aspect-ratio fixed right-4 top-1/2 z-10 -translate-y-1/2 transform ${isSelecting ? "pointer-events-none" : "pointer-events-auto"
+          }`}
+      >
+        <HistoryStack
+          onClear={() => {
+            store.RESET();
+            setPendingGenerations([]);
+          }}
+          imageBundles={historyImageBundles}
+          pendingPlaceholders={pendingGenerations}
+          blurredBackgroundUrl={baseImageUrl}
+          onImageSelect={(baseImage) => {
+            store.clearLineNodes();
+            store.setBaseImageInfo(baseImage);
+          }}
+          onImageRemove={(baseImage) => {
+            if (
+              pendingGenerations.length === 0 &&
+              store.historyImageBundles.length === 1 &&
+              store.historyImageBundles[0].images.length <= 1
+            ) {
+              store.RESET();
+            } else {
+              removeHistoryImage(baseImage);
+            }
+          }}
+          onNewImageBundle={(newBundle: ImageBundle) => {
+            addHistoryImageBundle(newBundle);
+          }}
+          onResolvePending={(id: string) =>
+            setPendingGenerations((prev) => prev.filter((p) => p.id !== id))
+          }
+          selectedImageToken={store.baseImageInfo?.mediaToken}
+        />
+      </div>
+      <div
+        className={`preserve-aspect-ratio fixed bottom-0 left-1/2 z-10 -translate-x-1/2 transform ${isSelecting ? "pointer-events-none" : "pointer-events-auto"
+          }`}
       >
         <PromptEditor
           initialPrompt=""
diff --git a/frontend/apps/artcraft/app/src/pages/PageDraw/stores/SceneState.ts b/frontend/apps/artcraft/app/src/pages/PageDraw/stores/SceneState.ts
index ef97308a6e..39c7a0a50a 100644
--- a/frontend/apps/artcraft/app/src/pages/PageDraw/stores/SceneState.ts
+++ b/frontend/apps/artcraft/app/src/pages/PageDraw/stores/SceneState.ts
@@ -1,6 +1,8 @@
 import { create } from "zustand";
 import { Node, NodeType } from "../Node";
 import { EnqueueImageBgRemoval } from "libs/tauri-api/src/lib/enqueue/EnqueueImageBgRemovalCommand";
+import { ImageBundle } from "~/pages/PageEdit/HistoryStack";
+import { BaseSelectorImage } from "~/pages/PageEdit/BaseImageSelector";
 
 // Add LineNode type
 export type LineNode = {
@@ -175,6 +177,7 @@ interface SceneState {
     dy: number,
     shouldSaveState?: boolean,
   ) => void;
+  clearLineNodes: () => void;
 
   // Add a specific method for file uploads
   createImageFromFile: (
@@ -253,6 +256,17 @@ interface SceneState {
   // Add aspect ratio action
   setAspectRatioType: (type: AspectRatioType) => void;
   getAspectRatioDimensions: () => { width: number; height: number };
+
+  // For the history stack
+  historyImageBundles: ImageBundle[];
+  clearHistoryImages: () => void;
+  addHistoryImageBundle: (bundle: ImageBundle) => void;
+  removeHistoryImage: (image: BaseSelectorImage) => void;
+
+  // Base image state
+  baseImageInfo: BaseSelectorImage | null;
+  baseImageBitmap: HTMLImageElement | null;
+  setBaseImageInfo: (image: BaseSelectorImage | null) => void;
 }
 
 export const generateId = (): string => {
@@ -293,6 +307,13 @@ export const useSceneStore = create<SceneState>((set, get, store) => ({
   cursorPosition: null,
   cursorVisible: false,
 
+  // Base image initial state
+  baseImageInfo: null,
+  baseImageBitmap: null,
+
+  // History stack state
+  historyImageBundles: [],
+
   // Add initial aspect ratio state
   aspectRatioType: AspectRatioType.NONE,
 
@@ -1554,4 +1575,60 @@ export const useSceneStore = create<SceneState>((set, get, store) => ({
         return { width: 1024, height: 683 }; // Default to landscape
     }
   },
+
+  setBaseImageInfo: (image: BaseSelectorImage | null) => {
+    set({ baseImageInfo: image });
+
+    if (!image) {
+      return;
+    }
+
+    // Load the image bitmap as well
+    const imgBitmap = new Image();
+    imgBitmap.onload = () => {
+      set({ baseImageBitmap: imgBitmap });
+    };
+    imgBitmap.onerror = (event) => {
+      console.error("Failed to load base image, discarding", event);
+      set({ baseImageInfo: null, baseImageBitmap: null });
+      imgBitmap.onload = null;
+      imgBitmap.onerror = null;
+    };
+    imgBitmap.crossOrigin = "anonymous";
+    const isDataUrl =
+      typeof image.url === "string" && image.url.startsWith("data:");
+    imgBitmap.src = isDataUrl ? image.url : image.url + "?basecanvasimg";
+  },
+
+  clearLineNodes() {
+    set({ lineNodes: [] });
+  },
+
+  clearHistoryImages: () => {
+    set({ historyImageBundles: [] });
+  },
+
+  addHistoryImageBundle(bundle) {
+    set((state) => ({
+      historyImageBundles: [...state.historyImageBundles, bundle],
+    }));
+  },
+
+  removeHistoryImage(image) {
+    console.log("Removing history image:", image);
+    set((state) => {
+      const updatedBundles = state.historyImageBundles
+        .map((bundle) => {
+          return {
+            ...bundle,
+            images: bundle.images.filter(
+              (img) => img.mediaToken !== image.mediaToken,
+            ),
+          };
+        })
+        .filter((bundle) => bundle.images.length > 0); // Remove empty bundles
+      console.log("Updated history image bundles:", updatedBundles);
+      return { historyImageBundles: updatedBundles };
+    });
+  },
 }));