diff --git a/frontend/apps/editor3d/app/src/pages/PageImage/TextToImage.tsx b/frontend/apps/editor3d/app/src/pages/PageImage/TextToImage.tsx index 537b686374..8c7e32b2b2 100644 --- a/frontend/apps/editor3d/app/src/pages/PageImage/TextToImage.tsx +++ b/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(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); }} /> @@ -184,7 +186,7 @@ const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => { {!showPromptAtBottom && }
- void; +} + +export const useClassyModelSelectorStore = create((set) => ({ + selectedModels: {}, + setSelectedModel: (page, model) => + set((state) => ({ + selectedModels: { + ...state.selectedModels, + [page]: model, + }, + })), +})); diff --git a/frontend/libs/components/model-selector/src/lib/classy-model-selector.tsx b/frontend/libs/components/model-selector/src/lib/classy-model-selector.tsx new file mode 100644 index 0000000000..9499ad7ff1 --- /dev/null +++ b/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[]; + 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 ( + + ); +} + +export default ClassyModelSelector; 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 3fd26cde05..fab1642b28 100644 --- a/frontend/libs/components/model-selector/src/lib/model-lists.tsx +++ b/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[]; @@ -53,11 +54,6 @@ const buildItems2 = ( * OLD MODEL LISTS */ -export const TEXT_TO_IMAGE_PAGE_MODEL_LIST : ModelList = buildItems( - TEXT_TO_IMAGE_PAGE_MODELS, - -); - export const CANVAS_2D_PAGE_MODEL_LIST : ModelList = buildItems( CANVAS_2D_PAGE_MODELS, @@ -84,3 +80,8 @@ export const IMAGE_TO_VIDEO_PAGE_MODEL_LIST : ModelList = buildItems2( VIDEO_MODELS, ); + +export const TEXT_TO_IMAGE_PAGE_MODEL_LIST : ModelList = buildItems2( + IMAGE_MODELS, + +); diff --git a/frontend/libs/components/model-selector/src/lib/model-selector.spec.tsx b/frontend/libs/components/model-selector/src/lib/model-selector.spec.tsx deleted file mode 100644 index 5648146f67..0000000000 --- a/frontend/libs/components/model-selector/src/lib/model-selector.spec.tsx +++ /dev/null @@ -1,12 +0,0 @@ -import { render } from '@testing-library/react'; - -import ModelSelector from './model-selector'; - -describe('ModelSelector', () => { - - it('should render successfully', () => { - const { baseElement } = render(); - expect(baseElement).toBeTruthy(); - }); - -}); diff --git a/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx b/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx index f9fca84e78..f18dd8cb96 100644 --- a/frontend/libs/components/promptbox/src/lib/PromptBoxImage.tsx +++ b/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; - 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", diff --git a/frontend/libs/model-list/src/index.ts b/frontend/libs/model-list/src/index.ts index f686c17dd9..25f68cc685 100644 --- a/frontend/libs/model-list/src/index.ts +++ b/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"; \ 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 390944181a..401059af07 100644 --- a/frontend/libs/model-list/src/lib/Models.ts +++ b/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; }; diff --git a/frontend/libs/model-list/src/lib/classes/ImageModel.ts b/frontend/libs/model-list/src/lib/classes/ImageModel.ts index 71933726de..1e59863219 100644 --- a/frontend/libs/model-list/src/lib/classes/ImageModel.ts +++ b/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; } } diff --git a/frontend/libs/model-list/src/lib/lists/AllModels.ts b/frontend/libs/model-list/src/lib/lists/AllModels.ts new file mode 100644 index 0000000000..fb6ae4b898 --- /dev/null +++ b/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]; diff --git a/frontend/libs/model-list/src/lib/lists/ImageModels.ts b/frontend/libs/model-list/src/lib/lists/ImageModels.ts new file mode 100644 index 0000000000..eb68d93d3a --- /dev/null +++ b/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, + }), +]; diff --git a/frontend/libs/model-list/src/lib/VideoModels.ts b/frontend/libs/model-list/src/lib/lists/VideoModels.ts similarity index 94% rename from frontend/libs/model-list/src/lib/VideoModels.ts rename to frontend/libs/model-list/src/lib/lists/VideoModels.ts index 08549480b2..9e4fdc9616 100644 --- a/frontend/libs/model-list/src/lib/VideoModels.ts +++ b/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({ diff --git a/frontend/libs/tauri-api/src/lib/enqueue/EnqueueTextToImage.ts b/frontend/libs/tauri-api/src/lib/enqueue/EnqueueTextToImage.ts index 0a201e7f8f..92c6c25f7c 100644 --- a/frontend/libs/tauri-api/src/lib/enqueue/EnqueueTextToImage.ts +++ b/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; }