Add diffusion dataset upload and training info endpoints
Training an image LoRA required knowing the Studio home layout and copying files onto the server by hand, which is the most confusing step of the whole flow. Two small endpoints fix that: - GET /api/train/diffusion/info reports the datasets and outputs roots plus every dataset folder that contains images (with image/caption counts), so the UI can offer a picker instead of a blind free-text path. - POST /api/train/diffusion/dataset uploads images and optional caption .txt / metadata.jsonl files into a named folder under the datasets root, creating it on first use and accumulating on repeat uploads so large sets can arrive in batches. Names are validated to a single path component and files stream to disk under the same per-upload size cap as LLM dataset uploads. The returned name is a valid data_dir for /diffusion/start.
This commit is contained in:
Daniel Han
parent
c5a0ad59cf
commit
cfde12451c
3 changed files with 253 additions and 1 deletions
|
|
@ -738,3 +738,34 @@ class DiffusionTrainingStatusResponse(BaseModel):
|
|||
lora_path: Optional[str] = None
|
||||
started_at: Optional[float] = None
|
||||
updated_at: Optional[float] = None
|
||||
|
||||
|
||||
class DiffusionDatasetSummary(BaseModel):
|
||||
"""One image-dataset folder under the Studio datasets root."""
|
||||
|
||||
name: str
|
||||
path: str
|
||||
image_count: int
|
||||
caption_count: int
|
||||
|
||||
|
||||
class DiffusionTrainingInfoResponse(BaseModel):
|
||||
"""Where diffusion training reads/writes on this Studio, plus usable datasets.
|
||||
|
||||
Lets the UI show real on-disk locations and offer existing dataset folders,
|
||||
instead of asking users to know the Studio home layout."""
|
||||
|
||||
datasets_root: str
|
||||
outputs_root: str
|
||||
datasets: List[DiffusionDatasetSummary]
|
||||
|
||||
|
||||
class DiffusionDatasetUploadResponse(BaseModel):
|
||||
"""Result of uploading images/captions into a named dataset folder. Counts are
|
||||
for the whole folder after the upload, so repeat uploads show the running total."""
|
||||
|
||||
name: str
|
||||
path: str
|
||||
image_count: int
|
||||
caption_count: int
|
||||
uploaded: int
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ Training API routes
|
|||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile
|
||||
from fastapi.responses import StreamingResponse
|
||||
from typing import Dict, Optional, Any
|
||||
import structlog
|
||||
|
|
@ -58,6 +58,9 @@ from models import (
|
|||
TrainingProgress,
|
||||
)
|
||||
from models.training import (
|
||||
DiffusionDatasetSummary,
|
||||
DiffusionDatasetUploadResponse,
|
||||
DiffusionTrainingInfoResponse,
|
||||
DiffusionTrainingStartRequest,
|
||||
DiffusionTrainingStartResponse,
|
||||
DiffusionTrainingStatusResponse,
|
||||
|
|
@ -1192,3 +1195,139 @@ async def diffusion_training_status(current_subject: str = Depends(get_current_s
|
|||
"""Poll the current diffusion training job's status/progress (JSON)."""
|
||||
from core.training.diffusion_training_service import get_diffusion_training_service
|
||||
return DiffusionTrainingStatusResponse(**get_diffusion_training_service().status())
|
||||
|
||||
|
||||
# Extensions accepted into an image-training dataset folder: images the trainer reads,
|
||||
# plus its caption sources (per-image sidecars and metadata/captions jsonl).
|
||||
_DIFFUSION_DATASET_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
|
||||
_DIFFUSION_DATASET_TEXT_EXTS = {".txt", ".caption", ".jsonl"}
|
||||
|
||||
|
||||
def _diffusion_dataset_summary(folder: Path) -> DiffusionDatasetSummary:
|
||||
images = captions = 0
|
||||
for f in folder.iterdir():
|
||||
if not f.is_file():
|
||||
continue
|
||||
ext = f.suffix.lower()
|
||||
if ext in _DIFFUSION_DATASET_IMAGE_EXTS:
|
||||
images += 1
|
||||
elif ext in (".txt", ".caption"):
|
||||
captions += 1
|
||||
return DiffusionDatasetSummary(
|
||||
name = folder.name, path = str(folder), image_count = images, caption_count = captions
|
||||
)
|
||||
|
||||
|
||||
@router.get("/diffusion/info", response_model = DiffusionTrainingInfoResponse)
|
||||
async def diffusion_training_info(current_subject: str = Depends(get_current_subject)):
|
||||
"""Describe where diffusion training reads/writes, and list usable dataset folders.
|
||||
|
||||
A dataset folder is any direct child of the datasets root that contains at least one
|
||||
image. The UI uses this to offer a picker instead of a blind free-text path."""
|
||||
from utils.paths import datasets_root, outputs_root
|
||||
|
||||
def scan() -> DiffusionTrainingInfoResponse:
|
||||
root = datasets_root()
|
||||
found: list[DiffusionDatasetSummary] = []
|
||||
try:
|
||||
children = sorted(p for p in root.iterdir() if p.is_dir())
|
||||
except OSError:
|
||||
children = []
|
||||
for child in children:
|
||||
try:
|
||||
summary = _diffusion_dataset_summary(child)
|
||||
except OSError:
|
||||
continue
|
||||
if summary.image_count > 0:
|
||||
found.append(summary)
|
||||
return DiffusionTrainingInfoResponse(
|
||||
datasets_root = str(root), outputs_root = str(outputs_root()), datasets = found
|
||||
)
|
||||
|
||||
return await asyncio.to_thread(scan)
|
||||
|
||||
|
||||
_DATASET_NAME_RE = None # compiled lazily; module keeps its import block torch-free
|
||||
|
||||
|
||||
def _clean_diffusion_dataset_name(name: str) -> str:
|
||||
"""Validate a dataset folder name: a single path component, no traversal, printable."""
|
||||
import re
|
||||
global _DATASET_NAME_RE
|
||||
if _DATASET_NAME_RE is None:
|
||||
_DATASET_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._ -]{0,127}$")
|
||||
cleaned = (name or "").strip()
|
||||
if not _DATASET_NAME_RE.fullmatch(cleaned) or ".." in cleaned:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = (
|
||||
"Dataset name must be a plain folder name (letters, numbers, dots, "
|
||||
"dashes, spaces; no slashes), e.g. 'my-style-photos'."
|
||||
),
|
||||
)
|
||||
return cleaned
|
||||
|
||||
|
||||
@router.post("/diffusion/dataset", response_model = DiffusionDatasetUploadResponse)
|
||||
async def upload_diffusion_dataset(
|
||||
name: str = Form(...),
|
||||
files: list[UploadFile] = File(...),
|
||||
current_subject: str = Depends(get_current_subject),
|
||||
):
|
||||
"""Upload training images (and optional caption .txt / metadata.jsonl files) into a
|
||||
named folder under the Studio datasets root, creating it if needed. Repeat uploads
|
||||
into the same name accumulate, so large datasets can arrive in batches. The returned
|
||||
name can be passed directly as ``data_dir`` to /diffusion/start."""
|
||||
from utils.paths import datasets_root
|
||||
from utils.upload_limits import get_upload_limit_bytes, get_upload_limit_label
|
||||
|
||||
cleaned = _clean_diffusion_dataset_name(name)
|
||||
folder = datasets_root() / cleaned
|
||||
folder.mkdir(parents = True, exist_ok = True)
|
||||
|
||||
limit_bytes = get_upload_limit_bytes()
|
||||
total_bytes = 0
|
||||
uploaded = 0
|
||||
allowed = _DIFFUSION_DATASET_IMAGE_EXTS | _DIFFUSION_DATASET_TEXT_EXTS
|
||||
for f in files:
|
||||
filename = Path(f.filename or "").name.strip().replace("\x00", "")
|
||||
ext = Path(filename).suffix.lower()
|
||||
if not filename or ext not in allowed:
|
||||
exts = ", ".join(sorted(allowed))
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = f"Unsupported file '{f.filename}'. Allowed: {exts}",
|
||||
)
|
||||
dest = folder / filename
|
||||
complete = False
|
||||
try:
|
||||
with open(dest, "wb") as out:
|
||||
while chunk := await f.read(1024 * 1024):
|
||||
total_bytes += len(chunk)
|
||||
if total_bytes > limit_bytes:
|
||||
raise HTTPException(
|
||||
status_code = 413,
|
||||
detail = (
|
||||
"Dataset upload too large. "
|
||||
f"Maximum is {get_upload_limit_label()} per upload; "
|
||||
"add the remaining images in another batch."
|
||||
),
|
||||
)
|
||||
out.write(chunk)
|
||||
complete = True
|
||||
finally:
|
||||
if not complete:
|
||||
try:
|
||||
dest.unlink(missing_ok = True)
|
||||
except OSError:
|
||||
pass
|
||||
uploaded += 1
|
||||
|
||||
summary = _diffusion_dataset_summary(folder)
|
||||
return DiffusionDatasetUploadResponse(
|
||||
name = cleaned,
|
||||
path = str(folder),
|
||||
image_count = summary.image_count,
|
||||
caption_count = summary.caption_count,
|
||||
uploaded = uploaded,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -377,3 +377,85 @@ def test_stale_pump_events_cannot_corrupt_new_job():
|
|||
# The current job's events still apply.
|
||||
svc._apply_event({"type": "progress", "step": 9}, proc = current)
|
||||
assert svc.status()["step"] == 9
|
||||
|
||||
|
||||
# ── /diffusion/info + /diffusion/dataset (dataset discovery + upload) ─────────
|
||||
@pytest.fixture
|
||||
def dataset_roots(client, monkeypatch, tmp_path):
|
||||
# The endpoints import these lazily per-request, so patching the package attr works.
|
||||
import utils.paths as up
|
||||
|
||||
ds_root = tmp_path / "assets" / "datasets"
|
||||
out_root = tmp_path / "outputs"
|
||||
ds_root.mkdir(parents = True)
|
||||
out_root.mkdir(parents = True)
|
||||
monkeypatch.setattr(up, "datasets_root", lambda: ds_root)
|
||||
monkeypatch.setattr(up, "outputs_root", lambda: out_root)
|
||||
return ds_root, out_root
|
||||
|
||||
|
||||
def test_diffusion_info_lists_image_dataset_folders(client, dataset_roots):
|
||||
ds_root, out_root = dataset_roots
|
||||
good = ds_root / "cat-photos"
|
||||
good.mkdir()
|
||||
(good / "a.png").write_bytes(b"x")
|
||||
(good / "b.jpg").write_bytes(b"x")
|
||||
(good / "a.txt").write_text("a cat")
|
||||
(ds_root / "empty-dir").mkdir() # no images -> not a dataset
|
||||
(ds_root / "stray.txt").write_text("not a folder")
|
||||
|
||||
r = client.get("/api/train/diffusion/info")
|
||||
assert r.status_code == 200, r.text
|
||||
body = r.json()
|
||||
assert body["datasets_root"] == str(ds_root)
|
||||
assert body["outputs_root"] == str(out_root)
|
||||
assert [d["name"] for d in body["datasets"]] == ["cat-photos"]
|
||||
assert body["datasets"][0]["image_count"] == 2
|
||||
assert body["datasets"][0]["caption_count"] == 1
|
||||
|
||||
|
||||
def test_diffusion_dataset_upload_accumulates(client, dataset_roots):
|
||||
ds_root, _ = dataset_roots
|
||||
files = [
|
||||
("files", ("a.png", b"png-bytes", "image/png")),
|
||||
("files", ("b.JPG", b"jpg-bytes", "image/jpeg")),
|
||||
("files", ("a.txt", b"a caption", "text/plain")),
|
||||
]
|
||||
r = client.post("/api/train/diffusion/dataset", data = {"name": "my style"}, files = files)
|
||||
assert r.status_code == 200, r.text
|
||||
body = r.json()
|
||||
assert body["name"] == "my style"
|
||||
assert body["uploaded"] == 3
|
||||
assert body["image_count"] == 2
|
||||
assert body["caption_count"] == 1
|
||||
assert (ds_root / "my style" / "a.png").read_bytes() == b"png-bytes"
|
||||
|
||||
# A second batch into the same name accumulates (large sets arrive in chunks).
|
||||
r = client.post(
|
||||
"/api/train/diffusion/dataset",
|
||||
data = {"name": "my style"},
|
||||
files = [("files", ("c.webp", b"w", "image/webp"))],
|
||||
)
|
||||
assert r.status_code == 200, r.text
|
||||
assert r.json()["uploaded"] == 1
|
||||
assert r.json()["image_count"] == 3
|
||||
|
||||
|
||||
def test_diffusion_dataset_upload_rejects_traversal_names(client, dataset_roots):
|
||||
for bad in ("../evil", "a/b", ".hidden", " "):
|
||||
r = client.post(
|
||||
"/api/train/diffusion/dataset",
|
||||
data = {"name": bad},
|
||||
files = [("files", ("a.png", b"x", "image/png"))],
|
||||
)
|
||||
assert r.status_code == 400, f"{bad!r}: {r.status_code}"
|
||||
|
||||
|
||||
def test_diffusion_dataset_upload_rejects_unsupported_files(client, dataset_roots):
|
||||
r = client.post(
|
||||
"/api/train/diffusion/dataset",
|
||||
data = {"name": "ok-name"},
|
||||
files = [("files", ("weights.exe", b"mz", "application/octet-stream"))],
|
||||
)
|
||||
assert r.status_code == 400
|
||||
assert "Unsupported file" in r.json()["detail"]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue