generation count update

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

@@ -134,40 +134,13 @@ export const PromptBox2D = ({
     },
   ]);
 
-  const [generationCountList, setGenerationCountList] = useState<PopoverItem[]>(
-    [],
-  );
-
-  // Update generation count options from selected model capabilities
   useEffect(() => {
-    const caps = getCapabilitiesForModel(selectedImageModel);
-    const items: PopoverItem[] = Array.from(
-      { length: caps.maxGenerationCount },
-      (_, i) => i + 1,
-    ).map((count) => ({
-      label: String(count),
-      selected: count === generationCount,
-    }));
-    setGenerationCountList(items);
-
-    // Clamp selection to allowed range
-    if (generationCount < 1 || generationCount > caps.maxGenerationCount) {
-      setGenerationCount(
-        Math.min(Math.max(1, generationCount), caps.maxGenerationCount),
-      );
+    if (selectedImageModel?.isValidGenerationCount(generationCount)) {
+      return;
     }
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [selectedImageModel]);
-
-  // Keep UI selection in sync when store value changes
-  useEffect(() => {
-    setGenerationCountList((prev) =>
-      prev.map((item) => ({
-        ...item,
-        selected: item.label === String(generationCount),
-      })),
-    );
-  }, [generationCount]);
+    const defaultGenerations = selectedImageModel?.defaultGenerationCount || 4;
+    setGenerationCount(defaultGenerations);
+  }, [selectedImageModel, generationCount]);
 
   useEffect(() => {
     setResolutionList((prev) =>
@@ -209,17 +182,6 @@ export const PromptBox2D = ({
 
   // Image prompt row replaces legacy upload/gallery UI
 
-  const handleGenerationCountSelect = (selectedItem: PopoverItem) => {
-    const count = parseInt(selectedItem.label, 10);
-    setGenerationCount(count);
-    setGenerationCountList((prev) =>
-      prev.map((item) => ({
-        ...item,
-        selected: item.label === selectedItem.label,
-      })),
-    );
-  };
-
   const handlePaste = (e: React.ClipboardEvent<HTMLTextAreaElement>) => {
     e.preventDefault();
     e.stopPropagation();
@@ -613,23 +575,13 @@ export const PromptBox2D = ({
                   </Button>
                 </Tooltip>
               )}
-              <Tooltip
-                content="Number of generations"
-                position="top"
-                className="z-50"
-                closeOnClick={true}
-              >
-                <PopoverMenu
-                  items={generationCountList}
-                  onSelect={handleGenerationCountSelect}
-                  mode="toggle"
-                  panelTitle="No. of images"
-                  triggerIcon={
-                    <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />
-                  }
-                  buttonClassName="h-9"
-                />
-              </Tooltip>
+              <GenerationCountPicker
+                currentModel={selectedImageModel}
+                currentCount={generationCount}
+                handleCountChange={(count) => {
+                  setGenerationCount(count);
+                }}
+              />
               <Button
                 className="flex items-center border-none bg-primary px-3 text-sm text-white disabled:cursor-not-allowed disabled:opacity-50"
                 icon={!isEnqueueing ? faSparkles : undefined}

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

@@ -31,6 +31,7 @@ import { ImagePromptRow, type UploadImageFn } from "./ImagePromptRow";
 import { RefImage, usePromptEditStore } from "./promptStore";
 import { GenerationProvider } from "@storyteller/api-enums";
 import { AspectRatioPicker } from "./common/AspectRatioPicker";
+import { GenerationCountPicker } from "./common/GenerationCountPicker";
 
 export interface PromptBoxEditProps {
   onModeChange?: (mode: string) => void;
@@ -82,9 +83,6 @@ export const PromptBoxEdit = ({
   const [prompt, setPrompt] = useState("");
   const [useSystemPrompt, setUseSystemPrompt] = useState(true);
   const [isFocused, setIsFocused] = useState(false);
-  const [generationCount, setGenerationCount] = useState<number>(
-    typeof generationCountProp === "number" ? generationCountProp : 1,
-  );
   const [internalEnqueueing, setInternalEnqueueing] = useState(false);
   const referenceImages = usePromptEditStore((s) => s.referenceImages);
   const setReferenceImages = usePromptEditStore((s) => s.setReferenceImages);
@@ -95,10 +93,6 @@ export const PromptBoxEdit = ({
 
   const [showImagePrompts, setShowImagePrompts] = useState(false);
 
-  const [generationCountList, setGenerationCountList] = useState<PopoverItem[]>(
-    [],
-  );
-
   const [aspectRatioList, setAspectRatioList] = useState<PopoverItem[]>([
     {
       label: "Auto",
@@ -142,6 +136,8 @@ export const PromptBoxEdit = ({
     },
   ]);
 
+  const generationCount = generationCountProp || 1;
+
   useEffect(() => {
     setAspectRatioList((prev) =>
       prev.map((item) => ({
@@ -193,42 +189,12 @@ export const PromptBoxEdit = ({
   });
 
   useEffect(() => {
-    if (typeof generationCountProp === "number") {
-      setGenerationCount(generationCountProp);
+    if (selectedImageModel?.isValidGenerationCount(generationCount)) {
+      return;
     }
-  }, [generationCountProp]);
-
-  useEffect(() => {
-    const caps = getCapabilitiesForModel(selectedImageModel);
-    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);
-
-    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
-  }, [selectedImageModel]);
-
-  useEffect(() => {
-    setGenerationCountList((prev) =>
-      prev.map((item) => ({
-        ...item,
-        selected: item.label === String(generationCount),
-      })),
-    );
-  }, [generationCount]);
+    const defaultGenerations = selectedImageModel?.defaultGenerationCount || 4;
+    onGenerationCountChange?.(defaultGenerations);
+  }, [selectedImageModel, generationCount]);
 
   const handlePaste = (e: React.ClipboardEvent<HTMLTextAreaElement>) => {
     e.preventDefault();
@@ -532,24 +498,13 @@ export const PromptBoxEdit = ({
                     </Button>
                   </Tooltip>
                 )}
-                <Tooltip
-                  content="Number of generations"
-                  position="top"
-                  className="z-50"
-                  closeOnClick={true}
-                >
-                  <PopoverMenu
-                    items={generationCountList}
-                    onSelect={handleGenerationCountSelect}
-                    mode="toggle"
-                    panelTitle="No. of images"
-                    triggerIcon={
-                      <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />
-                    }
-                    panelClassName="min-w-28"
-                    buttonClassName="h-9"
-                  />
-                </Tooltip>
+                <GenerationCountPicker
+                  currentModel={selectedImageModel}
+                  currentCount={generationCount}
+                  handleCountChange={(count) => {
+                    onGenerationCountChange?.(count);
+                  }}
+                />
                 <Button
                   className="flex items-center border-none bg-primary px-3 text-sm text-white disabled:cursor-not-allowed disabled:opacity-50"
                   icon={

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

@@ -177,9 +177,7 @@ export const PromptBoxImage = ({
     if (selectedModel?.isValidGenerationCount(generationCount)) {
       return;
     }
-
     const defaultGenerations = selectedModel?.defaultGenerationCount || 4;
-
     setGenerationCount(defaultGenerations);
   }, [selectedModel, generationCount]);