number of generations dropdown inpainting

frontend/apps/editor3d/app/src/pages/PageEdit/PageEdit.tsx

@@ -6,22 +6,22 @@ import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import { faClock, faImage } from "@fortawesome/pro-solid-svg-icons";
 import Konva from "konva"; // just for types
 
-import { setCanvasRenderBitmap } from "../../signals/canvasRenderBitmap"
+// import { setCanvasRenderBitmap } from "../../signals/canvasRenderBitmap"
 import {
   EnqueueImageInpaint,
   EnqueueImageInpaintModel,
 } from "@storyteller/tauri-api";
 import { ContextMenuContainer } from "../PageDraw/components/ui/ContextMenu";
-import SideToolbar from "../PageDraw/components/ui/SideToolbar";
+// import SideToolbar from "../PageDraw/components/ui/SideToolbar";
 import { useCopyPasteHotkeys } from "../PageDraw/hooks/useCopyPasteHotkeys";
 import { useDeleteHotkeys } from "../PageDraw/hooks/useDeleteHotkeys";
 import { useUndoRedoHotkeys } from "../PageDraw/hooks/useUndoRedoHotkeys";
-import { captureStageImageBitmap } from "../PageDraw/hooks/useUpdateSnapshot";
+// import { captureStageImageBitmap } from "../PageDraw/hooks/useUpdateSnapshot";
 import PromptEditor from "./PromptEditor/PromptEditor";
-import { AspectRatioType } from "../PageDraw/stores/SceneState";
+// import { AspectRatioType } from "../PageDraw/stores/SceneState";
 import { ActiveEditTool, useEditStore } from "./stores/EditState";
 import { EditPaintSurface } from "./EditPaintSurface";
-import { normalizeCanvas } from "~/Helpers/CanvasHelpers";
+import { normalizeCanvas } from "../../Helpers/CanvasHelpers";
 import { BaseImageSelector, BaseSelectorImage } from "./BaseImageSelector";
 import DrawToolControlBar from "./DrawToolControlBar";
 
