From 85c4e0edab39bdaf13af42e45ed4add7e5214fe2 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 20 Aug 2025 00:56:58 -0400 Subject: [PATCH] clean up model metadata and plumbing (#770) --- .../enqueue_image_inpaint_command.rs | 4 +- .../src/components/signaled/TopBar/TopBar.tsx | 25 +++++-- .../app/src/pages/PageDraw/PageDraw.tsx | 21 +++--- .../app/src/pages/PageEdit/PageEdit.tsx | 57 ++++++---------- .../app/src/pages/PageEnigma/PageEditor.tsx | 19 +++--- .../app/src/pages/PageImage/TextToImage.tsx | 16 +++-- .../app/src/pages/PageVideo/ImageToVideo.tsx | 16 +++-- .../components/model-selector/src/index.ts | 2 +- .../src/lib/model-categories.ts | 6 -- .../model-selector/src/lib/model-lists.tsx | 67 ++++++++++--------- .../model-selector/src/lib/model-pages.ts | 9 +++ .../src/lib/model-selector-store.ts | 10 +-- .../model-selector/src/lib/model-selector.tsx | 10 +-- .../promptbox/src/lib/PromptBoxEdit.tsx | 5 ++ frontend/libs/model-list/src/index.ts | 4 +- .../libs/model-list/src/lib/ModelConfig.ts | 22 ++++++ frontend/libs/model-list/src/lib/ModelTag.ts | 11 +++ frontend/libs/model-list/src/lib/Models.ts | 54 +++++++++------ .../libs/model-list/src/lib/ModelsByPage.ts | 35 ++++++++++ 19 files changed, 241 insertions(+), 152 deletions(-) delete mode 100644 frontend/libs/components/model-selector/src/lib/model-categories.ts create mode 100644 frontend/libs/components/model-selector/src/lib/model-pages.ts create mode 100644 frontend/libs/model-list/src/lib/ModelConfig.ts create mode 100644 frontend/libs/model-list/src/lib/ModelTag.ts create mode 100644 frontend/libs/model-list/src/lib/ModelsByPage.ts diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/enqueue_image_inpaint_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/enqueue_image_inpaint_command.rs index d825104b5f..d92ce410ca 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/enqueue_image_inpaint_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_inpaint/enqueue_image_inpaint_command.rs @@ -145,8 +145,8 @@ pub async fn enqueue_image_inpaint_command( sora_task_queue: State<'_, SoraTaskQueue>, ) -> ResponseOrErrorType { - info!("enqueue_image_inpaint_command called; model: {:?}, image_media_token: {:?}, prompt: {:?}", - &request.model, &request.image_media_token, &request.prompt); + info!("enqueue_image_inpaint_command called; model: {:?}, image_media_token: {:?}, prompt: {:?}, image_count: {:?}", + &request.model, &request.image_media_token, &request.prompt, &request.image_count); let result = handle_request( &request, diff --git a/frontend/apps/editor3d/app/src/components/signaled/TopBar/TopBar.tsx b/frontend/apps/editor3d/app/src/components/signaled/TopBar/TopBar.tsx index c2368f373e..a1869658f8 100644 --- a/frontend/apps/editor3d/app/src/components/signaled/TopBar/TopBar.tsx +++ b/frontend/apps/editor3d/app/src/components/signaled/TopBar/TopBar.tsx @@ -124,6 +124,25 @@ export const TopBar = ({ pageName, loginSignUpPressed }: Props) => { } }; + const getPageTitle = () : string => { + switch (tabStore.activeTabId) { + case "2D": + return "Canvas"; + case "3D": + return "3D Editor"; + case "IMAGE": + return "Text to Image"; + case "VIDEO": + return "Image to Video"; + case "EDIT": + return "Edit Image"; + default: + return "Artcraft"; + } + }; + + const pageTitle = getPageTitle(); + // const { toggleModal } = usePricingModalStore(); - Uncomment for pricing modal - BFlat return ( @@ -179,11 +198,7 @@ export const TopBar = ({ pageName, loginSignUpPressed }: Props) => { ) : (

- {tabStore.activeTabId === "2D" - ? "Canvas" - : tabStore.activeTabId === "VIDEO" - ? "Generate Video" - : "Generate Image"} + {pageTitle}

)} diff --git a/frontend/apps/editor3d/app/src/pages/PageDraw/PageDraw.tsx b/frontend/apps/editor3d/app/src/pages/PageDraw/PageDraw.tsx index 4f7b088b5f..11a91ae88a 100644 --- a/frontend/apps/editor3d/app/src/pages/PageDraw/PageDraw.tsx +++ b/frontend/apps/editor3d/app/src/pages/PageDraw/PageDraw.tsx @@ -7,24 +7,21 @@ import { AspectRatioType, useSceneStore } from "./stores/SceneState"; import { useUndoRedoHotkeys } from "./hooks/useUndoRedoHotkeys"; import { useDeleteHotkeys } from "./hooks/useDeleteHotkeys"; import { useCopyPasteHotkeys } from "./hooks/useCopyPasteHotkeys"; -import { PopoverItem, PopoverMenu } from "@storyteller/ui-popover"; -import { FontAwesomeIcon } from "@fortawesome/react-fontawesome"; -import { faClock, faImage } from "@fortawesome/pro-solid-svg-icons"; import Konva from "konva"; import { setCanvasRenderBitmap } from "../../signals/canvasRenderBitmap"; import { captureStageImageBitmap } from "./hooks/useUpdateSnapshot"; import { ContextMenuContainer } from "./components/ui/ContextMenu"; -import { FalBackgroundRemoval } from "@storyteller/tauri-api"; -import { EnqueueImageBgRemoval } from "@storyteller/tauri-api"; import { ModelInfo } from "@storyteller/model-list"; import { - instructiveImageEditModels, - ModelCategory, + CANVAS_2D_PAGE_MODEL_LIST, + ModelPage, ModelSelector, useModelSelectorStore, } from "@storyteller/ui-model-selector"; import { useCanvasBgRemovedEvent } from "@storyteller/tauri-api"; +const PAGE_ID : ModelPage = ModelPage.Canvas2D; + export const DecodeBase64ToImage = async ( base64String: string, ): Promise => { @@ -60,11 +57,11 @@ const PageDraw = () => { const store = useSceneStore(); const selectedModel = - selectedModels[ModelCategory.Canvas2D] || - instructiveImageEditModels[0]?.label; + selectedModels[PAGE_ID] || + CANVAS_2D_PAGE_MODEL_LIST[0]?.label; const selectedModelInfo: ModelInfo | undefined = - instructiveImageEditModels.find( + CANVAS_2D_PAGE_MODEL_LIST.find( (m) => m.label === selectedModel, )?.modelInfo; @@ -450,8 +447,8 @@ const PageDraw = () => {
{ //useStateSceneLoader(); + const { selectedModels } = useModelSelectorStore(); + + const selectedModel = + selectedModels[PAGE_ID] || + IMAGE_EDITOR_PAGE_MODEL_LIST[0]?.label; + + const selectedModelInfo: ModelInfo | undefined = + IMAGE_EDITOR_PAGE_MODEL_LIST.find( + (m) => m.label === selectedModel, + )?.modelInfo; // State for canvas dimensions const canvasWidth = useRef(1024); const canvasHeight = useRef(1024); + // Add new state to track if user is selecting const [isSelecting, setIsSelecting] = useState(false); - const [selectedModel, setSelectedModel] = useState("GPT-4o"); + // Create refs for stage and image const stageRef = useRef({} as Konva.Stage); const leftPanelRef = useRef({} as Konva.Layer); @@ -42,6 +47,7 @@ const PageEdit = () => { const [isEnqueuing, setIsEnqueuing] = useState(false); const [generationCount, setGenerationCount] = useState(1); + // Use the Zustand store const store = useEditStore(); @@ -115,27 +121,6 @@ const PageEdit = () => { }; }, [store]); - const modelList: PopoverItem[] = [ - { - label: "GPT-4o", - icon: , - selected: selectedModel === "GPT-4o", - description: "High quality model", - badges: [{ label: "2 min.", icon: }], - }, - { - label: "FLUX.1 Kontext", - icon: , - selected: selectedModel === "FLUX.1 Kontext", - description: "Fast and high-quality model", - badges: [{ label: "20 sec.", icon: }], - }, - ]; - - const handleModelSelect = (item: PopoverItem) => { - setSelectedModel(item.label); - }; - const onFitPressed = async () => { // Get the stage and its container dimensions const stage = stageRef.current; @@ -263,6 +248,7 @@ const PageEdit = () => { }`} > { store.setActiveTool(mode as ActiveEditTool); }} @@ -328,10 +314,9 @@ const PageEdit = () => {
- { useSignals(); const { triggerRecheck } = useLoginModalStore(); @@ -112,11 +111,11 @@ export const PageEditor = () => { }, []); const selectedModel = - selectedModels[ModelCategory.Editor3D] || - instructiveImageEditModels[0]?.label; + selectedModels[PAGE_ID] || + STAGE_3D_PAGE_MODEL_LIST[0]?.label; const selectedModelInfo: ModelInfo | undefined = - instructiveImageEditModels.find( + STAGE_3D_PAGE_MODEL_LIST.find( (m) => m.label === selectedModel, )?.modelInfo; @@ -550,8 +549,8 @@ export const PageEditor = () => {
{ const containerRef = useRef(null); const { selectedModels } = useModelSelectorStore(); const selectedModel = - selectedModels[ModelCategory.TextToImage] || - imageGenerationModels[0]?.label; + selectedModels[PAGE_ID] || + TEXT_TO_IMAGE_PAGE_MODEL_LIST[0]?.label; - const selectedModelInfo: ModelInfo | undefined = imageGenerationModels.find( + const selectedModelInfo: ModelInfo | undefined = TEXT_TO_IMAGE_PAGE_MODEL_LIST.find( (m) => m.label === selectedModel, )?.modelInfo; @@ -66,8 +68,8 @@ const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => {
{ const { selectedModels } = useModelSelectorStore(); const selectedModel = - selectedModels[ModelCategory.ImageToVideo] || - videoGenerationModels[0]?.label; + selectedModels[PAGE_ID] || + IMAGE_TO_VIDEO_PAGE_MODEL_LIST[0]?.label; - const selectedModelInfo: ModelInfo | undefined = videoGenerationModels.find( + const selectedModelInfo: ModelInfo | undefined = IMAGE_TO_VIDEO_PAGE_MODEL_LIST.find( (m) => m.label === selectedModel, )?.modelInfo; @@ -67,8 +69,8 @@ const ImageToVideo = ({ imageMediaId, imageUrl }: ImageToVideoProps) => {
[]; @@ -12,45 +16,42 @@ export type ModelList = Omit[]; const withIcon = (creatorIcon: any, fallback: any) => creatorIcon || fallback; const buildItems = ( - models: { - label: string; - description?: string; - badges?: { label: string }[]; - info: { creator: any }; - }[], + models: ModelConfig[], 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) => ({ + icon: withIcon(getCreatorIcon(m.info.creator), fallbackIcon), + description: m.description, + badges: m.badges?.map((b) => ({ label: b.label, icon: , })), - modelInfo: (m as any).info, + modelConfig: m, // Access to full object. + modelInfo: m.info, })); -const imageModels = getModelsByCategory("image"); -const videoModels = getModelsByCategory("video"); -const instructiveModels = getInstructiveImageEditModels(); +export const TEXT_TO_IMAGE_PAGE_MODEL_LIST : ModelList = buildItems( + TEXT_TO_IMAGE_PAGE_MODELS, + +); -export const allModels = { - video: buildItems( - videoModels as any, - - ), - image: buildItems( - imageModels as any, - - ), - instructiveImageEdits: buildItems( - instructiveModels as any, - - ), -}; +export const IMAGE_TO_VIDEO_PAGE_MODEL_LIST : ModelList = buildItems( + IMAGE_TO_VIDEO_PAGE_MODELS, + +); -export const videoGenerationModels: ModelList = allModels.video; -export const imageGenerationModels: ModelList = allModels.image; -export const instructiveImageEditModels: ModelList = - allModels.instructiveImageEdits; +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, + +); + +export const IMAGE_EDITOR_PAGE_MODEL_LIST : ModelList = buildItems( + IMAGE_EDITOR_PAGE_MODELS, + +); diff --git a/frontend/libs/components/model-selector/src/lib/model-pages.ts b/frontend/libs/components/model-selector/src/lib/model-pages.ts new file mode 100644 index 0000000000..351b48a26b --- /dev/null +++ b/frontend/libs/components/model-selector/src/lib/model-pages.ts @@ -0,0 +1,9 @@ +export enum ModelPage { + TextToImage = "text-to-image", + ImageToVideo = "image-to-video", + Canvas2D = "canvas-2d", + Stage3D = "stage-3d", + + // The "inpainting + outpainting + text" image editor + ImageEditor = "image-editor", +} diff --git a/frontend/libs/components/model-selector/src/lib/model-selector-store.ts b/frontend/libs/components/model-selector/src/lib/model-selector-store.ts index 0339e174cd..e0887c4a93 100644 --- a/frontend/libs/components/model-selector/src/lib/model-selector-store.ts +++ b/frontend/libs/components/model-selector/src/lib/model-selector-store.ts @@ -1,18 +1,18 @@ import { create } from "zustand"; -import { ModelCategory } from "./model-categories"; +import { ModelPage } from "./model-pages"; interface ModelSelectorState { - selectedModels: { [category in ModelCategory]?: string }; - setSelectedModel: (category: ModelCategory, modelLabel: string) => void; + selectedModels: { [page in ModelPage]?: string }; + setSelectedModel: (page: ModelPage, modelLabel: string) => void; } export const useModelSelectorStore = create((set) => ({ selectedModels: {}, - setSelectedModel: (category, modelLabel) => + setSelectedModel: (page, modelLabel) => set((state) => ({ selectedModels: { ...state.selectedModels, - [category]: modelLabel, + [page]: modelLabel, }, })), })); diff --git a/frontend/libs/components/model-selector/src/lib/model-selector.tsx b/frontend/libs/components/model-selector/src/lib/model-selector.tsx index 448c81961d..b11d8658df 100644 --- a/frontend/libs/components/model-selector/src/lib/model-selector.tsx +++ b/frontend/libs/components/model-selector/src/lib/model-selector.tsx @@ -1,11 +1,11 @@ import { PopoverMenu, type PopoverItem } from "@storyteller/ui-popover"; import { useModelSelectorStore } from "./model-selector-store"; import { useMemo } from "react"; -import { ModelCategory } from "./model-categories"; +import { ModelPage } from "./model-pages"; interface ModelSelectorProps { items: Omit[]; - category: ModelCategory; + page: ModelPage; mode?: "hoverSelect" | "default" | "toggle" | "button"; panelTitle?: string; buttonClassName?: string; @@ -16,14 +16,14 @@ interface ModelSelectorProps { export function ModelSelector({ items, - category, + page, ...popoverProps }: ModelSelectorProps) { const { selectedModels, setSelectedModel } = useModelSelectorStore(); - const selectedModel = selectedModels[category] || items[0]?.label; + const selectedModel = selectedModels[page] || items[0]?.label; const handleModelSelect = (item: PopoverItem) => { - setSelectedModel(category, item.label); + setSelectedModel(page, item.label); }; const modelList = useMemo( diff --git a/frontend/libs/components/promptbox/src/lib/PromptBoxEdit.tsx b/frontend/libs/components/promptbox/src/lib/PromptBoxEdit.tsx index 7465d9cbad..a7dacc3f77 100644 --- a/frontend/libs/components/promptbox/src/lib/PromptBoxEdit.tsx +++ b/frontend/libs/components/promptbox/src/lib/PromptBoxEdit.tsx @@ -37,9 +37,11 @@ export const PromptBoxEdit = ({ }: PromptBoxEditProps) => { const [prompt, setPrompt] = useState(""); const [useSystemPrompt, setUseSystemPrompt] = useState(true); + const [generationCount, setGenerationCount] = useState( typeof generationCountProp === "number" ? generationCountProp : 1 ); + const [generationCountList, setGenerationCountList] = useState( [] ); @@ -53,6 +55,8 @@ export const PromptBoxEdit = ({ } }); + console.log(">>> Model info", modelInfo); + // Sync internal state if a controlled prop is provided useEffect(() => { if (typeof generationCountProp === "number") { @@ -63,6 +67,7 @@ export const PromptBoxEdit = ({ // Build generation count options based on selected model useEffect(() => { const caps = getCapabilitiesForModel(modelInfo); + console.log(">>> caps", caps) const items: PopoverItem[] = Array.from( { length: caps.maxGenerationCount }, (_, i) => i + 1 diff --git a/frontend/libs/model-list/src/index.ts b/frontend/libs/model-list/src/index.ts index 015a1f98ea..fb88aa4e7a 100644 --- a/frontend/libs/model-list/src/index.ts +++ b/frontend/libs/model-list/src/index.ts @@ -1,5 +1,7 @@ -export * from "./lib/ModelInfo.js"; +export * from "./lib/ModelConfig.js"; export * from "./lib/ModelCreator.js"; export * from "./lib/ModelCreatorIcons.js"; +export * from "./lib/ModelInfo.js"; export * from "./lib/ModelMapping.js"; export * from "./lib/Models.js"; +export * from "./lib/ModelsByPage.js"; diff --git a/frontend/libs/model-list/src/lib/ModelConfig.ts b/frontend/libs/model-list/src/lib/ModelConfig.ts new file mode 100644 index 0000000000..7718c1e916 --- /dev/null +++ b/frontend/libs/model-list/src/lib/ModelConfig.ts @@ -0,0 +1,22 @@ +import { ModelInfo } from "./ModelInfo.js"; +import { ModelTag } from "./ModelTag.js"; + +export type ModelCategory = "image" | "video"; + +// Centralized model capability definition +// For setting options for the model +// TODO: add more capabilities here - BFlat +export interface ModelCapabilities { + maxGenerationCount: number; +} + +export interface ModelConfig { + id: string; + label: string; // UI label + description?: string; + badges?: { label: string }[]; + category: ModelCategory; + info: ModelInfo; + capabilities: ModelCapabilities; + tags?: (ModelTag|string)[]; // optional tags, e.g. ["instructiveEdit"] - for filtering +} diff --git a/frontend/libs/model-list/src/lib/ModelTag.ts b/frontend/libs/model-list/src/lib/ModelTag.ts new file mode 100644 index 0000000000..248990eb2a --- /dev/null +++ b/frontend/libs/model-list/src/lib/ModelTag.ts @@ -0,0 +1,11 @@ + +export enum ModelTag { + // Models that show up in the 2D and 3D editors + InstructiveEdit = "instructiveEdit", + + // Models that show up in the inpainting editor and that use masking. + MaskedInpainting = "maskedInpainting", + + // Models that show up in the inpainting editor, but that do not use masking. + NonMaskedInpainting = "nonMaskedInpainting", +} \ No newline at end of file diff --git a/frontend/libs/model-list/src/lib/Models.ts b/frontend/libs/model-list/src/lib/Models.ts index 63ff285491..aac23daf8e 100644 --- a/frontend/libs/model-list/src/lib/Models.ts +++ b/frontend/libs/model-list/src/lib/Models.ts @@ -1,28 +1,10 @@ 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 -} +import { ModelTag } from "./ModelTag.js"; +import { ModelConfig, ModelCapabilities, ModelCategory } from "./ModelConfig.js"; 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, }; @@ -129,8 +111,28 @@ export const ALL_MODELS: ModelConfig[] = [ }, description: "Fast and high-quality model", badges: [{ label: "20 sec." }], - capabilities: { maxGenerationCount: 4 }, - tags: ["instructiveEdit"], + capabilities: { + maxGenerationCount: 4, + }, + tags: [ + ModelTag.InstructiveEdit, + ModelTag.NonMaskedInpainting, + ], + }), + cfg({ + id: "flux_pro_inpaint", + category: "image", + info: { + name: "Flux Pro 1 (Inpainting)", + tauri_id: "flux_pro_1", + creator: mc.BlackForestLabs, + }, + description: "Fast and high-quality model", + badges: [{ label: "20 sec." }], + capabilities: { + maxGenerationCount: 1 // NB: For some reason Fal only supports ONE image! + }, + tags: [ModelTag.MaskedInpainting], }), ////////////////////////////// @@ -194,6 +196,9 @@ export const ALL_MODELS: ModelConfig[] = [ }), ]; +export const ALL_MODELS_BY_ID : Map = + new Map(ALL_MODELS.map(model => [model.id, model])); + export const getAllModels = (): ModelConfig[] => ALL_MODELS; export const getModelsByCategory = (category: ModelCategory): ModelConfig[] => @@ -204,6 +209,11 @@ export const getInstructiveImageEditModels = (): ModelConfig[] => (m) => m.category === "image" && m.tags?.includes("instructiveEdit") ); +export const getMaskedInpaintModels = (): ModelConfig[] => + ALL_MODELS.filter( + (m) => m.category === "image" && m.tags?.includes(ModelTag.MaskedInpainting) + ); + export const lookupModelByTauriId = ( tauriId: string ): ModelConfig | undefined => diff --git a/frontend/libs/model-list/src/lib/ModelsByPage.ts b/frontend/libs/model-list/src/lib/ModelsByPage.ts new file mode 100644 index 0000000000..0650b49d4a --- /dev/null +++ b/frontend/libs/model-list/src/lib/ModelsByPage.ts @@ -0,0 +1,35 @@ +import { ALL_MODELS_BY_ID } from "./Models.js"; +import { ModelConfig } from "./ModelConfig.js"; +import { ModelTag } from "./ModelTag.js"; +import { getModelsByCategory } from "./Models.js"; + +/** + * We might not want to always programatically build the list of models, because + * we likely want to control the order of models in the UI to prioritize + * more popular models. We may also want to selectively prune non-useful models. + */ + +export const TEXT_TO_IMAGE_PAGE_MODELS : ModelConfig[] = + getModelsByCategory("image") + // Remove inpainting and editing models. + .filter((m) => m.id !== "flux_pro_inpaint") + .filter((m) => m.id !== "flux_pro_kontext_max"); + +export const IMAGE_TO_VIDEO_PAGE_MODELS : ModelConfig[] = + getModelsByCategory("video"); + +export const CANVAS_2D_PAGE_MODELS : ModelConfig[] = + getModelsByCategory("image") + // Remove non-instructive models. + .filter((m) => m.tags?.includes(ModelTag.InstructiveEdit)) + +export const STAGE_3D_PAGE_MODELS : ModelConfig[] = + getModelsByCategory("image") + // Remove non-instructive models. + .filter((m) => m.tags?.includes(ModelTag.InstructiveEdit)) + +export const IMAGE_EDITOR_PAGE_MODELS : ModelConfig[] = [ + // Explicit list for now. + ALL_MODELS_BY_ID.get("flux_pro_inpaint")!, + ALL_MODELS_BY_ID.get("flux_pro_kontext_max")!, +];