Add PixlStash remote dataset source plugin
A reference plugin for the remote dataset source framework that browses and
imports training images from a PixlStash server.
- Browse characters and picture sets, import selected groups as a dataset.
- Captions are built from each picture's natural-language description and WD14
tags, read per-picture from GET /pictures/{id}/metadata (the listing
endpoints only return scalar grid fields).
- Supports SSL and configurable API paths.
Also un-ignores extensions/pixlstash in .gitignore so the plugin is tracked.
This commit is contained in:
Gaute Lindkvist
Gaute Lindkvist
parent
418fc82bdc
commit
349098a7cc
|
|
@ -179,6 +179,7 @@ cython_debug/
|
|||
!/output/.gitkeep
|
||||
/extensions/*
|
||||
!/extensions/example
|
||||
!/extensions/pixlstash
|
||||
/temp
|
||||
/wandb
|
||||
.vscode/settings.json
|
||||
|
|
|
|||
|
|
@ -0,0 +1,221 @@
|
|||
"""
|
||||
PixlStash dataset fetch process for AI-Toolkit.
|
||||
|
||||
Downloads images and captions from a PixlStash server into the AI-Toolkit
|
||||
datasets folder so the dataset appears automatically in the UI.
|
||||
|
||||
Config keys
|
||||
-----------
|
||||
pixlstash_url : str
|
||||
Base URL of the PixlStash server, e.g. "http://localhost:9537".
|
||||
pixlstash_token : str
|
||||
Personal API token (PixlStash → Settings → API Tokens).
|
||||
source_type : "character" | "picture_set"
|
||||
Whether to fetch by character or by picture set.
|
||||
source_id : int
|
||||
The integer ID of the character or picture set to fetch.
|
||||
caption_mode : "description" | "tags" | "both" (default: "description")
|
||||
Which PixlStash caption source to use.
|
||||
"description" — Florence-2 natural-language caption.
|
||||
"tags" — WD14 comma-separated tags.
|
||||
"both" — description first, then tags.
|
||||
trigger_word : str (optional)
|
||||
Token prepended to every caption, e.g. your LoRA trigger word.
|
||||
dataset_name : str (optional)
|
||||
Name used for the output subfolder inside the AI-Toolkit datasets root.
|
||||
Defaults to the character/set name returned by PixlStash.
|
||||
datasets_root : str (optional)
|
||||
Absolute path to the AI-Toolkit datasets folder.
|
||||
Defaults to "<toolkit_root>/datasets" (matches the UI default).
|
||||
overwrite : bool (default: false)
|
||||
If false, images that already exist on disk are skipped.
|
||||
If true, every image is re-downloaded and captions are rewritten.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections import OrderedDict
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
from jobs.process import BaseExtensionProcess
|
||||
from toolkit.paths import TOOLKIT_ROOT
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from jobs import ExtensionJob
|
||||
|
||||
|
||||
class PixlStashFetchProcess(BaseExtensionProcess):
|
||||
def __init__(
|
||||
self,
|
||||
process_id: int,
|
||||
job: "ExtensionJob",
|
||||
config: OrderedDict,
|
||||
) -> None:
|
||||
super().__init__(process_id, job, config)
|
||||
|
||||
self.pixlstash_url: str = self.get_conf("pixlstash_url", required=True)
|
||||
self.pixlstash_token: str = self.get_conf("pixlstash_token", required=True)
|
||||
self.source_type: str = self.get_conf("source_type", required=True)
|
||||
self.source_id: int = int(self.get_conf("source_id", required=True))
|
||||
self.caption_mode: str = self.get_conf("caption_mode", default="description")
|
||||
self.trigger_word: str = self.get_conf("trigger_word", default="")
|
||||
self.dataset_name: str | None = self.get_conf("dataset_name", default=None)
|
||||
self.overwrite: bool = self.get_conf("overwrite", default=False)
|
||||
self.min_score: int = int(self.get_conf("min_score", default=0))
|
||||
self.verify_ssl: bool = str(
|
||||
self.get_conf("verify_ssl", default="true")
|
||||
).lower() not in ("false", "0", "no", "off")
|
||||
|
||||
# Where to write the dataset. Defaults to the same root the UI watches.
|
||||
default_datasets_root = os.path.join(TOOLKIT_ROOT, "datasets")
|
||||
self.datasets_root: str = self.get_conf(
|
||||
"datasets_root", default=default_datasets_root
|
||||
)
|
||||
|
||||
if self.source_type not in ("character", "picture_set"):
|
||||
raise ValueError(
|
||||
f"source_type must be 'character' or 'picture_set', got '{self.source_type}'"
|
||||
)
|
||||
if self.caption_mode not in ("description", "tags", "both"):
|
||||
raise ValueError(
|
||||
f"caption_mode must be 'description', 'tags', or 'both', got '{self.caption_mode}'"
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def run(self) -> None:
|
||||
super().run()
|
||||
|
||||
# Import here so the module is only loaded when this process runs
|
||||
from extensions.pixlstash.pixlstash_client import PixlStashClient
|
||||
|
||||
print(f"\n[PixlStash] Connecting to {self.pixlstash_url} …")
|
||||
client = PixlStashClient(
|
||||
self.pixlstash_url, self.pixlstash_token, verify_ssl=self.verify_ssl
|
||||
)
|
||||
client.login()
|
||||
print("[PixlStash] Authenticated.")
|
||||
|
||||
# ---- resolve source name and picture list -------------------------
|
||||
if self.source_type == "character":
|
||||
source = client.get_character(self.source_id)
|
||||
source_label = f"character '{source['name']}' (id={self.source_id})"
|
||||
pictures = client.list_pictures_for_character(self.source_id)
|
||||
else:
|
||||
source = client.get_picture_set(self.source_id)
|
||||
source_label = f"picture set '{source['name']}' (id={self.source_id})"
|
||||
pictures = client.list_pictures_for_set(self.source_id)
|
||||
|
||||
total = len(pictures)
|
||||
print(
|
||||
f"[PixlStash] Fetched {source_label} — {total} picture(s) found in source.",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# ---- apply score filter ------------------------------------------
|
||||
if self.min_score > 0:
|
||||
before = total
|
||||
pictures = [p for p in pictures if (p.get("score") or 0) >= self.min_score]
|
||||
total = len(pictures)
|
||||
filtered_out = before - total
|
||||
print(
|
||||
f"[PixlStash] Score filter ≥{self.min_score}★: "
|
||||
f"{total} picture(s) kept, {filtered_out} filtered out.",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
print(
|
||||
f"[PixlStash] Downloading {source_label} — {total} picture(s) found.",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# ---- resolve output folder ----------------------------------------
|
||||
dataset_name = self.dataset_name or self._safe_folder_name(source["name"])
|
||||
output_folder = os.path.join(self.datasets_root, dataset_name)
|
||||
os.makedirs(output_folder, exist_ok=True)
|
||||
print(f"[PixlStash] Output folder: {output_folder}")
|
||||
|
||||
# ---- download loop -----------------------------------------------
|
||||
downloaded = 0
|
||||
skipped = 0
|
||||
errors = 0
|
||||
|
||||
for pic in tqdm(pictures, desc="Downloading", unit="img"):
|
||||
pic_id = pic["id"]
|
||||
img_filename = f"{pic_id}.jpg"
|
||||
txt_filename = f"{pic_id}.txt"
|
||||
img_path = os.path.join(output_folder, img_filename)
|
||||
txt_path = os.path.join(output_folder, txt_filename)
|
||||
|
||||
if (
|
||||
not self.overwrite
|
||||
and os.path.exists(img_path)
|
||||
and os.path.exists(txt_path)
|
||||
):
|
||||
skipped += 1
|
||||
print(f"PROGRESS:{downloaded + skipped}/{total}", flush=True)
|
||||
continue
|
||||
|
||||
try:
|
||||
# The listing/embed rows only carry scalar grid fields, so the
|
||||
# natural-language description and WD14 tags are read per-picture
|
||||
# from GET /pictures/{id}/metadata. Fetched here, after the
|
||||
# on-disk skip check, so we never query metadata for images we
|
||||
# are about to skip.
|
||||
meta = client.get_picture_metadata(pic_id)
|
||||
fmt = meta.get("format", "jpg") or "jpg"
|
||||
|
||||
# Build caption
|
||||
caption = client.build_caption(
|
||||
meta,
|
||||
mode=self.caption_mode,
|
||||
trigger=self.trigger_word,
|
||||
)
|
||||
|
||||
# Download image
|
||||
img_bytes = client.download_image_bytes(pic_id, fmt)
|
||||
|
||||
# Write image (always save as .jpg for maximum AI-Toolkit compat)
|
||||
if fmt.lower() in ("jpg", "jpeg"):
|
||||
with open(img_path, "wb") as f:
|
||||
f.write(img_bytes)
|
||||
else:
|
||||
# Convert to JPEG via PIL so AI-Toolkit doesn't have to
|
||||
import io
|
||||
from PIL import Image
|
||||
|
||||
pil_img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
|
||||
pil_img.save(img_path, format="JPEG", quality=95)
|
||||
|
||||
# Write caption
|
||||
with open(txt_path, "w", encoding="utf-8") as f:
|
||||
f.write(caption)
|
||||
|
||||
downloaded += 1
|
||||
print(f"PROGRESS:{downloaded + skipped}/{total}", flush=True)
|
||||
|
||||
except Exception as exc:
|
||||
print(
|
||||
f"\n[PixlStash] WARNING: Failed to fetch picture id={pic_id}: {exc}",
|
||||
flush=True,
|
||||
)
|
||||
errors += 1
|
||||
print(f"PROGRESS:{downloaded + skipped}/{total}", flush=True)
|
||||
|
||||
# ---- summary -----------------------------------------------------
|
||||
print(
|
||||
f"\n[PixlStash] Done — {downloaded} downloaded, "
|
||||
f"{skipped} skipped (already on disk), {errors} errors."
|
||||
)
|
||||
print(f"[PixlStash] Dataset '{dataset_name}' is ready in the AI-Toolkit UI.")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _safe_folder_name(name: str) -> str:
|
||||
"""Convert an arbitrary string into a safe directory name."""
|
||||
safe = "".join(c if c.isalnum() or c in " -_" else "_" for c in name)
|
||||
return safe.strip().replace(" ", "_")
|
||||
|
|
@ -0,0 +1,25 @@
|
|||
from toolkit.extension import Extension
|
||||
from toolkit.dataset_sources.registry import register_source
|
||||
|
||||
|
||||
class PixlStashFetchExtension(Extension):
|
||||
uid = "pixlstash_fetch"
|
||||
name = "PixlStash Dataset Fetch"
|
||||
|
||||
@classmethod
|
||||
def get_process(cls):
|
||||
from .PixlStashFetchProcess import PixlStashFetchProcess
|
||||
return PixlStashFetchProcess
|
||||
|
||||
|
||||
# Register the dataset source so the dataloader and UI can discover it
|
||||
def _register():
|
||||
from .pixlstash_source import PixlStashDatasetSource
|
||||
register_source(PixlStashDatasetSource)
|
||||
|
||||
|
||||
_register()
|
||||
|
||||
AI_TOOLKIT_EXTENSIONS = [
|
||||
PixlStashFetchExtension,
|
||||
]
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 204 KiB |
|
|
@ -0,0 +1,51 @@
|
|||
---
|
||||
# PixlStash Dataset Fetch — fetch by character
|
||||
#
|
||||
# This job downloads all images assigned to a character in PixlStash and writes
|
||||
# them as jpg + txt caption pairs into the AI-Toolkit datasets folder.
|
||||
# Once complete the dataset appears automatically in the AI-Toolkit UI.
|
||||
#
|
||||
# Run with:
|
||||
# python run.py config/pixlstash/fetch_character.yaml
|
||||
|
||||
job: extension
|
||||
config:
|
||||
name: "pixlstash_fetch"
|
||||
process:
|
||||
- type: "pixlstash_fetch"
|
||||
|
||||
# ---- PixlStash connection ------------------------------------------
|
||||
# URL of your running PixlStash instance
|
||||
pixlstash_url: "http://localhost:9537"
|
||||
|
||||
# Personal API token — create one in PixlStash → Settings → API Tokens
|
||||
pixlstash_token: "paste-your-api-token-here"
|
||||
|
||||
# ---- What to fetch ------------------------------------------------
|
||||
# "character" fetches all pictures assigned to a character via face detection
|
||||
source_type: "character"
|
||||
|
||||
# Integer ID of the character (visible in the PixlStash URL when browsing)
|
||||
source_id: 1
|
||||
|
||||
# ---- Caption options ----------------------------------------------
|
||||
# "description" — Florence-2 natural-language caption (best for FLUX / SD3)
|
||||
# "tags" — WD14 comma-separated tags (best for SDXL / SD1.5)
|
||||
# "both" — description first, then tags appended
|
||||
caption_mode: "description"
|
||||
|
||||
# Optional trigger word prepended to every caption.
|
||||
# Leave blank or remove if you don't want one.
|
||||
# trigger_word: "mycharacter"
|
||||
|
||||
# ---- Output -------------------------------------------------------
|
||||
# Subfolder name inside the AI-Toolkit datasets root.
|
||||
# Defaults to the character name from PixlStash if not set.
|
||||
# dataset_name: "my_character_lora"
|
||||
|
||||
# Absolute path to the datasets root.
|
||||
# Defaults to <ai-toolkit root>/datasets — the same folder the UI watches.
|
||||
# datasets_root: "/absolute/path/to/datasets"
|
||||
|
||||
# Set to true to re-download images that already exist on disk.
|
||||
overwrite: false
|
||||
|
|
@ -0,0 +1,51 @@
|
|||
---
|
||||
# PixlStash Dataset Fetch — fetch by picture set
|
||||
#
|
||||
# This job downloads all images belonging to a picture set in PixlStash and
|
||||
# writes them as jpg + txt caption pairs into the AI-Toolkit datasets folder.
|
||||
# Once complete the dataset appears automatically in the AI-Toolkit UI.
|
||||
#
|
||||
# Run with:
|
||||
# python run.py config/pixlstash/fetch_picture_set.yaml
|
||||
|
||||
job: extension
|
||||
config:
|
||||
name: "pixlstash_fetch"
|
||||
process:
|
||||
- type: "pixlstash_fetch"
|
||||
|
||||
# ---- PixlStash connection ------------------------------------------
|
||||
# URL of your running PixlStash instance
|
||||
pixlstash_url: "http://localhost:9537"
|
||||
|
||||
# Personal API token — create one in PixlStash → Settings → API Tokens
|
||||
pixlstash_token: "paste-your-api-token-here"
|
||||
|
||||
# ---- What to fetch ------------------------------------------------
|
||||
# "picture_set" fetches all pictures that are members of a picture set
|
||||
source_type: "picture_set"
|
||||
|
||||
# Integer ID of the picture set (visible in the PixlStash URL when browsing)
|
||||
source_id: 7
|
||||
|
||||
# ---- Caption options ----------------------------------------------
|
||||
# "description" — Florence-2 natural-language caption (best for FLUX / SD3)
|
||||
# "tags" — WD14 comma-separated tags (best for SDXL / SD1.5)
|
||||
# "both" — description first, then tags appended
|
||||
caption_mode: "tags"
|
||||
|
||||
# Optional trigger word prepended to every caption.
|
||||
# Leave blank or remove if you don't want one.
|
||||
# trigger_word: "mystyle"
|
||||
|
||||
# ---- Output -------------------------------------------------------
|
||||
# Subfolder name inside the AI-Toolkit datasets root.
|
||||
# Defaults to the picture set name from PixlStash if not set.
|
||||
# dataset_name: "my_style_lora"
|
||||
|
||||
# Absolute path to the datasets root.
|
||||
# Defaults to <ai-toolkit root>/datasets — the same folder the UI watches.
|
||||
# datasets_root: "/absolute/path/to/datasets"
|
||||
|
||||
# Set to true to re-download images that already exist on disk.
|
||||
overwrite: false
|
||||
|
|
@ -0,0 +1,194 @@
|
|||
"""
|
||||
PixlStash API client — corrected against live API (v1.0.0b3).
|
||||
|
||||
Key endpoint notes (verified from /redoc):
|
||||
- Base prefix: /api/v1
|
||||
- Picture sets: /api/v1/picture_sets/{id} (underscore, not hyphen)
|
||||
- Picture file: /api/v1/pictures/{id}.{ext}
|
||||
- Thumbnail: /api/v1/pictures/thumbnails/{id}.webp
|
||||
- Set members: /api/v1/picture_sets/{id}/members -> returns integer IDs only
|
||||
- Char pictures: /api/v1/pictures/list?character_id={id}
|
||||
- Auth: Authorization: Bearer <token> header on all requests
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import warnings
|
||||
from typing import List, Optional
|
||||
|
||||
import requests
|
||||
import urllib3
|
||||
|
||||
_EMPTY_TAG_SENTINEL = ""
|
||||
|
||||
|
||||
class PixlStashError(RuntimeError):
|
||||
"""Raised when the PixlStash server returns an unexpected response."""
|
||||
|
||||
|
||||
class PixlStashClient:
|
||||
"""Thin wrapper around the PixlStash REST API."""
|
||||
|
||||
def __init__(self, base_url: str, token: str, verify_ssl: bool = True) -> None:
|
||||
self.base_url = base_url.rstrip("/") + "/api/v1"
|
||||
self.token = token
|
||||
self._session = requests.Session()
|
||||
self._session.verify = verify_ssl
|
||||
self._session.headers.update({"Authorization": f"Bearer {token}"})
|
||||
if not verify_ssl:
|
||||
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
|
||||
|
||||
def login(self) -> None:
|
||||
"""No-op: authentication is handled via Bearer token header."""
|
||||
pass
|
||||
|
||||
def _get(self, path: str, **params) -> requests.Response:
|
||||
r = self._session.get(
|
||||
f"{self.base_url}{path}", params=params or None, timeout=60
|
||||
)
|
||||
if not r.ok:
|
||||
raise PixlStashError(f"GET {path} failed ({r.status_code}): {r.text[:200]}")
|
||||
return r
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Characters
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_character(self, character_id: int) -> dict:
|
||||
return self._get(f"/characters/{character_id}").json()
|
||||
|
||||
def list_characters(self, name: Optional[str] = None) -> List[dict]:
|
||||
params = {}
|
||||
if name:
|
||||
params["name"] = name
|
||||
return self._get("/characters", **params).json()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Picture sets
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_picture_set(self, set_id: int) -> dict:
|
||||
"""Return picture set metadata."""
|
||||
return self._get(f"/picture_sets/{set_id}").json()["set"]
|
||||
|
||||
def list_picture_sets(self) -> List[dict]:
|
||||
"""Return all non-reference picture sets."""
|
||||
all_sets = self._get("/picture_sets").json()
|
||||
# Reference sets are auto-created per character for face recognition;
|
||||
# they have a non-null `reference_character` field and are not useful
|
||||
# as training datasets.
|
||||
return [s for s in all_sets if s.get("reference_character") is None]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Picture listing
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def list_pictures_for_character(self, character_id: int) -> List[dict]:
|
||||
"""Return the listing rows for every picture this character appears in.
|
||||
|
||||
Rows carry only scalar grid fields (id, score, format) — enough to
|
||||
filter and download. The natural-language ``description`` and WD14
|
||||
``tags`` are projected out of the listing response, so they are read
|
||||
per-picture via :meth:`get_picture_metadata`.
|
||||
"""
|
||||
return self._get("/pictures", character_id=character_id).json()
|
||||
|
||||
def list_pictures_for_set(self, set_id: int) -> List[dict]:
|
||||
"""Return the picture rows for every member of a set.
|
||||
|
||||
Uses a single GET /picture_sets/{id} call which embeds all members.
|
||||
As with the character listing, ``description`` and ``tags`` must be
|
||||
fetched per-picture via :meth:`get_picture_metadata`.
|
||||
"""
|
||||
return self._get(f"/picture_sets/{set_id}").json()["pictures"]
|
||||
|
||||
def get_picture_metadata(self, pic_id: int) -> dict:
|
||||
"""Return full metadata for one picture, including description and tags.
|
||||
|
||||
The listing endpoints only return scalar grid fields — the
|
||||
natural-language caption and the WD14 tags are not included — so both
|
||||
are read from the per-picture metadata endpoint::
|
||||
|
||||
GET /pictures/{id}/metadata
|
||||
|
||||
``description`` comes back as a plain string and ``tags`` as a list of
|
||||
``{"id": int, "tag": str}`` objects, exactly what :meth:`build_caption`
|
||||
consumes.
|
||||
"""
|
||||
return self._get(f"/pictures/{pic_id}/metadata").json()
|
||||
|
||||
def download_image_bytes(self, pic_id: int, fmt: str = "jpg") -> bytes:
|
||||
"""Return raw image bytes. Endpoint: GET /pictures/{id}.{ext}"""
|
||||
r = self._session.get(
|
||||
f"{self.base_url}/pictures/{pic_id}.{fmt}",
|
||||
timeout=120,
|
||||
)
|
||||
if not r.ok:
|
||||
raise PixlStashError(
|
||||
f"Image download for id={pic_id} failed ({r.status_code})"
|
||||
)
|
||||
return r.content
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Thumbnail URLs (used by the UI browse modal)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def thumbnail_url(self, pic_id: int) -> str:
|
||||
"""Full URL for a picture's WebP thumbnail."""
|
||||
return f"{self.base_url}/pictures/thumbnails/{pic_id}.webp"
|
||||
|
||||
def picture_set_thumbnail_url(self, set_id: int) -> str:
|
||||
return f"{self.base_url}/picture_sets/{set_id}/thumbnail"
|
||||
|
||||
def character_thumbnail_url(self, character_id: int) -> str:
|
||||
return f"{self.base_url}/characters/{character_id}/thumbnail"
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Caption building
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def tags_to_string(meta: dict) -> str:
|
||||
return ", ".join(
|
||||
t["tag"]
|
||||
for t in (meta.get("tags") or [])
|
||||
if t.get("tag") and t["tag"] != _EMPTY_TAG_SENTINEL
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def build_caption(
|
||||
cls,
|
||||
meta: dict,
|
||||
mode: str = "description",
|
||||
trigger: str = "",
|
||||
) -> str:
|
||||
"""
|
||||
Build a caption string for one picture.
|
||||
|
||||
mode: "description" | "tags" | "both"
|
||||
trigger: optional token prepended to every caption.
|
||||
"""
|
||||
parts: List[str] = []
|
||||
|
||||
if trigger:
|
||||
parts.append(trigger)
|
||||
|
||||
if mode in ("description", "both"):
|
||||
desc = (meta.get("description") or "").strip()
|
||||
if desc:
|
||||
parts.append(desc)
|
||||
|
||||
if mode in ("tags", "both"):
|
||||
tag_str = cls.tags_to_string(meta)
|
||||
if tag_str:
|
||||
parts.append(tag_str)
|
||||
|
||||
content_count = len(parts) - (1 if trigger else 0)
|
||||
if content_count == 0:
|
||||
fallback = (meta.get("description") or "").strip() or cls.tags_to_string(
|
||||
meta
|
||||
)
|
||||
if fallback:
|
||||
parts.append(fallback)
|
||||
|
||||
return ", ".join(parts)
|
||||
|
|
@ -0,0 +1,194 @@
|
|||
"""
|
||||
PixlStash implementation of RemoteDatasetSource.
|
||||
|
||||
Registered under type_id = "pixlstash".
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import os
|
||||
from typing import List
|
||||
|
||||
from tqdm import tqdm
|
||||
|
||||
from toolkit.dataset_sources.base import (
|
||||
ImportField,
|
||||
RemoteDatasetSource,
|
||||
SettingField,
|
||||
SourceGroup,
|
||||
SourceItem,
|
||||
)
|
||||
|
||||
|
||||
class PixlStashDatasetSource(RemoteDatasetSource):
|
||||
type_id = "pixlstash"
|
||||
display_name = "PixlStash"
|
||||
icon_path = os.path.join(os.path.dirname(__file__), "assets", "logo.png")
|
||||
|
||||
# Settings keys stored in the AI-Toolkit DB
|
||||
SETTING_URL = "PIXLSTASH_URL"
|
||||
SETTING_TOKEN = "PIXLSTASH_TOKEN"
|
||||
|
||||
SETTING_VERIFY_SSL = "PIXLSTASH_VERIFY_SSL"
|
||||
|
||||
@classmethod
|
||||
def get_settings_schema(cls) -> List[SettingField]:
|
||||
return [
|
||||
SettingField(
|
||||
key=cls.SETTING_URL,
|
||||
label="PixlStash URL",
|
||||
input_type="text",
|
||||
description="Base URL of your PixlStash server.",
|
||||
placeholder="http://localhost:9537",
|
||||
),
|
||||
SettingField(
|
||||
key=cls.SETTING_TOKEN,
|
||||
label="PixlStash API Token",
|
||||
input_type="password",
|
||||
description="Create a token in PixlStash → Settings → API Tokens.",
|
||||
placeholder="Paste your API token here",
|
||||
),
|
||||
SettingField(
|
||||
key=cls.SETTING_VERIFY_SSL,
|
||||
label="Verify SSL Certificate",
|
||||
input_type="checkbox",
|
||||
description="Uncheck to allow self-signed certificates (e.g. local HTTPS).",
|
||||
placeholder="",
|
||||
required=False,
|
||||
),
|
||||
]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _make_client(self):
|
||||
from extensions.pixlstash.pixlstash_client import (
|
||||
PixlStashClient,
|
||||
PixlStashError,
|
||||
)
|
||||
|
||||
url = self.get_setting(self.SETTING_URL)
|
||||
token = self.get_setting(self.SETTING_TOKEN)
|
||||
if not url or not token:
|
||||
raise ValueError(
|
||||
"PixlStash URL and API token must be set in AI-Toolkit Settings "
|
||||
"before using a PixlStash dataset source."
|
||||
)
|
||||
verify_ssl = self.get_setting(
|
||||
self.SETTING_VERIFY_SSL, default="true"
|
||||
).lower() not in ("false", "0", "no", "off")
|
||||
client = PixlStashClient(url, token, verify_ssl=verify_ssl)
|
||||
client.login()
|
||||
return client
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Thumbnail
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_thumbnail(self, thumbnail_id: str, thumbnail_type: str) -> tuple:
|
||||
client = self._make_client()
|
||||
if thumbnail_type == "character":
|
||||
url = client.character_thumbnail_url(int(thumbnail_id))
|
||||
elif thumbnail_type == "set":
|
||||
url = client.picture_set_thumbnail_url(int(thumbnail_id))
|
||||
else:
|
||||
url = client.thumbnail_url(int(thumbnail_id))
|
||||
r = client._session.get(url, timeout=30)
|
||||
r.raise_for_status()
|
||||
return r.content, r.headers.get("content-type", "image/webp")
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Job config
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def build_job_config(self, params: dict) -> dict:
|
||||
verify_ssl_raw = self.get_setting(self.SETTING_VERIFY_SSL, default="true")
|
||||
verify_ssl = verify_ssl_raw.lower() not in ("false", "0", "no", "off")
|
||||
cfg = {
|
||||
"type": "pixlstash_fetch",
|
||||
"pixlstash_url": self.get_setting(self.SETTING_URL),
|
||||
"pixlstash_token": self.get_setting(self.SETTING_TOKEN),
|
||||
"verify_ssl": verify_ssl,
|
||||
"source_type": params["source_type"],
|
||||
"source_id": int(params["source_id"]),
|
||||
"caption_mode": params.get("caption_mode", "description"),
|
||||
"overwrite": bool(params.get("overwrite", False)),
|
||||
}
|
||||
if params.get("trigger_word"):
|
||||
cfg["trigger_word"] = params["trigger_word"]
|
||||
if params.get("dataset_name"):
|
||||
cfg["dataset_name"] = params["dataset_name"]
|
||||
score = int(params.get("min_score") or 0)
|
||||
if score > 0:
|
||||
cfg["min_score"] = score
|
||||
return cfg
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Browse — grouped items for the UI
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def browse(self) -> List[SourceGroup]:
|
||||
client = self._make_client()
|
||||
characters = client.list_characters()
|
||||
picture_sets = client.list_picture_sets()
|
||||
|
||||
char_items = [
|
||||
SourceItem(
|
||||
id=str(c["id"]),
|
||||
name=c["name"],
|
||||
picture_count=c.get("picture_count", -1) or -1,
|
||||
thumbnail_id=str(c["id"]),
|
||||
thumbnail_type="character",
|
||||
)
|
||||
for c in characters
|
||||
]
|
||||
|
||||
set_items = [
|
||||
SourceItem(
|
||||
id=str(s["id"]),
|
||||
name=s["name"],
|
||||
picture_count=s.get("member_count", -1) or -1,
|
||||
thumbnail_id=str(s["id"]),
|
||||
thumbnail_type="set",
|
||||
)
|
||||
for s in picture_sets
|
||||
]
|
||||
|
||||
return [
|
||||
SourceGroup(id="character", label="Characters", items=char_items),
|
||||
SourceGroup(id="picture_set", label="Picture Sets", items=set_items),
|
||||
]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Source-specific import fields
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@classmethod
|
||||
def get_import_fields(cls) -> List[ImportField]:
|
||||
return [
|
||||
ImportField(
|
||||
id="caption_mode",
|
||||
label="Caption Mode",
|
||||
field_type="select",
|
||||
options=[
|
||||
{"value": "description", "label": "Description"},
|
||||
{"value": "tags", "label": "Tags"},
|
||||
{"value": "both", "label": "Description + Tags"},
|
||||
],
|
||||
default="description",
|
||||
),
|
||||
ImportField(
|
||||
id="min_score",
|
||||
label="Minimum Star Rating",
|
||||
field_type="select",
|
||||
options=[
|
||||
{"value": 0, "label": "All"},
|
||||
{"value": 3, "label": "3★ and above"},
|
||||
{"value": 4, "label": "4★ and above"},
|
||||
{"value": 5, "label": "5★ only"},
|
||||
],
|
||||
default=0,
|
||||
),
|
||||
]
|
||||
Loading…
Reference in New Issue