From af02cc6525abffe0b9fe562be6a4a66e535cb5d1 Mon Sep 17 00:00:00 2001
From: bflatastic <bombayapirat@gmail.com>
Date: Tue, 29 Jul 2025 22:50:41 +0700
Subject: [PATCH] number of generation dropdown

---
 .../components/popover/src/lib/popover.tsx    |  2 +-
 .../promptbox/src/lib/PromptBoxImage.tsx      | 77 ++++++++++++++++++-
 .../promptbox/src/lib/promptStore.ts          |  4 +
 3 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/frontend/libs/components/popover/src/lib/popover.tsx b/frontend/libs/components/popover/src/lib/popover.tsx
index 7c2781bd10..bb66d9f894 100644
--- a/frontend/libs/components/popover/src/lib/popover.tsx
+++ b/frontend/libs/components/popover/src/lib/popover.tsx
@@ -369,7 +369,7 @@ export const PopoverMenu = ({
                         <div key={index}>
                           <Button
                             className={twMerge(
-                              "flex w-full items-center justify-between border-transparent bg-transparent px-1.5",
+                              "flex w-full items-center shadow-none justify-between border-transparent bg-transparent px-1.5",
                               "hover:bg-[#63636B]/60",
                               mode === "toggle" && item.selected
                                 ? "hover:bg-[#63636B]"
diff --git a/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx b/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx
index fceec50c5f..4e7ac4ce04 100644
--- a/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx
+++ b/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx
@@ -18,6 +18,7 @@ import {
   faTimes,
   faPlus,
   faImages,
+  faCopy,
 } from "@fortawesome/pro-solid-svg-icons";
 import {
   faRectangle,
@@ -72,6 +73,8 @@ export const PromptBoxImage = ({
   const setUseSystemPrompt = usePromptImageStore((s) => s.setUseSystemPrompt);
   const aspectRatio = usePromptImageStore((s) => s.aspectRatio);
   const setAspectRatio = usePromptImageStore((s) => s.setAspectRatio);
+  const generationCount = usePromptImageStore((s) => s.generationCount);
+  const setGenerationCount = usePromptImageStore((s) => s.setGenerationCount);
   const [isEnqueueing, setIsEnqueueing] = useState(false);
   const [selectedGalleryImages, setSelectedGalleryImages] = useState<string[]>(
     []
@@ -98,6 +101,30 @@ export const PromptBoxImage = ({
       icon: <FontAwesomeIcon icon={faSquare} className="h-4 w-4" />,
     },
   ]);
+  const [generationCountList, setGenerationCountList] = useState<PopoverItem[]>(
+    [
+      {
+        label: "1",
+        selected: generationCount === 1,
+        icon: <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />,
+      },
+      {
+        label: "2",
+        selected: generationCount === 2,
+        icon: <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />,
+      },
+      {
+        label: "3",
+        selected: generationCount === 3,
+        icon: <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />,
+      },
+      {
+        label: "4",
+        selected: generationCount === 4,
+        icon: <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />,
+      },
+    ]
+  );
 
   const fileInputRef = useRef<HTMLInputElement>(null);
   const textareaRef = useRef<HTMLTextAreaElement>(null);
@@ -121,6 +148,24 @@ export const PromptBoxImage = ({
     }
   }, [imageMediaId, url]);
 
+  useEffect(() => {
+    setAspectRatioList((prev) =>
+      prev.map((item) => ({
+        ...item,
+        selected: item.label === aspectRatio,
+      }))
+    );
+  }, [aspectRatio]);
+
+  useEffect(() => {
+    setGenerationCountList((prev) =>
+      prev.map((item) => ({
+        ...item,
+        selected: item.label === generationCount.toString(),
+      }))
+    );
+  }, [generationCount]);
+
   const handleAspectRatioSelect = (selectedItem: PopoverItem) => {
     setAspectRatio(selectedItem.label as any);
     setAspectRatioList((prev) =>
@@ -131,6 +176,17 @@ export const PromptBoxImage = ({
     );
   };
 
+  const handleGenerationCountSelect = (selectedItem: PopoverItem) => {
+    const count = parseInt(selectedItem.label, 10);
+    setGenerationCount(count);
+    setGenerationCountList((prev) =>
+      prev.map((item) => ({
+        ...item,
+        selected: item.label === selectedItem.label,
+      }))
+    );
+  };
+
   const handleRemoveReference = (id: string) => {
     setReferenceImages(referenceImages.filter((img) => img.id !== id));
     if (fileInputRef.current) {
@@ -197,8 +253,6 @@ export const PromptBoxImage = ({
   const handleEnqueue = async () => {
     setIsEnqueueing(true);
 
-    console.log("PromptBoxImage - Prompting with model", modelInfo);
-
     gtagEvent("enqueue_image");
 
     setTimeout(() => {
@@ -214,6 +268,7 @@ export const PromptBoxImage = ({
       prompt: prompt,
       model: modelInfo,
       aspect_ratio: aspectRatio,
+      //generation_count: generationCount, TODO: Backend doesn't support generation count yet?
     });
 
     console.log("PromptBoxImage - generateResponse", generateResponse);
@@ -382,6 +437,24 @@ 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>
               <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}
diff --git a/frontend/libs/components/promptbox/src/lib/promptStore.ts b/frontend/libs/components/promptbox/src/lib/promptStore.ts
index 3a7436e83c..da66193c13 100644
--- a/frontend/libs/components/promptbox/src/lib/promptStore.ts
+++ b/frontend/libs/components/promptbox/src/lib/promptStore.ts
@@ -59,10 +59,12 @@ interface PromptImageStore {
   aspectRatio: AspectRatio;
   useSystemPrompt: boolean;
   referenceImages: RefImage[];
+  generationCount: number;
   setPrompt: (prompt: string) => void;
   setAspectRatio: (ratio: AspectRatio) => void;
   setUseSystemPrompt: (value: boolean) => void;
   setReferenceImages: (images: RefImage[]) => void;
+  setGenerationCount: (count: number) => void;
 }
 
 export const usePromptImageStore = create<PromptImageStore>()((set) => ({
@@ -70,10 +72,12 @@ export const usePromptImageStore = create<PromptImageStore>()((set) => ({
   aspectRatio: "3:2",
   useSystemPrompt: true,
   referenceImages: [],
+  generationCount: 1,
   setPrompt: (prompt) => set({ prompt }),
   setAspectRatio: (aspectRatio) => set({ aspectRatio }),
   setUseSystemPrompt: (useSystemPrompt) => set({ useSystemPrompt }),
   setReferenceImages: (referenceImages) => set({ referenceImages }),
+  setGenerationCount: (generationCount) => set({ generationCount }),
 }));
 
 // ----- Video Prompt Box Store -----