Merge branch 'main' into feature/model-flags

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

@@ -20,9 +20,10 @@ import {
   ModelSelector,
   useModelSelectorStore,
 } from "@storyteller/ui-model-selector";
-import { ModelInfo } from "@storyteller/model-list";
+import { lookupModelByTauriId, ModelInfo } from "@storyteller/model-list";
 import { useImageEditCompleteEvent } from "@storyteller/tauri-events";
 import { HistoryStack, ImageBundle } from "./HistoryStack";
+import { ModelTag } from "libs/model-list/src/lib/ModelTag";
 import LeftResetStack from "./LeftResetStack";
 
 const PAGE_ID: ModelPage = ModelPage.ImageEditor;
@@ -31,13 +32,15 @@ const PageEdit = () => {
   //useStateSceneLoader();
   const { selectedModels } = useModelSelectorStore();
 
-  const selectedModel =
+  const selectedModelLabel =
     selectedModels[PAGE_ID] || IMAGE_EDITOR_PAGE_MODEL_LIST[0]?.label;
 
+  const selectedModel = IMAGE_EDITOR_PAGE_MODEL_LIST.find(
+    (m) => m.label === selectedModelLabel,
+  );
+
   const selectedModelInfo: ModelInfo | undefined =
-    IMAGE_EDITOR_PAGE_MODEL_LIST.find(
-      (m) => m.label === selectedModel,
-    )?.modelInfo;
+    selectedModel?.modelInfo;
 
   // State for canvas dimensions
   const canvasWidth = useRef<number>(1024);
@@ -267,6 +270,24 @@ const PageEdit = () => {
     ],
   );
 
+  const modelConfig = lookupModelByTauriId(selectedModelInfo!.tauri_id);
+  const supportsMaskedInpainting = modelConfig?.tags?.includes(ModelTag.MaskedInpainting) ?? false;
+
+  if (!supportsMaskedInpainting && (store.activeTool !== "select" || store.lineNodes.length > 0)) {
+    // TODO: Implement a new mode for unsupported masking and hide the nodes layer instead of clearing
+    store.setActiveTool("select");
+    store.clearLineNodes();
+  }
+
+  /*
+    *
+    * Only component logic below this
+    *
+    *
+    *
+    *
+    */
+
   // Display image selector on launch, otherwise hide it
   // Also show loading state if info is set but image is loading
   if (!store.baseImageInfo || !store.baseImageBitmap) {
@@ -342,6 +363,7 @@ const PageEdit = () => {
           isDisabled={isEnqueuing}
           generationCount={generationCount}
           onGenerationCountChange={setGenerationCount}
+          supportsMaskedInpainting={supportsMaskedInpainting}
         />
       </div>
       <div className="relative z-0">

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

@@ -10,6 +10,7 @@ const PromptEditor: React.FC<PromptBoxEditProps> = ({
   modelInfo,
   generationCount,
   onGenerationCountChange,
+  supportsMaskedInpainting
 }) => {
   return (
     <div className="mx-auto flex w-full max-w-3xl flex-col space-y-2">
@@ -22,6 +23,7 @@ const PromptEditor: React.FC<PromptBoxEditProps> = ({
         modelInfo={modelInfo}
         generationCount={generationCount}
         onGenerationCountChange={onGenerationCountChange}
+        supportsMaskedInpainting={supportsMaskedInpainting}
       />
     </div>
   );

frontend/apps/editor3d/app/src/pages/PageEdit/stores/EditState.ts

@@ -810,7 +810,14 @@ export const useEditStore = create<EditState>((set, get, store) => ({
   },
 
   // Toolbar actions
-  setActiveTool: (tool: ActiveEditTool) => set({ activeTool: tool }),
+  setActiveTool: (tool: ActiveEditTool) => {
+    // Avoid triggering state change if the tool is same
+    if (tool === get().activeTool) {
+      return;
+    }
+
+    set({ activeTool: tool })
+  },
   setEditOperation: (mode: EditOperation) => set({ editOperation: mode }),
   setBrushColor: (color: string) => set({ brushColor: color }),
   setBrushSize: (size: number) => set({ brushSize: size }),

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

@@ -1,6 +1,6 @@
 import { PopoverMenu, type PopoverItem } from "@storyteller/ui-popover";
 import { useModelSelectorStore } from "./model-selector-store";
-import { useMemo } from "react";
+import { useEffect, useMemo } from "react";
 import { ModelPage } from "./model-pages";
 
 interface ModelSelectorProps {
@@ -22,6 +22,14 @@ export function ModelSelector({
   const { selectedModels, setSelectedModel } = useModelSelectorStore();
   const selectedModel = selectedModels[page] || items[0]?.label;
 
+  // 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].label);
+    }
+  }, []);
+
   const handleModelSelect = (item: PopoverItem) => {
     setSelectedModel(page, item.label);
   };

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

@@ -26,6 +26,7 @@ export interface PromptBoxEditProps {
   modelInfo?: ModelInfo;
   generationCount?: number;
   onGenerationCountChange?: (count: number) => void;
+  supportsMaskedInpainting?: boolean;
 }
 
 export const PromptBoxEdit = ({
@@ -37,6 +38,7 @@ export const PromptBoxEdit = ({
   modelInfo,
   generationCount: generationCountProp,
   onGenerationCountChange,
+  supportsMaskedInpainting = false,
 }: PromptBoxEditProps) => {
   const [prompt, setPrompt] = useState("");
   const [useSystemPrompt, setUseSystemPrompt] = useState(true);
@@ -180,17 +182,19 @@ export const PromptBoxEdit = ({
               onChange={handleChange}
               onPaste={handlePaste}
               onKeyDown={handleKeyDown}
-              onFocus={() => {}}
-              onBlur={() => {}}
+              onFocus={() => { }}
+              onBlur={() => { }}
             />
           </div>
           <div className="mt-2 flex items-center justify-between gap-2">
             <div className="flex items-center gap-2">
-              <ButtonIconSelect
-                options={modes}
-                onOptionChange={onModeSelectionChange}
-                selectedOption={selectedMode}
-              />
+              {supportsMaskedInpainting && (
+                <ButtonIconSelect
+                  options={modes}
+                  onOptionChange={onModeSelectionChange}
+                  selectedOption={selectedMode}
+                />
+              )}
 
               <Tooltip
                 content={