promptbox refactor

frontend/apps/editor3d/app/src/pages/PageVideo/ImageToVideo.tsx

@@ -4,11 +4,10 @@ import { PromptBoxVideo } from "@storyteller/ui-promptbox";
 import BackgroundGallery from "./BackgroundGallery";
 import {
   ModelPage,
-  ModelSelector,
-  useModelSelectorStore,
   IMAGE_TO_VIDEO_PAGE_MODEL_LIST,
+  VideoModelSelector,
 } from "@storyteller/ui-model-selector";
-import { ModelInfo } from "@storyteller/model-list";
+import { useVideoModelSelectorStore } from "@storyteller/ui-model-selector";
 
 const PAGE_ID : ModelPage = ModelPage.ImageToVideo;
 
@@ -19,15 +18,7 @@ interface ImageToVideoProps {
 
 const ImageToVideo = ({ imageMediaId, imageUrl }: ImageToVideoProps) => {
   const containerRef = useRef<HTMLDivElement>(null);
-  const { selectedModels } = useModelSelectorStore();
-
-  const selectedModel =
-    selectedModels[PAGE_ID] ||
-    IMAGE_TO_VIDEO_PAGE_MODEL_LIST[0]?.label;
-
-  const selectedModelInfo: ModelInfo | undefined = IMAGE_TO_VIDEO_PAGE_MODEL_LIST.find(
-    (m) => m.label === selectedModel,
-  )?.modelInfo;
+  const { selectedModel } = useVideoModelSelectorStore();
 
   const jobContext: JobContextType = {
     jobTokens: [],
@@ -56,8 +47,7 @@ const ImageToVideo = ({ imageMediaId, imageUrl }: ImageToVideoProps) => {
               console.log("Using job context");
               return jobContext;
             }}
-            model={selectedModel}
-            modelInfo={selectedModelInfo}
+            selectedModel={selectedModel}
             imageMediaId={imageMediaId}
             url={imageUrl ?? undefined}
             onEnqueuePressed={async () => {
@@ -68,7 +58,7 @@ const ImageToVideo = ({ imageMediaId, imageUrl }: ImageToVideoProps) => {
           <BackgroundGallery />
 
           <div className="absolute bottom-6 left-6 z-20 flex items-center gap-2">
-            <ModelSelector
+            <VideoModelSelector
               items={IMAGE_TO_VIDEO_PAGE_MODEL_LIST}
               page={PAGE_ID}
               panelTitle="Select Model"

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

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

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

@@ -3,12 +3,13 @@ import { faClock, faFilm, faImage } from "@fortawesome/pro-solid-svg-icons";
 import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import {
   getCreatorIcon,
+  Model,
   ModelConfig,
   TEXT_TO_IMAGE_PAGE_MODELS,
-  IMAGE_TO_VIDEO_PAGE_MODELS,
   CANVAS_2D_PAGE_MODELS,
   STAGE_3D_PAGE_MODELS,
   IMAGE_EDITOR_PAGE_MODELS,
+  VIDEO_MODELS,
 } from "@storyteller/model-list";
 
 export type ModelList = Omit<PopoverItem, "selected">[];
@@ -31,16 +32,32 @@ const buildItems = (
     modelInfo: m.info,
   }));
 
+const buildItems2 = (
+  models: Model[],
+  fallbackIcon: any
+) =>
+  models.map((model : Model) => ({
+    label: model.selectorName,
+    icon: withIcon(getCreatorIcon(model.creator), fallbackIcon),
+    description: model.selectorDescription,
+    badges: model.toLegacyBadges()?.map((b) => ({
+      label: b.label,
+      icon: <FontAwesomeIcon icon={faClock} />,
+    })),
+    modelConfig: model.toLegacyModelConfig(), // Access to full object.
+    modelInfo: undefined, // NB: Dead.
+    model: model,
+  }));
+
+/**
+ * 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 IMAGE_TO_VIDEO_PAGE_MODEL_LIST : ModelList = buildItems(
-  IMAGE_TO_VIDEO_PAGE_MODELS,
-  <FontAwesomeIcon icon={faFilm} 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" />
@@ -55,3 +72,15 @@ export const IMAGE_EDITOR_PAGE_MODEL_LIST : ModelList = buildItems(
   IMAGE_EDITOR_PAGE_MODELS,
   <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
 );
+
+/**
+ * IN-PROGRESS MIGRATION (messy for now)
+ * We're gradually going to phase out ModelList, ModelConfig, etc.
+ * We won't index by name, but rather id, or simply will always have full 
+ * access to the object directly.
+ */ 
+
+export const IMAGE_TO_VIDEO_PAGE_MODEL_LIST : ModelList = buildItems2(
+  VIDEO_MODELS,
+  <FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
+);

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

@@ -0,0 +1,19 @@
+import { create } from "zustand";
+import { VideoModel } from "@storyteller/model-list";
+import { VIDEO_MODELS } from "@storyteller/model-list";
+
+/**
+ * TODO: This is temporary. We'll create a new "ModelSelectorState" that is typesafe and isn't just scoped to videos.
+ */
+interface VideoModelSelectorState {
+  selectedModel: VideoModel;
+  setSelectedModel: (model: VideoModel) => void;
+}
+
+export const useVideoModelSelectorStore = create<VideoModelSelectorState>((set) => ({
+  selectedModel: VIDEO_MODELS[0],
+  setSelectedModel: (model) =>
+    set((state) => ({
+      selectedModel: model
+    })),
+}));

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

@@ -0,0 +1,61 @@
+import { PopoverMenu, type PopoverItem } from "@storyteller/ui-popover";
+import { useEffect, useMemo } from "react";
+import { ModelPage } from "./model-pages";
+import { useVideoModelSelectorStore } from "./video-model-selector-store";
+import { VideoModel } from "@storyteller/model-list";
+
+/**
+ * TODO: This is temporary. We'll create a new "ModelSelectorState" that is typesafe and isn't just scoped to videos.
+ */
+interface VideoModelSelectorProps {
+  items: Omit<PopoverItem, "selected">[];
+  page: ModelPage;
+  mode?: "hoverSelect" | "default" | "toggle" | "button";
+  panelTitle?: string;
+  buttonClassName?: string;
+  panelClassName?: string;
+  showIconsInList?: boolean;
+  triggerLabel?: string;
+}
+
+export function VideoModelSelector({
+  items,
+  page,
+  ...popoverProps
+}: VideoModelSelectorProps) {
+  const { selectedModel, setSelectedModel } = useVideoModelSelectorStore();
+
+  //// 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].label);
+  //  }
+  //}, []);
+
+  const handleModelSelect = (item: PopoverItem) => {
+    if (!item.model) return;
+    if (!(item.model instanceof VideoModel)) return;
+    setSelectedModel(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 VideoModelSelector;

frontend/libs/components/popover/src/lib/popover.tsx

@@ -13,7 +13,7 @@ import {
   faChevronUp,
   faCircleCheck,
 } from "@fortawesome/pro-solid-svg-icons";
-import { ModelInfo } from "@storyteller/model-list";
+import { Model, ModelInfo } from "@storyteller/model-list";
 
 export interface PopoverItem {
   label: string;
@@ -28,6 +28,7 @@ export interface PopoverItem {
     icon?: ReactNode;
   }>;
   modelInfo?: ModelInfo;
+  model?: Model; // NB: Let's migrate to using this.
 }
 
 interface PopoverMenuProps {

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

@@ -21,15 +21,14 @@ import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
 import { IconProp } from "@fortawesome/fontawesome-svg-core";
 import { IsDesktopApp } from "@storyteller/tauri-utils";
 import { GalleryItem, GalleryModal } from "@storyteller/ui-gallery-modal";
-import { ModelInfo } from "@storyteller/model-list";
+import { Model } from "@storyteller/model-list";
 import { usePromptVideoStore, RefImage } from "./promptStore";
 import { gtagEvent } from "@storyteller/google-analytics";
 
 interface PromptBoxVideoProps {
   useJobContext: () => JobContextType;
   onEnqueuePressed?: () => void | Promise<void>;
-  model: string;
-  modelInfo?: ModelInfo;
+  selectedModel?: Model;
   imageMediaId?: string;
   url?: string;
 }
@@ -37,8 +36,7 @@ interface PromptBoxVideoProps {
 export const PromptBoxVideo = ({
   useJobContext,
   onEnqueuePressed,
-  model,
-  modelInfo,
+  selectedModel,
   imageMediaId,
   url,
 }: PromptBoxVideoProps) => {
@@ -194,7 +192,7 @@ export const PromptBoxVideo = ({
   const handleEnqueue = async () => {
     setIsEnqueueing(true);
 
-    console.log("PromptBoxVideo - Prompting with model", modelInfo);
+    console.log("PromptBoxVideo - Prompting with model", selectedModel);
 
     gtagEvent("enqueue_video");
 
@@ -206,7 +204,7 @@ export const PromptBoxVideo = ({
     }, 10000);
 
     const generateResponse = await EnqueueImageToVideo({
-      model: modelInfo,
+      model: selectedModel,
       image_media_token: referenceImages[0].mediaToken,
       prompt: prompt,
     });

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

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

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

@@ -0,0 +1,65 @@
+import { VideoModel } from "./classes/VideoModel.js";
+import { ModelCreator } from "./ModelCreator.js";
+
+export const VIDEO_MODELS: VideoModel[] = [
+  new VideoModel({
+    id: "kling_1_6_pro",
+    tauriId: "kling_1.6_pro",
+    fullName: "Kling 1.6 Pro",
+    category: "video",
+    creator: ModelCreator.Kling,
+    selectorName: "Kling 1.6 Pro",
+    selectorDescription: "Fast video model",
+    selectorBadges: ["2 min."],
+    startFrame: true,
+    endFrame: true,
+  }),
+  new VideoModel({
+    id: "kling_2_1_pro",
+    tauriId: "kling_2.1_pro",
+    fullName: "Kling 2.1 Pro",
+    category: "video",
+    creator: ModelCreator.Kling,
+    selectorName: "Kling 2.1 Pro",
+    selectorDescription: "Fast video model",
+    selectorBadges: ["2 min."],
+    startFrame: true,
+    endFrame: true,
+  }),
+  new VideoModel({
+    id: "kling_2_1_master",
+    tauriId: "kling_2.1_master",
+    fullName: "Kling 2.1 Master",
+    category: "video",
+    creator: ModelCreator.Kling,
+    selectorName: "Kling 2.1 Master",
+    selectorDescription: "Fast video model",
+    selectorBadges: ["2 min."],
+    startFrame: true,
+    endFrame: false,
+  }),
+  new VideoModel({
+    id: "seedance_1_0_lite",
+    tauriId: "seedance_1.0_lite",
+    fullName: "Seedance 1.0 Lite",
+    category: "video",
+    creator: ModelCreator.Bytedance,
+    selectorName: "Seedance 1.0 Lite",
+    selectorDescription: "Fast video model",
+    selectorBadges: ["2 min."],
+    startFrame: true,
+    endFrame: false,
+  }),
+  new VideoModel({
+    id: "veo_2",
+    tauriId: "veo_2",
+    fullName: "Google Veo 2",
+    category: "video",
+    creator: ModelCreator.Google,
+    selectorName: "Google Veo 2",
+    selectorDescription: "Fast video model",
+    selectorBadges: ["2 min."],
+    startFrame: true,
+    endFrame: false,
+  }),
+];

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

@@ -1,5 +1,5 @@
 import { ModelCreator } from "../ModelCreator.js";
-import { ModelCategory } from "../ModelConfig.js";
+import { ModelCategory, ModelConfig } from "../ModelConfig.js";
 
 // NB: Do not create instances of this class directly, use subclasses.
 export class Model {
@@ -48,4 +48,29 @@ export class Model {
     this.selectorDescription = args.selectorDescription;
     this.selectorBadges = args.selectorBadges;
   }
+
+  toLegacyBadges() : { label: string }[] {
+    return this.selectorBadges.map((b) => ({ label: b }));
+  }
+
+  // TODO: This is a method to support migration. Kill it after we no longer need it.
+  toLegacyModelConfig(): ModelConfig {
+    return {
+      id: this.id,
+      label: this.selectorName,
+      description: this.selectorDescription,
+      badges: this.toLegacyBadges(),
+      category: this.category,
+      info: {
+        name: this.fullName,
+        tauri_id: this.tauriId,
+        creator: this.creator,
+      },
+      capabilities: {
+        maxGenerationCount: 9, // NB: Sentinel value to detect continued use
+        defaultGenerationCount: 9, // NB: Sentinel value to detect continued use
+      },
+      tags: [],
+    };
+  }
 }

frontend/libs/tauri-api/src/lib/enqueue/EnqueueImageToVideo.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 } from "@storyteller/model-list";
 
 export enum EnqueueImageToVideoErrorType {
   /// Caller didn't specify a model
@@ -15,7 +15,7 @@ export enum EnqueueImageToVideoErrorType {
 
 export interface EnqueueImageToVideoRequest {
   // Required. The model to use.
-  model?: ModelInfo;
+  model?: Model;
 
   // Required. Starting frame.
   image_media_token?: string;
@@ -49,10 +49,9 @@ export interface EnqueueImageToVideoSuccess extends CommandResult {
 export type EnqueueImageToVideoResult = EnqueueImageToVideoSuccess | EnqueueImageToVideoError;
 
 export const EnqueueImageToVideo = async (request: EnqueueImageToVideoRequest) : Promise<EnqueueImageToVideoResult> => {
-  const modelName = request.model?.tauri_id;
-
+  
   let mutableRequest : RawEnqueueImageToVideoRequest = {
-    model: modelName,
+    model: request.model?.tauriId,
     image_media_token: request.image_media_token,
   };