generation count update

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

@@ -41,6 +41,7 @@ import { ImageModel } from "@storyteller/model-list";
 import { ImagePromptRow } from "./ImagePromptRow";
 import { twMerge } from "tailwind-merge";
 import { GenerationProvider } from "@storyteller/api-enums";
+import { GenerationCountPicker } from "./common/GenerationCountPicker";
 
 export type AspectRatio = "wide" | "tall" | "square";
 
@@ -133,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) =>
@@ -208,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();
@@ -612,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

@@ -34,6 +34,7 @@ import { twMerge } from "tailwind-merge";
 import { ImagePromptRow } from "./ImagePromptRow";
 import { GenerationProvider } from "@storyteller/api-enums";
 import { AspectRatioPicker } from "./common/AspectRatioPicker";
+import { GenerationCountPicker } from "./common/GenerationCountPicker";
 
 interface PromptBoxImageProps {
   useJobContext: () => JobContextType;
@@ -150,9 +151,6 @@ export const PromptBoxImage = ({
       icon: <FontAwesomeIcon icon={faExpand} className="h-4 w-4" />,
     },
   ]);
-  const [generationCountList, setGenerationCountList] = useState<PopoverItem[]>(
-    [],
-  );
 
   const textareaRef = useRef<HTMLTextAreaElement>(null);
 
