From f91f1965a49f9cdc9ff7a0aa91b5b4adae8939fc Mon Sep 17 00:00:00 2001
From: BFlat <bombayapirat@gmail.com>
Date: Thu, 14 Aug 2025 21:54:04 +0700
Subject: [PATCH] centralized models refactor (#760)

- simplified adding models (use **Models.ts** to add more models)
---
 .../model-selector/src/lib/model-lists.tsx    | 172 ++++----------
 .../promptbox/src/lib/PromptBox2D.tsx         |  74 ++++++-
 .../promptbox/src/lib/PromptBoxImage.tsx      |  65 +++---
 .../promptbox/src/lib/promptStore.ts          |   4 +
 frontend/libs/model-list/src/index.ts         |   6 +-
 .../libs/model-list/src/lib/ImageModels.ts    |  35 ---
 .../{ModelTypeMapping.ts => ModelMapping.ts}  |   3 +-
 frontend/libs/model-list/src/lib/Models.ts    | 209 ++++++++++++++++++
 .../libs/model-list/src/lib/VideoModels.ts    |  30 ---
 .../model-list/src/lib/model-list.spec.ts     |   7 -
 .../libs/model-list/src/lib/model-list.ts     |   3 -
 11 files changed, 360 insertions(+), 248 deletions(-)
 delete mode 100644 frontend/libs/model-list/src/lib/ImageModels.ts
 rename frontend/libs/model-list/src/lib/{ModelTypeMapping.ts => ModelMapping.ts} (95%)
 create mode 100644 frontend/libs/model-list/src/lib/Models.ts
 delete mode 100644 frontend/libs/model-list/src/lib/VideoModels.ts
 delete mode 100644 frontend/libs/model-list/src/lib/model-list.spec.ts
 delete mode 100644 frontend/libs/model-list/src/lib/model-list.ts

diff --git a/frontend/libs/components/model-selector/src/lib/model-lists.tsx b/frontend/libs/components/model-selector/src/lib/model-lists.tsx
index faadfb20ef..e4f303bb04 100644
--- a/frontend/libs/components/model-selector/src/lib/model-lists.tsx
+++ b/frontend/libs/components/model-selector/src/lib/model-lists.tsx
@@ -2,144 +2,52 @@ import type { PopoverItem } from "@storyteller/ui-popover";
 import { faClock, faFilm, faImage } from "@fortawesome/pro-solid-svg-icons";
 import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import {
-  IMAGE_MODELS,
-  VIDEO_MODELS,
   getCreatorIcon,
-  ModelCreator,
+  getModelsByCategory,
+  getInstructiveImageEditModels,
 } from "@storyteller/model-list";
 
 export type ModelList = Omit<PopoverItem, "selected">[];
 
+const withIcon = (creatorIcon: any, fallback: any) => creatorIcon || fallback;
+
+const buildItems = (
+  models: {
+    label: string;
+    description?: string;
+    badges?: { label: string }[];
+    info: { creator: any };
+  }[],
+  fallbackIcon: any
+) =>
+  models.map((m) => ({
+    label: m.label,
+    icon: withIcon(getCreatorIcon((m as any).info.creator), fallbackIcon),
+    description: (m as any).description,
+    badges: (m as any).badges?.map((b: any) => ({
+      label: b.label,
+      icon: <FontAwesomeIcon icon={faClock} />,
+    })),
+    modelInfo: (m as any).info,
+  }));
+
+const imageModels = getModelsByCategory("image");
+const videoModels = getModelsByCategory("video");
+const instructiveModels = getInstructiveImageEditModels();
+
 export const allModels = {
-  video: [
-    {
-      label: "Kling 1.6 Pro",
-      icon: getCreatorIcon(ModelCreator.Kling) || (
-        <FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
-      ),
-      description: "Good quality model",
-      badges: [{ label: "2 min.", icon: <FontAwesomeIcon icon={faClock} /> }],
-      modelInfo: VIDEO_MODELS.kling_1_6_pro,
-    },
-    {
-      label: "Kling 2.1 Pro",
-      icon: getCreatorIcon(ModelCreator.Kling) || (
-        <FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
-      ),
-      description: "High quality model",
-      badges: [{ label: "2 min.", icon: <FontAwesomeIcon icon={faClock} /> }],
-      modelInfo: VIDEO_MODELS.kling_2_1_pro,
-    },
-    {
-      label: "Kling 2.1 Master",
-      icon: getCreatorIcon(ModelCreator.Kling) || (
-        <FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
-      ),
-      description: "Master quality model ($$)",
-      badges: [{ label: "2 min.", icon: <FontAwesomeIcon icon={faClock} /> }],
-      modelInfo: VIDEO_MODELS.kling_2_1_master,
-    },
-    {
-      label: "Seedance 1.0 Lite",
-      icon: getCreatorIcon(ModelCreator.Bytedance) || (
-        <FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
-      ),
-      description: "Fast and high-quality model",
-      badges: [{ label: "2 min.", icon: <FontAwesomeIcon icon={faClock} /> }],
-      modelInfo: VIDEO_MODELS.seedance_1_0_lite,
-    },
-    {
-      label: "Google Veo 2",
-      icon: getCreatorIcon(ModelCreator.Google) || (
-        <FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
-      ),
-      description: "Fast and high-quality model",
-      badges: [{ label: "2 min.", icon: <FontAwesomeIcon icon={faClock} /> }],
-      modelInfo: VIDEO_MODELS.veo_2,
-    },
-  ],
-  image: [
-    {
-      label: "Flux Pro 1.1 Ultra",
-      icon: getCreatorIcon(ModelCreator.BlackForestLabs) || (
-        <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
-      ),
-      description: "High quality model",
-      badges: [{ label: "15 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
-      modelInfo: IMAGE_MODELS.flux_pro_1_1_ultra,
-    },
-    {
-      label: "Flux Pro 1.1",
-      icon: getCreatorIcon(ModelCreator.BlackForestLabs) || (
-        <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
-      ),
-      description: "High quality model",
-      badges: [{ label: "15 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
-      modelInfo: IMAGE_MODELS.flux_pro_1_1,
-    },
-    {
-      label: "Flux 1 Dev",
-      icon: getCreatorIcon(ModelCreator.BlackForestLabs) || (
-        <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
-      ),
-      description: "High quality model",
-      badges: [{ label: "15 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
-      modelInfo: IMAGE_MODELS.flux_1_dev,
-    },
-    {
-      label: "Flux 1 Schnell",
-      icon: getCreatorIcon(ModelCreator.BlackForestLabs) || (
-        <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
-      ),
-      description: "High quality model",
-      badges: [{ label: "15 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
-      modelInfo: IMAGE_MODELS.flux_1_schnell,
-    },
-    {
-      label: "GPT Image 1 (GPT-4o)",
-      icon: getCreatorIcon(ModelCreator.OpenAi) || (
-        <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
-      ),
-      description: "Slow, ultra instructive model",
-      badges: [{ label: "45 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
-      modelInfo: IMAGE_MODELS.gpt_image_1,
-    },
-    //{
-    //  id: "todo",
-    //  label: "Recraft 3",
-    //  icon: <FontAwesomeIcon icon={faImage} className="h-4 w-4" />,
-    //  description: "Fast and high-quality model",
-    //  badges: [{ label: "15 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
-    //},
-    //{
-    //  id: "todo",
-    //  label: "Flux.1 Kontext",
-    //  icon: <FontAwesomeIcon icon={faImage} className="h-4 w-4" />,
-    //  description: "New model with advanced context",
-    //  badges: [{ label: "25 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
-    //},
-  ],
-  // These models are used for 2D and 3D editing and scene building.
-  instructiveImageEdits: [
-    {
-      label: "GPT-4o",
-      icon: getCreatorIcon(ModelCreator.OpenAi) || (
-        <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
-      ),
-      description: "High quality model",
-      badges: [{ label: "2 min.", icon: <FontAwesomeIcon icon={faClock} /> }],
-      modelInfo: IMAGE_MODELS.gpt_image_1,
-    },
-    {
-      label: "FLUX.1 Pro Kontext Max",
-      icon: getCreatorIcon(ModelCreator.BlackForestLabs) || (
-        <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
-      ),
-      description: "Fast and high-quality model",
-      badges: [{ label: "20 sec.", icon: <FontAwesomeIcon icon={faClock} /> }],
-      modelInfo: IMAGE_MODELS.flux_pro_kontext_max,
-    },
-  ],
+  video: buildItems(
+    videoModels as any,
+    <FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
+  ),
+  image: buildItems(
+    imageModels as any,
+    <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
+  ),
+  instructiveImageEdits: buildItems(
+    instructiveModels as any,
+    <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
+  ),
 };
 
 export const videoGenerationModels: ModelList = allModels.video;
diff --git a/frontend/libs/components/promptbox/src/lib/PromptBox2D.tsx b/frontend/libs/components/promptbox/src/lib/PromptBox2D.tsx
index 84ca95ae76..24e7b87551 100644
--- a/frontend/libs/components/promptbox/src/lib/PromptBox2D.tsx
+++ b/frontend/libs/components/promptbox/src/lib/PromptBox2D.tsx
@@ -22,6 +22,7 @@ import {
   faUpload,
   faImages,
   faFrame,
+  faCopy,
 } from "@fortawesome/pro-solid-svg-icons";
 import {
   faRectangleVertical,
@@ -46,7 +47,7 @@ import {
 // import { invoke } from "@tauri-apps/api/core";
 import { usePrompt2DStore, RefImage } from "./promptStore";
 import { gtagEvent } from "@storyteller/google-analytics";
-import { ModelInfo } from "@storyteller/model-list";
+import { ModelInfo, getCapabilitiesForModel } from "@storyteller/model-list";
 
 export type AspectRatio = "1:1" | "3:2" | "2:3";
 
@@ -92,6 +93,8 @@ export const PromptBox2D = ({
   const setUseSystemPrompt = usePrompt2DStore((s) => s.setUseSystemPrompt);
   const aspectRatio = usePrompt2DStore((s) => s.aspectRatio);
   const setAspectRatio = usePrompt2DStore((s) => s.setAspectRatio);
+  const generationCount = usePrompt2DStore((s) => s.generationCount);
+  const setGenerationCount = usePrompt2DStore((s) => s.setGenerationCount);
 
   const [isEnqueueing, setIsEnqueueing] = useState(false);
   const [selectedGalleryImages, setSelectedGalleryImages] = useState<string[]>(
@@ -120,6 +123,41 @@ export const PromptBox2D = ({
     },
   ]);
 
+  const [generationCountList, setGenerationCountList] = useState<PopoverItem[]>(
+    []
+  );
+
+  // Update generation count options from selected model capabilities
+  useEffect(() => {
+    const caps = getCapabilitiesForModel(selectedModelInfo);
+    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)
+      );
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [selectedModelInfo]);
+
+  // Keep UI selection in sync when store value changes
+  useEffect(() => {
+    setGenerationCountList((prev) =>
+      prev.map((item) => ({
+        ...item,
+        selected: item.label === String(generationCount),
+      }))
+    );
+  }, [generationCount]);
+
   const fileInputRef = useRef<HTMLInputElement>(null);
   const textareaRef = useRef<HTMLTextAreaElement>(null);
 
@@ -249,6 +287,17 @@ export const PromptBox2D = ({
     }
   };
 
+  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();
@@ -327,7 +376,7 @@ export const PromptBox2D = ({
       image_media_tokens: referenceImages.map((image) => image.mediaToken),
       disable_system_prompt: !useSystemPrompt,
       prompt: prompt,
-      image_count: 1,
+      image_count: generationCount,
       aspect_ratio: aspectRatio,
     });
 
@@ -564,8 +613,8 @@ export const PromptBox2D = ({
               onChange={handleChange}
               onPaste={handlePaste}
               onKeyDown={handleKeyDown}
-              onFocus={() => { }}
-              onBlur={() => { }}
+              onFocus={() => {}}
+              onBlur={() => {}}
             />
           </div>
           <div className="mt-2 flex items-center justify-between gap-2">
@@ -626,6 +675,23 @@ export const PromptBox2D = ({
               </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"
+                  panelTitle="No. of images"
+                  triggerIcon={
+                    <FontAwesomeIcon icon={faCopy} className="h-4 w-4" />
+                  }
+                  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/PromptBoxImage.tsx b/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx
index d009b4d144..6c8b5f3fe9 100644
--- a/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx
+++ b/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx
@@ -28,7 +28,7 @@ import {
 import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import { IconProp } from "@fortawesome/fontawesome-svg-core";
 import { GalleryItem, GalleryModal } from "@storyteller/ui-gallery-modal";
-import { ModelInfo } from "@storyteller/model-list";
+import { ModelInfo, getCapabilitiesForModel } from "@storyteller/model-list";
 import { usePromptImageStore, RefImage } from "./promptStore";
 import { gtagEvent } from "@storyteller/google-analytics";
 
@@ -102,28 +102,7 @@ export const PromptBoxImage = ({
     },
   ]);
   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);
@@ -148,6 +127,37 @@ export const PromptBoxImage = ({
     }
   }, [imageMediaId, url]);
 
+  useEffect(() => {
+    // Build generation count options based on selected model
+    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) {
+      setGenerationCount(
+        Math.min(Math.max(1, generationCount), caps.maxGenerationCount)
+      );
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [modelInfo]);
+
+  useEffect(() => {
+    setGenerationCountList((prev) =>
+      prev.map((item) => ({
+        ...item,
+        selected: item.label === String(generationCount),
+      }))
+    );
+  }, [generationCount]);
+
   useEffect(() => {
     setAspectRatioList((prev) =>
       prev.map((item) => ({
@@ -157,15 +167,6 @@ export const PromptBoxImage = ({
     );
   }, [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) =>
diff --git a/frontend/libs/components/promptbox/src/lib/promptStore.ts b/frontend/libs/components/promptbox/src/lib/promptStore.ts
index da66193c13..a41757b7d0 100644
--- a/frontend/libs/components/promptbox/src/lib/promptStore.ts
+++ b/frontend/libs/components/promptbox/src/lib/promptStore.ts
@@ -15,10 +15,12 @@ interface Prompt2DStore {
   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 usePrompt2DStore = create<Prompt2DStore>()((set) => ({
@@ -26,10 +28,12 @@ export const usePrompt2DStore = create<Prompt2DStore>()((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 }),
 }));
 
 export { usePrompt2DStore as usePromptStore };
diff --git a/frontend/libs/model-list/src/index.ts b/frontend/libs/model-list/src/index.ts
index 94eea6b676..015a1f98ea 100644
--- a/frontend/libs/model-list/src/index.ts
+++ b/frontend/libs/model-list/src/index.ts
@@ -1,7 +1,5 @@
-export * from "./lib/model-list.js";
 export * from "./lib/ModelInfo.js";
 export * from "./lib/ModelCreator.js";
-export * from "./lib/ImageModels.js";
-export * from "./lib/VideoModels.js";
 export * from "./lib/ModelCreatorIcons.js";
-export * from "./lib/ModelTypeMapping.js";
+export * from "./lib/ModelMapping.js";
+export * from "./lib/Models.js";
diff --git a/frontend/libs/model-list/src/lib/ImageModels.ts b/frontend/libs/model-list/src/lib/ImageModels.ts
deleted file mode 100644
index 40c8efd8e8..0000000000
--- a/frontend/libs/model-list/src/lib/ImageModels.ts
+++ /dev/null
@@ -1,35 +0,0 @@
-import { ModelCreator } from "./ModelCreator.js";
-import { ModelInfo } from "./ModelInfo.js";
-
-export const IMAGE_MODELS: Record<string, ModelInfo> = {
-  flux_pro_1_1_ultra: {
-    name: "Flux Pro 1.1 Ultra",
-    tauri_id: "flux_pro_11_ultra",
-    creator: ModelCreator.BlackForestLabs,
-  },
-  flux_pro_1_1: {
-    name: "Flux Pro 1.1",
-    tauri_id: "flux_pro_11",
-    creator: ModelCreator.BlackForestLabs,
-  },
-  flux_1_dev: {
-    name: "Flux 1 Dev",
-    tauri_id: "flux_1_dev",
-    creator: ModelCreator.BlackForestLabs,
-  },
-  flux_1_schnell: {
-    name: "Flux 1 Schnell",
-    tauri_id: "flux_1_schnell",
-    creator: ModelCreator.BlackForestLabs,
-  },
-  gpt_image_1: {
-    name: "GPT Image 1 (GPT-4o)",
-    tauri_id: "gpt_image_1",
-    creator: ModelCreator.OpenAi,
-  },
-  flux_pro_kontext_max: {
-    name: "Flux Pro Kontext Max",
-    tauri_id: "flux_pro_kontext_max",
-    creator: ModelCreator.BlackForestLabs,
-  },
-};
diff --git a/frontend/libs/model-list/src/lib/ModelTypeMapping.ts b/frontend/libs/model-list/src/lib/ModelMapping.ts
similarity index 95%
rename from frontend/libs/model-list/src/lib/ModelTypeMapping.ts
rename to frontend/libs/model-list/src/lib/ModelMapping.ts
index 4a7050a4dc..b9caa43bbb 100644
--- a/frontend/libs/model-list/src/lib/ModelTypeMapping.ts
+++ b/frontend/libs/model-list/src/lib/ModelMapping.ts
@@ -94,12 +94,13 @@ export const getModelDisplayName = (modelType: string): string => {
     hunyuan_3d_2p1: "Hunyuan 3D 2.1",
     hunyuan_3d: "Hunyuan 3D",
     flux_pro_kontext_max: "Flux Pro Kontext Max",
+    // TODO: add more models here - BFlat
   };
 
   return displayNames[modelType] || modelType;
 };
 
-// Convert provider string to human-readable display name
+// Convert provider string to human-readable display name (this is for the provider priority in settings)
 export const getProviderDisplayName = (provider: string): string => {
   const displayNames: Record<string, string> = {
     artcraft: "ArtCraft",
diff --git a/frontend/libs/model-list/src/lib/Models.ts b/frontend/libs/model-list/src/lib/Models.ts
new file mode 100644
index 0000000000..d2eed07bb0
--- /dev/null
+++ b/frontend/libs/model-list/src/lib/Models.ts
@@ -0,0 +1,209 @@
+import { ModelCreator } from "./ModelCreator.js";
+import { ModelInfo } from "./ModelInfo.js";
+
+export type ModelCategory = "image" | "video";
+
+export interface ModelConfig {
+  id: string;
+  label: string; // UI label
+  description?: string;
+  badges?: { label: string }[];
+  category: ModelCategory;
+  info: ModelInfo;
+  capabilities: ModelCapabilities;
+  tags?: string[]; // optional tags, e.g. ["instructiveEdit"] - for filtering
+}
+
+const mc = ModelCreator;
+
+// Centralized model capability definition
+// For setting options for the model
+// TODO: add more capabilities here - BFlat
+export interface ModelCapabilities {
+  maxGenerationCount: number;
+}
+
+const DEFAULT_CAPABILITIES: ModelCapabilities = {
+  maxGenerationCount: 1,
+};
+
+const cfg = (
+  m: Partial<ModelConfig> & {
+    id: string;
+    category: ModelCategory;
+    info: ModelInfo;
+  }
+): ModelConfig => ({
+  label: m.label ?? m.info.name,
+  description: m.description,
+  badges: m.badges,
+  capabilities: m.capabilities ?? DEFAULT_CAPABILITIES,
+  tags: m.tags ?? [],
+  ...m,
+});
+
+export const ALL_MODELS: ModelConfig[] = [
+  //////////////////////////////
+  // Image models
+  //////////////////////////////
+  cfg({
+    id: "flux_pro_1_1_ultra",
+    category: "image",
+    info: {
+      name: "Flux Pro 1.1 Ultra",
+      tauri_id: "flux_pro_11_ultra",
+      creator: mc.BlackForestLabs,
+    },
+    description: "High quality model",
+    badges: [{ label: "15 sec." }],
+    capabilities: { maxGenerationCount: 4 },
+    tags: ["instructiveEdit"],
+  }),
+  cfg({
+    id: "flux_pro_1_1",
+    category: "image",
+    info: {
+      name: "Flux Pro 1.1",
+      tauri_id: "flux_pro_11",
+      creator: mc.BlackForestLabs,
+    },
+    description: "High quality model",
+    badges: [{ label: "15 sec." }],
+    capabilities: { maxGenerationCount: 4 },
+    tags: ["instructiveEdit"],
+  }),
+  cfg({
+    id: "flux_1_dev",
+    category: "image",
+    info: {
+      name: "Flux 1 Dev",
+      tauri_id: "flux_1_dev",
+      creator: mc.BlackForestLabs,
+    },
+    description: "High quality model",
+    badges: [{ label: "15 sec." }],
+    capabilities: { maxGenerationCount: 4 },
+  }),
+  cfg({
+    id: "flux_1_schnell",
+    category: "image",
+    info: {
+      name: "Flux 1 Schnell",
+      tauri_id: "flux_1_schnell",
+      creator: mc.BlackForestLabs,
+    },
+    description: "High quality model",
+    badges: [{ label: "15 sec." }],
+    capabilities: { maxGenerationCount: 4 },
+  }),
+  cfg({
+    id: "gpt_image_1",
+    category: "image",
+    info: {
+      name: "GPT Image 1 (GPT-4o)",
+      tauri_id: "gpt_image_1",
+      creator: mc.OpenAi,
+    },
+    description: "Slow, ultra instructive model",
+    badges: [{ label: "45 sec." }],
+    capabilities: { maxGenerationCount: 1 },
+    tags: ["instructiveEdit"],
+  }),
+  cfg({
+    id: "flux_pro_kontext_max",
+    category: "image",
+    info: {
+      name: "Flux Pro Kontext Max",
+      tauri_id: "flux_pro_kontext_max",
+      creator: mc.BlackForestLabs,
+    },
+    description: "Fast and high-quality model",
+    badges: [{ label: "20 sec." }],
+    capabilities: { maxGenerationCount: 4 },
+    tags: ["instructiveEdit"],
+  }),
+
+  //////////////////////////////
+  // Video models
+  //////////////////////////////
+  cfg({
+    id: "kling_1_6_pro",
+    category: "video",
+    info: {
+      name: "Kling 1.6 Pro",
+      tauri_id: "kling_1.6_pro",
+      creator: mc.Kling,
+    },
+    description: "Good quality model",
+    badges: [{ label: "2 min." }],
+    capabilities: { maxGenerationCount: 1 },
+  }),
+  cfg({
+    id: "kling_2_1_pro",
+    category: "video",
+    info: {
+      name: "Kling 2.1 Pro",
+      tauri_id: "kling_2.1_pro",
+      creator: mc.Kling,
+    },
+    description: "High quality model",
+    badges: [{ label: "2 min." }],
+    capabilities: { maxGenerationCount: 1 },
+  }),
+  cfg({
+    id: "kling_2_1_master",
+    category: "video",
+    info: {
+      name: "Kling 2.1 Master",
+      tauri_id: "kling_2.1_master",
+      creator: mc.Kling,
+    },
+    description: "Master quality model ($$)",
+    badges: [{ label: "2 min." }],
+    capabilities: { maxGenerationCount: 1 },
+  }),
+  cfg({
+    id: "seedance_1_0_lite",
+    category: "video",
+    info: {
+      name: "Seedance 1.0 Lite",
+      tauri_id: "seedance_1.0_lite",
+      creator: mc.Bytedance,
+    },
+    description: "Fast and high-quality model",
+    badges: [{ label: "2 min." }],
+    capabilities: { maxGenerationCount: 1 },
+  }),
+  cfg({
+    id: "veo_2",
+    category: "video",
+    info: { name: "Google Veo 2", tauri_id: "veo_2", creator: mc.Google },
+    description: "Fast and high-quality model",
+    badges: [{ label: "2 min." }],
+    capabilities: { maxGenerationCount: 1 },
+  }),
+];
+
+export const getAllModels = (): ModelConfig[] => ALL_MODELS;
+
+export const getModelsByCategory = (category: ModelCategory): ModelConfig[] =>
+  ALL_MODELS.filter((m) => m.category === category);
+
+export const getInstructiveImageEditModels = (): ModelConfig[] =>
+  ALL_MODELS.filter(
+    (m) => m.category === "image" && m.tags?.includes("instructiveEdit")
+  );
+
+export const lookupModelByTauriId = (
+  tauriId: string
+): ModelConfig | undefined =>
+  ALL_MODELS.find((m) => m.info.tauri_id === tauriId);
+
+// Single exported capability resolver so callers never need to touch anything else
+export const getCapabilitiesForModel = (
+  model?: ModelInfo
+): ModelCapabilities => {
+  if (!model) return DEFAULT_CAPABILITIES;
+  const cfg = lookupModelByTauriId(model.tauri_id);
+  return cfg?.capabilities ?? DEFAULT_CAPABILITIES;
+};
diff --git a/frontend/libs/model-list/src/lib/VideoModels.ts b/frontend/libs/model-list/src/lib/VideoModels.ts
deleted file mode 100644
index ee1a934c88..0000000000
--- a/frontend/libs/model-list/src/lib/VideoModels.ts
+++ /dev/null
@@ -1,30 +0,0 @@
-import { ModelCreator } from "./ModelCreator.js";
-import { ModelInfo } from "./ModelInfo.js";
-
-export const VIDEO_MODELS: Record<string, ModelInfo> = {
-  kling_1_6_pro: {
-    name: "Kling 1.6 Pro",
-    tauri_id: "kling_1.6_pro",
-    creator: ModelCreator.Kling,
-  },
-  kling_2_1_pro: {
-    name: "Kling 2.1 Pro",
-    tauri_id: "kling_2.1_pro",
-    creator: ModelCreator.Kling,
-  },
-  kling_2_1_master: {
-    name: "Kling 2.1 Master",
-    tauri_id: "kling_2.1_master",
-    creator: ModelCreator.Kling,
-  },
-  seedance_1_0_lite: {
-    name: "Seedance 1.0 Lite",
-    tauri_id: "seedance_1.0_lite",
-    creator: ModelCreator.Bytedance,
-  },
-  veo_2: {
-    name: "Google Veo 2",
-    tauri_id: "veo_2",
-    creator: ModelCreator.Google,
-  },
-};
diff --git a/frontend/libs/model-list/src/lib/model-list.spec.ts b/frontend/libs/model-list/src/lib/model-list.spec.ts
deleted file mode 100644
index 4b1e5101c7..0000000000
--- a/frontend/libs/model-list/src/lib/model-list.spec.ts
+++ /dev/null
@@ -1,7 +0,0 @@
-import { modelList } from './model-list.js';
-
-describe('modelList', () => {
-  it('should work', () => {
-    expect(modelList()).toEqual('model-list');
-  })
-})
diff --git a/frontend/libs/model-list/src/lib/model-list.ts b/frontend/libs/model-list/src/lib/model-list.ts
deleted file mode 100644
index a46bc7f2cd..0000000000
--- a/frontend/libs/model-list/src/lib/model-list.ts
+++ /dev/null
@@ -1,3 +0,0 @@
-export function modelList(): string {
-  return "model-list";
-}