diff --git a/frontend/apps/editor3d/app/src/pages/PageEdit/PageEdit.tsx b/frontend/apps/editor3d/app/src/pages/PageEdit/PageEdit.tsx index 808352a53b..33d7a3138f 100644 --- a/frontend/apps/editor3d/app/src/pages/PageEdit/PageEdit.tsx +++ b/frontend/apps/editor3d/app/src/pages/PageEdit/PageEdit.tsx @@ -24,16 +24,30 @@ import { EditPaintSurface } from "./EditPaintSurface"; import { normalizeCanvas } from "../../Helpers/CanvasHelpers"; import { BaseImageSelector, BaseSelectorImage } from "./BaseImageSelector"; import DrawToolControlBar from "./DrawToolControlBar"; +import { imageEditorModels, ModelCategory, ModelSelector, useModelSelectorStore } from "@storyteller/ui-model-selector"; +import { ModelInfo } from "@storyteller/model-list"; const PageEdit = () => { //useStateSceneLoader(); + const { selectedModels } = useModelSelectorStore(); + + const selectedModel = + selectedModels[ModelCategory.Canvas2D] || + imageEditorModels[0]?.label; + + const selectedModelInfo: ModelInfo | undefined = + imageEditorModels.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"); + //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 +56,7 @@ const PageEdit = () => { const [isEnqueuing, setIsEnqueuing] = useState(false); const [generationCount, setGenerationCount] = useState(1); + // Use the Zustand store const store = useEditStore(); @@ -115,27 +130,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 +257,7 @@ const PageEdit = () => { }`} > { store.setActiveTool(mode as ActiveEditTool); }} @@ -328,10 +323,9 @@ const PageEdit = () => {
- []; @@ -35,6 +36,9 @@ const imageModels = getModelsByCategory("image"); const videoModels = getModelsByCategory("video"); const instructiveModels = getInstructiveImageEditModels(); +// Models for the editor page (not 2d canvas, the inpaint/outpaint editor) +const imageEditModels = getEditorModels(); + export const allModels = { video: buildItems( videoModels as any, @@ -48,9 +52,15 @@ export const allModels = { instructiveModels as any, ), + imageEditorModels: buildItems( + imageEditModels as any, + + ), }; export const videoGenerationModels: ModelList = allModels.video; export const imageGenerationModels: ModelList = allModels.image; export const instructiveImageEditModels: ModelList = allModels.instructiveImageEdits; + +export const imageEditorModels: ModelList = allModels.imageEditorModels; diff --git a/frontend/libs/components/promptbox/src/lib/PromptBoxEdit.tsx b/frontend/libs/components/promptbox/src/lib/PromptBoxEdit.tsx index 7465d9cbad..37ab4d87cf 100644 --- a/frontend/libs/components/promptbox/src/lib/PromptBoxEdit.tsx +++ b/frontend/libs/components/promptbox/src/lib/PromptBoxEdit.tsx @@ -63,6 +63,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/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..d3314cf702 100644 --- a/frontend/libs/model-list/src/lib/Models.ts +++ b/frontend/libs/model-list/src/lib/Models.ts @@ -1,5 +1,6 @@ import { ModelCreator } from "./ModelCreator.js"; import { ModelInfo } from "./ModelInfo.js"; +import { ModelTag } from "./ModelTag.js"; export type ModelCategory = "image" | "video"; @@ -11,7 +12,7 @@ export interface ModelConfig { category: ModelCategory; info: ModelInfo; capabilities: ModelCapabilities; - tags?: string[]; // optional tags, e.g. ["instructiveEdit"] - for filtering + tags?: (ModelTag|string)[]; // optional tags, e.g. ["instructiveEdit"] - for filtering } const mc = ModelCreator; @@ -130,7 +131,23 @@ export const ALL_MODELS: ModelConfig[] = [ description: "Fast and high-quality model", badges: [{ label: "20 sec." }], capabilities: { maxGenerationCount: 4 }, - tags: ["instructiveEdit"], + 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: 4 }, + tags: [ModelTag.MaskedInpainting], }), ////////////////////////////// @@ -204,6 +221,20 @@ 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) + ); + +// Models for the editor page +export const getEditorModels = (): ModelConfig[] => + ALL_MODELS.filter( + (m) => m.category === "image" && ( + m.tags?.includes(ModelTag.MaskedInpainting) + || m.tags?.includes(ModelTag.NonMaskedInpainting) + ) + ); + export const lookupModelByTauriId = ( tauriId: string ): ModelConfig | undefined =>