@@ -176,33 +174,12 @@ export const PromptBoxImage = ({
   }, [imageMediaId, url]);
 
   useEffect(() => {
-    const caps = getCapabilitiesForModel(selectedModel);
-    const defaultCount = Math.min(
-      Math.max(1, caps.defaultGenerationCount ?? 1),
-      caps.maxGenerationCount,
-    );
-
-    setGenerationCount(defaultCount);
-
-    const items: PopoverItem[] = Array.from(
-      { length: caps.maxGenerationCount },
-      (_, i) => i + 1,
-    ).map((count) => ({
-      label: String(count),
-      selected: count === defaultCount,
-      icon: <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />,
-    }));
-    setGenerationCountList(items);
-  }, [selectedModel]);
-
-  useEffect(() => {
-    setGenerationCountList((prev) =>
-      prev.map((item) => ({
-        ...item,
-        selected: item.label === String(generationCount),
-      })),
-    );
-  }, [generationCount]);
+    if (selectedModel?.isValidGenerationCount(generationCount)) {
+      return;
+    }
+    const defaultGenerations = selectedModel?.defaultGenerationCount || 4;
+    setGenerationCount(defaultGenerations);
+  }, [selectedModel, generationCount]);
 
   useEffect(() => {
     setAspectRatioList((prev) =>
@@ -236,17 +213,6 @@ export const PromptBoxImage = ({
     setResolution(selectedItem.label.toLowerCase() as any);
   };
 
-  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();
     const pastedText = e.clipboardData.getData("text");
@@ -544,24 +510,13 @@ export const PromptBoxImage = ({
               </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>
+              <GenerationCountPicker
+                currentModel={selectedModel}
+                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/common/GenerationCountPicker.tsx

@@ -0,0 +1,89 @@
+import { faBookCopy, faCopy } from "@fortawesome/pro-solid-svg-icons"
+import { FontAwesomeIcon } from "@fortawesome/react-fontawesome"
+import { PopoverMenu, PopoverItem } from "@storyteller/ui-popover";
+import { Tooltip } from "@storyteller/ui-tooltip";
+import { ImageModel } from "@storyteller/model-list";
+
+const DEFAULT_GENERATION_COUNT : number = 4;
+
+interface GenerationCountPickerProps {
+  currentModel?: ImageModel,
+  currentCount: number,
+  handleCountChange: (count: number) => void,
+}
+
+export const GenerationCountPicker = ({ 
+  currentModel,
+  currentCount,
+  handleCountChange
+}: GenerationCountPickerProps) => {
+
+  const maxGenerationCount = currentModel?.maxGenerationCount || DEFAULT_GENERATION_COUNT;
+  const hasPredefinedOptions = !!currentModel?.predefinedGenerationCounts;
+
+  let generationCountOptions: PopoverItem[];
+
+  // Count pickers either have a "[1,2, ... max]" set of options, or a pre-defined list of options.
+  if (hasPredefinedOptions) {
+    generationCountOptions = buildPredefinedCountOptions(currentModel?.predefinedGenerationCounts || [], currentCount);
+  } else {
+    generationCountOptions = buildSequentialCountOptions(maxGenerationCount, currentCount);
+  }
+
+  const onSelect = (item: PopoverItem) => {
+    let count = parseInt(item.label, 10);
+    if (isNaN(count)) {
+      return;
+    }
+    if (count < 1 || count > maxGenerationCount) {
+      // Clamp to valid range
+      count = Math.min(Math.max(1, count), maxGenerationCount);
+    }
+    handleCountChange(count);
+  };
+
+  return (
+    <>
+      <Tooltip
+        content="Number of generations"
+        position="top"
+        className="z-50"
+        closeOnClick={true}
+      >
+        <PopoverMenu
+          items={generationCountOptions}
+          onSelect={onSelect}
+          mode="toggle"
+          panelTitle="No. of images"
+          triggerIcon={
+            <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />
+          }
+          buttonClassName="h-9"
+        />
+      </Tooltip>
+    </>
+  )
+}
+
+const buildSequentialCountOptions = (maxCount: number, currentCount: number): PopoverItem[] => {
+  const options = [];
+  for (let i = 0; i < maxCount; i++) {
+    const count = i + 1;
+    options.push({
+      label: String(count),
+      selected: count === currentCount,
+    });
+  }
+  return options;
+}
+
+const buildPredefinedCountOptions = (options: number[], currentCount: number): PopoverItem[] => {
+  const result: PopoverItem[] = [];
+  options.forEach(option => {
+    result.push({
+      label: String(option),
+      selected: option === currentCount,
+    });
+  });
+  return result;
+};

frontend/libs/model-list/src/lib/classes/ImageModel.ts

@@ -17,6 +17,9 @@ export class ImageModel extends Model {
   // Default number of images that can be generated at once
   readonly defaultGenerationCount: number;
 
+  // If set, use this instead of `maxGenerationCount`.
+  readonly predefinedGenerationCounts?: number[];
+
   // Signals image editing models that focus on editing a single image.
   readonly canEditImages: boolean;
 
@@ -63,6 +66,7 @@ export class ImageModel extends Model {
     selectorBadges: string[];
     maxGenerationCount: number;
     defaultGenerationCount: number;
+    predefinedGenerationCounts?: number[];
     canEditImages?: boolean;
     usesInpaintingMask?: boolean;
     editingIsInpainting?: boolean;
@@ -91,6 +95,7 @@ export class ImageModel extends Model {
     super(args);
     this.maxGenerationCount = args.maxGenerationCount;
     this.defaultGenerationCount = args.defaultGenerationCount;
+    this.predefinedGenerationCounts = args.predefinedGenerationCounts;
     this.canEditImages = args.canEditImages ?? false;
     this.usesInpaintingMask = args.usesInpaintingMask ?? false;
     this.editingIsInpainting = args.editingIsInpainting ?? false;
@@ -125,4 +130,13 @@ export class ImageModel extends Model {
       && this.aspectRatios 
       && this.aspectRatios.length > 0;
   }
+
+  // Return whether the count of generations is valid for this model
+  isValidGenerationCount(count: number): boolean {
+    if (this.predefinedGenerationCounts) {
+      return this.predefinedGenerationCounts.includes(count);
+    } else {
+      return count >= 1 && count <= this.maxGenerationCount;
+    }
+  }
 }

frontend/libs/model-list/src/lib/lists/ImageModels.ts

@@ -19,8 +19,10 @@ export const IMAGE_MODELS: ImageModel[] = [
     selectorName: "Grok",
     selectorDescription: "Fast af",
     selectorBadges: ["10 sec."],
-    defaultGenerationCount: 4,
+    // NB: Grok only supports 6 images at a time
+    defaultGenerationCount: 6,
     maxGenerationCount: 6,
+    predefinedGenerationCounts: [6],
     canUseImagePrompt: false,
     maxImagePromptCount: 6,
     progressBarTime: 10000,
@@ -44,8 +46,10 @@ export const IMAGE_MODELS: ImageModel[] = [
     selectorName: "Midjourney",
     selectorDescription: "Stunning style and quality",
     selectorBadges: ["45 sec."],
+    // NB: Midjourney only supports 4 images at a time
     maxGenerationCount: 4,
     defaultGenerationCount: 4,
+    predefinedGenerationCounts: [4],
     canUseImagePrompt: false,
     maxImagePromptCount: 6,
     progressBarTime: 45000,