From dc9c60e1749ecb80180734007041dc718b7e5b8d Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 23 Mar 2024 14:46:55 -0400 Subject: [PATCH] gallery optimizations --- .eslintrc.json | 7 ++- CHANGELOG.md | 6 +-- javascript/gallery.js | 91 +++++++++++++++++++++++++++++++++++---- javascript/imageViewer.js | 8 +++- javascript/indexdb.js | 83 +++++++++++++++++++++++++++++++++++ modules/api/gallery.py | 41 +++++++++++++----- modules/api/middleware.py | 2 +- modules/images.py | 11 ++--- modules/shared.py | 3 +- requirements.txt | 4 +- 10 files changed, 221 insertions(+), 35 deletions(-) create mode 100644 javascript/indexdb.js diff --git a/.eslintrc.json b/.eslintrc.json index 6bd030605..e0fe61fba 100644 --- a/.eslintrc.json +++ b/.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", diff --git a/CHANGELOG.md b/CHANGELOG.md index 34517bb1c..58347bfe1 100644 --- a/CHANGELOG.md +++ b/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* diff --git a/javascript/gallery.js b/javascript/gallery.js index 0e140f79a..2e5eef383 100644 --- a/javascript/gallery.js +++ b/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('
', '\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'); diff --git a/javascript/imageViewer.js b/javascript/imageViewer.js index fd72ddfa4..1e3536716 100644 --- a/javascript/imageViewer.js +++ b/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 = `Image ${el.src} Size ${el.naturalWidth}x${el.naturalHeight}
`; let html = ''; let params; diff --git a/javascript/indexdb.js b/javascript/indexdb.js new file mode 100644 index 000000000..789a15b46 --- /dev/null +++ b/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); diff --git a/modules/api/gallery.py b/modules/api/gallery.py index d6113d57c..ce717ddb8 100644 --- a/modules/api/gallery.py +++ b/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) diff --git a/modules/api/middleware.py b/modules/api/middleware.py index 4c004ca08..540275bea 100644 --- a/modules/api/middleware.py +++ b/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, diff --git a/modules/images.py b/modules/images.py index 8ff381bf5..ee7ed6f53 100644 --- a/modules/images.py +++ b/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 diff --git a/modules/shared.py b/modules/shared.py index 8546c5ca4..ac87984dd 100644 --- a/modules/shared.py +++ b/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("

Image Browser

", "", gr.HTML), + "image_sep_browser": OptionInfo("

Image Gallery

", "", 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"), diff --git a/requirements.txt b/requirements.txt index f15b40793..8ba4cd7be 100644 --- a/requirements.txt +++ b/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