add mask tab

web_app/src/components/FileManager.tsx

@@ -50,6 +50,7 @@ const SORT_BY_MODIFIED_TIME = "Modified time"
 
 const IMAGE_TAB = "input"
 const OUTPUT_TAB = "output"
+export const MASK_TAB = "mask"
 
 const SortByMap = {
   [SortBy.NAME]: SORT_BY_NAME,
@@ -264,6 +265,7 @@ export default function FileManager(props: Props) {
             <TabsList aria-label="Manage your account">
               <TabsTrigger value={IMAGE_TAB}>Image Directory</TabsTrigger>
               <TabsTrigger value={OUTPUT_TAB}>Output Directory</TabsTrigger>
+              <TabsTrigger value={MASK_TAB}>Mask Directory</TabsTrigger>
             </TabsList>
           </Tabs>
 

web_app/src/components/Header.tsx

@@ -7,8 +7,8 @@ import { useImage } from "@/hooks/useImage"
 import { Popover, PopoverContent, PopoverTrigger } from "./ui/popover"
 import PromptInput from "./PromptInput"
 import { RotateCw, Image, Upload } from "lucide-react"
-import FileManager from "./FileManager"
-import { getMediaFile } from "@/lib/api"
+import FileManager, { MASK_TAB } from "./FileManager"
+import { getMediaBlob, getMediaFile } from "@/lib/api"
 import { useStore } from "@/lib/states"
 import SettingsDialog from "./Settings"
 import { cn, fileToImage } from "@/lib/utils"
@@ -31,6 +31,7 @@ const Header = () => {
     hidePrevMask,
     imageHeight,
     imageWidth,
+    handleFileManagerMaskSelect,
   ] = useStore((state) => [
     state.file,
     state.customMask,
@@ -46,6 +47,7 @@ const Header = () => {
     state.hidePrevMask,
     state.imageHeight,
     state.imageWidth,
+    state.handleFileManagerMaskSelect,
   ])
 
   const { toast } = useToast()
@@ -64,25 +66,29 @@ const Header = () => {
     hidePrevMask()
   }
 
+  const handleOnPhotoClick = async (tab: string, filename: string) => {
+    try {
+      if (tab === MASK_TAB) {
+        const maskBlob = await getMediaBlob(tab, filename)
+        handleFileManagerMaskSelect(maskBlob)
+      } else {
+        const newFile = await getMediaFile(tab, filename)
+        setFile(newFile)
+      }
+    } catch (e: any) {
+      toast({
+        variant: "destructive",
+        description: e.message ? e.message : e.toString(),
+      })
+      return
+    }
+  }
+
   return (
     <header className="h-[60px] px-6 py-4 absolute top-[0] flex justify-between items-center w-full z-20 border-b backdrop-filter backdrop-blur-md bg-background/70">
       <div className="flex items-center gap-1">
         {serverConfig.enableFileManager ? (
-          <FileManager
-            photoWidth={512}
-            onPhotoClick={async (tab: string, filename: string) => {
-              try {
-                const newFile = await getMediaFile(tab, filename)
-                setFile(newFile)
-              } catch (e: any) {
-                toast({
-                  variant: "destructive",
-                  description: e.message ? e.message : e.toString(),
-                })
-                return
-              }
-            }}
-          />
+          <FileManager photoWidth={512} onPhotoClick={handleOnPhotoClick} />
         ) : (
           <></>
         )}

web_app/src/lib/api.ts

@@ -135,7 +135,7 @@ export async function runPlugin(
     body: JSON.stringify({
       name,
       image: imageBase64,
-      scale:upscale,
+      scale: upscale,
       clicks,
     }),
   })
@@ -167,6 +167,23 @@ export async function getMediaFile(tab: string, filename: string) {
   throw new Error(errMsg.errors)
 }
 
+export async function getMediaBlob(tab: string, filename: string) {
+  const res = await fetch(
+    `${API_ENDPOINT}/media_file?tab=${tab}&filename=${encodeURIComponent(
+      filename
+    )}`,
+    {
+      method: "GET",
+    }
+  )
+  if (res.ok) {
+    const blob = await res.blob()
+    return blob
+  }
+  const errMsg = await res.json()
+  throw new Error(errMsg.errors)
+}
+
 export async function getMedias(tab: string): Promise<Filename[]> {
   const res = await api.get(`medias`, { params: { tab } })
   return res.data

web_app/src/lib/states.ts

@@ -207,6 +207,7 @@ type AppAction = {
   updateInteractiveSegState: (newState: Partial<InteractiveSegState>) => void
   resetInteractiveSegState: () => void
   handleInteractiveSegAccept: () => void
+  handleFileManagerMaskSelect: (blob: Blob) => Promise<void>
   showPromptInput: () => boolean
 
   runInpainting: () => Promise<void>
@@ -903,6 +904,16 @@ export const useStore = createWithEqualityFn<AppState & AppAction>()(
         })
       },
 
+      handleFileManagerMaskSelect: async (blob: Blob) => {
+        const newMask = new Image()
+
+        await loadImage(newMask, URL.createObjectURL(blob))
+        set((state) => {
+          state.editorState.extraMasks.push(castDraft(newMask))
+        })
+        get().runInpainting()
+      },
+
       setIsInpainting: (newValue: boolean) =>
         set((state) => {
           state.isInpainting = newValue