continued rework of the prompt box

frontend/apps/editor3d/app/src/pages/PageImage/TextToImage.tsx

@@ -5,15 +5,15 @@ import BackgroundGallery from "./BackgroundGallery";
 import {
   TEXT_TO_IMAGE_PAGE_MODEL_LIST,
   ModelPage,
-  ModelSelector,
-  useModelSelectorStore,
+  ClassyModelSelector,
+  useClassyModelSelectorStore,
 } from "@storyteller/ui-model-selector";
-
+import { ImageModel } from "@storyteller/model-list";
 interface TextToImageProps {
   imageMediaId?: string;
   imageUrl?: string;
 }
-import { ModelInfo } from "@storyteller/model-list";
+import { Model, ModelInfo } from "@storyteller/model-list";
 import { useTextToImageGenerationCompleteEvent } from "@storyteller/tauri-events";
 import { useTextToImageStore } from "./TextToImageStore";
 import { animated, useSpring } from "@react-spring/web";
@@ -27,19 +27,21 @@ const PAGE_ID: ModelPage = ModelPage.TextToImage;
 
 const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => {
   const containerRef = useRef<HTMLDivElement>(null);
-  const { selectedModels } = useModelSelectorStore();
+  //const { selectedModels } = useModelSelectorStore();
   const batches = useTextToImageStore((s) => s.batches);
   const startBatch = useTextToImageStore((s) => s.startBatch);
   const completeBatch = useTextToImageStore((s) => s.completeBatch);
   const resetBatches = useTextToImageStore((s) => s.reset);
 
-  const selectedModel =
-    selectedModels[PAGE_ID] || TEXT_TO_IMAGE_PAGE_MODEL_LIST[0]?.label;
+  const { selectedModels } = useClassyModelSelectorStore();
+
+  const selectedModel : Model | undefined =
+    selectedModels[PAGE_ID] || TEXT_TO_IMAGE_PAGE_MODEL_LIST[0]?.model;
+
+  const selectedImageModel = (selectedModel instanceof ImageModel) ? selectedModel : undefined;
 
   const selectedModelInfo: ModelInfo | undefined =
-    TEXT_TO_IMAGE_PAGE_MODEL_LIST.find(
-      (m) => m.label === selectedModel,
-    )?.modelInfo;
+    selectedModel?.toLegacyModelConfig().info;
 
   const jobContext: JobContextType = {
     jobTokens: [],
@@ -171,12 +173,12 @@ const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => {
               useJobContext={() => {
                 return jobContext;
               }}
-              model={selectedModel}
-              modelInfo={selectedModelInfo}
+              selectedModel={selectedImageModel}
               imageMediaId={imageMediaId}
               url={imageUrl ?? undefined}
               onEnqueuePressed={async (prompt, count, subscriberId) => {
-                startBatch(prompt, count, selectedModel, subscriberId);
+                const modelLabel = selectedModelInfo?.name ?? "";
+                startBatch(prompt, count, modelLabel, subscriberId);
               }}
             />
           </animated.div>
@@ -184,7 +186,7 @@ const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => {
           {!showPromptAtBottom && <BackgroundGallery />}
 
           <div className="absolute bottom-6 left-6 z-20 flex items-center gap-2">
-            <ModelSelector
+            <ClassyModelSelector
               items={TEXT_TO_IMAGE_PAGE_MODEL_LIST}
               page={PAGE_ID}
               mode="hoverSelect"

frontend/libs/components/model-selector/src/index.ts

@@ -4,4 +4,6 @@ export * from "./lib/model-lists";
 export * from "./lib/model-pages";
 export * from "./lib/video-model-selector-store";
 export * from "./lib/video-model-selector";
+export * from "./lib/classy-model-selector";
+export * from "./lib/classy-model-selector-store";
 

frontend/libs/components/model-selector/src/lib/classy-model-selector-store.ts

@@ -0,0 +1,19 @@
+import { create } from "zustand";
+import { ModelPage } from "./model-pages";
+import { Model } from "@storyteller/model-list";
+
+interface ClassyModelSelectorState {
+  selectedModels: { [page in ModelPage]?: Model };
+  setSelectedModel: (page: ModelPage, model: Model) => void;
+}
+
+export const useClassyModelSelectorStore = create<ClassyModelSelectorState>((set) => ({
+  selectedModels: {},
+  setSelectedModel: (page, model) =>
+    set((state) => ({
+      selectedModels: {
+        ...state.selectedModels,
+        [page]: model,
+      },
+    })),
+}));

frontend/libs/components/model-selector/src/lib/classy-model-selector.tsx

@@ -0,0 +1,56 @@
+import { PopoverMenu, type PopoverItem } from "@storyteller/ui-popover";
+import { useClassyModelSelectorStore } from "./classy-model-selector-store";
+import { useEffect, useMemo } from "react";
+import { ModelPage } from "./model-pages";
+
+interface ClassyModelSelectorProps {
+  items: Omit<PopoverItem, "selected">[];
+  page: ModelPage;
+  mode?: "hoverSelect" | "default" | "toggle" | "button";
+  panelTitle?: string;
+  buttonClassName?: string;
+  panelClassName?: string;
+  showIconsInList?: boolean;
+  triggerLabel?: string;
+}
+
+export function ClassyModelSelector({
+  items,
+  page,
+  ...popoverProps
+}: ClassyModelSelectorProps) {
+  const { selectedModels, setSelectedModel } = useClassyModelSelectorStore();
+  const selectedModel = selectedModels[page] || items[0]?.model;
+
+  // For the first mount, make sure the selected model is set for other components to listen
+  useEffect(() => {
+    // Initialize selected model if not set
+    if (!selectedModels[page] && items[0]) {
+      setSelectedModel(page, items[0].model!);
+    }
+  }, []);
+
+  const handleModelSelect = (item: PopoverItem) => {
+    setSelectedModel(page, item.model!);
+  };
+
+  const modelList = useMemo(
+    () =>
+      items.map((item) => ({
+        ...item,
+        selected: item.model === selectedModel,
+      })),
+    [items, selectedModel]
+  );
+
+  return (
+    <PopoverMenu
+      items={modelList}
+      onSelect={handleModelSelect}
+      mode="hoverSelect"
+      {...popoverProps}
+    />
+  );
+}
+
+export default ClassyModelSelector;

frontend/libs/components/model-selector/src/lib/model-lists.tsx

@@ -10,6 +10,7 @@ import {
   STAGE_3D_PAGE_MODELS,
   IMAGE_EDITOR_PAGE_MODELS,
   VIDEO_MODELS,
+  IMAGE_MODELS,
 } from "@storyteller/model-list";
 
 export type ModelList = Omit<PopoverItem, "selected">[];
@@ -53,11 +54,6 @@ const buildItems2 = (
  * OLD MODEL LISTS
  */ 
 
-export const TEXT_TO_IMAGE_PAGE_MODEL_LIST : ModelList = buildItems(
-  TEXT_TO_IMAGE_PAGE_MODELS,
-  <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
-);
-
 export const CANVAS_2D_PAGE_MODEL_LIST : ModelList = buildItems(
   CANVAS_2D_PAGE_MODELS,
   <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
@@ -84,3 +80,8 @@ export const IMAGE_TO_VIDEO_PAGE_MODEL_LIST : ModelList = buildItems2(
   VIDEO_MODELS,
   <FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
 );
+
+export const TEXT_TO_IMAGE_PAGE_MODEL_LIST : ModelList = buildItems2(
+  IMAGE_MODELS,
+  <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
+);

frontend/libs/components/model-selector/src/lib/model-selector.spec.tsx

@@ -1,12 +0,0 @@
-import { render } from '@testing-library/react';
-
-import ModelSelector from './model-selector';
-
-describe('ModelSelector', () => {
-  
-  it('should render successfully', () => {
-    const { baseElement } = render(<ModelSelector />);
-    expect(baseElement).toBeTruthy();
-  });
-  
-});

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, getCapabilitiesForModel } from "@storyteller/model-list";
+import { ImageModel, ModelInfo, getCapabilitiesForModel } from "@storyteller/model-list";
 import { usePromptImageStore, RefImage } from "./promptStore";
 import { gtagEvent } from "@storyteller/google-analytics";
 
@@ -39,8 +39,7 @@ interface PromptBoxImageProps {
     count: number,
     subscriberId: string
   ) => void | Promise<void>;
-  model: string;
-  modelInfo?: ModelInfo;
+  selectedModel?: ImageModel;
   imageMediaId?: string;
   url?: string;
 }
@@ -48,8 +47,7 @@ interface PromptBoxImageProps {
 export const PromptBoxImage = ({
   useJobContext,
   onEnqueuePressed,
-  model,
-  modelInfo,
+  selectedModel,
   imageMediaId,
   url,
 }: PromptBoxImageProps) => {
@@ -133,7 +131,7 @@ export const PromptBoxImage = ({
 
   useEffect(() => {
     // Build generation count options based on selected model
-    const caps = getCapabilitiesForModel(modelInfo);
+    const caps = getCapabilitiesForModel(selectedModel);
     const defaultCount = Math.min(
       Math.max(1, caps.defaultGenerationCount ?? 1),
       caps.maxGenerationCount
@@ -151,7 +149,7 @@ export const PromptBoxImage = ({
     }));
     setGenerationCountList(items);
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [modelInfo]);
+  }, [selectedModel]);
 
   useEffect(() => {
     setGenerationCountList((prev) =>
@@ -285,7 +283,7 @@ export const PromptBoxImage = ({
 
     const generateResponse = await EnqueueTextToImage({
       prompt: prompt,
-      model: modelInfo,
+      model: selectedModel,
       aspect_ratio: aspectRatio,
       number_images: generationCount,
       frontend_caller: "text_to_image",

frontend/libs/model-list/src/index.ts

@@ -5,7 +5,9 @@ export * from "./lib/ModelInfo.js";
 export * from "./lib/ModelMapping.js";
 export * from "./lib/Models.js";
 export * from "./lib/ModelsByPage.js";
-export * from "./lib/VideoModels.js";
 export * from "./lib/classes/ImageModel.js";
 export * from "./lib/classes/Model.js";
 export * from "./lib/classes/VideoModel.js";
+export * from "./lib/lists/AllModels.js";
+export * from "./lib/lists/ImageModels.js";
+export * from "./lib/lists/VideoModels.js";

frontend/libs/model-list/src/lib/Models.ts

@@ -6,6 +6,7 @@ import {
   ModelCapabilities,
   ModelCategory,
 } from "./ModelConfig.js";
+import { Model } from "./classes/Model.js";
 
 const mc = ModelCreator;
 
@@ -239,9 +240,10 @@ export const lookupModelByTauriId = (
 
 // Single exported capability resolver so callers never need to touch anything else
 export const getCapabilitiesForModel = (
-  model?: ModelInfo
+  model?: Model | ModelInfo
 ): ModelCapabilities => {
   if (!model) return DEFAULT_CAPABILITIES;
-  const cfg = lookupModelByTauriId(model.tauri_id);
+  const tauriId = model instanceof Model ? model.tauriId : model.tauri_id;
+  const cfg = lookupModelByTauriId(tauriId);
   return cfg?.capabilities ?? DEFAULT_CAPABILITIES;
 };

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

@@ -3,6 +3,12 @@ import { Model } from "./Model.js";
 import { ModelCategory } from "../ModelConfig.js";
 
 export class ImageModel extends Model {
+  readonly maxGenerationCount: number;
+  readonly defaultGenerationCount: number;
+
+  // For inpainting models, does it require sending a mask?
+  readonly usesInpaintingMask: boolean;
+
   constructor(args: {
     id: string;
     tauriId: string;
@@ -12,7 +18,22 @@ export class ImageModel extends Model {
     selectorName: string;
     selectorDescription: string;
     selectorBadges: string[];
+    maxGenerationCount: number;
+    defaultGenerationCount: number;
+    usesInpaintingMask?: boolean;
   }) {
+    if (args.maxGenerationCount < 1) {
+      throw new Error("maxGenerationCount must be at least 1");
+    }
+    if (args.defaultGenerationCount < 1) {
+      throw new Error("defaultGenerationCount must be at least 1");
+    }
+    if (args.defaultGenerationCount > args.maxGenerationCount) {
+      throw new Error("defaultGenerationCount must be less than or equal to maxGenerationCount");
+    }
     super(args);
+    this.maxGenerationCount = args.maxGenerationCount;
+    this.defaultGenerationCount = args.defaultGenerationCount;
+    this.usesInpaintingMask = args.usesInpaintingMask ?? false; 
   } 
 }

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

@@ -0,0 +1,5 @@
+import { IMAGE_MODELS } from "./ImageModels.js";
+import { VIDEO_MODELS } from "./VideoModels.js";
+import { Model } from "../classes/Model.js";
+
+export const ALL_MODELS_LIST: Model[] = [...IMAGE_MODELS, ...VIDEO_MODELS];

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

@@ -0,0 +1,115 @@
+import { ImageModel } from "../classes/ImageModel.js";
+import { ModelCreator } from "../ModelCreator.js";
+
+export const IMAGE_MODELS: ImageModel[] = [
+  new ImageModel({
+    id: "midjourney",
+    tauriId: "midjourney",
+    fullName: "Midjourney",
+    category: "image",
+    creator: ModelCreator.Midjourney,
+    selectorName: "Midjourney",
+    selectorDescription: "Incredible style and quality",
+    selectorBadges: ["15 sec."],
+    maxGenerationCount: 4,
+    defaultGenerationCount: 4,
+  }),
+  new ImageModel({
+    id: "flux_pro_1_1_ultra",
+    tauriId: "flux_pro_11_ultra",
+    fullName: "Flux Pro 1.1 Ultra",
+    category: "image",
+    creator: ModelCreator.BlackForestLabs,
+    selectorName: "Flux Pro 1.1 Ultra",
+    selectorDescription: "High quality model",
+    selectorBadges: ["15 sec."],
+    maxGenerationCount: 4,
+    defaultGenerationCount: 4,
+  }),
+  new ImageModel({
+    id: "flux_pro_1_1",
+    tauriId: "flux_pro_11",
+    fullName: "Flux Pro 1.1",
+    category: "image",
+    creator: ModelCreator.BlackForestLabs,
+    selectorName: "Flux Pro 1.1",
+    selectorDescription: "High quality model",
+    selectorBadges: ["15 sec."],
+    maxGenerationCount: 4,
+    defaultGenerationCount: 4,
+  }),
+  new ImageModel({
+    id: "flux_1_dev",
+    tauriId: "flux_1_dev",
+    fullName: "Flux 1 Dev",
+    category: "image",
+    creator: ModelCreator.BlackForestLabs,
+    selectorName: "Flux 1 Dev",
+    selectorDescription: "High quality model",
+    selectorBadges: ["15 sec."],
+    maxGenerationCount: 4,
+    defaultGenerationCount: 4,
+  }),
+  new ImageModel({
+    id: "flux_1_schnell",
+    tauriId: "flux_1_schnell",
+    fullName: "Flux 1 Schnell",
+    category: "image",
+    creator: ModelCreator.BlackForestLabs,
+    selectorName: "Flux 1 Schnell",
+    selectorDescription: "High quality model",
+    selectorBadges: ["15 sec."],
+    maxGenerationCount: 4,
+    defaultGenerationCount: 4,
+  }),
+  new ImageModel({
+    id: "gpt_image_1",
+    tauriId: "gpt_image_1",
+    fullName: "GPT Image 1 (GPT-4o)",
+    category: "image",
+    creator: ModelCreator.OpenAi,
+    selectorName: "GPT Image 1 (GPT-4o)",
+    selectorDescription: "Slow, ultra instructive model",
+    selectorBadges: ["45 sec."],
+    maxGenerationCount: 1,
+    defaultGenerationCount: 1,
+  }),
+  new ImageModel({
+    id: "flux_pro_kontext_max",
+    tauriId: "flux_pro_kontext_max",
+    fullName: "Flux Pro Kontext Max",
+    category: "image",
+    creator: ModelCreator.BlackForestLabs,
+    selectorName: "Flux Pro Kontext Max",
+    selectorDescription: "Fast and high-quality model",
+    selectorBadges: ["20 sec."],
+    maxGenerationCount: 4,
+    defaultGenerationCount: 4,
+  }),
+  new ImageModel({
+    id: "flux_pro_inpaint",
+    tauriId: "flux_pro_inpaint",
+    fullName: "Flux Pro Inpaint",
+    category: "image",
+    creator: ModelCreator.BlackForestLabs,
+    selectorName: "Flux Pro Inpaint",
+    selectorDescription: "Fast and high-quality model",
+    selectorBadges: ["20 sec."],
+    maxGenerationCount: 1,
+    defaultGenerationCount: 1,
+    usesInpaintingMask: true,
+  }),
+  new ImageModel({
+    id: "flux_dev_juggernaut_inpaint",
+    tauriId: "flux_dev_juggernaut_inpaint",
+    fullName: "Flux Dev Juggernaut Inpaint",
+    category: "image",
+    creator: ModelCreator.BlackForestLabs,
+    selectorName: "Flux Dev Juggernaut Inpaint",
+    selectorDescription: "Fast and high-quality model",
+    selectorBadges: ["20 sec."],
+    maxGenerationCount: 4,
+    defaultGenerationCount: 4,
+    usesInpaintingMask: true,
+  }),
+];

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

@@ -1,5 +1,5 @@
-import { VideoModel } from "./classes/VideoModel.js";
-import { ModelCreator } from "./ModelCreator.js";
+import { VideoModel } from "../classes/VideoModel.js";
+import { ModelCreator } from "../ModelCreator.js";
 
 export const VIDEO_MODELS: VideoModel[] = [
   new VideoModel({

frontend/libs/tauri-api/src/lib/enqueue/EnqueueTextToImage.ts

@@ -1,6 +1,6 @@
 import { invoke } from "@tauri-apps/api/core";
 import { CommandResult } from "../common/CommandStatus";
-import { ModelInfo } from "@storyteller/model-list";
+import { Model, ModelInfo } from "@storyteller/model-list";
 
 export enum EnqueueTextToImageErrorType {
   /// Caller didn't specify a model
@@ -15,7 +15,7 @@ export enum EnqueueTextToImageErrorType {
 
 export interface EnqueueTextToImageRequest {
   prompt?: string;
-  model?: ModelInfo | EnqueueTextToImageModel;
+  model?: Model | ModelInfo | EnqueueTextToImageModel;
   aspect_ratio?: EnqueueTextToImageSize;
   number_images?: number;
 
@@ -73,6 +73,8 @@ export const EnqueueTextToImage = async (request: EnqueueTextToImageRequest) : P
   if (!!request.model) {
     if (typeof request.model === "string") {
       modelName = request.model;
+    } else if (request.model instanceof Model) {
+      modelName = request.model.tauriId;
     } else {
       modelName = request.model.tauri_id;
     }