lots of update 2

lama_cleaner/app/package.json

@@ -18,12 +18,14 @@
     "@testing-library/jest-dom": "^5.14.1",
     "@testing-library/react": "^12.1.2",
     "@testing-library/user-event": "^13.5.0",
+    "@types/flexsearch": "^0.7.3",
     "@types/jest": "^27.0.2",
     "@types/lodash": "^4.14.182",
     "@types/node": "^16.11.1",
     "@types/react": "^17.0.30",
     "@types/react-dom": "^17.0.9",
     "cross-env": "7.x",
+    "flexsearch": "0.7.21",
     "hacktimer": "^1.1.3",
     "lodash": "^4.17.21",
     "mitt": "^3.0.0",

lama_cleaner/app/src/App.tsx

@@ -1,5 +1,5 @@
 import React, { useCallback, useEffect, useMemo } from 'react'
-import { useRecoilState } from 'recoil'
+import { useRecoilState, useSetRecoilState } from 'recoil'
 import { nanoid } from 'nanoid'
 import useInputImage from './hooks/useInputImage'
 import { themeState } from './components/Header/ThemeChanger'
@@ -30,14 +30,10 @@ const SUPPORTED_FILE_TYPE = [
 function App() {
   const [file, setFile] = useRecoilState(fileState)
   const [theme, setTheme] = useRecoilState(themeState)
-  const [toastVal, setToastState] = useRecoilState(toastState)
+  const setToastState = useSetRecoilState(toastState)
   const userInputImage = useInputImage()
-  const [isDisableModelSwitch, setIsDisableModelSwitch] = useRecoilState(
-    isDisableModelSwitchState
-  )
-  const [enableFileManager, setEnableFileManager] = useRecoilState(
-    enableFileManagerState
-  )
+  const setIsDisableModelSwitch = useSetRecoilState(isDisableModelSwitchState)
+  const setEnableFileManager = useSetRecoilState(enableFileManagerState)
 
   // Set Input Image
   useEffect(() => {
@@ -70,7 +66,7 @@ function App() {
       setEnableFileManager(isEnabled === 'true')
     }
     fetchData2()
-  }, [])
+  }, [setEnableFileManager, setIsDisableModelSwitch])
 
   // Dark Mode Hotkey
   useHotKey(
@@ -118,35 +114,38 @@ function App() {
     setIsDragging(false)
   }, [])
 
