gallery optimizations

.eslintrc.json

@@ -97,7 +97,12 @@
     "removeSplash": "readonly",
     // nvml.js
     "initNVML": "readonly",
-    "disableNVML": "readonly"
+    "disableNVML": "readonly",
+    // indexdb.js
+    "idbGet": "readonly",
+    "idbPut": "readonly",
+    "idbDel": "readonly",
+    "idbAdd": "readonly"
   },
   "ignorePatterns": [
     "node_modules",

CHANGELOG.md

@@ -5,16 +5,14 @@
 - Include reference styles
 - Quick apply style
 - Add refine workflow in img2img
-- Gallery client-side caching? thumbnails, metadata
-- Gallery search by metadata
 - Control API/CLI
 
 ## Update for 2024-03-23
 
 - **Features**:
   - **Gallery**:
-      implemented as infinite-scroll with lazy-loading while being fully async and non-blocking  
-      search or sort by path, name, size, width, height, mtime with extended syntax like *width > 1000*  
+      implemented as infinite-scroll with client-side-caching and lazy-loading while being fully async and non-blocking  
+      search or sort by path, name, size, width, height, mtime or any image metadata item, also with extended syntax like *width > 1000*  
       *settings*: optional additional user-defined folders, thumbnails in fixed or variable aspect-ratio  
 - **Changes**:
   - Removed built-in extensions: *ControlNet* and *Image-Browser*  

javascript/gallery.js

@@ -3,6 +3,8 @@
 let ws;
 let url;
 let currentImage;
+let pruneImagesTimer;
+let outstanding = 0;
 const el = {
   folders: undefined,
   files: undefined,
@@ -49,6 +51,32 @@ class GalleryFolder extends HTMLElement {
   }
 }
 
+async function createThumb(img) {
+  const height = opts.extra_networks_card_size;
+  const width = opts.browser_fixed_width ? opts.extra_networks_card_size : 0;
+  const canvas = document.createElement('canvas');
+  const scaleY = height / img.height;
+  const scaleX = width > 0 ? width / img.width : scaleY;
+  const scale = Math.min(scaleX, scaleY);
+  const scaledWidth = img.width * scale;
+  const scaledHeight = img.height * scale;
+  canvas.width = scaledWidth;
+  canvas.height = scaledHeight;
+  const ctx = canvas.getContext('2d');
+  ctx.drawImage(img, 0, 0, scaledWidth, scaledHeight);
+  const dataURL = canvas.toDataURL('image/jpeg', 0.5);
+  return dataURL;
+}
+
+async function delayFetchThumb(fn) {
+  while (outstanding > 10) await new Promise((resolve) => setTimeout(resolve, 50)); // eslint-disable-line no-promise-executor-return
+  outstanding++;
+  const res = await fetch(`/sdapi/v1/browser/thumb?file=${fn}`, { priority: 'low' });
+  const json = await res.json();
+  outstanding--;
+  return json;
+}
+
 class GalleryFile extends HTMLElement {
   constructor({ folder, file, size, mtime }) {
     super();
@@ -60,6 +88,7 @@ class GalleryFile extends HTMLElement {
     this.exif = '';
     this.width = 0;
     this.height = 0;
+    this.src = `${this.folder}/${this.name}`;
     this.shadow = this.attachShadow({ mode: 'open' });
   }
 
@@ -81,22 +110,54 @@ class GalleryFile extends HTMLElement {
       }
     `;
 
+    const cache = opts.browser_cache ? await idbGet(this.hash) : undefined;
     this.shadow.appendChild(style);
     const img = document.createElement('img');
     img.className = 'gallery-file';
     img.loading = 'lazy';
     img.title = `Folder: ${this.folder}\nFile: ${this.name}\nSize: ${this.size.toLocaleString()} bytes\nModified: ${this.mtime.toLocaleString()}`;
     img.onload = async () => {
-      this.width = img.naturalWidth;
-      this.height = img.naturalHeight;
       img.title += `\nResolution: ${this.width} x ${this.height}`;
       this.title = img.title;
-      // let exif = await getExif(img);
-      // if (exif) this.exif = exif.replaceAll('<br>', '\n').replace(/<\/?[^>]+(>|$)/g, "");
+      if (!cache && opts.browser_cache) {
+        if ((this.width === 0) || (this.height === 0)) { // fetch thumb failed so we use actual image
+          this.width = img.naturalWidth;
+          this.height = img.naturalHeight;
+        }
+      }
     };
-    img.src = `file=${this.folder}/${this.name}`;
+    if (cache) {
+      img.src = cache.img;
+      this.exif = cache.exif;
+      this.width = cache.width;
+      this.height = cache.height;
+    } else {
+      try {
+        const json = await delayFetchThumb(this.src);
+        img.src = json.data;
+        this.exif = json.exif;
+        this.width = json.width;
+        this.height = json.height;
+        await idbAdd({
+          hash: this.hash,
+          folder: this.folder,
+          file: this.name,
+          size: this.size,
+          mtime: this.mtime,
+          width: this.width,
+          height: this.height,
+          src: this.src,
+          exif: this.exif,
+          img: img.src,
+          // exif: await getExif(img), // alternative client-side exif
+          // img: await createThumb(img), // alternative client-side thumb
+        });
+      } catch (err) { // thumb fetch failed so assign actual image
+        img.src = `file=${this.src}`;
+      }
+    }
     img.onclick = () => {
-      currentImage = `${this.folder}/${this.name}`;
+      currentImage = this.src;
       el.btnSend.click();
     };
     this.title = img.title;
@@ -147,7 +208,6 @@ async function gallerySearch(evt) {
         const op = match[2].trim();
         let val = match[3].trim();
         if (key === 'mtime') val = new Date(val);
-        console.log('HERE', key, op, val, f[key]);
         if (((op === '=') && (f[key] === val)) || ((op === '>') && (f[key] > val)) || ((op === '<') && (f[key] < val))) {
           f.style.display = 'unset';
           numFound++;
@@ -229,6 +289,7 @@ async function fetchFiles(evt) { // fetch file-by-file list over websockets
   const t0 = performance.now();
   let t1 = performance.now();
   let lastDir;
+  let fragment = document.createDocumentFragment();
   ws.onmessage = (event) => { // time is 20% list 80% create item
     numFiles++;
     t1 = performance.now();
@@ -246,11 +307,16 @@ async function fetchFiles(evt) { // fetch file-by-file list over websockets
         el.files.appendChild(sep);
       }
       const file = new GalleryFile(json);
-      el.files.appendChild(file);
+      fragment.appendChild(file);
+      if (numFiles % 100 === 0) {
+        el.files.appendChild(fragment);
+        fragment = document.createDocumentFragment();
+      }
       el.status.innerText = `Folder | ${evt.target.name} | ${numFiles.toLocaleString()} images | ${Math.floor(t1 - t0).toLocaleString()}ms`;
     }
   };
   ws.onclose = (event) => {
+    el.files.appendChild(fragment);
     // log('gallery ws file enum', event);
   };
   ws.onerror = (event) => {
@@ -259,6 +325,10 @@ async function fetchFiles(evt) { // fetch file-by-file list over websockets
   ws.send(evt.target.name);
 }
 
+async function pruneImages() {
+  // TODO replace img.src with placeholder for images that are not visible
+}
+
 async function galleryVisible() {
   // if (el.folders.children.length > 0) return;
   const res = await fetch('/sdapi/v1/browser/folders');
@@ -270,9 +340,12 @@ async function galleryVisible() {
     const f = new GalleryFolder(folder);
     el.folders.appendChild(f);
   }
+  pruneImagesTimer = setInterval(pruneImages, 1000);
 }
 
-async function galleryHidden() { /**/ }
+async function galleryHidden() {
+  if (pruneImagesTimer) clearInterval(pruneImagesTimer);
+}
 
 async function galleryObserve() { // triggered on gradio change to monitor when ui gets sufficiently constructed
   log('initBrowser');

javascript/imageViewer.js

@@ -55,7 +55,13 @@ function modalKeyHandler(event) {
 }
 
 async function getExif(el) {
-  const exif = await window.exifr.parse(el, { userComment: true });
+  let exif = '';
+  try {
+    exif = await window.exifr.parse(el, { userComment: true });
+  } catch (e) {
+    log('getExif', el, e);
+    return exif;
+  }
   // let html = `<b>Image</b> <a href="${el.src}" target="_blank">${el.src}</a> <b>Size</b> ${el.naturalWidth}x${el.naturalHeight}<br>`;
   let html = '';
   let params;

javascript/indexdb.js

@@ -0,0 +1,83 @@
+let db;
+
+async function initIndexDB() {
+  async function createDB() {
+    return new Promise((resolve, reject) => {
+      const request = indexedDB.open('SDNext');
+      request.onerror = (evt) => reject(evt);
+      request.onsuccess = (evt) => {
+        db = evt.target.result;
+        const countAll = db
+          .transaction(['thumbs'], 'readwrite')
+          .objectStore('thumbs')
+          .count();
+        countAll.onsuccess = () => log('initIndexDB', countAll.result);
+        resolve();
+      };
+      request.onupgradeneeded = (evt) => {
+        db = evt.target.result;
+        const store = db.createObjectStore('thumbs', { keyPath: 'hash' });
+        store.createIndex('hash', 'hash', { unique: true });
+        const index = store.index('hash');
+        resolve();
+      };
+    });
+  }
+
+  if (!db) await createDB();
+}
+
+async function add(record) {
+  if (!db) return null;
+  return new Promise((resolve, reject) => {
+    const request = db
+      .transaction(['thumbs'], 'readwrite')
+      .objectStore('thumbs')
+      .add(record);
+    request.onsuccess = (evt) => resolve(evt);
+    request.onerror = (evt) => reject(evt);
+  });
+}
+
+async function del(hash) {
+  if (!db) return null;
+  return new Promise((resolve, reject) => {
+    const request = db
+      .transaction(['thumbs'], 'readwrite')
+      .objectStore('thumbs')
+      .delete(hash);
+    request.onsuccess = (evt) => resolve(evt);
+    request.onerror = (evt) => reject(evt);
+  });
+}
+
+async function get(hash) {
+  if (!db) return null;
+  return new Promise((resolve, reject) => {
+    const request = db
+      .transaction(['thumbs'], 'readwrite')
+      .objectStore('thumbs')
+      .get(hash);
+    request.onsuccess = () => resolve(request.result);
+    request.onerror = (evt) => reject(evt);
+  });
+}
+
+async function put(record) {
+  if (!db) return null;
+  return new Promise((resolve, reject) => {
+    const request = db
+      .transaction(['thumbs'], 'readwrite')
+      .objectStore('thumbs')
+      .put(record);
+    request.onsuccess = (evt) => resolve(evt);
+    request.onerror = (evt) => reject(evt);
+  });
+}
+
+window.idbAdd = add;
+window.idbDel = del;
+window.idbGet = get;
+window.idbPut = put;
+
+onUiLoaded(initIndexDB);

modules/api/gallery.py

@@ -1,10 +1,14 @@
+import io
 import os
+import time
+import base64
 from typing import List
 from fastapi import FastAPI
 from fastapi.responses import JSONResponse
 from starlette.websockets import WebSocket, WebSocketState, WebSocketDisconnect
 from pydantic import BaseModel, Field # pylint: disable=no-name-in-module
-from modules import shared, files_cache
+from PIL import Image
+from modules import shared, images, files_cache
 
 
 debug = shared.log.debug if os.environ.get('SD_BROWSER_DEBUG', None) is not None else lambda *args, **kwargs: None
@@ -83,13 +87,35 @@ def register_api(app: FastAPI): # register api
         debug(f'Browser folders: {folders}')
         return JSONResponse(content=folders)
 
+    @app.get("/sdapi/v1/browser/thumb", response_model=dict)
+    async def get_thumb(file: str):
+        image = Image.open(file)
+        geninfo, _items = images.read_info_from_image(image)
+        h = shared.opts.extra_networks_card_size
+        w = shared.opts.extra_networks_card_size if shared.opts.browser_fixed_width else image.width * h // image.height
+        width, height = image.width, image.height
+        image.thumbnail((w, h), Image.Resampling.HAMMING)
+        buffered = io.BytesIO()
+        image.save(buffered, format='jpeg')
+        data_url = f'data:image/jpeg;base64,{base64.b64encode(buffered.getvalue()).decode("ascii")}'
+        image.close()
+        content = {
+            'exif': geninfo,
+            'data': data_url,
+            'width': width,
+            'height': height,
+        }
+        return JSONResponse(content=content)
+
     @app.websocket("/sdapi/v1/browser/files")
     async def ws_files(ws: WebSocket):
         try:
             await manager.connect(ws)
             folder = await ws.receive_text()
-            debug(f'Browser WS folder: {folder}')
+            t0 = time.time()
+            numFiles = 0
             for f in files_cache.directory_files(folder, recursive=True):
+                numFiles += 1
                 file = os.path.relpath(f, folder)
                 stat = os.stat(f)
                 dct = {
@@ -100,15 +126,8 @@ def register_api(app: FastAPI): # register api
                 }
                 await manager.send(ws, dct)
             await manager.send(ws, '#END#')
+            t1 = time.time()
+            shared.log.debug(f'Gallery: folder={folder} files={numFiles} time={t1-t0:.3f}')
         except WebSocketDisconnect:
             debug('Browser WS unexpected disconnect')
         manager.disconnect(ws)
-
-    @app.websocket("/sdapi/v1/browser/file/{file}")
-    async def ws_file(ws: WebSocket, file: str):
-        try:
-            await manager.connect(ws)
-            with open(file, 'rb') as f:  # noqa: ASYNC101
-                await manager.send(ws, f.read())
-        except WebSocketDisconnect:
-            manager.disconnect(ws)

modules/api/middleware.py

@@ -43,7 +43,7 @@ def setup_middleware(app: FastAPI, cmd_opts):
             endpoint = req.scope.get('path', 'err')
             token = req.cookies.get("access-token") or req.cookies.get("access-token-unsecure")
             if (cmd_opts.api_log or cmd_opts.api_only) and endpoint.startswith('/sdapi'):
-                if '/sdapi/v1/log' in endpoint:
+                if '/sdapi/v1/log' or '/sdapi/v1/browser' in endpoint:
                     return res
                 log.info('API {user} {code} {prot}/{ver} {method} {endpoint} {cli} {duration}'.format( # pylint: disable=consider-using-f-string, logging-format-interpolation
                     user = app.tokens.get(token) if hasattr(app, 'tokens') else None,

modules/images.py

@@ -741,7 +741,7 @@ def safe_decode_string(s: bytes):
     return None
 
 
-def read_info_from_image(image: Image):
+def read_info_from_image(image: Image, watermark: bool = False):
     items = image.info or {}
     geninfo = items.pop('parameters', None)
     if geninfo is None:
@@ -772,10 +772,11 @@ def read_info_from_image(image: Image):
                             items[ExifTags.TAGS[key]] = val
                     elif val is not None and key in ExifTags.GPSTAGS:
                         items[ExifTags.GPSTAGS[key]] = val
-    wm = get_watermark(image)
-    if wm != '':
-        # geninfo += f' Watermark: {wm}'
-        items['watermark'] = wm
+    if watermark:
+        wm = get_watermark(image)
+        if wm != '':
+            # geninfo += f' Watermark: {wm}'
+            items['watermark'] = wm
 
     for key, val in items.items():
         if isinstance(val, bytes): # decode bytestring

modules/shared.py

@@ -529,7 +529,8 @@ options_templates.update(options_section(('saving-images', "Image Options"), {
     "font": OptionInfo("", "Font file"),
     "font_color": OptionInfo("#FFFFFF", "Font color", gr.ColorPicker, {}),
 
-    "image_sep_browser": OptionInfo("<h2>Image Browser</h2>", "", gr.HTML),
+    "image_sep_browser": OptionInfo("<h2>Image Gallery</h2>", "", gr.HTML),
+    "browser_cache": OptionInfo(True, "Use image gallery cache"),
     "browser_folders": OptionInfo("", "Additional image browser folders"),
     "browser_fixed_width": OptionInfo(False, "Use fixed with thumbnails"),
     "viewer_show_metadata": OptionInfo(True, "Show metadata in full screen image browser"),

requirements.txt

@@ -44,7 +44,7 @@ requests==2.31.0
 tqdm==4.66.1
 accelerate==0.28.0
 opencv-contrib-python-headless==4.9.0.80
-diffusers==0.27.0
+diffusers==0.27.2
 einops==0.4.1
 gradio==3.43.2
 huggingface_hub==0.21.4
@@ -55,7 +55,7 @@ pandas
 protobuf==3.20.3
 pytorch_lightning==1.9.4
 tokenizers==0.15.2
-transformers==4.38.2
+transformers==4.39.1
 tomesd==0.1.3
 urllib3==1.26.18
 Pillow==10.2.0