diff --git a/toolkit/dataset_sources/__init__.py b/toolkit/dataset_sources/__init__.py new file mode 100644 index 00000000..00ae4bd4 --- /dev/null +++ b/toolkit/dataset_sources/__init__.py @@ -0,0 +1,12 @@ +from .base import RemoteDatasetSource, SettingField, SourceItem +from .registry import register_source, get_source, get_all_sources, resolve_dataset_source + +__all__ = [ + "RemoteDatasetSource", + "SettingField", + "SourceItem", + "register_source", + "get_source", + "get_all_sources", + "resolve_dataset_source", +] diff --git a/toolkit/dataset_sources/base.py b/toolkit/dataset_sources/base.py new file mode 100644 index 00000000..9b3f59f4 --- /dev/null +++ b/toolkit/dataset_sources/base.py @@ -0,0 +1,180 @@ +""" +Base class and supporting types for remote dataset sources (data-source plugins). + +A RemoteDatasetSource plugin knows how to: + 1. Declare what settings it needs from the user (URL, token, etc.) + 2. Return grouped, browseable items available on the remote + 3. Declare source-specific import form fields (e.g. caption mode, score filter) + +Download logic lives in each extension's FetchProcess (e.g. PixlStashFetchProcess), +triggered explicitly by the user via the UI — never automatically during training. + +Implementations live in extensions, e.g. extensions/pixlstash/. +They register themselves via toolkit.dataset_sources.registry.register_source(). +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any, List, Optional + + +@dataclass +class SettingField: + """Describes one user-configurable setting required by a source.""" + + # Key stored in the AI-Toolkit settings DB (e.g. "PIXLSTASH_URL") + key: str + # Human-readable label shown in the UI + label: str + # Input type hint for the UI: "text" | "password" + input_type: str = "text" + # Shown below the label in the UI + description: str = "" + # Placeholder text inside the input + placeholder: str = "" + # Whether this setting must be non-empty for the plugin to be considered configured + required: bool = True + + +@dataclass +class SourceItem: + """One item inside a SourceGroup (e.g. a single character or album).""" + + id: str # string to handle both ints and UUIDs + name: str + picture_count: int = -1 # -1 = unknown + # ID and type passed to the thumbnail proxy route + thumbnail_id: str = "" + thumbnail_type: str = "" + + +@dataclass +class SourceGroup: + """A labelled collection of SourceItems shown as one tab in the browse modal.""" + + id: str # e.g. "character", "picture_set", "person", "album" + label: str # e.g. "Characters", "Albums" + items: List[SourceItem] = field(default_factory=list) + + +@dataclass +class ImportField: + """Describes one source-specific field shown in the import form.""" + + id: str + label: str + # "select" | "text" | "checkbox" + field_type: str = "text" + # For "select": list of {"value": ..., "label": ...} dicts + options: List[dict] = field(default_factory=list) + default: Any = None + required: bool = False + + +class RemoteDatasetSource(ABC): + """ + Abstract base class for remote dataset source plugins. + + Subclasses must set ``type_id`` to a unique string identifier. + """ + + # Unique plugin identifier, e.g. "pixlstash" + type_id: str = None + + # Human-readable name shown in the UI, e.g. "PixlStash" + display_name: str = "" + + # Optional absolute path to an icon image (PNG/SVG) shown in the UI + icon_path: Optional[str] = None + + def __init__(self, settings: dict) -> None: + """ + Parameters + ---------- + settings: + Key/value pairs loaded from the AI-Toolkit settings DB. + """ + self.settings = settings + + # ------------------------------------------------------------------ + # Schema — settings this plugin needs the user to configure + # ------------------------------------------------------------------ + + @classmethod + @abstractmethod + def get_settings_schema(cls) -> List[SettingField]: + """Return the list of settings fields this source requires.""" + + # ------------------------------------------------------------------ + # Thumbnail — serve a thumbnail image to the UI + # ------------------------------------------------------------------ + + @abstractmethod + def get_thumbnail(self, thumbnail_id: str, thumbnail_type: str) -> tuple: + """ + Fetch a thumbnail image for a SourceItem. + + Returns + ------- + (image_bytes: bytes, content_type: str) + """ + + # ------------------------------------------------------------------ + # Job config — describe how to run the import job + # ------------------------------------------------------------------ + + @abstractmethod + def build_job_config(self, params: dict) -> dict: + """ + Build and return the process-config dict for run.py. + + ``params`` mirrors the POST body from the import route: + source_type, source_id, trigger_word, dataset_name, overwrite, + plus any source-specific fields from get_import_fields(). + + The returned dict is placed inside: + { job: 'extension', config: { name: '...', process: [] } } + """ + + # ------------------------------------------------------------------ + # Browse — return grouped items the user can pick from + # ------------------------------------------------------------------ + + @abstractmethod + def browse(self) -> List[SourceGroup]: + """ + Return a list of SourceGroups, each representing one browseable + category (e.g. Characters, Albums). Items within each group are + displayed as thumbnails in the UI. + """ + + # ------------------------------------------------------------------ + # Import form — extra fields beyond trigger_word / dataset_name + # ------------------------------------------------------------------ + + @classmethod + def get_import_fields(cls) -> List[ImportField]: + """ + Return source-specific fields to show in the import form. + The base implementation returns an empty list (no extra fields). + Override in subclasses to add e.g. caption_mode or min_score. + """ + return [] + + # ------------------------------------------------------------------ + # Helpers available to subclasses + # ------------------------------------------------------------------ + + def get_setting(self, key: str, default: str = "") -> str: + return self.settings.get(key, default) + + @classmethod + def is_configured(cls, settings: dict) -> bool: + """Return True if all required settings have non-empty values.""" + return all( + settings.get(f.key, "").strip() + for f in cls.get_settings_schema() + if f.required + ) diff --git a/toolkit/dataset_sources/cli.py b/toolkit/dataset_sources/cli.py new file mode 100644 index 00000000..a26b0819 --- /dev/null +++ b/toolkit/dataset_sources/cli.py @@ -0,0 +1,232 @@ +""" +CLI bridge between the Next.js API routes and the Python data-source plugin system. + +Commands +-------- +plugins Print JSON array of { id, display_name, settings_schema } + for every registered source plugin. + +browse SOURCE_TYPE_ID Print a JSON object: + { groups: [...], import_fields: [...] } + where groups contains all browseable items for the + given plugin (e.g. characters + picture sets for + PixlStash). + +thumbnail SOURCE_ID THUMBNAIL_ID THUMBNAIL_TYPE + Fetch a thumbnail from the plugin and print a JSON + object: { content_type: str, data: } + +job-config SOURCE_ID Read import params JSON from stdin, print the + process-config JSON dict for run.py to stdout. + +Usage +----- + python -m toolkit.dataset_sources.cli plugins + python -m toolkit.dataset_sources.cli browse pixlstash + python -m toolkit.dataset_sources.cli thumbnail pixlstash 42 character + echo '{"source_type":"character","source_id":"1"}' | \\ + python -m toolkit.dataset_sources.cli job-config pixlstash +""" + +from __future__ import annotations + +import json +import os +import sys + + +def _load_extensions() -> None: + """Auto-discover and import all extensions so they register their sources.""" + ext_dir = os.path.join(os.path.dirname(__file__), "..", "..", "extensions") + if not os.path.isdir(ext_dir): + return + for name in sorted(os.listdir(ext_dir)): + pkg = os.path.join(ext_dir, name, "__init__.py") + if os.path.isfile(pkg): + try: + __import__(f"extensions.{name}") + except Exception: + pass # extension may have missing optional deps — skip silently + + +def cmd_plugins() -> None: + import base64 + import dataclasses + from toolkit.dataset_sources.registry import get_all_sources, load_settings_from_db + from toolkit.paths import TOOLKIT_ROOT + + db_path = os.path.join(TOOLKIT_ROOT, "aitk_db.db") + settings = load_settings_from_db(db_path) + + def _icon_data_url(cls) -> str | None: + if not cls.icon_path: + return None + try: + with open(cls.icon_path, "rb") as f: + data = base64.b64encode(f.read()).decode() + ext = os.path.splitext(cls.icon_path)[1].lower() + mime = {'.png': 'image/png', '.svg': 'image/svg+xml', '.jpg': 'image/jpeg'}.get(ext, 'image/png') + return f"data:{mime};base64,{data}" + except Exception: + return None + + out = [ + { + "id": cls.type_id, + "display_name": cls.display_name, + "icon": _icon_data_url(cls), + "settings_schema": [ + dataclasses.asdict(f) for f in cls.get_settings_schema() + ], + } + for cls in get_all_sources() + if cls.is_configured(settings) + ] + print(json.dumps(out)) + + +def cmd_browse(type_id: str) -> None: + from toolkit.dataset_sources.registry import get_source, load_settings_from_db + from toolkit.paths import TOOLKIT_ROOT + + SourceClass = get_source(type_id) + if SourceClass is None: + print(json.dumps({"error": f"Unknown plugin: {type_id}"})) + sys.exit(1) + + db_path = os.path.join(TOOLKIT_ROOT, "aitk_db.db") + settings = load_settings_from_db(db_path) + source = SourceClass(settings) + + try: + groups = source.browse() + import_fields = SourceClass.get_import_fields() + except Exception as exc: + print(json.dumps({"error": str(exc)})) + sys.exit(1) + + def _item(i): + return { + "id": i.id, + "name": i.name, + "picture_count": i.picture_count, + "thumbnail_id": i.thumbnail_id, + "thumbnail_type": i.thumbnail_type, + } + + def _group(g): + return {"id": g.id, "label": g.label, "items": [_item(i) for i in g.items]} + + def _field(f): + return { + "id": f.id, + "label": f.label, + "field_type": f.field_type, + "options": f.options, + "default": f.default, + "required": f.required, + } + + print( + json.dumps( + { + "groups": [_group(g) for g in groups], + "import_fields": [_field(f) for f in import_fields], + } + ) + ) + + +def cmd_thumbnail(source_id: str, thumbnail_id: str, thumbnail_type: str) -> None: + import base64 + from toolkit.dataset_sources.registry import get_source, load_settings_from_db + from toolkit.paths import TOOLKIT_ROOT + + SourceClass = get_source(source_id) + if SourceClass is None: + print(json.dumps({"error": f"Unknown plugin: {source_id}"})) + sys.exit(1) + + db_path = os.path.join(TOOLKIT_ROOT, "aitk_db.db") + settings = load_settings_from_db(db_path) + source = SourceClass(settings) + + try: + image_bytes, content_type = source.get_thumbnail(thumbnail_id, thumbnail_type) + print( + json.dumps( + { + "content_type": content_type, + "data": base64.b64encode(image_bytes).decode("ascii"), + } + ) + ) + except Exception as exc: + print(json.dumps({"error": str(exc)})) + sys.exit(1) + + +def cmd_job_config(source_id: str) -> None: + from toolkit.dataset_sources.registry import get_source, load_settings_from_db + from toolkit.paths import TOOLKIT_ROOT + + SourceClass = get_source(source_id) + if SourceClass is None: + print(json.dumps({"error": f"Unknown plugin: {source_id}"})) + sys.exit(1) + + db_path = os.path.join(TOOLKIT_ROOT, "aitk_db.db") + settings = load_settings_from_db(db_path) + source = SourceClass(settings) + + params = json.loads(sys.stdin.read()) + + try: + config = source.build_job_config(params) + print(json.dumps(config)) + except Exception as exc: + print(json.dumps({"error": str(exc)})) + sys.exit(1) + + +if __name__ == "__main__": + _load_extensions() + + if len(sys.argv) < 2: + print( + "Usage: python -m toolkit.dataset_sources.cli ...", + file=sys.stderr, + ) + sys.exit(1) + + cmd = sys.argv[1] + + if cmd == "plugins": + cmd_plugins() + elif cmd == "browse": + if len(sys.argv) < 3: + print( + "Usage: python -m toolkit.dataset_sources.cli browse ", + file=sys.stderr, + ) + sys.exit(1) + cmd_browse(sys.argv[2]) + elif cmd == "thumbnail": + if len(sys.argv) < 5: + print( + "Usage: python -m toolkit.dataset_sources.cli thumbnail ", + file=sys.stderr, + ) + sys.exit(1) + cmd_thumbnail(sys.argv[2], sys.argv[3], sys.argv[4]) + elif cmd == "job-config": + if len(sys.argv) < 3: + print( + "Usage: python -m toolkit.dataset_sources.cli job-config ", + file=sys.stderr, + ) + sys.exit(1) + cmd_job_config(sys.argv[2]) + else: + print(f"Unknown command: {cmd}", file=sys.stderr) + sys.exit(1) diff --git a/toolkit/dataset_sources/registry.py b/toolkit/dataset_sources/registry.py new file mode 100644 index 00000000..7569e4e6 --- /dev/null +++ b/toolkit/dataset_sources/registry.py @@ -0,0 +1,86 @@ +""" +Registry for RemoteDatasetSource implementations. + +Sources register themselves at startup by calling register_source(). +Extensions do this in their __init__.py so they are only loaded on demand. +""" + +from __future__ import annotations + +import os +from typing import Dict, List, Optional, Type, TYPE_CHECKING + +if TYPE_CHECKING: + from toolkit.dataset_sources.base import RemoteDatasetSource + +_registry: Dict[str, Type["RemoteDatasetSource"]] = {} + + +def register_source(source_class: Type["RemoteDatasetSource"]) -> None: + """Register a RemoteDatasetSource subclass by its type_id.""" + if not source_class.type_id: + raise ValueError(f"{source_class.__name__} must set type_id") + _registry[source_class.type_id] = source_class + + +def get_source(type_id: str) -> Optional[Type["RemoteDatasetSource"]]: + """Return the source class for *type_id*, or None if not registered.""" + return _registry.get(type_id) + + +def get_all_sources() -> List[Type["RemoteDatasetSource"]]: + """Return all registered source classes.""" + return list(_registry.values()) + + +def resolve_dataset_source(source_config: dict, settings: dict, cache_dir: str) -> str: + """ + Convenience function called by the dataloader. + + Looks up the correct RemoteDatasetSource by type_id, instantiates it + with the current settings, and calls resolve(). + + Parameters + ---------- + source_config: + The ``dataset_source`` dict from the YAML config. + settings: + Key/value settings dict loaded from the AI-Toolkit settings DB. + cache_dir: + Root directory for cached dataset folders. + + Returns + ------- + str + Absolute local folder path ready for AiToolkitDataset. + """ + type_id = source_config.get("type") + if not type_id: + raise ValueError("dataset_source config is missing required 'type' field") + + source_class = get_source(type_id) + if source_class is None: + raise ValueError( + f"No RemoteDatasetSource registered for type '{type_id}'. " + f"Registered types: {list(_registry.keys())}" + ) + + source = source_class(settings) + return source.resolve(source_config, cache_dir) + + +def load_settings_from_db(db_path: str) -> dict: + """ + Load the AI-Toolkit settings from the SQLite database. + Returns an empty dict if the DB does not exist or cannot be read. + """ + try: + import sqlite3 + conn = sqlite3.connect(db_path) + cursor = conn.cursor() + cursor.execute("SELECT key, value FROM Settings") + rows = cursor.fetchall() + conn.close() + return {row[0]: row[1] for row in rows} + except Exception: + return {} diff --git a/ui/src/app/api/datasets/remote/[source]/browse/route.ts b/ui/src/app/api/datasets/remote/[source]/browse/route.ts new file mode 100644 index 00000000..dbdea846 --- /dev/null +++ b/ui/src/app/api/datasets/remote/[source]/browse/route.ts @@ -0,0 +1,81 @@ +/** + * GET /api/datasets/remote/[source]/browse + * + * Calls the Python CLI to run the registered plugin's browse() method and + * returns normalized groups + import_fields for the generic browse modal. + */ +import { NextResponse } from 'next/server'; +import { spawn } from 'child_process'; +import path from 'path'; +import { existsSync } from 'fs'; +import { TOOLKIT_ROOT } from '@/paths'; + +function getPythonPath(): string { + for (const dir of ['.venv', 'venv']) { + const p = path.join(TOOLKIT_ROOT, dir, 'bin', 'python'); + if (existsSync(p)) return p; + const pw = path.join(TOOLKIT_ROOT, dir, 'Scripts', 'python.exe'); + if (existsSync(pw)) return pw; + } + return 'python'; +} + +// Short-lived in-memory cache to avoid spawning a Python process on every +// modal open or React StrictMode double-invocation. +const browseCache = new Map(); +const BROWSE_CACHE_TTL_MS = 30_000; // 30 seconds + +export async function GET( + _req: Request, + { params }: { params: { source: string } }, +) { + const { source } = params; + + const cached = browseCache.get(source); + if (cached && cached.expires > Date.now()) { + return NextResponse.json(cached.data); + } + + return new Promise(resolve => { + const child = spawn( + getPythonPath(), + ['-m', 'toolkit.dataset_sources.cli', 'browse', source], + { cwd: TOOLKIT_ROOT, stdio: ['ignore', 'pipe', 'pipe'] }, + ); + + let out = ''; + let err = ''; + child.stdout.on('data', (d: Buffer) => (out += d.toString())); + child.stderr.on('data', (d: Buffer) => (err += d.toString())); + + child.on('close', code => { + if (code !== 0) { + // CLI may write {"error": "..."} to stdout even on non-zero exit + let pluginError: string | undefined; + try { + const parsed = JSON.parse(out); + if (parsed?.error) pluginError = parsed.error; + } catch { /* ignore */ } + const msg = pluginError || err.trim() || 'unknown error'; + resolve( + NextResponse.json( + { error: `Browse failed for "${source}": ${msg}` }, + { status: 502 }, + ), + ); + return; + } + try { + const data = JSON.parse(out); + if (data.error) { + resolve(NextResponse.json({ error: data.error }, { status: 400 })); + } else { + browseCache.set(source, { data, expires: Date.now() + BROWSE_CACHE_TTL_MS }); + resolve(NextResponse.json(data)); + } + } catch { + resolve(NextResponse.json({ error: 'Invalid response from plugin' }, { status: 500 })); + } + }); + }); +} diff --git a/ui/src/app/api/datasets/remote/[source]/import/route.ts b/ui/src/app/api/datasets/remote/[source]/import/route.ts new file mode 100644 index 00000000..cf5c3a42 --- /dev/null +++ b/ui/src/app/api/datasets/remote/[source]/import/route.ts @@ -0,0 +1,175 @@ +/** + * POST /api/datasets/remote/[source]/import + * + * Generic import endpoint for any registered data-source plugin. + * Asks the Python plugin for its job config via the `job-config` CLI command, + * then spawns run.py and streams progress back as Server-Sent Events. + * + * Expected body: + * source_type string — group id from browse() (e.g. "character", "person") + * source_id string — item id from browse() + * trigger_word string — written into caption .txt files + * dataset_name string — output folder name under datasets/ + * overwrite bool — re-download even if already cached + * ... any source-specific fields declared by get_import_fields() + * + * SSE event types: + * total { count: number } + * progress { done: number, total: number } + * complete { downloaded: number } + * error { message: string } + */ +import { NextRequest, NextResponse } from 'next/server'; +import { spawn } from 'child_process'; +import { writeFileSync, mkdirSync, existsSync } from 'fs'; +import * as readline from 'readline'; +import path from 'path'; +import { TOOLKIT_ROOT } from '@/paths'; + +const isWindows = process.platform === 'win32'; + +function getPythonPath(): string { + for (const dir of ['.venv', 'venv']) { + const p = path.join(TOOLKIT_ROOT, dir, isWindows ? 'Scripts/python.exe' : 'bin/python'); + if (existsSync(p)) return p; + } + return 'python'; +} + +/** Ask the Python plugin for its process config by writing params to stdin. */ +function getJobConfig( + source: string, + params: Record, +): Promise> { + return new Promise((resolve, reject) => { + const child = spawn( + getPythonPath(), + ['-m', 'toolkit.dataset_sources.cli', 'job-config', source], + { cwd: TOOLKIT_ROOT, stdio: ['pipe', 'pipe', 'pipe'] }, + ); + + let out = ''; + let err = ''; + child.stdout.on('data', (d: Buffer) => (out += d.toString())); + child.stderr.on('data', (d: Buffer) => (err += d.toString())); + + child.stdin.write(JSON.stringify(params)); + child.stdin.end(); + + child.on('close', code => { + try { + const data = JSON.parse(out); + if (data.error || code !== 0) { + reject(new Error(data.error || err.trim() || 'job-config failed')); + } else { + resolve(data); + } + } catch { + reject(new Error(`Invalid job-config response: ${out}`)); + } + }); + }); +} + +export async function POST( + request: NextRequest, + { params }: { params: { source: string } }, +) { + const { source } = params; + const body = await request.json() as Record; + const { source_type, source_id } = body; + + if (!source_type || source_id == null) { + return NextResponse.json({ error: 'source_type and source_id are required' }, { status: 400 }); + } + + let processConfig: Record; + try { + processConfig = await getJobConfig(source, body); + } catch (e: unknown) { + return NextResponse.json({ error: String(e) }, { status: 400 }); + } + + const jobConfig = { + job: 'extension', + config: { + name: `${source}_import_${source_type}_${source_id}`, + process: [processConfig], + }, + }; + + const tmpDir = path.join(TOOLKIT_ROOT, 'output', `.${source}_tmp`); + mkdirSync(tmpDir, { recursive: true }); + const configPath = path.join( + tmpDir, + `fetch_${source_type}_${source_id}_${Date.now()}.json`, + ); + writeFileSync(configPath, JSON.stringify(jobConfig, null, 2)); + + const pythonPath = getPythonPath(); + const runFilePath = path.join(TOOLKIT_ROOT, 'run.py'); + + if (!existsSync(runFilePath)) { + return NextResponse.json({ error: 'run.py not found' }, { status: 500 }); + } + + const encoder = new TextEncoder(); + const sseEvent = (type: string, data: object) => + encoder.encode(`event: ${type}\ndata: ${JSON.stringify(data)}\n\n`); + + const stream = new ReadableStream({ + start(controller) { + const child = spawn(pythonPath, [runFilePath, configPath], { + stdio: ['ignore', 'pipe', 'pipe'], + cwd: TOOLKIT_ROOT, + }); + + let currentTotal = 0; + const errorLines: string[] = []; + + const processLine = (line: string, isStderr: boolean) => { + if (isStderr) errorLines.push(line); + const progressMatch = line.match(/^PROGRESS:(\d+)\/(\d+)$/); + if (progressMatch) { + currentTotal = parseInt(progressMatch[2]); + controller.enqueue( + sseEvent('progress', { + done: parseInt(progressMatch[1]), + total: currentTotal, + }), + ); + return; + } + const totalMatch = line.match(/— (\d+) (?:picture|image)\(s\) (?:found|downloaded)/); + if (totalMatch) { + currentTotal = parseInt(totalMatch[1]); + controller.enqueue(sseEvent('total', { count: currentTotal })); + return; + } + const doneMatch = line.match(/Done — (\d+) downloaded/); + if (doneMatch) { + controller.enqueue(sseEvent('complete', { downloaded: parseInt(doneMatch[1]) })); + } + }; + + readline.createInterface({ input: child.stdout! }).on('line', line => processLine(line, false)); + readline.createInterface({ input: child.stderr! }).on('line', line => processLine(line, true)); + + child.on('close', code => { + if (code !== 0) { + const errMsg = errorLines.slice(-10).join('\n').trim() || `Process exited with code ${code}`; + controller.enqueue(sseEvent('error', { message: errMsg })); + } + controller.close(); + }); + }, + }); + + return new Response(stream, { + headers: { + 'Content-Type': 'text/event-stream', + 'Cache-Control': 'no-cache', + Connection: 'keep-alive', + }, + }); +} diff --git a/ui/src/app/api/datasets/remote/[source]/thumbnail/[id]/route.ts b/ui/src/app/api/datasets/remote/[source]/thumbnail/[id]/route.ts new file mode 100644 index 00000000..7aa74ec7 --- /dev/null +++ b/ui/src/app/api/datasets/remote/[source]/thumbnail/[id]/route.ts @@ -0,0 +1,67 @@ +/** + * GET /api/datasets/remote/[source]/thumbnail/[id]?type= + * + * Delegates entirely to the Python plugin via the CLI `thumbnail` command. + * The core app has no knowledge of how individual plugins authenticate or + * construct thumbnail URLs. + */ +import { NextRequest, NextResponse } from 'next/server'; +import { spawn } from 'child_process'; +import path from 'path'; +import { existsSync } from 'fs'; +import { TOOLKIT_ROOT } from '@/paths'; + +function getPythonPath(): string { + for (const dir of ['.venv', 'venv']) { + const p = path.join(TOOLKIT_ROOT, dir, 'bin', 'python'); + if (existsSync(p)) return p; + const pw = path.join(TOOLKIT_ROOT, dir, 'Scripts', 'python.exe'); + if (existsSync(pw)) return pw; + } + return 'python'; +} + +export async function GET( + request: NextRequest, + { params }: { params: { source: string; id: string } }, +) { + const { source, id } = params; + const type = request.nextUrl.searchParams.get('type') || ''; + + return new Promise(resolve => { + const child = spawn( + getPythonPath(), + [ + '-m', 'toolkit.dataset_sources.cli', + 'thumbnail', source, decodeURIComponent(id), type, + ], + { cwd: TOOLKIT_ROOT, stdio: ['ignore', 'pipe', 'pipe'] }, + ); + + let out = ''; + child.stdout.on('data', (d: Buffer) => (out += d.toString())); + + child.on('close', code => { + try { + const data = JSON.parse(out); + if (data.error || code !== 0) { + resolve(new NextResponse(data.error || 'Thumbnail fetch failed', { status: 502 })); + return; + } + const imageBuffer = Buffer.from(data.data as string, 'base64'); + resolve( + new NextResponse(imageBuffer, { + status: 200, + headers: { + 'Content-Type': data.content_type as string, + 'Cache-Control': 'public, max-age=86400', + }, + }), + ); + } catch { + resolve(new NextResponse('Invalid response from plugin', { status: 500 })); + } + }); + }); +} + diff --git a/ui/src/app/api/datasets/remote/plugins/route.ts b/ui/src/app/api/datasets/remote/plugins/route.ts new file mode 100644 index 00000000..1a7924a5 --- /dev/null +++ b/ui/src/app/api/datasets/remote/plugins/route.ts @@ -0,0 +1,45 @@ +/** + * GET /api/datasets/remote/plugins + * + * Returns the list of registered data-source plugins by calling the Python CLI. + * The datasets page uses this to dynamically render "Browse " buttons. + */ +import { NextResponse } from 'next/server'; +import { spawn } from 'child_process'; +import path from 'path'; +import { existsSync } from 'fs'; +import { TOOLKIT_ROOT } from '@/paths'; + +function getPythonPath(): string { + for (const dir of ['.venv', 'venv']) { + const p = path.join(TOOLKIT_ROOT, dir, 'bin', 'python'); + if (existsSync(p)) return p; + const pw = path.join(TOOLKIT_ROOT, dir, 'Scripts', 'python.exe'); + if (existsSync(pw)) return pw; + } + return 'python'; +} + +export async function GET() { + return new Promise(resolve => { + const child = spawn( + getPythonPath(), + ['-m', 'toolkit.dataset_sources.cli', 'plugins'], + { cwd: TOOLKIT_ROOT, stdio: ['ignore', 'pipe', 'pipe'] }, + ); + + let out = ''; + child.stdout.on('data', (d: Buffer) => (out += d.toString())); + child.on('close', code => { + if (code !== 0) { + resolve(NextResponse.json({ error: 'Failed to list plugins' }, { status: 502 })); + return; + } + try { + resolve(NextResponse.json(JSON.parse(out))); + } catch { + resolve(NextResponse.json({ error: 'Invalid response from CLI' }, { status: 500 })); + } + }); + }); +} diff --git a/ui/src/app/api/settings/route.ts b/ui/src/app/api/settings/route.ts index 46cf47ee..ce7df8d9 100644 --- a/ui/src/app/api/settings/route.ts +++ b/ui/src/app/api/settings/route.ts @@ -34,31 +34,16 @@ export async function GET() { export async function POST(request: Request) { try { const body = await request.json(); - const { HF_TOKEN, TRAINING_FOLDER, DATASETS_FOLDER, MODELS_PATH } = body; - // Upsert both settings - await Promise.all([ - prisma.settings.upsert({ - where: { key: 'HF_TOKEN' }, - update: { value: HF_TOKEN }, - create: { key: 'HF_TOKEN', value: HF_TOKEN }, - }), - prisma.settings.upsert({ - where: { key: 'TRAINING_FOLDER' }, - update: { value: TRAINING_FOLDER }, - create: { key: 'TRAINING_FOLDER', value: TRAINING_FOLDER }, - }), - prisma.settings.upsert({ - where: { key: 'DATASETS_FOLDER' }, - update: { value: DATASETS_FOLDER }, - create: { key: 'DATASETS_FOLDER', value: DATASETS_FOLDER }, - }), - prisma.settings.upsert({ - where: { key: 'MODELS_PATH' }, - update: { value: MODELS_PATH }, - create: { key: 'MODELS_PATH', value: MODELS_PATH }, - }), - ]); + await Promise.all( + Object.entries(body as Record).map(([key, value]) => + prisma.settings.upsert({ + where: { key }, + update: { value: value ?? '' }, + create: { key, value: value ?? '' }, + }), + ), + ); flushCache(); diff --git a/ui/src/app/datasets/[datasetName]/page.tsx b/ui/src/app/datasets/[datasetName]/page.tsx index 94a0700f..ff6e4b06 100644 --- a/ui/src/app/datasets/[datasetName]/page.tsx +++ b/ui/src/app/datasets/[datasetName]/page.tsx @@ -20,7 +20,7 @@ export default function DatasetPage({ params }: { params: { datasetName: string const [imgList, setImgList] = useState<{ img_path: string }[]>([]); const [isAutoCaptioning, setIsAutoCaptioning] = useState(false); const usableParams = use(params as any) as { datasetName: string }; - const datasetName = usableParams.datasetName; + const datasetName = decodeURIComponent(usableParams.datasetName); const [status, setStatus] = useState<'idle' | 'loading' | 'success' | 'error'>('idle'); const { settings, isSettingsLoaded } = useSettings(); const [selectedImgPath, setSelectedImgPath] = useState(null); diff --git a/ui/src/app/datasets/page.tsx b/ui/src/app/datasets/page.tsx index 9003dabe..54ffd468 100644 --- a/ui/src/app/datasets/page.tsx +++ b/ui/src/app/datasets/page.tsx @@ -1,10 +1,11 @@ 'use client'; -import { useState } from 'react'; +import { useState, useEffect } from 'react'; import { Modal } from '@/components/Modal'; import Link from 'next/link'; import { TextInput } from '@/components/formInputs'; import useDatasetList from '@/hooks/useDatasetList'; +import RemoteSourceBrowseModal from '@/components/RemoteSourceBrowseModal'; import { Button } from '@headlessui/react'; import { FaRegTrashAlt } from 'react-icons/fa'; import { openConfirm } from '@/components/ConfirmModal'; @@ -13,12 +14,30 @@ import UniversalTable, { TableColumn } from '@/components/UniversalTable'; import { apiClient } from '@/utils/api'; import { useRouter } from 'next/navigation'; +interface RemotePlugin { + id: string; + display_name: string; + icon?: string; +} + export default function Datasets() { const router = useRouter(); const { datasets, status, refreshDatasets } = useDatasetList(); const [newDatasetName, setNewDatasetName] = useState(''); const [isNewDatasetModalOpen, setIsNewDatasetModalOpen] = useState(false); + // Remote data-source plugins + const [plugins, setPlugins] = useState([]); + const [openPluginId, setOpenPluginId] = useState(null); + + // Load registered data-source plugins once on mount + useEffect(() => { + apiClient + .get('/api/datasets/remote/plugins') + .then(res => setPlugins(res.data ?? [])) + .catch(() => setPlugins([])); + }, []); + // Transform datasets array into rows with objects const tableRows = datasets.map(dataset => ({ name: dataset, @@ -117,7 +136,19 @@ export default function Datasets() {

Datasets

-
+
+ {plugins.map(plugin => ( + + ))}
+ + {plugins.map(plugin => + plugin.settings_schema.length === 0 ? null : ( +
+
+

{plugin.display_name}

+
+ {plugin.settings_schema.map(field => ( +
+ {field.input_type === 'checkbox' ? ( + + ) : ( + <> + + + + )} +
+ ))} +
+
+ ), + )}
@@ -147,3 +221,4 @@ export default function Settings() { ); } + diff --git a/ui/src/components/RemoteSourceBrowseModal.tsx b/ui/src/components/RemoteSourceBrowseModal.tsx new file mode 100644 index 00000000..c4c65e83 --- /dev/null +++ b/ui/src/components/RemoteSourceBrowseModal.tsx @@ -0,0 +1,567 @@ +'use client'; + +import React, { useState, useEffect } from 'react'; +import { useRouter } from 'next/navigation'; +import { Modal } from '@/components/Modal'; +import { apiClient } from '@/utils/api'; + +// ─── Types (mirror Python dataclasses) ─────────────────────────────────────── + +interface SourceItem { + id: string; + name: string; + picture_count: number; + thumbnail_id: string; + thumbnail_type: string; +} + +interface SourceGroup { + id: string; + label: string; + items: SourceItem[]; +} + +interface ImportFieldOption { + value: string | number; + label: string; +} + +interface ImportField { + id: string; + label: string; + field_type: 'select' | 'text' | 'checkbox'; + options: ImportFieldOption[]; + default: string | number | boolean | null; + required: boolean; +} + +interface BrowseData { + groups: SourceGroup[]; + import_fields: ImportField[]; +} + +// ─── Thumbnail image ────────────────────────────────────────────────────────── + +function ThumbnailImage({ + sourceId, + thumbnailId, + thumbnailType, + alt, +}: { + sourceId: string; + thumbnailId: string; + thumbnailType: string; + alt: string; +}) { + const [errored, setErrored] = useState(false); + const src = `/api/datasets/remote/${sourceId}/thumbnail/${encodeURIComponent(thumbnailId)}?type=${thumbnailType}`; + if (!thumbnailId || errored) { + return ( +
?
+ ); + } + return ( + {alt} setErrored(true)} + /> + ); +} + +// ─── Progress bar ───────────────────────────────────────────────────────────── + +function ProgressBar({ done, total }: { done: number; total: number }) { + const pct = total > 0 ? Math.round((done / total) * 100) : 0; + return ( +
+
+ + {done} / {total} images + + {pct}% +
+
+
+
+
+ ); +} + +// ─── Import form ────────────────────────────────────────────────────────────── + +type ImportPhase = 'idle' | 'connecting' | 'downloading' | 'complete' | 'error'; + +interface ImportFormProps { + sourcePluginId: string; + groupId: string; + itemId: string; + itemName: string; + importFields: ImportField[]; + onSuccess: (datasetName: string) => void; + onCancel: () => void; +} + +function ImportForm({ + sourcePluginId, + groupId, + itemId, + itemName, + importFields, + onSuccess, + onCancel, +}: ImportFormProps) { + const [triggerWord, setTriggerWord] = useState(''); + const [datasetName, setDatasetName] = useState(itemName); + const [overwrite, setOverwrite] = useState(false); + const [phase, setPhase] = useState('idle'); + const [progress, setProgress] = useState({ done: 0, total: 0 }); + const [downloaded, setDownloaded] = useState(0); + const [errorMsg, setErrorMsg] = useState(null); + + // Initialise source-specific field values from their declared defaults + const [fieldValues, setFieldValues] = useState>(() => { + const init: Record = {}; + for (const f of importFields) { + init[f.id] = + f.default !== null && f.default !== undefined + ? f.default + : f.field_type === 'checkbox' + ? false + : (f.options[0]?.value ?? ''); + } + return init; + }); + + const handleImport = async () => { + setPhase('connecting'); + setErrorMsg(null); + setProgress({ done: 0, total: 0 }); + + try { + const token = + typeof window !== 'undefined' ? localStorage.getItem('AI_TOOLKIT_AUTH') : null; + const res = await fetch(`/api/datasets/remote/${sourcePluginId}/import`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + ...(token ? { Authorization: `Bearer ${token}` } : {}), + }, + body: JSON.stringify({ + source_type: groupId, + source_id: itemId, + trigger_word: triggerWord || undefined, + dataset_name: datasetName || undefined, + overwrite, + ...fieldValues, + }), + }); + + if (!res.ok || !res.body) { + const errBody = await res.json().catch(() => ({ error: 'Import failed' })); + setPhase('error'); + setErrorMsg(errBody.error || 'Import failed'); + return; + } + + const reader = res.body.getReader(); + const decoder = new TextDecoder(); + let buffer = ''; + + while (true) { + const { done, value } = await reader.read(); + if (done) break; + buffer += decoder.decode(value, { stream: true }); + + const parts = buffer.split('\n\n'); + buffer = parts.pop() ?? ''; + + for (const part of parts) { + const eventLine = part.split('\n').find(l => l.startsWith('event:')); + const dataLine = part.split('\n').find(l => l.startsWith('data:')); + if (!dataLine) continue; + + const eventType = eventLine ? eventLine.slice(7).trim() : 'message'; + const data = JSON.parse(dataLine.slice(5).trim()); + + if (eventType === 'total') { + setPhase('downloading'); + setProgress({ done: 0, total: data.count }); + } else if (eventType === 'progress') { + setPhase('downloading'); + setProgress({ done: data.done, total: data.total }); + } else if (eventType === 'complete') { + setDownloaded(data.downloaded); + setPhase('complete'); + onSuccess(datasetName || itemName); + } else if (eventType === 'error') { + setPhase('error'); + setErrorMsg(data.message || 'Download failed'); + } + } + } + } catch (err: unknown) { + setPhase('error'); + setErrorMsg(err instanceof Error ? err.message : 'Unexpected error'); + } + }; + + const inputClass = + 'w-full rounded bg-gray-700 border border-gray-600 px-3 py-2 text-gray-100 text-sm focus:outline-none focus:ring-2 focus:ring-blue-500'; + const labelClass = 'block text-sm text-gray-400 mb-1'; + const isRunning = phase === 'connecting' || phase === 'downloading'; + + if (phase === 'complete') { + return ( +
+
+ ✓ Downloaded {downloaded} image{downloaded !== 1 ? 's' : ''} to “ + {datasetName || itemName}” +
+ + +
+ ); + } + + return ( +
+
+ Downloading: {itemName} +
+ + {!isRunning && ( + <> +
+ + setDatasetName(e.target.value)} + placeholder={itemName} + /> +

+ Folder name created inside your datasets directory. +

+
+ +
+ + setTriggerWord(e.target.value)} + placeholder="e.g. myperson" + /> +

+ Prepended to every caption for LoRA training. +

+
+ + {/* Source-specific fields from plugin's get_import_fields() */} + {importFields.map(f => ( +
+ + {f.field_type === 'select' && ( + + )} + {f.field_type === 'text' && ( + setFieldValues(v => ({ ...v, [f.id]: e.target.value }))} + /> + )} + {f.field_type === 'checkbox' && ( + + )} +
+ ))} + + + + )} + + {isRunning && ( +
+
+ {phase === 'connecting' + ? 'Connecting…' + : `Downloading to "${datasetName || itemName}"…`} +
+ {progress.total > 0 ? ( + + ) : ( +
+
+
+ )} +
+ )} + + {errorMsg && ( +
+ {errorMsg} +
+ )} + +
+ + +
+
+ ); +} + +// ─── Main modal ─────────────────────────────────────────────────────────────── + +interface RemoteSourceBrowseModalProps { + /** Plugin id, e.g. "pixlstash" */ + sourceId: string; + /** Human-readable name shown in the modal title, e.g. "PixlStash" */ + sourceName: string; + isOpen: boolean; + onClose: () => void; + onImportStarted?: () => void; +} + +export default function RemoteSourceBrowseModal({ + sourceId, + sourceName, + isOpen, + onClose, + onImportStarted, +}: RemoteSourceBrowseModalProps) { + const router = useRouter(); + const [browseData, setBrowseData] = useState(null); + const [browseError, setBrowseError] = useState(null); + const [browseLoading, setBrowseLoading] = useState(false); + const [activeGroup, setActiveGroup] = useState(''); + const [selectedItem, setSelectedItem] = useState<{ + id: string; + name: string; + groupId: string; + } | null>(null); + const [retryCount, setRetryCount] = useState(0); + + useEffect(() => { + if (!isOpen) return; + setSelectedItem(null); + + let aborted = false; + setBrowseLoading(true); + setBrowseError(null); + + apiClient + .get(`/api/datasets/remote/${sourceId}/browse`) + .then(res => { + if (aborted) return; + const data: BrowseData = res.data; + setBrowseData(data); + if (data.groups?.length > 0) setActiveGroup(data.groups[0].id); + }) + .catch(err => { + if (aborted) return; + setBrowseError( + (err as any)?.response?.data?.error || + (err instanceof Error ? err.message : 'Failed to connect'), + ); + }) + .finally(() => { + if (!aborted) setBrowseLoading(false); + }); + + return () => { + aborted = true; + }; + }, [isOpen, sourceId, retryCount]); + + const handleImportSuccess = (datasetName: string) => { + onImportStarted?.(); + onClose(); + router.push(`/datasets/${encodeURIComponent(datasetName)}`); + }; + + const tabClass = (groupId: string) => + `px-4 py-2 text-sm font-medium rounded-t border-b-2 transition-colors ${ + activeGroup === groupId + ? 'border-blue-500 text-blue-400' + : 'border-transparent text-gray-400 hover:text-gray-200' + }`; + + const currentGroup = browseData?.groups.find(g => g.id === activeGroup); + + return ( + +
+ {/* Group tabs */} + {browseData && browseData.groups.length > 1 && ( +
+ {browseData.groups.map(g => ( + + ))} +
+ )} + + {/* Browse list */} + {!selectedItem && ( + <> + {browseLoading && ( +
+ Loading {sourceName} data… +
+ )} + {browseError && ( +
+
{browseError}
+ +
+ )} + {currentGroup && ( +
+ {currentGroup.items.length === 0 ? ( +
+ No {currentGroup.label.toLowerCase()} found. +
+ ) : ( +
    + {currentGroup.items.map(item => ( +
  • + +
    • + ))} +
+ )} +
+ )} + + )} + + {/* Import form */} + {selectedItem && ( +
+ + setSelectedItem(null)} + /> +
+ )} +
+ + ); +} diff --git a/ui/src/server/settings.ts b/ui/src/server/settings.ts index 7a0a6750..8ea0918f 100644 --- a/ui/src/server/settings.ts +++ b/ui/src/server/settings.ts @@ -66,6 +66,19 @@ export const getHFToken = async () => { return token; }; +/** + * Load multiple settings keys in a single DB query. + * Returns a map of key → value (empty string if not set). + */ +export const getSettingsByKeys = async (keys: string[]): Promise> => { + const rows = await prisma.settings.findMany({ where: { key: { in: keys } } }); + const result: Record = Object.fromEntries(keys.map(k => [k, ''])); + for (const row of rows) { + if (row.value) result[row.key] = row.value; + } + return result; +}; + export const getDataRoot = async () => { const key = 'DATA_ROOT'; let dataRoot = myCache.get(key) as string;