diff --git a/frontend/apps/editor3d/app/src/pages/PageDraw/PageDraw.tsx b/frontend/apps/editor3d/app/src/pages/PageDraw/PageDraw.tsx index 11a91ae88a..5ae04a22f1 100644 --- a/frontend/apps/editor3d/app/src/pages/PageDraw/PageDraw.tsx +++ b/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(1024); const canvasHeight = useRef(1024); - const { selectedModels } = useModelSelectorStore(); + //const { selectedModels } = useModelSelectorStore(); const [isSelecting, setIsSelecting] = useState(false); const stageRef = useRef({} 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} /> {
- = ({ onImageStyleChange, onEnqueuePressed, onAspectRatioChange, onFitPressed, - selectedModelInfo, + //selectedModelInfo, + selectedImageModel, }) => { const [images, setImages] = useState([]); @@ -44,7 +46,8 @@ const PromptEditor: React.FC = ({ onEnqueuePressed={onEnqueuePressed} onAspectRatioChange={onAspectRatioChange} onFitPressed={onFitPressed} - selectedModelInfo={selectedModelInfo} + //selectedModelInfo={selectedModelInfo} + selectedImageModel={selectedImageModel} /> diff --git a/frontend/apps/editor3d/app/src/pages/PageDraw/PromptEditor/types.ts b/frontend/apps/editor3d/app/src/pages/PageDraw/PromptEditor/types.ts index 001a917442..4978a35e12 100644 --- a/frontend/apps/editor3d/app/src/pages/PageDraw/PromptEditor/types.ts +++ b/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"; 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 1f7542e85b..a1b2c0d15b 100644 --- a/frontend/libs/components/model-selector/src/lib/model-lists.tsx +++ b/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[]; @@ -53,11 +54,6 @@ const buildItems2 = ( * OLD MODEL LISTS */ -export const CANVAS_2D_PAGE_MODEL_LIST : ModelList = buildItems( - CANVAS_2D_PAGE_MODELS, - -); - export const STAGE_3D_PAGE_MODEL_LIST : ModelList = buildItems( STAGE_3D_PAGE_MODELS, @@ -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)), + +); + + export const IMAGE_TO_VIDEO_PAGE_MODEL_LIST : ModelList = buildItems2( VIDEO_MODELS, diff --git a/frontend/libs/components/promptbox/src/lib/PromptBox2D.tsx b/frontend/libs/components/promptbox/src/lib/PromptBox2D.tsx index fb83ad8457..0592ac563d 100644 --- a/frontend/libs/components/promptbox/src/lib/PromptBox2D.tsx +++ b/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; - selectedModelInfo?: ModelInfo; + //selectedModelInfo?: ModelInfo; + selectedImageModel?: ImageModel; getCanvasRenderBitmap: () => MaybeCanvasRenderBitmapType; EncodeImageBitmapToBase64: (imageBitmap: ImageBitmap) => Promise; 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, diff --git a/frontend/libs/model-list/src/lib/classes/ImageModel.ts b/frontend/libs/model-list/src/lib/classes/ImageModel.ts index 1e59863219..beeb42897e 100644 --- a/frontend/libs/model-list/src/lib/classes/ImageModel.ts +++ b/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"); diff --git a/frontend/libs/model-list/src/lib/classes/Model.ts b/frontend/libs/model-list/src/lib/classes/Model.ts index 6d8892d55d..b58d4bd1f4 100644 --- a/frontend/libs/model-list/src/lib/classes/Model.ts +++ b/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 }[] { diff --git a/frontend/libs/model-list/src/lib/classes/VideoModel.ts b/frontend/libs/model-list/src/lib/classes/VideoModel.ts index 219f59ea87..9814f11b92 100644 --- a/frontend/libs/model-list/src/lib/classes/VideoModel.ts +++ b/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; diff --git a/frontend/libs/model-list/src/lib/lists/ImageModels.ts b/frontend/libs/model-list/src/lib/lists/ImageModels.ts index eb68d93d3a..a10ffb5143 100644 --- a/frontend/libs/model-list/src/lib/lists/ImageModels.ts +++ b/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", diff --git a/frontend/libs/tauri-api/src/lib/enqueue/EnqueueContextualEditImage.ts b/frontend/libs/tauri-api/src/lib/enqueue/EnqueueContextualEditImage.ts index 848d927447..809bdf14b3 100644 --- a/frontend/libs/tauri-api/src/lib/enqueue/EnqueueContextualEditImage.ts +++ b/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; }