fix flask file manager cache

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

@@ -109,6 +109,9 @@ export default function FileManager(props: Props) {
   )
 
   useEffect(() => {
+    if (!show) {
+      return
+    }
     const fetchData = async () => {
       try {
         const filenames = await getMedias(tab)
@@ -141,7 +144,15 @@ export default function FileManager(props: Props) {
       }
     }
     fetchData()
-  }, [setToastState, tab, debouncedSearchText, sortBy, sortOrder, photoWidth])
+  }, [
+    setToastState,
+    tab,
+    debouncedSearchText,
+    sortBy,
+    sortOrder,
+    photoWidth,
+    show,
+  ])
 
   const onScroll = (event: SyntheticEvent) => {
     setScrollTop(event.currentTarget.scrollTop)

lama_cleaner/file_manager/file_manager.py

@@ -1,13 +1,17 @@
 # Copy from https://github.com/silentsokolov/flask-thumbnails/blob/master/flask_thumbnails/thumbnail.py
 import os
-from cachetools import TTLCache, cached
+from datetime import datetime
+
 import cv2
 import time
 from io import BytesIO
 from pathlib import Path
 import numpy as np
+from watchdog.events import FileSystemEventHandler
+from watchdog.observers import Observer
 
 from PIL import Image, ImageOps, PngImagePlugin
+from loguru import logger
 
 LARGE_ENOUGH_NUMBER = 100
 PngImagePlugin.MAX_TEXT_CHUNK = LARGE_ENOUGH_NUMBER * (1024 ** 2)
@@ -15,7 +19,7 @@ from .storage_backends import FilesystemStorageBackend
 from .utils import aspect_to_string, generate_filename, glob_img
 
 
-class FileManager:
+class FileManager(FileSystemEventHandler):
     def __init__(self, app=None):
         self.app = app
         self._default_root_directory = "media"
@@ -28,6 +32,43 @@ class FileManager:
         if app is not None:
             self.init_app(app)
 
+        self.image_dir_filenames = []
+        self.output_dir_filenames = []
+
+        self.image_dir_observer = None
+        self.output_dir_observer = None
+
+        self.modified_time = {
+            "image": datetime.utcnow(),
+            "output": datetime.utcnow(),
+        }
+
+    def start(self):
+        self.image_dir_filenames = self._media_names(self.root_directory)
+        self.output_dir_filenames = self._media_names(self.output_dir)
+
+        logger.info(f"Start watching image directory: {self.root_directory}")
+        self.image_dir_observer = Observer()
+        self.image_dir_observer.schedule(self, self.root_directory, recursive=False)
+        self.image_dir_observer.start()
+
+        logger.info(f"Start watching output directory: {self.output_dir}")
+        self.output_dir_observer = Observer()
+        self.output_dir_observer.schedule(self, self.output_dir, recursive=False)
+        self.output_dir_observer.start()
+
+    def on_modified(self, event):
+        if not os.path.isdir(event.src_path):
+            return
+        if event.src_path == str(self.root_directory):
+            logger.info(f"Image directory {event.src_path} modified")
+            self.image_dir_filenames = self._media_names(self.root_directory)
+            self.modified_time['image'] = datetime.utcnow()
+        elif event.src_path == str(self.output_dir):
+            logger.info(f"Output directory {event.src_path} modified")
+            self.output_dir_filenames = self._media_names(self.output_dir)
+            self.modified_time['output'] = datetime.utcnow()
+
     def init_app(self, app):
         if self.app is None:
             self.app = app
@@ -80,14 +121,12 @@ class FileManager:
         return self.app.config["THUMBNAIL_MEDIA_URL"]
 
     @property
-    @cached(cache=TTLCache(maxsize=1024, ttl=30))
     def media_names(self):
-        return self._media_names(self.root_directory)
+        return self.image_dir_filenames
 
     @property
-    @cached(cache=TTLCache(maxsize=1024, ttl=30))
     def output_media_names(self):
-        return self._media_names(self.output_dir)
+        return self.output_dir_filenames
 
     @staticmethod
     def _media_names(directory: Path):

lama_cleaner/server.py

@@ -16,6 +16,7 @@ import cv2
 import torch
 import numpy as np
 from loguru import logger
+from watchdog.events import FileSystemEventHandler
 
 from lama_cleaner.interactive_seg import InteractiveSeg, Click
 from lama_cleaner.model_manager import ModelManager
@@ -71,7 +72,7 @@ app.config["JSON_AS_ASCII"] = False
 CORS(app, expose_headers=["Content-Disposition"])
 
 model: ModelManager = None
-thumb = FileManager(app)
+thumb: FileManager = None
 interactive_seg_model: InteractiveSeg = None
 device = None
 input_image_path: str = None
@@ -105,9 +106,14 @@ def save_image():
 @app.route("/medias/<tab>")
 def medias(tab):
     if tab == 'image':
-        # all images in input folder
-        return jsonify(thumb.media_names), 200
-    return jsonify(thumb.output_media_names), 200
+        response = make_response(jsonify(thumb.media_names), 200)
+    else:
+        response = make_response(jsonify(thumb.output_media_names), 200)
+    response.last_modified = thumb.modified_time[tab]
+    # response.cache_control.no_cache = True
+    # response.cache_control.max_age = 0
+    response.make_conditional(request)
+    return response
 
 
 @app.route('/media/<tab>/<filename>')
@@ -340,6 +346,11 @@ def set_input_photo():
         return "No Input Image"
 
 
+class FSHandler(FileSystemEventHandler):
+    def on_modified(self, event):
+        print("File modified: %s" % event.src_path)
+
+
 def main(args):
     global model
     global interactive_seg_model
@@ -348,6 +359,7 @@ def main(args):
     global is_disable_model_switch
     global is_enable_file_manager
     global is_desktop
+    global thumb
 
     device = torch.device(args.device)
     is_disable_model_switch = args.disable_model_switch
@@ -356,10 +368,22 @@ def main(args):
         logger.info(f"Start with --disable-model-switch, model switch on frontend is disable")
 
     if os.path.isdir(args.input):
+        logger.info(f"Initialize file manager")
+        thumb = FileManager(app)
+        is_enable_file_manager = True
         app.config["THUMBNAIL_MEDIA_ROOT"] = args.input
         app.config["THUMBNAIL_MEDIA_THUMBNAIL_ROOT"] = os.path.join(args.output_dir, 'lama_cleaner_thumbnails')
-        is_enable_file_manager = True
         thumb.output_dir = Path(args.output_dir)
+        thumb.start()
+        # try:
+        #     while True:
+        #         time.sleep(1)
+        # finally:
+        #     thumb.image_dir_observer.stop()
+        #     thumb.image_dir_observer.join()
+        #     thumb.output_dir_observer.stop()
+        #     thumb.output_dir_observer.join()
+
     else:
         input_image_path = args.input
 
@@ -388,5 +412,4 @@ def main(args):
         )
         ui.run()
     else:
-        # TODO: socketio
         app.run(host=args.host, port=args.port, debug=args.debug)

requirements.txt

@@ -12,4 +12,4 @@ markupsafe==2.0.1
 scikit-image==0.19.3
 diffusers[torch]==0.10.2
 transformers>=4.25.1
-cachetools==5.2.0
+watchdog==2.2.1