cleanup wip

frontend/apps/editor3d/app/src/pages/PageDraw/PageDraw.tsx

@@ -11,14 +11,16 @@ import Konva from "konva";
 import { setCanvasRenderBitmap } from "../../signals/canvasRenderBitmap";
 import { captureStageImageBitmap } from "./hooks/useUpdateSnapshot";
 import { ContextMenuContainer } from "./components/ui/ContextMenu";
-import { ModelInfo } from "@storyteller/model-list";
+import { ImageModel, ModelInfo } from "@storyteller/model-list";
 import {
   CANVAS_2D_PAGE_MODEL_LIST,
+  ClassyModelSelector,
   ModelPage,
   ModelSelector,
   useModelSelectorStore,
 } from "@storyteller/ui-model-selector";
 import { useCanvasBgRemovedEvent } from "@storyteller/tauri-api";
+import { getSelectedImageModel } from "@storyteller/ui-model-selector";
 
 const PAGE_ID : ModelPage = ModelPage.Canvas2D;
 
@@ -50,20 +52,22 @@ export const DecodeBase64ToImage = async (
 const PageDraw = () => {
   const canvasWidth = useRef<number>(1024);
   const canvasHeight = useRef<number>(1024);
-  const { selectedModels } = useModelSelectorStore();
+  //const { selectedModels } = useModelSelectorStore();
   const [isSelecting, setIsSelecting] = useState<boolean>(false);
   const stageRef = useRef<Konva.Stage>({} as Konva.Stage);
   const transformerRefs = useRef<{ [key: string]: Konva.Transformer }>({});
   const store = useSceneStore();
 
-  const selectedModel =
-    selectedModels[PAGE_ID] ||
-    CANVAS_2D_PAGE_MODEL_LIST[0]?.label;
+  //const selectedModel =
+  //  selectedModels[PAGE_ID] ||
+  //  CANVAS_2D_PAGE_MODEL_LIST[0]?.label;
 
-  const selectedModelInfo: ModelInfo | undefined =
-    CANVAS_2D_PAGE_MODEL_LIST.find(
-      (m) => m.label === selectedModel,
-    )?.modelInfo;
+  //const selectedModelInfo: ModelInfo | undefined =
+  //  CANVAS_2D_PAGE_MODEL_LIST.find(
+  //    (m) => m.label === selectedModel,
+  //  )?.modelInfo;
+
+  const selectedImageModel : ImageModel | undefined = getSelectedImageModel(PAGE_ID);
 
   useDeleteHotkeys({ onDelete: store.deleteSelectedItems });
   useUndoRedoHotkeys({ undo: store.undo, redo: store.redo });
@@ -305,7 +309,8 @@ const PageDraw = () => {
           }}
           onEnqueuePressed={onEnqueuedPressed}
           onFitPressed={onFitPressed}
-          selectedModelInfo={selectedModelInfo}
+          //selectedModelInfo={selectedModelInfo}
+          selectedImageModel={selectedImageModel}
         />
       </div>
       <SideToolbar
@@ -446,7 +451,7 @@ const PageDraw = () => {
         </ContextMenuContainer>
       </div>
       <div className="absolute bottom-6 left-6 z-20 flex items-center gap-2">
-        <ModelSelector
+        <ClassyModelSelector
           items={CANVAS_2D_PAGE_MODEL_LIST}
           page={PAGE_ID}
           panelTitle="Select Model"

frontend/apps/editor3d/app/src/pages/PageDraw/PromptEditor/PromptEditor.tsx

@@ -9,12 +9,14 @@ import { JobProvider, useJobContext } from "../JobContext";
 
 // Set this value on when enqueue is pressed nasty global variable.
 import { getCanvasRenderBitmap } from "../../../signals/canvasRenderBitmap";
+
 const PromptEditor: React.FC<PromptEditorProps> = ({
   onImageStyleChange,
   onEnqueuePressed,
   onAspectRatioChange,
   onFitPressed,
-  selectedModelInfo,
+  //selectedModelInfo,
+  selectedImageModel,
 }) => {
   const [images, setImages] = useState<ImageStyle[]>([]);
 
@@ -44,7 +46,8 @@ const PromptEditor: React.FC<PromptEditorProps> = ({
           onEnqueuePressed={onEnqueuePressed}
           onAspectRatioChange={onAspectRatioChange}
           onFitPressed={onFitPressed}
-          selectedModelInfo={selectedModelInfo}
+          //selectedModelInfo={selectedModelInfo}
+          selectedImageModel={selectedImageModel}
         />
       </JobProvider>
 

frontend/apps/editor3d/app/src/pages/PageDraw/PromptEditor/types.ts

@@ -1,3 +1,4 @@
+import { ImageModel } from "libs/model-list/src/lib/classes/ImageModel";
 import { ModelInfo } from "libs/model-list/src/lib/ModelInfo";
 
 export interface PromptEditorProps {
@@ -11,6 +12,7 @@ export interface PromptEditorProps {
   onEnqueuePressed?: () => void;
   onFitPressed?: () => void;
   selectedModelInfo?: ModelInfo;
+  selectedImageModel?: ImageModel;
 }
 
 export type AspectRatio = "1:1" | "3:2" | "2:3";

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

@@ -11,6 +11,7 @@ import {
   VIDEO_MODELS,
   IMAGE_MODELS,
 } from "@storyteller/model-list";
+import { ModelTag } from "libs/model-list/src/lib/ModelTag";
 
 export type ModelList = Omit<PopoverItem, "selected">[];
 
@@ -53,11 +54,6 @@ const buildItems2 = (
  * OLD MODEL LISTS
  */ 
 
-export const CANVAS_2D_PAGE_MODEL_LIST : ModelList = buildItems(
-  CANVAS_2D_PAGE_MODELS,
-  <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
-);
-
 export const STAGE_3D_PAGE_MODEL_LIST : ModelList = buildItems(
   STAGE_3D_PAGE_MODELS,
   <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
@@ -75,6 +71,12 @@ export const IMAGE_EDITOR_PAGE_MODEL_LIST : ModelList = buildItems(
  * access to the object directly.
  */ 
 
+export const CANVAS_2D_PAGE_MODEL_LIST : ModelList = buildItems2(
+  IMAGE_MODELS.filter((m) => m.tags?.includes(ModelTag.InstructiveEdit)),
+  <FontAwesomeIcon icon={faImage} className="h-4 w-4" />
+);
+
+
 export const IMAGE_TO_VIDEO_PAGE_MODEL_LIST : ModelList = buildItems2(
   VIDEO_MODELS,
   <FontAwesomeIcon icon={faFilm} className="h-4 w-4" />

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

@@ -48,6 +48,7 @@ import {
 import { usePrompt2DStore, RefImage } from "./promptStore";
 import { gtagEvent } from "@storyteller/google-analytics";
 import { ModelInfo, getCapabilitiesForModel } from "@storyteller/model-list";
+import { ImageModel } from "@storyteller/model-list";
 
 export type AspectRatio = "1:1" | "3:2" | "2:3";
 
@@ -61,7 +62,8 @@ interface PromptBox2DProps {
     assetFile: File;
     progressCallback: (newState: UploaderState) => void;
   }) => Promise<void>;
-  selectedModelInfo?: ModelInfo;
+  //selectedModelInfo?: ModelInfo;
+  selectedImageModel?: ImageModel;
   getCanvasRenderBitmap: () => MaybeCanvasRenderBitmapType;
   EncodeImageBitmapToBase64: (imageBitmap: ImageBitmap) => Promise<string>;
   useJobContext: () => JobContextType;
@@ -77,7 +79,8 @@ export const PromptBox2D = ({
   useJobContext,
   onEnqueuePressed,
   onAspectRatioChange,
-  selectedModelInfo,
+  //selectedModelInfo,
+  selectedImageModel,
   onFitPressed,
 }: PromptBox2DProps) => {
   useSignals();
@@ -129,7 +132,7 @@ export const PromptBox2D = ({
 
   // Update generation count options from selected model capabilities
   useEffect(() => {
-    const caps = getCapabilitiesForModel(selectedModelInfo);
+    const caps = getCapabilitiesForModel(selectedImageModel);
     const items: PopoverItem[] = Array.from(
       { length: caps.maxGenerationCount },
       (_, i) => i + 1
@@ -146,7 +149,7 @@ export const PromptBox2D = ({
       );
     }
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [selectedModelInfo]);
+  }, [selectedImageModel]);
 
   // Keep UI selection in sync when store value changes
   useEffect(() => {
@@ -379,7 +382,7 @@ export const PromptBox2D = ({
     const aspectRatio = getCurrentAspectRatio();
 
     const generateResponse = await EnqueueContextualEditImage({
-      model: selectedModelInfo,
+      model: selectedImageModel,
       scene_image_media_token: snapshotMediaToken.data!,
       image_media_tokens: referenceImages.map((image) => image.mediaToken),
       disable_system_prompt: !useSystemPrompt,

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

@@ -1,6 +1,7 @@
 import { ModelCreator } from "src/index.js";
 import { Model } from "./Model.js";
 import { ModelCategory } from "../ModelConfig.js";
+import { ModelTag } from "../ModelTag.js";
 
 export class ImageModel extends Model {
   readonly maxGenerationCount: number;
@@ -21,6 +22,7 @@ export class ImageModel extends Model {
     maxGenerationCount: number;
     defaultGenerationCount: number;
     usesInpaintingMask?: boolean;
+    tags?: ModelTag[];
   }) {
     if (args.maxGenerationCount < 1) {
       throw new Error("maxGenerationCount must be at least 1");

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

@@ -1,5 +1,6 @@
 import { ModelCreator } from "../ModelCreator.js";
 import { ModelCategory, ModelConfig } from "../ModelConfig.js";
+import { ModelTag } from "../ModelTag.js";
 
 // NB: Do not create instances of this class directly, use subclasses.
 export class Model {
@@ -29,6 +30,9 @@ export class Model {
   // Labels for the selector
   readonly selectorBadges: string[];
 
+  // A list of filterable "capabilities" that can be used to filter models.
+  readonly tags: ModelTag[];
+
   protected constructor(args: {
     id: string;
     tauriId: string;
@@ -38,6 +42,7 @@ export class Model {
     selectorName: string;
     selectorDescription: string;
     selectorBadges: string[];
+    tags?: ModelTag[];
   }) {
     this.id = args.id;
     this.tauriId = args.tauriId;
@@ -47,6 +52,7 @@ export class Model {
     this.selectorName = args.selectorName;
     this.selectorDescription = args.selectorDescription;
     this.selectorBadges = args.selectorBadges;
+    this.tags = args.tags ?? [];
   }
 
   toLegacyBadges() : { label: string }[] {

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

@@ -1,6 +1,7 @@
 import { ModelCreator } from "src/index.js";
 import { Model } from "./Model.js";
 import { ModelCategory } from "../ModelConfig.js";
+import { ModelTag } from "../ModelTag.js";
 
 export class VideoModel extends Model {
   // Whether the model supports image starting frames
@@ -20,6 +21,7 @@ export class VideoModel extends Model {
     selectorBadges: string[];
     startFrame: boolean;
     endFrame: boolean;
+    tags?: ModelTag[];
   }) {
     super(args);
     this.startFrame = args.startFrame;

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

@@ -1,5 +1,6 @@
 import { ImageModel } from "../classes/ImageModel.js";
 import { ModelCreator } from "../ModelCreator.js";
+import { ModelTag } from "../ModelTag.js";
 
 export const IMAGE_MODELS: ImageModel[] = [
   new ImageModel({
@@ -73,6 +74,7 @@ export const IMAGE_MODELS: ImageModel[] = [
     selectorBadges: ["45 sec."],
     maxGenerationCount: 1,
     defaultGenerationCount: 1,
+    tags: [ModelTag.InstructiveEdit],
   }),
   new ImageModel({
     id: "flux_pro_kontext_max",
@@ -85,6 +87,7 @@ export const IMAGE_MODELS: ImageModel[] = [
     selectorBadges: ["20 sec."],
     maxGenerationCount: 4,
     defaultGenerationCount: 4,
+    tags: [ModelTag.InstructiveEdit],
   }),
   new ImageModel({
     id: "flux_pro_inpaint",

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

@@ -1,11 +1,11 @@
 import { invoke } from "@tauri-apps/api/core";
 import { CommandResult } from "../common/CommandStatus";
-import { ModelInfo } from "@storyteller/model-list";
+import { ImageModel, ModelInfo } from "@storyteller/model-list";
 
 export interface EnqueueContextualEditImageRequest {
 
   /// The model to use.
-  model?: ModelInfo | EnqueueContextualEditImageModel;
+  model?: ImageModel | ModelInfo | EnqueueContextualEditImageModel;
 
   /// If set, this is the first image in the contextual image set.
   /// This gets submitted along with `image_media_tokens` and will 
@@ -95,6 +95,8 @@ export const EnqueueContextualEditImage = async (request: EnqueueContextualEditI
   if (!!request.model) {
     if (typeof request.model === "string") {
       modelName = request.model;
+    } else if (request.model instanceof ImageModel) {
+      modelName = request.model.tauriId;
     } else {
       modelName = request.model.tauri_id;
     }