diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py new file mode 100644 index 0000000000..d0b3d3220e --- /dev/null +++ b/studio/backend/core/training/diffusion_training_service.py @@ -0,0 +1,275 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""A small job service for diffusion LoRA training. + +This is deliberately separate from the LLM ``TrainingBackend``: that backend's +lifecycle (LLM config build, per-run SQLite rows, matplotlib plots, transfer-to-chat- +inference) is specific to text training and would mis-handle a diffusion run. This +service does only what a diffusion job needs -- spawn the trainer subprocess, pump its +events (``model_load_*`` / ``progress`` / ``complete`` / ``error``) into an in-memory +status snapshot, and support stop -- and is polled over JSON by the route layer. + +The subprocess context, target, and queues are injectable so the service can be unit +tested without real multiprocessing or torch: tests pass a fake context whose Process +runs a scripted target on a thread. +""" + +from __future__ import annotations + +import multiprocessing as mp +import threading +import time +import uuid +from typing import Any, Callable, Optional + +# Spawn (not fork): a fresh interpreter, matching the LLM training worker, so CUDA/torch +# state from the parent never leaks into the trainer. +_CTX = mp.get_context("spawn") + +# Terminal event types after which the pump stops. +_TERMINAL = ("complete", "error") + + +def _run_diffusion_child(*, event_queue: Any, stop_queue: Any, config: dict) -> None: + # Imported lazily so this module (and the route layer) stays torch-free at import. + from .diffusion_lora_trainer import run_diffusion_training_process + run_diffusion_training_process(event_queue = event_queue, stop_queue = stop_queue, config = config) + + +def _default_target(*, event_queue: Any, stop_queue: Any, config: dict) -> None: + # First thing in the spawned child (before torch is imported): bind to the parent's + # death on Linux and scrub the native path lease secret, exactly like the inference / + # export / LLM-training workers. multiprocessing children cannot be given a + # parent-set preexec_fn, so the child must self-bind; otherwise a Studio crash or + # kill leaves this trainer holding the GPU. Tests inject their own target, so this + # binding only runs for the real production spawn. + from utils.native_path_leases import run_without_native_path_secret + run_without_native_path_secret( + _run_diffusion_child, event_queue = event_queue, stop_queue = stop_queue, config = config + ) + + +def _idle_state() -> dict[str, Any]: + return { + "active": False, + "job_id": None, + "status": "idle", + "message": "", + "step": 0, + "total_steps": 0, + "loss": None, + "avg_loss": None, + "learning_rate": None, + "num_images": None, + "in_model_load": False, + "output_dir": None, + "lora_path": None, + "started_at": None, + "updated_at": None, + } + + +class DiffusionTrainingService: + """One diffusion LoRA training job at a time, spawned as a subprocess.""" + + def __init__( + self, + *, + ctx: Any = None, + target: Optional[Callable[..., None]] = None, + ) -> None: + self._ctx = ctx if ctx is not None else _CTX + self._target = target if target is not None else _default_target + self._lock = threading.Lock() + self._proc: Any = None + self._stop_queue: Any = None + self._pump: Optional[threading.Thread] = None + self._state: dict[str, Any] = _idle_state() + + # ── lifecycle ──────────────────────────────────────────────────────────── + def is_active(self) -> bool: + with self._lock: + return self._proc is not None and self._proc.is_alive() + + def start(self, config: dict) -> str: + """Validate ``config``, spawn the trainer, and start pumping its events. + + Raises ValueError for an unusable config (before any spawn) and RuntimeError if a + job is already running. Returns the new job id.""" + # Validate cheaply BEFORE spawning so a bad request fails fast with a clear error. + from .diffusion_lora_trainer import _config_from_dict + + _config_from_dict(config).normalized() + + # Join a finished job's pump OUTSIDE the lock: its final state writes take this + # lock (via _apply_event / the exit handler), so joining under it would stall + # the start for the whole timeout and then let the stale pump overwrite the new + # job's state once the lock was released. + with self._lock: + if self._proc is not None and self._proc.is_alive(): + raise RuntimeError("A diffusion training job is already running.") + pump = self._pump + if pump is not None and pump.is_alive(): + pump.join(timeout = 5.0) + + with self._lock: + # Re-check: another start() may have won the race while we joined. + if self._proc is not None and self._proc.is_alive(): + raise RuntimeError("A diffusion training job is already running.") + + job_id = uuid.uuid4().hex + event_queue = self._ctx.Queue() + self._stop_queue = self._ctx.Queue() + self._proc = self._ctx.Process( + target = self._target, + kwargs = { + "event_queue": event_queue, + "stop_queue": self._stop_queue, + "config": config, + }, + daemon = True, + ) + self._proc.start() + try: + from utils.process_lifetime import adopt_pid + adopt_pid(self._proc.pid) # bind to parent lifetime (no zombie on exit) + except Exception: # noqa: BLE001 -- lifetime binding is best-effort + pass + + now = time.time() + self._state = _idle_state() + self._state.update( + active = True, + job_id = job_id, + status = "running", + message = "Starting diffusion LoRA training...", + started_at = now, + updated_at = now, + ) + self._pump = threading.Thread( + target = self._pump_loop, args = (event_queue, self._proc), daemon = True + ) + self._pump.start() + return job_id + + def stop(self) -> bool: + """Request a clean stop (the trainer finishes the current step and saves a partial + adapter). Returns True if a stop was signalled, False if nothing was running.""" + with self._lock: + if self._proc is None or not self._proc.is_alive() or self._stop_queue is None: + return False + try: + self._stop_queue.put(True) + except Exception: # noqa: BLE001 + return False + self._state["message"] = "Stop requested; finishing the current step..." + self._state["updated_at"] = time.time() + return True + + def status(self) -> dict[str, Any]: + with self._lock: + snap = dict(self._state) + # Keep ``active`` honest even if the process died between events. + snap["active"] = self._proc is not None and self._proc.is_alive() + return snap + + # ── event pump ─────────────────────────────────────────────────────────── + def _pump_loop(self, event_queue: Any, proc: Any) -> None: + while True: + try: + ev = event_queue.get(timeout = 1.0) + except Exception: # noqa: BLE001 -- Empty (timeout) or a closed queue + if not proc.is_alive(): + # Drain anything buffered, then decide if it exited cleanly. + drained = False + while True: + try: + self._apply_event(event_queue.get_nowait(), proc = proc) + drained = True + except Exception: # noqa: BLE001 + break + with self._lock: + if self._proc is not proc: + return # superseded by a newer job; don't touch its state + if self._state.get("status") not in ("completed", "stopped", "error"): + self._state.update( + active = False, + status = "error", + message = "Training process exited unexpectedly.", + updated_at = time.time(), + ) + _ = drained + return + continue + self._apply_event(ev, proc = proc) + if ev.get("type") in _TERMINAL: + return + + def _apply_event( + self, + ev: dict[str, Any], + proc: Any = None, + ) -> None: + """Fold one trainer event into the status snapshot. Pure state update -- unit + tested by feeding events directly. ``proc`` (when given) fences a stale pump: + an event from a superseded job's process must not touch the current job's + state.""" + etype = ev.get("type") + with self._lock: + if proc is not None and self._proc is not proc: + return + s = self._state + s["updated_at"] = time.time() + if etype == "model_load_started": + s.update(in_model_load = True, status = "running", message = "Loading base model...") + if ev.get("num_images") is not None: + s["num_images"] = ev.get("num_images") + elif etype == "model_load_completed": + s.update(in_model_load = False, message = "Training...") + elif etype == "progress": + s.update( + status = "running", + step = ev.get("step", s["step"]), + total_steps = ev.get("total_steps", s["total_steps"]), + loss = ev.get("loss", s["loss"]), + avg_loss = ev.get("avg_loss", s["avg_loss"]), + learning_rate = ev.get("learning_rate", s["learning_rate"]), + message = "Training...", + ) + elif etype == "complete": + # Reset in_model_load: a stop during model load emits complete without a + # preceding model_load_completed, which would otherwise leave a stale + # loading indicator after the job ended. + s.update( + active = False, + in_model_load = False, + status = "stopped" if ev.get("stopped") else "completed", + output_dir = ev.get("output_dir"), + lora_path = ev.get("lora_path"), + message = "Stopped (partial adapter saved)." + if ev.get("stopped") + else "Training complete.", + ) + elif etype == "error": + # Reset in_model_load too: an error raised during model loading has no + # model_load_completed, so the terminal state must clear it explicitly. + s.update( + active = False, + in_model_load = False, + status = "error", + message = str(ev.get("message", "error")), + ) + + +_service: Optional[DiffusionTrainingService] = None +_service_lock = threading.Lock() + + +def get_diffusion_training_service() -> DiffusionTrainingService: + """Process-wide singleton used by the route layer.""" + global _service + with _service_lock: + if _service is None: + _service = DiffusionTrainingService() + return _service diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index ff815a2fa9..ec8900c4a3 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -665,3 +665,107 @@ class TrainingRunDeleteResponse(BaseModel): status: str message: str + + +class DiffusionTrainingStartRequest(BaseModel): + """Request to start a diffusion (SDXL) LoRA training job. + + Field names mirror ``core.training.diffusion_lora_trainer.DiffusionLoraConfig`` so the + service can pass ``model_dump()`` straight through. Only the paths are required; the + rest carry the trainer's defaults. + """ + + model_config = ConfigDict(protected_namespaces = ()) + + base_model: str = Field(..., description = "HF repo id or local path to an SDXL pipeline") + data_dir: str = Field(..., description = "Folder of training images (+ captions)") + output_dir: str = Field(..., description = "Directory to write the LoRA .safetensors into") + instance_prompt: Optional[str] = Field( + None, description = "Dreambooth caption applied to images without their own caption" + ) + resolution: int = Field( + 1024, ge = 64, le = 2048, description = "Square training resolution (multiple of 8)" + ) + train_steps: int = Field(500, ge = 1, le = 100000) + learning_rate: float = Field(1e-4, gt = 0) + train_batch_size: int = Field(1, ge = 1, le = 64) + gradient_accumulation_steps: int = Field(1, ge = 1, le = 256) + lora_rank: int = Field(16, ge = 1, le = 320) + lora_alpha: Optional[int] = Field(None, ge = 1, le = 640, description = "Defaults to lora_rank") + lora_dropout: float = Field(0.0, ge = 0.0, le = 1.0) + # Mirror the remaining training-affecting knobs of DiffusionLoraConfig so a client that + # sets them is not silently trained with defaults. Default the target list to the SDXL + # attention projections (the trainer's DEFAULT_LORA_TARGETS) so it is never None. + lora_target_modules: List[str] = Field( + default_factory = lambda: ["to_k", "to_q", "to_v", "to_out.0"], + description = "U-Net modules to attach LoRA to", + ) + max_grad_norm: float = Field(1.0, gt = 0, description = "Gradient clipping max-norm") + seed: int = Field(42) + mixed_precision: Literal["bf16", "fp16", "no"] = Field("bf16") + snr_gamma: Optional[float] = Field(5.0, description = "Min-SNR loss weighting; null disables") + gradient_checkpointing: bool = Field(True) + lr_scheduler: str = Field("constant") + lr_warmup_steps: int = Field(0, ge = 0) + center_crop: bool = Field(False) + random_flip: bool = Field(True) + caption_column: str = Field("text") + hf_token: Optional[str] = Field(None) + + +class DiffusionTrainingStartResponse(BaseModel): + """Response for starting a diffusion training job.""" + + job_id: str + status: str + + +class DiffusionTrainingStatusResponse(BaseModel): + """A snapshot of the current diffusion training job (or idle).""" + + active: bool + job_id: Optional[str] = None + status: str + message: str = "" + step: int = 0 + total_steps: int = 0 + loss: Optional[float] = None + avg_loss: Optional[float] = None + learning_rate: Optional[float] = None + num_images: Optional[int] = None + in_model_load: bool = False + output_dir: Optional[str] = None + 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 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index b8a970a106..5f2169f86a 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2826,12 +2826,26 @@ def _guard_chat_load_against_training( from routes.training_vram import can_load_chat_during_training try: - if not get_training_backend().is_training_active(): - return + llm_active = get_training_backend().is_training_active() except Exception as e: logger.warning("Could not check training state for chat-load guard: %s", e) return + if not llm_active: + # An SDXL LoRA trainer runs in its own subprocess and its VRAM can't be cheaply + # fit-checked here, so refuse the chat load outright while one is active rather + # than risk OOMing the run. Symmetric with the image-load guard. + if _diffusion_training_active(): + raise HTTPException( + status_code = 409, + detail = ( + "Can't load this model while diffusion (Images) training is running: " + "its GPU memory use can't be verified against the trainer, so the load " + "was refused to protect the run. Try again after training finishes." + ), + ) + return + is_gguf = bool(getattr(config, "is_gguf", False)) required_override_gb = ( _estimate_gguf_required_gb( @@ -11024,6 +11038,16 @@ async def _openai_passthrough_non_streaming( # ────────────────────────────────────────────────────────────────────────── +def _diffusion_training_active() -> bool: + """Whether a diffusion (SDXL) LoRA job is running. Best-effort so a load is never + blocked just because the training service could not be imported/read.""" + try: + from core.training.diffusion_training_service import get_diffusion_training_service + return get_diffusion_training_service().is_active() + except Exception: # noqa: BLE001 + return False + + def _guard_diffusion_load_against_training() -> None: """Refuse loading an image model while a training run is active. Unlike chat, a diffusion pipeline's VRAM can't be cheaply estimated before the load, so the @@ -11032,11 +11056,15 @@ def _guard_diffusion_load_against_training() -> None: from core.training import get_training_backend try: - if not get_training_backend().is_training_active(): - return + llm_active = get_training_backend().is_training_active() except Exception as e: logger.warning("Could not check training state for image-load guard: %s", e) return + # An SDXL LoRA trainer runs in its own subprocess on the same GPU, so an image + # load must be refused while one is active too -- otherwise the resident pipeline + # competes with the trainer for VRAM. Symmetric with the diffusion-start interlock. + if not llm_active and not _diffusion_training_active(): + return raise HTTPException( status_code = 409, detail = ( diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index d235a2e843..c56a5e9bc9 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -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 @@ -57,6 +57,14 @@ from models import ( TrainingStatus, TrainingProgress, ) +from models.training import ( + DiffusionDatasetSummary, + DiffusionDatasetUploadResponse, + DiffusionTrainingInfoResponse, + DiffusionTrainingStartRequest, + DiffusionTrainingStartResponse, + DiffusionTrainingStatusResponse, +) from models.responses import TrainingStopResponse, TrainingMetricsResponse from pydantic import BaseModel as PydanticBaseModel @@ -172,6 +180,20 @@ async def start_training( error = "Training already active", ) + # A diffusion (SDXL) LoRA job runs in its own subprocess on the same GPU, so an + # LLM start must also refuse while one is active -- otherwise the two trainers + # contend for VRAM and both fail. Symmetric with the check in start_diffusion_training. + if _diffusion_training_active(): + return TrainingJobResponse( + job_id = "", + status = "error", + message = ( + "A diffusion (Images) LoRA training job is already running. " + "Stop it before starting an LLM training run." + ), + error = "Diffusion training already active", + ) + # Job ID; start_training() sets it on the backend only after the old # pump thread is dead. job_id = f"job_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{_uuid.uuid4().hex[:8]}" @@ -1040,3 +1062,283 @@ async def stream_training_progress( "X-Accel-Buffering": "no", }, ) + + +# ── Diffusion (SDXL) LoRA training ──────────────────────────────────────────── +# A separate, lightweight job path from the LLM training endpoints above: diffusion +# runs are driven by DiffusionTrainingService (its own subprocess + event pump), not +# the LLM TrainingBackend, so the two never contend and diffusion never triggers LLM +# lifecycle (DB run rows, plots, transfer-to-chat-inference). + + +def _diffusion_training_active() -> bool: + """Whether a diffusion (SDXL) LoRA job is currently running. Best-effort so the + interlock never blocks a start just because the service could not be imported.""" + try: + from core.training.diffusion_training_service import get_diffusion_training_service + return get_diffusion_training_service().is_active() + except Exception: # noqa: BLE001 + return False + + +def _free_gpu_for_diffusion_training() -> None: + """Free GPU residents before the diffusion trainer spawns its own SDXL pipeline. + + The trainer subprocess loads a full SDXL pipeline; an export worker, a resident + Images pipeline, or loaded chat models would otherwise keep their VRAM allocated and + OOM the run. Mirrors the LLM start path's pre-spawn cleanup (export + diffusion + pipeline + chat). Best-effort: a failure to free one resident never blocks the start.""" + try: + from core.export import get_export_backend + exp_backend = get_export_backend() + if exp_backend.current_checkpoint or exp_backend.is_export_active(): + logger.info("Shutting down export subprocess to free GPU memory for diffusion training") + exp_backend._shutdown_subprocess() + exp_backend.current_checkpoint = None + exp_backend.is_vision = False + exp_backend.is_peft = False + except Exception as e: # noqa: BLE001 + logger.warning("Could not shut down export subprocess: %s", e) + + try: + from core.inference import gpu_arbiter + from core.inference.diffusion import get_diffusion_backend + + diffusion = get_diffusion_backend() + if diffusion.is_loaded: + logger.info("Unloading resident Images pipeline to free GPU memory for training") + diffusion.unload() # no-op when nothing is loaded; also preempts an in-flight load + gpu_arbiter.release(gpu_arbiter.DIFFUSION) + except Exception as e: # noqa: BLE001 + logger.warning("Could not unload Images pipeline for diffusion training: %s", e) + + try: + # The SDXL trainer's footprint can't be cheaply sized against a resident chat + # model, so free chat unconditionally (same conservative choice the LLM path + # makes for an in-flight chat load) rather than risk an OOM. + from routes.training_vram import free_chat_models_for_training, summarize_resident_chat + if summarize_resident_chat()["any"]: + freed = free_chat_models_for_training(reason = "diffusion training starting") + logger.info("Freed chat model(s) for diffusion training: %s", freed) + except Exception as e: # noqa: BLE001 + logger.warning("Could not free chat models for diffusion training: %s", e) + + +@router.post("/diffusion/start", response_model = DiffusionTrainingStartResponse) +async def start_diffusion_training( + body: DiffusionTrainingStartRequest, current_subject: str = Depends(get_current_subject) +): + """Start an SDXL LoRA training job from an image + caption dataset.""" + from core.training.diffusion_training_service import get_diffusion_training_service + + # Interlock: refuse while an LLM training run holds the GPU (symmetric with the + # diffusion check in start_training), so the two trainers never contend for VRAM. + try: + if get_training_backend().is_training_active(): + raise HTTPException( + status_code = 409, + detail = ( + "An LLM training job is already running. " + "Stop it before starting diffusion (Images) training." + ), + ) + except HTTPException: + raise + except Exception: # noqa: BLE001 -- backend import/health issue must not block a start + pass + + # Resolve + contain the dataset and output paths BEFORE spawning, so Studio-relative + # names ("uploads/my-images") work and absolute paths stay under a Studio root -- the + # trainer subprocess otherwise resolves them relative to its own cwd. + config = body.model_dump() + try: + from utils.paths import resolve_dataset_path, resolve_output_dir + config["data_dir"] = str(resolve_dataset_path(config["data_dir"])) + config["output_dir"] = str(resolve_output_dir(config["output_dir"])) + except ValueError as e: + raise HTTPException(status_code = 400, detail = str(e)) + + # Validate the config BEFORE freeing resident GPU workloads, so a start that is + # then refused (bad numbers, a non-SDXL base model) never tears down the user's + # loaded chat/Images model. service.start() re-runs this cheaply before spawn. + from core.training.diffusion_lora_trainer import _config_from_dict + + try: + _config_from_dict(config).normalized() + except ValueError as e: + raise HTTPException(status_code = 400, detail = str(e)) + + # Free resident GPU workloads (export / Images pipeline / chat) before the trainer + # loads its own SDXL pipeline. + _free_gpu_for_diffusion_training() + + service = get_diffusion_training_service() + try: + job_id = service.start(config) + except ValueError as e: + raise HTTPException(status_code = 400, detail = str(e)) + except RuntimeError as e: + # A job is already running. + raise HTTPException(status_code = 409, detail = str(e)) + except Exception as e: + raise log_and_http_error( + e, + 500, + "Failed to start diffusion training", + event = "diffusion_training.start_failed", + log = logger, + ) + return DiffusionTrainingStartResponse(job_id = job_id, status = "running") + + +@router.post("/diffusion/stop") +async def stop_diffusion_training(current_subject: str = Depends(get_current_subject)): + """Request a clean stop of the running diffusion training job (partial adapter saved).""" + from core.training.diffusion_training_service import get_diffusion_training_service + + stopped = get_diffusion_training_service().stop() + return {"status": "stopping" if stopped else "idle"} + + +@router.get("/diffusion/status", response_model = DiffusionTrainingStatusResponse) +async def diffusion_training_status(current_subject: str = Depends(get_current_subject)): + """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, + ) diff --git a/studio/backend/tests/test_diffusion_lora_trainer.py b/studio/backend/tests/test_diffusion_lora_trainer.py index fd17b29500..e238fdaff7 100644 --- a/studio/backend/tests/test_diffusion_lora_trainer.py +++ b/studio/backend/tests/test_diffusion_lora_trainer.py @@ -208,7 +208,9 @@ def test_config_rejects_known_non_sdxl_base_models(): "z-image-turbo-Q4_K_M.gguf", ): with pytest.raises(ValueError, match = "SDXL"): - DiffusionLoraConfig(base_model = bad, data_dir = "d", output_dir = "o").normalized() + DiffusionLoraConfig( + base_model = bad, data_dir = "d", output_dir = "o" + ).normalized() def test_config_accepts_sdxl_and_unknown_base_models(): diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py new file mode 100644 index 0000000000..9c4452df7f --- /dev/null +++ b/studio/backend/tests/test_diffusion_training.py @@ -0,0 +1,477 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for the diffusion LoRA training service + routes. + +The service's subprocess context and target are injected with in-thread fakes, so the +full start -> event-pump -> status -> complete path is exercised without real +multiprocessing or torch. The routes are hit with the FastAPI TestClient and a mocked +service, so wiring / validation / error mapping are covered without a GPU. +""" + +from __future__ import annotations + +import queue as _queue +import threading +import time + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from auth.authentication import get_current_subject +from core.training.diffusion_training_service import DiffusionTrainingService +from routes.training import router as training_router + + +# ── fake spawn context (runs the "process" target on a thread) ──────────────── +class _FakeQueue: + def __init__(self) -> None: + self._q: _queue.Queue = _queue.Queue() + + def put(self, x): + self._q.put(x) + + def get(self, timeout = None): + return self._q.get(timeout = timeout) # raises queue.Empty on timeout + + def get_nowait(self): + return self._q.get_nowait() + + def empty(self): + return self._q.empty() + + +class _FakeProc: + def __init__(self, target, kwargs, daemon): + self._target = target + self._kwargs = kwargs + self._thread: threading.Thread | None = None + self.pid = 4321 + + def start(self): + self._thread = threading.Thread(target = self._target, kwargs = self._kwargs, daemon = True) + self._thread.start() + + def is_alive(self): + return self._thread is not None and self._thread.is_alive() + + +class _FakeCtx: + def Queue(self): + return _FakeQueue() + + def Process(self, target, kwargs, daemon): + return _FakeProc(target, kwargs, daemon) + + +def _happy_target(*, event_queue, stop_queue, config): + event_queue.put({"type": "model_load_started", "num_images": 3}) + event_queue.put({"type": "model_load_completed"}) + event_queue.put( + { + "type": "progress", + "step": 1, + "total_steps": 2, + "loss": 0.5, + "avg_loss": 0.5, + "learning_rate": 1e-4, + } + ) + event_queue.put( + { + "type": "progress", + "step": 2, + "total_steps": 2, + "loss": 0.4, + "avg_loss": 0.45, + "learning_rate": 1e-4, + } + ) + event_queue.put( + { + "type": "complete", + "output_dir": config["output_dir"], + "lora_path": config["output_dir"] + "/pytorch_lora_weights.safetensors", + "stopped": False, + } + ) + + +def _stoppable_target(*, event_queue, stop_queue, config): + event_queue.put({"type": "model_load_completed"}) + stop_queue.get(timeout = 5.0) # block until stop() signals + event_queue.put( + {"type": "complete", "output_dir": config["output_dir"], "lora_path": "x", "stopped": True} + ) + + +def _crashing_target(*, event_queue, stop_queue, config): + event_queue.put({"type": "model_load_started"}) + # Exits without a terminal event -> the pump must mark it as an error. + + +_CFG = {"base_model": "b", "data_dir": "d", "output_dir": "/tmp/out", "train_steps": 2} + + +def _wait_status( + svc, + *terminal, + timeout = 3.0, +): + end = time.time() + timeout + while time.time() < end: + st = svc.status() + if st["status"] in terminal: + return st + time.sleep(0.02) + return svc.status() + + +def test_service_happy_path(): + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) + job_id = svc.start(dict(_CFG)) + assert job_id + st = _wait_status(svc, "completed") + assert st["status"] == "completed" + assert st["step"] == 2 and st["total_steps"] == 2 + assert st["num_images"] == 3 + assert st["loss"] == 0.4 and st["avg_loss"] == 0.45 + assert st["lora_path"].endswith("pytorch_lora_weights.safetensors") + assert st["active"] is False + + +def test_service_rejects_bad_config_before_spawn(): + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) + with pytest.raises(ValueError): + svc.start({**_CFG, "train_steps": 0}) + # Nothing was spawned; still idle. + assert svc.status()["status"] == "idle" + + +def test_service_rejects_second_concurrent_job(): + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _stoppable_target) + svc.start(dict(_CFG)) + _wait_status(svc, "running") + with pytest.raises(RuntimeError): + svc.start(dict(_CFG)) + assert svc.stop() is True + _wait_status(svc, "stopped") + + +def test_service_stop_marks_stopped(): + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _stoppable_target) + svc.start(dict(_CFG)) + _wait_status(svc, "running") + assert svc.stop() is True + st = _wait_status(svc, "stopped") + assert st["status"] == "stopped" + assert st["active"] is False + # Stopping again when idle is a no-op. + assert svc.stop() is False + + +def test_service_crash_without_terminal_event_is_error(): + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _crashing_target) + svc.start(dict(_CFG)) + st = _wait_status(svc, "error") + assert st["status"] == "error" + assert "unexpectedly" in st["message"] + + +def test_apply_event_transitions(): + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) + svc._apply_event({"type": "model_load_started", "num_images": 5}) + assert svc.status()["in_model_load"] is True and svc.status()["num_images"] == 5 + svc._apply_event({"type": "model_load_completed"}) + assert svc.status()["in_model_load"] is False + svc._apply_event({"type": "error", "message": "boom"}) + assert svc.status()["status"] == "error" and svc.status()["message"] == "boom" + + +def test_terminal_events_clear_model_load_flag(): + # A stop or error during model load emits complete/error WITHOUT a preceding + # model_load_completed, so the terminal update must reset in_model_load or the + # client shows a stale loading indicator after the job ended. + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) + svc._apply_event({"type": "model_load_started"}) + assert svc.status()["in_model_load"] is True + svc._apply_event({"type": "complete", "stopped": True}) + assert svc.status()["in_model_load"] is False and svc.status()["status"] == "stopped" + + svc2 = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) + svc2._apply_event({"type": "model_load_started"}) + svc2._apply_event({"type": "error", "message": "load failed"}) + assert svc2.status()["in_model_load"] is False and svc2.status()["status"] == "error" + + +# ── route wiring (mocked service) ───────────────────────────────────────────── +class _FakeService: + def __init__(self): + self._running = False + self.started_with = None + + def start(self, config): + self.started_with = config + self._running = True + return "job-123" + + def stop(self): + was = self._running + self._running = False + return was + + def status(self): + return { + "active": self._running, + "job_id": "job-123" if self._running else None, + "status": "running" if self._running else "idle", + "message": "", + "step": 1, + "total_steps": 2, + "loss": 0.5, + "avg_loss": 0.5, + "learning_rate": 1e-4, + "num_images": 3, + "in_model_load": False, + "output_dir": None, + "lora_path": None, + "started_at": None, + "updated_at": None, + } + + +class _FakeLLMBackend: + def __init__(self, active = False): + self._active = active + + def is_training_active(self): + return self._active + + +@pytest.fixture +def client(monkeypatch): + fake = _FakeService() + monkeypatch.setattr( + "core.training.diffusion_training_service.get_diffusion_training_service", lambda: fake + ) + # Neutralize the LLM interlock + GPU-free for the wiring tests (their own tests below + # exercise those behaviors). The route imports get_training_backend at module scope. + import routes.training as tr + + monkeypatch.setattr(tr, "get_training_backend", lambda: _FakeLLMBackend(active = False)) + monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None) + app = FastAPI() + app.include_router(training_router, prefix = "/api/train") + app.dependency_overrides[get_current_subject] = lambda: "test-user" + c = TestClient(app) + c._fake = fake # type: ignore[attr-defined] + return c + + +# Studio-relative paths: the route resolves/contains them before spawn. +_BODY = { + "base_model": "stabilityai/sdxl-turbo", + "data_dir": "uploads/my-images", + "output_dir": "my-lora-run", + "train_steps": 10, +} + + +def test_route_start_ok(client): + r = client.post("/api/train/diffusion/start", json = _BODY) + assert r.status_code == 200, r.text + assert r.json() == {"job_id": "job-123", "status": "running"} + assert client._fake.started_with["base_model"] == "stabilityai/sdxl-turbo" + # Paths were resolved to absolute Studio-contained locations before spawn. + from pathlib import Path + + assert Path(client._fake.started_with["data_dir"]).is_absolute() + assert Path(client._fake.started_with["output_dir"]).is_absolute() + + +def test_route_start_forwards_extra_training_knobs(client): + # max_grad_norm and lora_target_modules must reach the service, not be silently dropped. + body = {**_BODY, "max_grad_norm": 0.5, "lora_target_modules": ["to_q", "to_v"]} + r = client.post("/api/train/diffusion/start", json = body) + assert r.status_code == 200, r.text + assert client._fake.started_with["max_grad_norm"] == 0.5 + assert client._fake.started_with["lora_target_modules"] == ["to_q", "to_v"] + + +def test_route_start_rejects_uncontained_paths(client): + # An absolute path outside the Studio dataset roots is a 400, not silently accepted. + r = client.post("/api/train/diffusion/start", json = {**_BODY, "data_dir": "/etc"}) + assert r.status_code == 400 + + +def test_route_start_blocked_by_active_llm_training(client, monkeypatch): + import routes.training as tr + + monkeypatch.setattr(tr, "get_training_backend", lambda: _FakeLLMBackend(active = True)) + r = client.post("/api/train/diffusion/start", json = _BODY) + assert r.status_code == 409 + assert "LLM training" in r.json()["detail"] + + +def test_route_start_missing_required_is_422(client): + r = client.post( + "/api/train/diffusion/start", json = {"base_model": "x"} + ) # no data_dir/output_dir + assert r.status_code == 422 + + +def test_route_start_bad_config_maps_to_400(client, monkeypatch): + def _raise(_cfg): + raise ValueError("resolution must be a multiple of 8") + + client._fake.start = _raise # type: ignore[assignment] + r = client.post("/api/train/diffusion/start", json = _BODY) + assert r.status_code == 400 + assert "multiple of 8" in r.json()["detail"] + + +def test_route_start_conflict_maps_to_409(client): + def _raise(_cfg): + raise RuntimeError("A diffusion training job is already running.") + + client._fake.start = _raise # type: ignore[assignment] + r = client.post("/api/train/diffusion/start", json = _BODY) + assert r.status_code == 409 + + +def test_route_status_and_stop(client): + client.post("/api/train/diffusion/start", json = _BODY) + s = client.get("/api/train/diffusion/status") + assert s.status_code == 200 and s.json()["status"] == "running" + st = client.post("/api/train/diffusion/stop") + assert st.status_code == 200 and st.json()["status"] == "stopping" + # After stopping, a stop with nothing running reports idle. + st2 = client.post("/api/train/diffusion/stop") + assert st2.json()["status"] == "idle" + + +def test_service_restart_after_completion(): + # A finished job's pump is joined OUTSIDE the lock (it needs the lock for its + # final state writes), so a second start neither stalls nor deadlocks. + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) + svc.start(dict(_CFG)) + _wait_status(svc, "completed") + t0 = time.time() + job2 = svc.start(dict(_CFG)) + assert job2 + assert time.time() - t0 < 4.0 # no 5s join-under-lock stall + st = _wait_status(svc, "completed") + assert st["status"] == "completed" + + +def test_stale_pump_events_cannot_corrupt_new_job(): + # An event carrying a superseded job's proc identity must be dropped, so a + # straggler pump can never overwrite the state of a newly started job. + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) + svc.start(dict(_CFG)) + _wait_status(svc, "completed") + current = svc._proc + svc._apply_event({"type": "error", "message": "stale boom"}, proc = object()) + assert svc.status()["message"] != "stale boom" + # 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"] + + +def test_route_start_refuses_non_sdxl_base_without_freeing_gpu(client, monkeypatch): + # A doomed start (non-SDXL base) must 400 BEFORE resident GPU workloads are freed, + # so a bad pick never unloads the user's working chat/Images model. + import routes.training as tr + + freed = [] + monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: freed.append(1)) + r = client.post( + "/api/train/diffusion/start", json = {**_BODY, "base_model": "unsloth/FLUX.1-dev-GGUF"} + ) + assert r.status_code == 400 + assert "SDXL" in r.json()["detail"] + assert freed == [] + assert client._fake.started_with is None