initial version of native gallery

modules/api/api.py

@@ -5,7 +5,7 @@ from fastapi import FastAPI, APIRouter, Depends, Request
 from fastapi.security import HTTPBasic, HTTPBasicCredentials
 from fastapi.exceptions import HTTPException
 from modules import errors, shared, postprocessing
-from modules.api import models, endpoints, script, helpers, server, nvml, generate, process, control
+from modules.api import models, endpoints, script, helpers, server, nvml, generate, process, control, gallery
 
 
 errors.install()
@@ -83,6 +83,9 @@ class Api:
         self.add_api_route("/sdapi/v1/reload-checkpoint", endpoints.post_reload_checkpoint, methods=["POST"])
         self.add_api_route("/sdapi/v1/refresh-vae", endpoints.post_refresh_vae, methods=["POST"])
 
+        # gallery api
+        gallery.register_api(app)
+
     def add_api_route(self, path: str, endpoint, **kwargs):
         if (shared.cmd_opts.auth or shared.cmd_opts.auth_file) and shared.cmd_opts.api_only:
             return self.app.add_api_route(path, endpoint, dependencies=[Depends(self.auth)], **kwargs)

modules/api/gallery.py

@@ -0,0 +1,114 @@
+import os
+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
+
+
+debug = shared.log.debug if os.environ.get('SD_BROWSER_DEBUG', None) is not None else lambda *args, **kwargs: None
+
+
+OPTS_FOLDERS = [
+    "outdir_samples",
+    "outdir_txt2img_samples",
+    "outdir_img2img_samples",
+    "outdir_control_samples",
+    "outdir_extras_samples",
+    "outdir_save"
+    "outdir_video",
+    "outdir_init_images",
+    "outdir_grids",
+    "outdir_txt2img_grids",
+    "outdir_img2img_grids",
+    "outdir_control_grids",
+]
+
+### class definitions
+
+class ReqFiles(BaseModel):
+    folder: str = Field(title="Folder")
+
+### ws connection manager
+
+class ConnectionManager:
+    def __init__(self):
+        self.active: list[WebSocket] = []
+
+    async def connect(self, ws: WebSocket):
+        await ws.accept()
+        agent = ws._headers.get("user-agent", "") # pylint: disable=protected-access
+        debug(f'Browser WS connect: client={ws.client.host} agent="{agent}"')
+        self.active.append(ws)
+
+    def disconnect(self, ws: WebSocket):
+        debug(f'Browser WS disconnect: client={ws.client.host}')
+        self.active.remove(ws)
+
+    async def send(self, ws: WebSocket, data: str|dict|bytes):
+        # debug(f'Browser WS send: client={ws.client.host} data={type(data)}')
+        if ws.client_state != WebSocketState.CONNECTED:
+            return
+        if isinstance(data, bytes):
+            await ws.send_bytes(data)
+        elif isinstance(data, dict):
+            await ws.send_json(data)
+        elif isinstance(data, str):
+            await ws.send_text(data)
+        else:
+            debug(f'Browser WS send: client={ws.client.host} data={type(data)} unknown')
+
+    async def broadcast(self, data: str|dict|bytes):
+        for ws in self.active:
+            await self.send(ws, data)
+
+### api definitions
+
+def register_api(app: FastAPI): # register api
+    manager = ConnectionManager()
+
+    @app.get('/sdapi/v1/browser/folders', response_model=List[str])
+    def get_folders():
+        folders = [shared.opts.data.get(f, '') for f in OPTS_FOLDERS]
+        folders += list(shared.opts.browser_folders.split(','))
+        folders = [f.strip() for f in folders if f != '']
+        folders = list(dict.fromkeys(folders)) # filter duplicates
+        folders = [f for f in folders if os.path.isdir(f)]
+        if shared.demo is not None:
+            for f in folders:
+                if os.path.isabs(f) and f not in shared.demo.allowed_paths:
+                    debug(f'Browser folders allow: {f}')
+                    shared.demo.allowed_paths.append(f)
+        debug(f'Browser folders: {folders}')
+        return JSONResponse(content=folders)
+
+    @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}')
+            for f in files_cache.directory_files(folder, recursive=True):
+                file = os.path.relpath(f, folder)
+                stat = os.stat(f)
+                dct = {
+                    'folder': folder,
+                    'file': file,
+                    'size': stat.st_size,
+                    'mtime': stat.st_mtime,
+                }
+                await manager.send(ws, dct)
+            await manager.send(ws, '#END#')
+        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)