-  const handleDrop = React.useCallback(event => {
-    event.preventDefault()
-    event.stopPropagation()
-    setIsDragging(false)
-    if (event.dataTransfer.files && event.dataTransfer.files.length > 0) {
-      if (event.dataTransfer.files.length > 1) {
-        setToastState({
-          open: true,
-          desc: 'Please drag and drop only one file',
-          state: 'error',
-          duration: 3000,
-        })
-      } else {
-        const dragFile = event.dataTransfer.files[0]
-        const fileType = dragFile.type
-        if (SUPPORTED_FILE_TYPE.includes(fileType)) {
-          setFile(dragFile)
-        } else {
+  const handleDrop = React.useCallback(
+    event => {
+      event.preventDefault()
+      event.stopPropagation()
+      setIsDragging(false)
+      if (event.dataTransfer.files && event.dataTransfer.files.length > 0) {
+        if (event.dataTransfer.files.length > 1) {
           setToastState({
             open: true,
-            desc: 'Please drag and drop an image file',
+            desc: 'Please drag and drop only one file',
             state: 'error',
             duration: 3000,
           })
+        } else {
+          const dragFile = event.dataTransfer.files[0]
+          const fileType = dragFile.type
+          if (SUPPORTED_FILE_TYPE.includes(fileType)) {
+            setFile(dragFile)
+          } else {
+            setToastState({
+              open: true,
+              desc: 'Please drag and drop an image file',
+              state: 'error',
+              duration: 3000,
+            })
+          }
         }
+        event.dataTransfer.clearData()
       }
-      event.dataTransfer.clearData()
-    }
-  }, [])
+    },
+    [setToastState, setFile]
+  )
 
   const onPaste = useCallback((event: any) => {
     // TODO: when sd side panel open, ctrl+v not work

lama_cleaner/app/src/components/FileManager/FileManager.scss

@@ -13,6 +13,9 @@
 }
 
 .react-photo-album--photo {
+  -moz-user-select: none;
+  -webkit-user-select: none;
+  user-select: none;
   border-radius: 8px;
   border: 1px solid transparent;
 
@@ -82,3 +85,17 @@
 .ScrollAreaCorner {
   background: var(--blackA8);
 }
+
+.file-search-input {
+  width: 250px;
+  padding-left: 30px;
+  height: 32px;
+  border: 1px solid var(--border-color);
+  border-radius: 8px;
+}
+
+.sort-btn-inactive {
+  svg {
+    opacity: 0.5;
+  }
+}

lama_cleaner/app/src/components/FileManager/FileManager.tsx

@@ -4,13 +4,25 @@ import React, {
   useMemo,
   useState,
   useCallback,
+  useRef,
+  FormEvent,
 } from 'react'
-import { useRecoilState } from 'recoil'
-import PhotoAlbum, { RenderPhoto } from 'react-photo-album'
+import _ from 'lodash'
+import { useSetRecoilState } from 'recoil'
+import PhotoAlbum from 'react-photo-album'
+import { BarsArrowDownIcon, BarsArrowUpIcon } from '@heroicons/react/24/outline'
+import { MagnifyingGlassIcon } from '@radix-ui/react-icons'
+import { useDebounce } from 'react-use'
+import { Id, Index, IndexSearchResult } from 'flexsearch'
 import * as ScrollArea from '@radix-ui/react-scroll-area'
 import Modal from '../shared/Modal'
+import Flex from '../shared/Layout'
 import { toastState } from '../../store/Atoms'
 import { getMedias } from '../../adapters/inpainting'
+import Selector from '../shared/Selector'
+import Button from '../shared/Button'
+import TextInput from '../shared/Input'
+import { useAsyncMemo } from '../../hooks/useAsyncMemo'
 
 interface Photo {
   src: string
@@ -22,26 +34,26 @@ interface Filename {
   name: string
   height: number
   width: number
+  ctime: number
 }
 
-const renderPhoto: RenderPhoto = ({
-  layout,
-  layoutOptions,
-  imageProps: { alt, style, ...restImageProps },
-}) => (
-  <div
-    style={{
-      boxSizing: 'content-box',
-      alignItems: 'center',
-    }}
-  >
-    <img
-      alt={alt}
-      style={{ ...style, width: '100%', padding: 0 }}
-      {...restImageProps}
-    />
-  </div>
-)
+enum SortOrder {
+  DESCENDING = 'desc',
+  ASCENDING = 'asc',
+}
+
+enum SortBy {
+  NAME = 'name',
+  CTIME = 'ctime',
+}
+
+const SORT_BY_NAME = 'Name'
+const SORT_BY_CREATED_TIME = 'Created time'
+
+const SortByMap = {
+  [SortBy.NAME]: SORT_BY_NAME,
+  [SortBy.CTIME]: SORT_BY_CREATED_TIME,
+}
 
 interface Props {
   show: boolean
@@ -55,7 +67,20 @@ export default function FileManager(props: Props) {
   const [filenames, setFileNames] = useState<Filename[]>([])
   const [scrollTop, setScrollTop] = useState(0)
   const [closeScrollTop, setCloseScrollTop] = useState(0)
-  const [toastVal, setToastState] = useRecoilState(toastState)
+  const setToastState = useSetRecoilState(toastState)
+  const [sortBy, setSortBy] = useState<SortBy>(SortBy.CTIME)
+  const [sortOrder, setSortOrder] = useState<SortOrder>(SortOrder.DESCENDING)
+  const ref = useRef(null)
+  const [searchText, setSearchText] = useState('')
+  const [debouncedSearchText, setDebouncedSearchText] = useState('')
+
+  const [, cancel] = useDebounce(
+    () => {
+      setDebouncedSearchText(searchText)
+    },
+    500,
+    [searchText]
+  )
 
   useEffect(() => {
     if (!show) {
@@ -98,20 +123,37 @@ export default function FileManager(props: Props) {
     if (show) {
       fetchData()
     }
-  }, [show])
+  }, [show, setToastState])
 
   const onScroll = (event: SyntheticEvent) => {
     setScrollTop(event.currentTarget.scrollTop)
   }
 
-  const photos = useMemo(() => {
-    return filenames.map((filename: Filename) => {
-      const width = photoWidth
-      const height = filename.height * (width / filename.width)
-      const src = `/media_thumbnail/${filename.name}?width=${width}&height=${height}`
-      return { src, height, width }
-    })
-  }, [filenames])
+  const filteredFilenames: Filename[] | undefined = useAsyncMemo(async () => {
+    if (!debouncedSearchText) {
+      return filenames
+    }
+
+    const index = new Index()
+    filenames.forEach((filename: Filename, id: number) =>
+      index.add(id, filename.name)
+    )
+    const results: IndexSearchResult = await index.searchAsync(
+      debouncedSearchText
+    )
+    return results.map((id: Id) => filenames[id as number])
+  }, [filenames, debouncedSearchText])
+
+  const photos: Photo[] = useMemo(() => {
+    return _.orderBy(filteredFilenames, sortBy, sortOrder).map(
+      (filename: Filename) => {
+        const width = photoWidth
+        const height = filename.height * (width / filename.width)
+        const src = `/media_thumbnail/${filename.name}?width=${width}&height=${height}`
+        return { src, height, width }
+      }
+    )
+  }, [filteredFilenames, photoWidth, sortBy, sortOrder])
 
   return (
     <Modal
@@ -120,6 +162,66 @@ export default function FileManager(props: Props) {
       className="file-manager-modal"
       show={show}
     >
+      <Flex style={{ justifyContent: 'end', gap: 8 }}>
+        <Flex
+          style={{
+            position: 'relative',
+            justifyContent: 'start',
+          }}
+        >
+          <MagnifyingGlassIcon style={{ position: 'absolute', left: 8 }} />
+          <TextInput
+            ref={ref}
+            value={searchText}
+            className="file-search-input"
+            tabIndex={-1}
+            onInput={(evt: FormEvent<HTMLInputElement>) => {
+              evt.preventDefault()
+              evt.stopPropagation()
+              const target = evt.target as HTMLInputElement
+              setSearchText(target.value)
+            }}
+            placeholder="Search by file name"
+          />
+        </Flex>
+        <Flex style={{ gap: 8 }}>
+          <Selector
+            width={130}
+            value={SortByMap[sortBy]}
+            options={Object.values(SortByMap)}
+            onChange={val => {
+              if (val === SORT_BY_CREATED_TIME) {
+                setSortBy(SortBy.CTIME)
+              } else {
+                setSortBy(SortBy.NAME)
+              }
+            }}
+            chevronDirection="down"
+          />
+          <Button
+            icon={<BarsArrowDownIcon />}
+            toolTip="Descending order"
+            tooltipPosition="bottom"
+            onClick={() => {
+              setSortOrder(SortOrder.DESCENDING)
+            }}
+            className={
+              sortOrder !== SortOrder.DESCENDING ? 'sort-btn-inactive' : ''
+            }
+          />
+          <Button
+            icon={<BarsArrowUpIcon />}
+            toolTip="Ascending order"
+            tooltipPosition="bottom"
+            onClick={() => {
+              setSortOrder(SortOrder.ASCENDING)
+            }}
+            className={
+              sortOrder !== SortOrder.ASCENDING ? 'sort-btn-inactive' : ''
+            }
+          />
+        </Flex>
+      </Flex>
       <ScrollArea.Root className="ScrollAreaRoot">
         <ScrollArea.Viewport
           className="ScrollAreaViewport"
@@ -129,9 +231,8 @@ export default function FileManager(props: Props) {
           <PhotoAlbum
             layout="masonry"
             photos={photos}
-            renderPhoto={renderPhoto}
             spacing={8}
-            padding={8}
+            padding={0}
             onClick={onClick}
           />
         </ScrollArea.Viewport>
@@ -141,12 +242,12 @@ export default function FileManager(props: Props) {
         >
           <ScrollArea.Thumb className="ScrollAreaThumb" />
         </ScrollArea.Scrollbar>
-        <ScrollArea.Scrollbar
+        {/* <ScrollArea.Scrollbar
           className="ScrollAreaScrollbar"
           orientation="horizontal"
         >
           <ScrollArea.Thumb className="ScrollAreaThumb" />
-        </ScrollArea.Scrollbar>
+        </ScrollArea.Scrollbar> */}
         <ScrollArea.Corner className="ScrollAreaCorner" />
       </ScrollArea.Root>
     </Modal>

lama_cleaner/app/src/components/Workspace.tsx

@@ -1,5 +1,5 @@
-import React, { useEffect, useRef } from 'react'
-import { useRecoilState, useRecoilValue } from 'recoil'
+import React, { useEffect } from 'react'
+import { useRecoilState, useRecoilValue, useSetRecoilState } from 'recoil'
 import Editor from './Editor/Editor'
 import ShortcutsModal from './Shortcuts/ShortcutsModal'
 import SettingModal from './Settings/SettingsModal'
@@ -24,7 +24,7 @@ import PESidePanel from './SidePanel/PESidePanel'
 import FileManager from './FileManager/FileManager'
 
 const Workspace = () => {
-  const [file, setFile] = useRecoilState(fileState)
+  const setFile = useSetRecoilState(fileState)
   const [settings, setSettingState] = useRecoilState(settingState)
   const [toastVal, setToastState] = useRecoilState(toastState)
   const isSD = useRecoilValue(isSDState)

lama_cleaner/app/src/components/shared/Layout.tsx

@@ -0,0 +1,27 @@
+import React, { ReactNode } from 'react'
+
+interface Props {
+  children: ReactNode
+  className?: string
+  style?: React.CSSProperties
+}
+
+const Flex: React.FC<Props> = props => {
+  const { children, className, style } = props
+
+  return (
+    <div
+      style={{
+        display: 'flex',
+        justifyContent: 'center',
+        alignItems: 'center',
+        ...style,
+      }}
+      className={className}
+    >
+      {children}
+    </div>
+  )
+}
+
+export default Flex

lama_cleaner/app/src/hooks/useAsyncMemo.tsx

@@ -0,0 +1,33 @@
+import { DependencyList, useEffect, useState } from 'react'
+
+export function useAsyncMemo<T>(
+  factory: () => Promise<T> | undefined | null,
+  deps: DependencyList
+): T | undefined
+export function useAsyncMemo<T>(
+  factory: () => Promise<T> | undefined | null,
+  deps: DependencyList,
+  initial: T
+): T
+export function useAsyncMemo<T>(
+  factory: () => Promise<T> | undefined | null,
+  deps: DependencyList,
+  initial?: T
+) {
+  const [val, setVal] = useState<T | undefined>(initial)
+
+  useEffect(() => {
+    let cancel = false
+    const promise = factory()
+    if (promise === undefined || promise === null) return
+    promise.then(v => {
+      if (!cancel) {
+        setVal(v)
+      }
+    })
+    return () => {
+      cancel = true
+    }
+  }, deps)
+  return val
+}

lama_cleaner/app/yarn.lock

@@ -3526,6 +3526,11 @@
     "@types/qs" "*"
     "@types/serve-static" "*"
 
+"@types/flexsearch@^0.7.3":
+  version "0.7.3"
+  resolved "https://registry.npmmirror.com/@types/flexsearch/-/flexsearch-0.7.3.tgz#ee79b1618035c82284278e05652e91116765b634"
+  integrity sha512-HXwADeHEP4exXkCIwy2n1+i0f1ilP1ETQOH5KDOugjkTFZPntWo0Gr8stZOaebkxsdx+k0X/K6obU/+it07ocg==
+
 "@types/graceful-fs@^4.1.2":
   version "4.1.5"
   resolved "https://registry.npmjs.org/@types/graceful-fs/-/graceful-fs-4.1.5.tgz"
@@ -6411,6 +6416,11 @@ flatted@^3.1.0:
   resolved "https://registry.npmjs.org/flatted/-/flatted-3.2.2.tgz"
   integrity sha512-JaTY/wtrcSyvXJl4IMFHPKyFur1sE9AUqc0QnhOaJ0CxHtAoIV8pYDzeEfAaNEtGkOfq4gr3LBFmdXW5mOQFnA==
 
+flexsearch@0.7.21:
+  version "0.7.21"
+  resolved "https://registry.npmmirror.com/flexsearch/-/flexsearch-0.7.21.tgz#0f5ede3f2aae67ddc351efbe3b24b69d29e9d48b"
+  integrity sha512-W7cHV7Hrwjid6lWmy0IhsWDFQboWSng25U3VVywpHOTJnnAZNPScog67G+cVpeX9f7yDD21ih0WDrMMT+JoaYg==
+
 follow-redirects@^1.0.0:
   version "1.14.4"
   resolved "https://registry.npmjs.org/follow-redirects/-/follow-redirects-1.14.4.tgz"

lama_cleaner/file_manager/file_manager.py

@@ -85,8 +85,9 @@ class FileManager:
         names = sorted([it.name for it in glob_img(self.root_directory)])
         res = []
         for name in names:
-            img = Image.open(os.path.join(self.root_directory, name))
-            res.append({"name": name, "height": img.height, "width": img.width})
+            path = os.path.join(self.root_directory, name)
+            img = Image.open(path)
+            res.append({"name": name, "height": img.height, "width": img.width, "ctime": os.path.getctime(path)})
         return res
 
     @property

lama_cleaner/file_manager/utils.py

@@ -48,7 +48,7 @@ def aspect_to_string(size):
     return "x".join(map(str, size))
 
 
-IMG_SUFFIX = {'.jpg', '.jpeg', '.png'}
+IMG_SUFFIX = {'.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG'}
 
 
 def glob_img(p: Union[Path, str], recursive: bool = False):

lama_cleaner/model/sd.py

@@ -62,6 +62,7 @@ class SD(InpaintModel):
 
         if kwargs.get('cpu_offload', False) and torch.cuda.is_available():
             # TODO: gpu_id
+            logger.info("Enable sequential cpu offload")
             self.model.enable_sequential_cpu_offload(gpu_id=0)
         else:
             if kwargs['sd_cpu_textencoder']:

lama_cleaner/tests/test_paint_by_example.py

@@ -78,3 +78,18 @@ def test_paint_by_example_cpu_offload(strategy):
         fy=0.9,
         fx=1.3
     )
+
+
+@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
+def test_paint_by_example_cpu_offload_cpu_device(strategy):
+    model = ModelManager(name="paint_by_example", device = torch.device('cpu'), cpu_offload=True)
+    cfg = get_config(strategy, paint_by_example_steps=1, sd_scale=0.85)
+    assert_equal(
+        model,
+        cfg,
+        f"paint_by_example_{strategy.capitalize()}_cpu_offload_cpu_device.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+        fy=0.9,
+        fx=1.3
+    )

lama_cleaner/tests/test_sd_model.py

@@ -181,3 +181,28 @@ def test_runway_sd_1_5_cpu_offload(sd_device, strategy, sampler):
         img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
         mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
     )
+
+
+@pytest.mark.parametrize("sd_device", ['cpu'])
+@pytest.mark.parametrize("strategy", [HDStrategy.ORIGINAL])
+@pytest.mark.parametrize("sampler", [SDSampler.k_euler_a])
+def test_runway_sd_1_5_cpu_offload_cpu_device(sd_device, strategy, sampler):
+    model = ModelManager(name="sd1.5",
+                         device=torch.device(sd_device),
+                         hf_access_token="",
+                         sd_run_local=True,
+                         sd_disable_nsfw=False,
+                         sd_cpu_textencoder=False,
+                         cpu_offload=True)
+    cfg = get_config(strategy, prompt='a fox sitting on a bench', sd_steps=1, sd_scale=0.85)
+    cfg.sd_sampler = sampler
+
+    name = f"device_{sd_device}_{sampler}"
+
+    assert_equal(
+        model,
+        cfg,
+        f"runway_sd_{strategy.capitalize()}_{name}_cpu_offload_cpu_device.png",
+        img_p=current_dir / "overture-creations-5sI6fQgYIuo.png",
+        mask_p=current_dir / "overture-creations-5sI6fQgYIuo_mask.png",
+    )