@@ -40,6 +40,7 @@ const PageEdit = () => {
   const baseImageKonvaRef = useRef<Konva.Image>({} as Konva.Image);
   const transformerRefs = useRef<{ [key: string]: Konva.Transformer }>({});
   const [isEnqueuing, setIsEnqueuing] = useState<boolean>(false);
+  const [generationCount, setGenerationCount] = useState<number>(1);
 
   // Use the Zustand store
   const store = useEditStore();
@@ -150,13 +151,17 @@ const PageEdit = () => {
     // Calculate position to center canvas in container
     stage.position({
       x: (containerWidth - canvasW) / 2,
-      y: (containerHeight - canvasH) / 2
+      y: (containerHeight - canvasH) / 2,
     });
-  }
+  };
 
   // Create a function to use the left layer ref and download the bitmap from it
   const getMaskArrayBuffer = async (): Promise<Uint8Array> => {
-    if (!stageRef.current || !leftPanelRef.current || !baseImageKonvaRef.current) {
+    if (
+      !stageRef.current ||
+      !leftPanelRef.current ||
+      !baseImageKonvaRef.current
+    ) {
       console.error("Stage or left panel ref is not available");
       throw new Error("Stage or left panel or base image ref is not available");
     }
@@ -175,9 +180,13 @@ const PageEdit = () => {
 
     // Using the pixelRatio scaling may result in off-by-one rounding errors,
     // So we re-fit the image to a canvas of precise size.
-    const fittedCanvas = normalizeCanvas(layerCrop, rect.width(), rect.height());
+    const fittedCanvas = normalizeCanvas(
+      layerCrop,
+      rect.width(),
+      rect.height(),
+    );
 
-    const blob = await fittedCanvas.convertToBlob({ type: 'image/png' });
+    const blob = await fittedCanvas.convertToBlob({ type: "image/png" });
     const arrayBuffer = await blob.arrayBuffer();
 
     return new Uint8Array(arrayBuffer);
@@ -201,12 +210,12 @@ const PageEdit = () => {
     // TODO: Call inference API here
     const arrayBuffer = await getMaskArrayBuffer();
 
-    const response = await EnqueueImageInpaint({
+    await EnqueueImageInpaint({
       model: EnqueueImageInpaintModel.FluxPro1,
       image_media_token: editedImageToken,
       mask_image_raw_bytes: arrayBuffer,
       prompt: prompt,
-      image_count: 1,
+      image_count: generationCount,
     });
 
     setIsEnqueuing(false);
@@ -216,22 +225,29 @@ const PageEdit = () => {
   // Also show loading state if info is set but image is loading
   if (!store.baseImageInfo || !store.baseImageBitmap) {
     return (
-      <div className={"h-screen w-full flex items-center justify-center bg-ui-panel"}>
+      <div
+        className={
+          "flex h-screen w-full items-center justify-center bg-ui-panel"
+        }
+      >
         <BaseImageSelector
           onImageSelect={(image: BaseSelectorImage) => {
             store.setBaseImageInfo(image);
           }}
-          showLoading={store.baseImageInfo !== null && store.baseImageInfo === null}
+          showLoading={
+            store.baseImageInfo !== null && store.baseImageInfo === null
+          }
         />
       </div>
-    )
+    );
   }
 
   return (
     <>
       <div
-        className={`preserve-aspect-ratio fixed top-0 left-1/2 z-10 -translate-x-1/2 transform ${isSelecting ? "pointer-events-none" : "pointer-events-auto"
-          }`}
+        className={`preserve-aspect-ratio fixed left-1/2 top-0 z-10 -translate-x-1/2 transform ${
+          isSelecting ? "pointer-events-none" : "pointer-events-auto"
+        }`}
         style={{ display: store.activeTool === "edit" ? "block" : "none" }}
       >
         <DrawToolControlBar
@@ -242,14 +258,19 @@ const PageEdit = () => {
         />
       </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"
-          }`}
+        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
-          onModeChange={(mode: string) => { store.setActiveTool(mode as ActiveEditTool) }}
+          onModeChange={(mode: string) => {
+            store.setActiveTool(mode as ActiveEditTool);
+          }}
           selectedMode={store.activeTool}
           onGenerateClick={handleGenerate}
           isDisabled={isEnqueuing}
+          generationCount={generationCount}
+          onGenerationCountChange={setGenerationCount}
         />
       </div>
       <div className="relative z-0">
@@ -260,30 +281,30 @@ const PageEdit = () => {
               if (hasSelection) {
                 console.log("An item is selected.");
                 // You can add additional actions here based on the selection
-                return true
+                return true;
               } else {
                 console.log("No item is selected.");
-                return false
+                return false;
               }
             }
-            return false
+            return false;
           }}
           onMenuAction={async (action) => {
             switch (action) {
-              case 'DUPLICATE':
-                store.copySelectedItems()
-                store.pasteItems()
+              case "DUPLICATE":
+                store.copySelectedItems();
+                store.pasteItems();
                 break;
-              case 'DELETE':
-                store.deleteSelectedItems()
+              case "DELETE":
+                store.deleteSelectedItems();
                 break;
               default:
               // No action needed for unhandled cases
             }
           }}
-          isLocked={store.selectedNodeIds.some(id => {
-            const node = store.nodes.find(n => n.id === id);
-            const lineNode = store.lineNodes.find(n => n.id === id);
+          isLocked={store.selectedNodeIds.some((id) => {
+            const node = store.nodes.find((n) => n.id === id);
+            const lineNode = store.lineNodes.find((n) => n.id === id);
             return (node?.locked || lineNode?.locked) ?? false;
           })}
         >
@@ -318,7 +339,7 @@ const PageEdit = () => {
           triggerLabel="Model"
         />
         <button
-          className="bg-transparent p-2 text-lg text-white hover:text-white hover:bg-white/20 rounded transition"
+          className="rounded bg-transparent p-2 text-lg text-white transition hover:bg-white/20 hover:text-white"
           onClick={onFitPressed}
         >
           Fit

frontend/apps/editor3d/app/src/pages/PageEdit/PromptEditor/PromptEditor.tsx

@@ -5,7 +5,8 @@ const PromptEditor: React.FC<PromptBoxEditProps> = ({
   onModeChange,
   selectedMode,
   onGenerateClick,
-  isDisabled
+  isDisabled,
+  ...rest
 }) => {
   return (
     <div className="mx-auto flex w-full max-w-3xl flex-col space-y-2">
@@ -14,6 +15,7 @@ const PromptEditor: React.FC<PromptBoxEditProps> = ({
         selectedMode={selectedMode}
         onGenerateClick={onGenerateClick}
         isDisabled={isDisabled}
+        {...rest}
       />
     </div>
   );

frontend/libs/components/promptbox/src/lib/PromptBoxEdit.tsx

@@ -6,18 +6,24 @@ import {
   faMousePointer,
   faSparkles,
   faSpinnerThird,
+  faCopy,
 } from "@fortawesome/pro-solid-svg-icons";
 import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import { Button, ToggleButton } from "@storyteller/ui-button";
 import { ButtonIconSelect } from "@storyteller/ui-button-icon-select";
+import { PopoverMenu, PopoverItem } from "@storyteller/ui-popover";
 import { Tooltip } from "@storyteller/ui-tooltip";
 import { useEffect, useRef, useState } from "react";
+import { ModelInfo, getCapabilitiesForModel } from "@storyteller/model-list";
 
 export interface PromptBoxEditProps {
   onModeChange?: (mode: string) => void;
   selectedMode?: string;
   onGenerateClick: (prompt: string) => void;
   isDisabled?: boolean;
+  modelInfo?: ModelInfo;
+  generationCount?: number;
+  onGenerationCountChange?: (count: number) => void;
 }
 
 export const PromptBoxEdit = ({
@@ -25,9 +31,18 @@ export const PromptBoxEdit = ({
   selectedMode,
   onGenerateClick,
   isDisabled,
+  modelInfo,
+  generationCount: generationCountProp,
+  onGenerationCountChange,
 }: PromptBoxEditProps) => {
   const [prompt, setPrompt] = useState("");
   const [useSystemPrompt, setUseSystemPrompt] = useState(true);
+  const [generationCount, setGenerationCount] = useState<number>(
+    typeof generationCountProp === "number" ? generationCountProp : 1
+  );
+  const [generationCountList, setGenerationCountList] = useState<PopoverItem[]>(
+    []
+  );
 
   const textareaRef = useRef<HTMLTextAreaElement>(null);
 
@@ -38,6 +53,48 @@ export const PromptBoxEdit = ({
     }
   });
 
+  // Sync internal state if a controlled prop is provided
+  useEffect(() => {
+    if (typeof generationCountProp === "number") {
+      setGenerationCount(generationCountProp);
+    }
+  }, [generationCountProp]);
+
+  // Build generation count options based on selected model
+  useEffect(() => {
+    const caps = getCapabilitiesForModel(modelInfo);
+    const items: PopoverItem[] = Array.from(
+      { length: caps.maxGenerationCount },
+      (_, i) => i + 1
+    ).map((count) => ({
+      label: String(count),
+      selected: count === generationCount,
+      icon: <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />,
+    }));
+    setGenerationCountList(items);
+
+    // Clamp selection
+    if (generationCount < 1 || generationCount > caps.maxGenerationCount) {
+      const clamped = Math.min(
+        Math.max(1, generationCount),
+        caps.maxGenerationCount
+      );
+      setGenerationCount(clamped);
+      onGenerationCountChange?.(clamped);
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [modelInfo]);
+
+  // Keep UI selection in sync when store/value changes
+  useEffect(() => {
+    setGenerationCountList((prev) =>
+      prev.map((item) => ({
+        ...item,
+        selected: item.label === String(generationCount),
+      }))
+    );
+  }, [generationCount]);
+
   const handlePaste = (e: React.ClipboardEvent<HTMLTextAreaElement>) => {
     e.preventDefault();
     e.stopPropagation();
@@ -69,6 +126,18 @@ export const PromptBoxEdit = ({
     }
   };
 
+  const handleGenerationCountSelect = (selectedItem: PopoverItem) => {
+    const count = parseInt(selectedItem.label, 10);
+    setGenerationCount(count);
+    onGenerationCountChange?.(count);
+    setGenerationCountList((prev) =>
+      prev.map((item) => ({
+        ...item,
+        selected: item.label === selectedItem.label,
+      }))
+    );
+  };
+
   const modes = [
     {
       value: "select",
@@ -134,6 +203,24 @@ export const PromptBoxEdit = ({
               </Tooltip>
             </div>
             <div className="flex items-center gap-2">
+              <Tooltip
+                content="Number of generations"
+                position="top"
+                className="z-50"
+                closeOnClick={true}
+              >
+                <PopoverMenu
+                  items={generationCountList}
+                  onSelect={handleGenerationCountSelect}
+                  mode="toggle"
+                  triggerIcon={
+                    <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />
+                  }
+                  panelClassName="min-w-28"
+                  panelTitle="No. of images"
+                  buttonClassName="h-9"
+                />
+              </Tooltip>
               <Button
                 className="flex items-center border-none bg-primary px-3 text-sm text-white disabled:cursor-not-allowed disabled:opacity-50"
                 icon={!isDisabled ? faSparkles : undefined}