Add support for remote dataset source plugins
Introduce a plugin system that lets datasets be browsed and imported from external services directly in the UI. - toolkit/dataset_sources/: base class, registry, and a CLI bridge that the Next.js API routes call to list plugins, browse items, fetch thumbnails, and build import job configs. Plugins live in extensions/ and self-register from their __init__.py, so they are only loaded on demand. - ui: a RemoteSourceBrowseModal plus /api/datasets/remote/* routes for listing plugins, browsing, thumbnails, and importing, with the settings page driven by each plugin's settings_schema. The framework is plugin-agnostic; no plugin is bundled here.
This commit is contained in:
Gaute Lindkvist
Gaute Lindkvist
parent
685ce37a8d
commit
418fc82bdc
|
|
@ -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",
|
||||
]
|
||||
|
|
@ -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: [<returned dict>] } }
|
||||
"""
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 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
|
||||
)
|
||||
|
|
@ -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: <base64> }
|
||||
|
||||
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 <plugins|browse|thumbnail|job-config> ...",
|
||||
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 <source_type_id>",
|
||||
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 <source_id> <thumbnail_id> <thumbnail_type>",
|
||||
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 <source_id>",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
cmd_job_config(sys.argv[2])
|
||||
else:
|
||||
print(f"Unknown command: {cmd}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
|
@ -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 {}
|
||||
|
|
@ -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<string, { data: unknown; expires: number }>();
|
||||
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<NextResponse>(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 }));
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
|
@ -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<string, unknown>,
|
||||
): Promise<Record<string, unknown>> {
|
||||
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<string, unknown>;
|
||||
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<string, unknown>;
|
||||
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',
|
||||
},
|
||||
});
|
||||
}
|
||||
|
|
@ -0,0 +1,67 @@
|
|||
/**
|
||||
* GET /api/datasets/remote/[source]/thumbnail/[id]?type=<thumbnail_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<NextResponse>(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 }));
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
|
|
@ -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 <Name>" 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<NextResponse>(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 }));
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
|
@ -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<string, string>).map(([key, value]) =>
|
||||
prisma.settings.upsert({
|
||||
where: { key },
|
||||
update: { value: value ?? '' },
|
||||
create: { key, value: value ?? '' },
|
||||
}),
|
||||
),
|
||||
);
|
||||
|
||||
flushCache();
|
||||
|
||||
|
|
|
|||
|
|
@ -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<string | null>(null);
|
||||
|
|
|
|||
|
|
@ -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<RemotePlugin[]>([]);
|
||||
const [openPluginId, setOpenPluginId] = useState<string | null>(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() {
|
|||
<h1 className="text-base sm:text-lg">Datasets</h1>
|
||||
</div>
|
||||
<div className="flex-1"></div>
|
||||
<div>
|
||||
<div className="flex items-center gap-2">
|
||||
{plugins.map(plugin => (
|
||||
<Button
|
||||
key={plugin.id}
|
||||
className="text-gray-200 bg-indigo-700 px-4 py-2 rounded-md hover:bg-indigo-600 transition-colors flex items-center gap-2"
|
||||
onClick={() => setOpenPluginId(plugin.id)}
|
||||
>
|
||||
{plugin.icon && (
|
||||
<img src={plugin.icon} alt="" className="w-4 h-4 object-contain flex-shrink-0" />
|
||||
)}
|
||||
Browse {plugin.display_name}
|
||||
</Button>
|
||||
))}
|
||||
<Button
|
||||
className="text-white bg-slate-600 px-2 sm:px-3 py-1 rounded-md hover:bg-slate-500 transition-colors text-sm sm:text-base whitespace-nowrap"
|
||||
onClick={() => openNewDatasetModal()}
|
||||
|
|
@ -137,6 +168,17 @@ export default function Datasets() {
|
|||
/>
|
||||
</MainContent>
|
||||
|
||||
{plugins.map(plugin => (
|
||||
<RemoteSourceBrowseModal
|
||||
key={plugin.id}
|
||||
sourceId={plugin.id}
|
||||
sourceName={plugin.display_name}
|
||||
isOpen={openPluginId === plugin.id}
|
||||
onClose={() => setOpenPluginId(null)}
|
||||
onImportStarted={() => refreshDatasets()}
|
||||
/>
|
||||
))}
|
||||
|
||||
<Modal
|
||||
isOpen={isNewDatasetModalOpen}
|
||||
onClose={() => setIsNewDatasetModalOpen(false)}
|
||||
|
|
|
|||
|
|
@ -1,35 +1,58 @@
|
|||
'use client';
|
||||
|
||||
import { useEffect, useState } from 'react';
|
||||
import useSettings from '@/hooks/useSettings';
|
||||
import { TopBar, MainContent } from '@/components/layout';
|
||||
import { apiClient } from '@/utils/api';
|
||||
|
||||
interface SettingField {
|
||||
key: string;
|
||||
label: string;
|
||||
input_type: string;
|
||||
description: string;
|
||||
placeholder: string;
|
||||
}
|
||||
|
||||
interface Plugin {
|
||||
id: string;
|
||||
display_name: string;
|
||||
settings_schema: SettingField[];
|
||||
}
|
||||
|
||||
export default function Settings() {
|
||||
const { settings, setSettings } = useSettings();
|
||||
const [values, setValues] = useState<Record<string, string>>({
|
||||
HF_TOKEN: '',
|
||||
TRAINING_FOLDER: '',
|
||||
DATASETS_FOLDER: '',
|
||||
});
|
||||
const [plugins, setPlugins] = useState<Plugin[]>([]);
|
||||
const [status, setStatus] = useState<'idle' | 'saving' | 'success' | 'error'>('idle');
|
||||
|
||||
useEffect(() => {
|
||||
apiClient
|
||||
.get('/api/settings')
|
||||
.then(res => setValues(res.data as Record<string, string>))
|
||||
.catch(err => console.error('Error fetching settings:', err));
|
||||
|
||||
apiClient
|
||||
.get('/api/datasets/remote/plugins')
|
||||
.then(res => setPlugins(res.data as Plugin[]))
|
||||
.catch(err => console.error('Error fetching plugins:', err));
|
||||
}, []);
|
||||
|
||||
const handleChange = (e: React.ChangeEvent<HTMLInputElement>) => {
|
||||
const { name, type, value, checked } = e.target;
|
||||
setValues(prev => ({ ...prev, [name]: type === 'checkbox' ? (checked ? 'true' : 'false') : value }));
|
||||
};
|
||||
|
||||
const handleSubmit = async (e: React.FormEvent) => {
|
||||
e.preventDefault();
|
||||
setStatus('saving');
|
||||
|
||||
apiClient
|
||||
.post('/api/settings', settings)
|
||||
.then(() => {
|
||||
setStatus('success');
|
||||
})
|
||||
.catch(error => {
|
||||
console.error('Error saving settings:', error);
|
||||
setStatus('error');
|
||||
})
|
||||
.finally(() => {
|
||||
setTimeout(() => setStatus('idle'), 2000);
|
||||
});
|
||||
};
|
||||
|
||||
const handleChange = (e: React.ChangeEvent<HTMLInputElement>) => {
|
||||
const { name, value } = e.target;
|
||||
setSettings(prev => ({ ...prev, [name]: value }));
|
||||
.post('/api/settings', values)
|
||||
.then(() => setStatus('success'))
|
||||
.catch(() => setStatus('error'))
|
||||
.finally(() => setTimeout(() => setStatus('idle'), 2000));
|
||||
};
|
||||
|
||||
return (
|
||||
|
|
@ -61,7 +84,7 @@ export default function Settings() {
|
|||
type="password"
|
||||
id="HF_TOKEN"
|
||||
name="HF_TOKEN"
|
||||
value={settings.HF_TOKEN}
|
||||
value={values['HF_TOKEN'] ?? ''}
|
||||
onChange={handleChange}
|
||||
className="w-full px-4 py-2 bg-gray-800 border border-gray-700 rounded-lg focus:ring-2 focus:ring-gray-600 focus:border-transparent"
|
||||
placeholder="Enter your Hugging Face token"
|
||||
|
|
@ -80,7 +103,7 @@ export default function Settings() {
|
|||
type="text"
|
||||
id="TRAINING_FOLDER"
|
||||
name="TRAINING_FOLDER"
|
||||
value={settings.TRAINING_FOLDER}
|
||||
value={values['TRAINING_FOLDER'] ?? ''}
|
||||
onChange={handleChange}
|
||||
className="w-full px-4 py-2 bg-gray-800 border border-gray-700 rounded-lg focus:ring-2 focus:ring-gray-600 focus:border-transparent"
|
||||
placeholder="Enter training folder path"
|
||||
|
|
@ -102,7 +125,7 @@ export default function Settings() {
|
|||
type="text"
|
||||
id="DATASETS_FOLDER"
|
||||
name="DATASETS_FOLDER"
|
||||
value={settings.DATASETS_FOLDER}
|
||||
value={values['DATASETS_FOLDER'] ?? ''}
|
||||
onChange={handleChange}
|
||||
className="w-full px-4 py-2 bg-gray-800 border border-gray-700 rounded-lg focus:ring-2 focus:ring-gray-600 focus:border-transparent"
|
||||
placeholder="Enter datasets folder path"
|
||||
|
|
@ -122,12 +145,63 @@ export default function Settings() {
|
|||
type="text"
|
||||
id="MODELS_PATH"
|
||||
name="MODELS_PATH"
|
||||
value={settings.MODELS_PATH}
|
||||
value={values['MODELS_PATH'] ?? ''}
|
||||
onChange={handleChange}
|
||||
className="w-full px-4 py-2 bg-gray-800 border border-gray-700 rounded-lg focus:ring-2 focus:ring-gray-600 focus:border-transparent"
|
||||
placeholder="Enter models folder path"
|
||||
/>
|
||||
</div>
|
||||
|
||||
{plugins.map(plugin =>
|
||||
plugin.settings_schema.length === 0 ? null : (
|
||||
<div key={plugin.id}>
|
||||
<hr className="border-gray-700 my-2" />
|
||||
<h2 className="text-sm font-semibold mb-3">{plugin.display_name}</h2>
|
||||
<div className="space-y-4">
|
||||
{plugin.settings_schema.map(field => (
|
||||
<div key={field.key}>
|
||||
{field.input_type === 'checkbox' ? (
|
||||
<label className="flex items-center gap-3 text-sm font-medium cursor-pointer">
|
||||
<input
|
||||
type="checkbox"
|
||||
id={field.key}
|
||||
name={field.key}
|
||||
checked={(values[field.key] ?? 'true') !== 'false'}
|
||||
onChange={handleChange}
|
||||
className="w-4 h-4 rounded border-gray-700 bg-gray-800 accent-blue-500"
|
||||
/>
|
||||
<span>
|
||||
{field.label}
|
||||
{field.description && (
|
||||
<div className="text-gray-500 text-sm">{field.description}</div>
|
||||
)}
|
||||
</span>
|
||||
</label>
|
||||
) : (
|
||||
<>
|
||||
<label htmlFor={field.key} className="block text-sm font-medium mb-2">
|
||||
{field.label}
|
||||
{field.description && (
|
||||
<div className="text-gray-500 text-sm ml-1">{field.description}</div>
|
||||
)}
|
||||
</label>
|
||||
<input
|
||||
type={field.input_type}
|
||||
id={field.key}
|
||||
name={field.key}
|
||||
value={values[field.key] ?? ''}
|
||||
onChange={handleChange}
|
||||
className="w-full px-4 py-2 bg-gray-800 border border-gray-700 rounded-lg focus:ring-2 focus:ring-gray-600 focus:border-transparent"
|
||||
placeholder={field.placeholder}
|
||||
/>
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
),
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
|
@ -147,3 +221,4 @@ export default function Settings() {
|
|||
</>
|
||||
);
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
<div className="w-16 h-16 rounded bg-gray-700 flex items-center justify-center text-gray-500 text-xs">?</div>
|
||||
);
|
||||
}
|
||||
return (
|
||||
<img
|
||||
src={src}
|
||||
alt={alt}
|
||||
className="w-16 h-16 rounded object-cover flex-shrink-0"
|
||||
onError={() => setErrored(true)}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
// ─── Progress bar ─────────────────────────────────────────────────────────────
|
||||
|
||||
function ProgressBar({ done, total }: { done: number; total: number }) {
|
||||
const pct = total > 0 ? Math.round((done / total) * 100) : 0;
|
||||
return (
|
||||
<div className="space-y-1">
|
||||
<div className="flex justify-between text-xs text-gray-400">
|
||||
<span>
|
||||
{done} / {total} images
|
||||
</span>
|
||||
<span>{pct}%</span>
|
||||
</div>
|
||||
<div className="w-full bg-gray-700 rounded-full h-2 overflow-hidden">
|
||||
<div
|
||||
className="bg-blue-500 h-2 rounded-full transition-all duration-300"
|
||||
style={{ width: `${pct}%` }}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
// ─── 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<ImportPhase>('idle');
|
||||
const [progress, setProgress] = useState({ done: 0, total: 0 });
|
||||
const [downloaded, setDownloaded] = useState(0);
|
||||
const [errorMsg, setErrorMsg] = useState<string | null>(null);
|
||||
|
||||
// Initialise source-specific field values from their declared defaults
|
||||
const [fieldValues, setFieldValues] = useState<Record<string, string | number | boolean>>(() => {
|
||||
const init: Record<string, string | number | boolean> = {};
|
||||
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 (
|
||||
<div className="p-4 border border-green-800 rounded-lg bg-green-950/40 space-y-3">
|
||||
<div className="text-sm text-green-300 font-medium">
|
||||
✓ Downloaded {downloaded} image{downloaded !== 1 ? 's' : ''} to “
|
||||
{datasetName || itemName}”
|
||||
</div>
|
||||
<ProgressBar done={downloaded} total={downloaded} />
|
||||
<button
|
||||
className="w-full rounded-md bg-gray-700 px-4 py-2 text-sm text-gray-200 hover:bg-gray-600"
|
||||
onClick={onCancel}
|
||||
>
|
||||
Close
|
||||
</button>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="space-y-4 p-4 border border-gray-700 rounded-lg bg-gray-900">
|
||||
<div className="text-sm font-medium text-gray-200">
|
||||
Downloading: <span className="text-blue-400">{itemName}</span>
|
||||
</div>
|
||||
|
||||
{!isRunning && (
|
||||
<>
|
||||
<div>
|
||||
<label className={labelClass}>Dataset Name</label>
|
||||
<input
|
||||
type="text"
|
||||
className={inputClass}
|
||||
value={datasetName}
|
||||
onChange={e => setDatasetName(e.target.value)}
|
||||
placeholder={itemName}
|
||||
/>
|
||||
<p className="text-xs text-gray-500 mt-1">
|
||||
Folder name created inside your datasets directory.
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<div>
|
||||
<label className={labelClass}>Trigger Word (optional)</label>
|
||||
<input
|
||||
type="text"
|
||||
className={inputClass}
|
||||
value={triggerWord}
|
||||
onChange={e => setTriggerWord(e.target.value)}
|
||||
placeholder="e.g. myperson"
|
||||
/>
|
||||
<p className="text-xs text-gray-500 mt-1">
|
||||
Prepended to every caption for LoRA training.
|
||||
</p>
|
||||
</div>
|
||||
|
||||
{/* Source-specific fields from plugin's get_import_fields() */}
|
||||
{importFields.map(f => (
|
||||
<div key={f.id}>
|
||||
<label className={labelClass}>{f.label}</label>
|
||||
{f.field_type === 'select' && (
|
||||
<select
|
||||
className={inputClass}
|
||||
value={String(fieldValues[f.id] ?? f.default ?? '')}
|
||||
onChange={e =>
|
||||
setFieldValues(v => ({
|
||||
...v,
|
||||
[f.id]: isNaN(Number(e.target.value)) ? e.target.value : Number(e.target.value),
|
||||
}))
|
||||
}
|
||||
>
|
||||
{f.options.map(opt => (
|
||||
<option key={String(opt.value)} value={String(opt.value)}>
|
||||
{opt.label}
|
||||
</option>
|
||||
))}
|
||||
</select>
|
||||
)}
|
||||
{f.field_type === 'text' && (
|
||||
<input
|
||||
type="text"
|
||||
className={inputClass}
|
||||
value={String(fieldValues[f.id] ?? '')}
|
||||
onChange={e => setFieldValues(v => ({ ...v, [f.id]: e.target.value }))}
|
||||
/>
|
||||
)}
|
||||
{f.field_type === 'checkbox' && (
|
||||
<label className="flex items-center gap-2 cursor-pointer">
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={Boolean(fieldValues[f.id])}
|
||||
onChange={e => setFieldValues(v => ({ ...v, [f.id]: e.target.checked }))}
|
||||
className="rounded"
|
||||
/>
|
||||
<span className="text-sm text-gray-300">Enabled</span>
|
||||
</label>
|
||||
)}
|
||||
</div>
|
||||
))}
|
||||
|
||||
<label className="flex items-center gap-2 cursor-pointer">
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={overwrite}
|
||||
onChange={e => setOverwrite(e.target.checked)}
|
||||
className="rounded"
|
||||
/>
|
||||
<span className="text-sm text-gray-300">Overwrite existing files</span>
|
||||
</label>
|
||||
</>
|
||||
)}
|
||||
|
||||
{isRunning && (
|
||||
<div className="space-y-3 py-2">
|
||||
<div className="text-sm text-gray-300">
|
||||
{phase === 'connecting'
|
||||
? 'Connecting…'
|
||||
: `Downloading to "${datasetName || itemName}"…`}
|
||||
</div>
|
||||
{progress.total > 0 ? (
|
||||
<ProgressBar done={progress.done} total={progress.total} />
|
||||
) : (
|
||||
<div className="w-full bg-gray-700 rounded-full h-2 overflow-hidden">
|
||||
<div className="bg-blue-500 h-2 rounded-full animate-pulse w-1/3" />
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{errorMsg && (
|
||||
<div className="text-sm text-red-400 bg-red-950 border border-red-800 rounded px-3 py-2">
|
||||
{errorMsg}
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="flex justify-end gap-3 pt-1">
|
||||
<button
|
||||
type="button"
|
||||
className="rounded-md bg-gray-700 px-4 py-2 text-sm text-gray-200 hover:bg-gray-600 disabled:opacity-40"
|
||||
onClick={onCancel}
|
||||
disabled={isRunning}
|
||||
>
|
||||
Cancel
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
className="rounded-md bg-blue-600 px-4 py-2 text-sm text-white hover:bg-blue-700 disabled:opacity-50"
|
||||
onClick={handleImport}
|
||||
disabled={isRunning}
|
||||
>
|
||||
{phase === 'connecting'
|
||||
? 'Connecting…'
|
||||
: phase === 'downloading'
|
||||
? 'Downloading…'
|
||||
: 'Download to Dataset'}
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
// ─── 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<BrowseData | null>(null);
|
||||
const [browseError, setBrowseError] = useState<string | null>(null);
|
||||
const [browseLoading, setBrowseLoading] = useState(false);
|
||||
const [activeGroup, setActiveGroup] = useState<string>('');
|
||||
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 (
|
||||
<Modal isOpen={isOpen} onClose={onClose} title={`Browse ${sourceName}`} size="lg">
|
||||
<div className="flex flex-col" style={{ minHeight: '500px' }}>
|
||||
{/* Group tabs */}
|
||||
{browseData && browseData.groups.length > 1 && (
|
||||
<div className="flex border-b border-gray-700 mb-3">
|
||||
{browseData.groups.map(g => (
|
||||
<button
|
||||
key={g.id}
|
||||
className={tabClass(g.id)}
|
||||
onClick={() => {
|
||||
setActiveGroup(g.id);
|
||||
setSelectedItem(null);
|
||||
}}
|
||||
>
|
||||
{g.label}
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
|
||||
{/* Browse list */}
|
||||
{!selectedItem && (
|
||||
<>
|
||||
{browseLoading && (
|
||||
<div className="py-12 text-center text-gray-400 text-sm">
|
||||
Loading {sourceName} data…
|
||||
</div>
|
||||
)}
|
||||
{browseError && (
|
||||
<div className="py-8 text-center">
|
||||
<div className="text-red-400 text-sm mb-3">{browseError}</div>
|
||||
<button
|
||||
className="text-xs text-blue-400 hover:text-blue-300 underline"
|
||||
onClick={() => setRetryCount(c => c + 1)}
|
||||
>
|
||||
Retry
|
||||
</button>
|
||||
</div>
|
||||
)}
|
||||
{currentGroup && (
|
||||
<div className="flex-1 overflow-y-auto max-h-96">
|
||||
{currentGroup.items.length === 0 ? (
|
||||
<div className="py-12 text-center text-gray-500 text-sm">
|
||||
No {currentGroup.label.toLowerCase()} found.
|
||||
</div>
|
||||
) : (
|
||||
<ul className="divide-y divide-gray-700">
|
||||
{currentGroup.items.map(item => (
|
||||
<li key={item.id}>
|
||||
<button
|
||||
className="flex items-center gap-4 w-full px-4 py-3 text-left hover:bg-gray-700 transition-colors"
|
||||
onClick={() =>
|
||||
setSelectedItem({
|
||||
id: item.id,
|
||||
name: item.name,
|
||||
groupId: currentGroup.id,
|
||||
})
|
||||
}
|
||||
>
|
||||
<ThumbnailImage
|
||||
sourceId={sourceId}
|
||||
thumbnailId={item.thumbnail_id}
|
||||
thumbnailType={item.thumbnail_type}
|
||||
alt={item.name}
|
||||
/>
|
||||
<div className="flex-1 min-w-0">
|
||||
<div className="flex items-center gap-2">
|
||||
<span className="font-medium text-gray-100 truncate">
|
||||
{item.name}
|
||||
</span>
|
||||
{item.picture_count > 0 && (
|
||||
<span className="text-xs text-gray-500 flex-shrink-0">
|
||||
{item.picture_count} images
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
<span className="text-gray-500 text-sm flex-shrink-0">▶</span>
|
||||
</button>
|
||||
</li>
|
||||
))}
|
||||
</ul>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</>
|
||||
)}
|
||||
|
||||
{/* Import form */}
|
||||
{selectedItem && (
|
||||
<div>
|
||||
<button
|
||||
className="text-xs text-blue-400 hover:text-blue-300 mb-3 flex items-center gap-1"
|
||||
onClick={() => setSelectedItem(null)}
|
||||
>
|
||||
← Back to list
|
||||
</button>
|
||||
<ImportForm
|
||||
sourcePluginId={sourceId}
|
||||
groupId={selectedItem.groupId}
|
||||
itemId={selectedItem.id}
|
||||
itemName={selectedItem.name}
|
||||
importFields={browseData?.import_fields ?? []}
|
||||
onSuccess={handleImportSuccess}
|
||||
onCancel={() => setSelectedItem(null)}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</Modal>
|
||||
);
|
||||
}
|
||||
|
|
@ -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<Record<string, string>> => {
|
||||
const rows = await prisma.settings.findMany({ where: { key: { in: keys } } });
|
||||
const result: Record<string, string> = 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;
|
||||
|
|
|
|||
Loading…
Reference in New Issue