a091a862df
Two evict/corruption fixes surfaced by review of the diffusion training path: - krea-2 sets force_bf16 in the DiT trainer spec, but the route-level _FORCE_BF16_FAMILIES preflight listed only qwen-image and z-image, so a krea-2 start with mixed_precision=fp16 passed the route check, reserved training and evicted resident GPU models, and only the child trainer then raised. Add krea-2 to the set and a drift-guard test asserting it equals the trainer specs whose force_bf16 is set. - The dataset-upload same-stem duplicate check compared stems case-sensitively, so on case-insensitive filesystems (Windows / default macOS) sample.png and Sample.jpg both passed even though their caption sidecars sample.txt / Sample.txt resolve to the same file, silently sharing and corrupting one caption. Compare stems and the same-name guard with casefold at both the on-disk and in-batch sites.
2267 lines
100 KiB
Python
2267 lines
100 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""
|
|
Training API routes
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile
|
|
from fastapi.responses import StreamingResponse
|
|
from typing import Dict, Optional, Any
|
|
import structlog
|
|
from loggers import get_logger
|
|
import asyncio
|
|
from datetime import datetime
|
|
import uuid as _uuid
|
|
|
|
# Add backend directory to path.
|
|
backend_path = Path(__file__).parent.parent.parent
|
|
if str(backend_path) not in sys.path:
|
|
sys.path.insert(0, str(backend_path))
|
|
|
|
try:
|
|
from core.training import get_training_backend
|
|
from core.training.resume import (
|
|
can_resume_run,
|
|
get_resume_checkpoint_path,
|
|
normalize_resume_output_dir,
|
|
)
|
|
from storage.studio_db import get_resumable_run_by_output_dir
|
|
from utils.models.model_config import load_model_defaults
|
|
from utils.paths import resolve_dataset_path
|
|
except ImportError:
|
|
# Fallback: parent directory.
|
|
parent_backend = backend_path.parent / "backend"
|
|
if str(parent_backend) not in sys.path:
|
|
sys.path.insert(0, str(parent_backend))
|
|
from core.training import get_training_backend
|
|
from core.training.resume import (
|
|
can_resume_run,
|
|
get_resume_checkpoint_path,
|
|
normalize_resume_output_dir,
|
|
)
|
|
from storage.studio_db import get_resumable_run_by_output_dir
|
|
from utils.models.model_config import load_model_defaults
|
|
from utils.paths import resolve_dataset_path
|
|
|
|
# Auth
|
|
from auth.authentication import authenticated_via_api_key, get_current_subject
|
|
|
|
from utils.utils import log_and_http_error
|
|
|
|
from models import (
|
|
TrainingStartRequest,
|
|
TrainingJobResponse,
|
|
TrainingStatus,
|
|
TrainingProgress,
|
|
)
|
|
from models.training import (
|
|
DiffusionCaptionUpdateRequest,
|
|
DiffusionDatasetExample,
|
|
DiffusionDatasetExamplesResponse,
|
|
DiffusionDatasetImageRecord,
|
|
DiffusionDatasetImagesResponse,
|
|
DiffusionDatasetImportRequest,
|
|
DiffusionDatasetImportResponse,
|
|
DiffusionDatasetSummary,
|
|
DiffusionDatasetUploadResponse,
|
|
DiffusionMetricHistory,
|
|
DiffusionTrainableFamily,
|
|
DiffusionTrainingInfoResponse,
|
|
DiffusionTrainingRunDetail,
|
|
DiffusionTrainingRunsResponse,
|
|
DiffusionTrainingRunSummary,
|
|
DiffusionTrainingStartRequest,
|
|
DiffusionTrainingStartResponse,
|
|
DiffusionTrainingStatusResponse,
|
|
DiffusionTrainingStopRequest,
|
|
)
|
|
from models.responses import TrainingStopResponse, TrainingMetricsResponse
|
|
from pydantic import BaseModel as PydanticBaseModel
|
|
from pydantic import ValidationError
|
|
|
|
|
|
class TrainingStopRequest(PydanticBaseModel):
|
|
save: bool = True
|
|
|
|
|
|
router = APIRouter()
|
|
logger = get_logger(__name__)
|
|
|
|
# Consecutive 1s polls without a step update that count as a stall. Applied only
|
|
# once stepping: the pre-first-step phase (model load + tokenization) can take far
|
|
# longer, and timing out there made a healthy long-prep run look frozen.
|
|
_PROGRESS_STALL_TIMEOUT_POLLS = 1800 # ~30 min at 1 poll/sec
|
|
|
|
|
|
def _validate_local_dataset_paths(paths: list[str], label: str = "Local dataset") -> list[str]:
|
|
"""Resolve and validate a list of local dataset paths. Returns validated absolute paths."""
|
|
validated = []
|
|
missing = []
|
|
for dataset_path in paths:
|
|
dataset_file = resolve_dataset_path(dataset_path)
|
|
if not dataset_file.exists():
|
|
missing.append(f"{dataset_path} (resolved: {dataset_file})")
|
|
continue
|
|
logger.info(f"Found {label.lower()} file: {dataset_file}")
|
|
validated.append(str(dataset_file))
|
|
|
|
if missing:
|
|
missing_detail = "; ".join(missing[:3])
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = f"{label} not found: {missing_detail}",
|
|
)
|
|
return validated
|
|
|
|
|
|
@router.get("/hardware")
|
|
async def get_hardware_utilization(current_subject: str = Depends(get_current_subject)):
|
|
"""
|
|
Live snapshot of GPU hardware utilization for the active backend.
|
|
|
|
Polled by the frontend during training.
|
|
"""
|
|
from utils.hardware import get_gpu_utilization
|
|
return get_gpu_utilization()
|
|
|
|
|
|
@router.get("/hardware/visible")
|
|
async def get_visible_hardware_utilization(current_subject: str = Depends(get_current_subject)):
|
|
from utils.hardware import get_visible_gpu_utilization
|
|
return get_visible_gpu_utilization()
|
|
|
|
|
|
@router.post("/start")
|
|
async def start_training(
|
|
request: TrainingStartRequest,
|
|
current_subject: str = Depends(get_current_subject),
|
|
via_api_key: bool = Depends(authenticated_via_api_key),
|
|
):
|
|
"""
|
|
Start a training job.
|
|
|
|
Initiates training in the background and returns immediately. Use /status
|
|
to check progress.
|
|
"""
|
|
try:
|
|
logger.info(f"Starting training job with model: {request.model_name}")
|
|
|
|
# When Studio is driven as an inference API (API-key auth), refuse to start
|
|
# training while a request is in flight: training frees VRAM by unloading
|
|
# the chat model, which would kill the stream. The Studio UI (session auth)
|
|
# still starts training and coexists/frees VRAM as before. (A mixed UI+API
|
|
# session is not yet special-cased.)
|
|
if via_api_key is True:
|
|
from core.inference.llama_keepwarm import other_inference_request_count
|
|
if other_inference_request_count(current_request_counted = False) > 0:
|
|
raise HTTPException(
|
|
status_code = 409,
|
|
detail = (
|
|
"Cannot start training over the API while an inference request is in "
|
|
"progress. Wait for it to finish, or start training from the Studio UI."
|
|
),
|
|
)
|
|
|
|
# No in-process ensure_transformers_version(): the subprocess
|
|
# (worker.py) activates the correct version before importing ML libs.
|
|
|
|
backend = get_training_backend()
|
|
|
|
# S3 dataset loading needs the optional boto3 dependency. Reject early
|
|
# with a clear message so credentials are never accepted and then
|
|
# silently dropped on a host without boto3 installed.
|
|
if request.s3_config is not None:
|
|
from core.training.s3_dataset import boto3_available
|
|
if not boto3_available():
|
|
raise HTTPException(
|
|
status_code = 501,
|
|
detail = "S3 dataset loading requires boto3. Install it with: pip install boto3",
|
|
)
|
|
|
|
# Check before mutating state.
|
|
if backend.is_training_active():
|
|
existing_job_id: Optional[str] = getattr(backend, "current_job_id", "")
|
|
return TrainingJobResponse(
|
|
job_id = existing_job_id or "",
|
|
status = "error",
|
|
message = (
|
|
"Training is already in progress. "
|
|
"Stop current training before starting a new one."
|
|
),
|
|
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]}"
|
|
|
|
# Validate dataset paths if provided.
|
|
if request.local_datasets:
|
|
request.local_datasets = _validate_local_dataset_paths(
|
|
request.local_datasets, "Local dataset"
|
|
)
|
|
if request.local_eval_datasets and request.eval_steps > 0:
|
|
request.local_eval_datasets = _validate_local_dataset_paths(
|
|
request.local_eval_datasets, "Local eval dataset"
|
|
)
|
|
resume_output_dir: Optional[str] = None
|
|
if request.resume_from_checkpoint:
|
|
try:
|
|
resume_output_dir = normalize_resume_output_dir(request.resume_from_checkpoint)
|
|
except ValueError as e:
|
|
# Deliberate user-facing validation message.
|
|
validation_message = str(e)
|
|
raise HTTPException(status_code = 400, detail = validation_message)
|
|
|
|
resume_run = get_resumable_run_by_output_dir(resume_output_dir)
|
|
if not resume_run or not can_resume_run(resume_run):
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = "Resume checkpoint must belong to a stopped run with saved trainer state.",
|
|
)
|
|
resume_checkpoint = get_resume_checkpoint_path(resume_output_dir)
|
|
if not resume_checkpoint:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = "Resume checkpoint must include saved trainer state.",
|
|
)
|
|
request.resume_from_checkpoint = resume_checkpoint
|
|
|
|
# Validate streaming-mode compatibility before any expensive work.
|
|
# Streaming is supported only for Hugging Face text datasets.
|
|
if request.dataset_streaming:
|
|
if not request.hf_dataset:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = "dataset_streaming requires hf_dataset; streaming is not supported for local datasets.",
|
|
)
|
|
if request.is_dataset_image or request.is_dataset_audio:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = "dataset_streaming is not supported for vision or audio datasets.",
|
|
)
|
|
if request.is_embedding:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = "dataset_streaming is not supported for embedding training; the embedding loader needs the full dataset.",
|
|
)
|
|
from utils.hardware import hardware as _hw
|
|
|
|
if _hw.DEVICE == _hw.DeviceType.MLX:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = "dataset_streaming is not yet supported on Apple Silicon (MLX); the MLX loader materializes the full dataset.",
|
|
)
|
|
if request.max_steps is None or request.max_steps <= 0:
|
|
raise HTTPException(
|
|
status_code = 422,
|
|
detail = "dataset_streaming requires max_steps > 0 because streaming datasets have no known length.",
|
|
)
|
|
if request.train_on_completions:
|
|
raise HTTPException(
|
|
status_code = 422,
|
|
detail = "dataset_streaming is not supported with train_on_completions yet.",
|
|
)
|
|
if request.eval_steps > 0:
|
|
train_split = request.train_split or "train"
|
|
if not request.eval_split or request.eval_split == train_split:
|
|
raise HTTPException(
|
|
status_code = 422,
|
|
detail = "dataset_streaming with evaluation requires a separate eval_split.",
|
|
)
|
|
# Streaming is HF-only: reject when the request also carries a local
|
|
# dataset path or an S3 config; those sources cannot be streamed via
|
|
# HF's streaming loader.
|
|
if request.local_datasets:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = (
|
|
"dataset_streaming is HF-only; remove local_datasets / S3 source. "
|
|
"Streaming is not supported with local file paths."
|
|
),
|
|
)
|
|
if request.s3_config is not None:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = (
|
|
"dataset_streaming is HF-only; remove local_datasets / S3 source. "
|
|
"Streaming is not supported with S3 datasets."
|
|
),
|
|
)
|
|
|
|
# Convert request to backend kwargs.
|
|
training_kwargs = {
|
|
"model_name": request.model_name,
|
|
"project_name": request.project_name,
|
|
"training_type": request.training_type,
|
|
"hf_token": request.hf_token or "",
|
|
"load_in_4bit": request.load_in_4bit,
|
|
"max_seq_length": request.max_seq_length,
|
|
"vision_image_size": request.vision_image_size,
|
|
"hf_dataset": request.hf_dataset or "",
|
|
"local_datasets": request.local_datasets,
|
|
"local_eval_datasets": request.local_eval_datasets,
|
|
"format_type": request.format_type,
|
|
"subset": request.subset,
|
|
"train_split": request.train_split,
|
|
"dataset_streaming": request.dataset_streaming,
|
|
"eval_split": request.eval_split,
|
|
"eval_steps": request.eval_steps,
|
|
"dataset_slice_start": request.dataset_slice_start,
|
|
"dataset_slice_end": request.dataset_slice_end,
|
|
"custom_format_mapping": request.custom_format_mapping,
|
|
"num_epochs": request.num_epochs,
|
|
"learning_rate": request.learning_rate,
|
|
"embedding_learning_rate": request.embedding_learning_rate,
|
|
"batch_size": request.batch_size,
|
|
"gradient_accumulation_steps": request.gradient_accumulation_steps,
|
|
"warmup_steps": request.warmup_steps,
|
|
"warmup_ratio": request.warmup_ratio,
|
|
"max_steps": request.max_steps,
|
|
"save_steps": request.save_steps,
|
|
"weight_decay": request.weight_decay,
|
|
"max_grad_norm": request.max_grad_norm,
|
|
"max_grad_value": request.max_grad_value,
|
|
"max_grad_leaf_norm": request.max_grad_leaf_norm,
|
|
"cast_norm_output_to_input_dtype": request.cast_norm_output_to_input_dtype,
|
|
"random_seed": request.random_seed,
|
|
"packing": request.packing,
|
|
"optim": request.optim,
|
|
"lr_scheduler_type": request.lr_scheduler_type,
|
|
"use_lora": request.use_lora,
|
|
"lora_r": request.lora_r,
|
|
"lora_alpha": request.lora_alpha,
|
|
"lora_dropout": request.lora_dropout,
|
|
"target_modules": request.target_modules if request.target_modules else None,
|
|
"gradient_checkpointing": request.gradient_checkpointing.strip()
|
|
if request.gradient_checkpointing and request.gradient_checkpointing.strip()
|
|
else "unsloth",
|
|
"use_rslora": request.use_rslora,
|
|
"use_loftq": request.use_loftq,
|
|
"train_on_completions": request.train_on_completions,
|
|
"finetune_vision_layers": request.finetune_vision_layers,
|
|
"finetune_language_layers": request.finetune_language_layers,
|
|
"finetune_attention_modules": request.finetune_attention_modules,
|
|
"finetune_mlp_modules": request.finetune_mlp_modules,
|
|
"is_dataset_image": request.is_dataset_image,
|
|
"is_dataset_audio": request.is_dataset_audio,
|
|
"is_embedding": request.is_embedding,
|
|
"enable_wandb": request.enable_wandb,
|
|
"wandb_token": request.wandb_token or "",
|
|
"wandb_project": request.wandb_project or "",
|
|
"enable_tensorboard": request.enable_tensorboard,
|
|
"tensorboard_dir": request.tensorboard_dir or "",
|
|
"output_dir": resume_output_dir,
|
|
"resume_from_checkpoint": request.resume_from_checkpoint,
|
|
"trust_remote_code": request.trust_remote_code,
|
|
"approved_remote_code_fingerprint": request.approved_remote_code_fingerprint,
|
|
"subject": current_subject,
|
|
"gpu_ids": request.gpu_ids,
|
|
"s3_config": request.s3_config.model_dump() if request.s3_config else None,
|
|
}
|
|
|
|
# Training page has no trust_remote_code toggle, so honor the YAML default
|
|
# -- but only for genuine first-party (unsloth/nvidia) Hub repos, never a
|
|
# local path or a name merely starting with "unsloth/".
|
|
if not training_kwargs["trust_remote_code"]:
|
|
from utils.security.trusted_org import is_trusted_org_repo
|
|
|
|
model_defaults = load_model_defaults(request.model_name)
|
|
yaml_trust = model_defaults.get("training", {}).get("trust_remote_code", False)
|
|
if yaml_trust and is_trusted_org_repo(
|
|
request.model_name, hf_token = request.hf_token or None
|
|
):
|
|
logger.info(f"YAML config sets trust_remote_code=True for {request.model_name}")
|
|
training_kwargs["trust_remote_code"] = True
|
|
elif yaml_trust:
|
|
logger.warning(
|
|
"YAML sets trust_remote_code=True for %s but it is not a trusted "
|
|
"first-party repo; leaving disabled (user can opt in explicitly).",
|
|
request.model_name,
|
|
)
|
|
|
|
# Free VRAM for training: stop export, unload chat unless it can coexist.
|
|
# A before_spawn hook -> runs only after start_training's guards pass, so
|
|
# we never tear down chat/export VRAM for a start that is then refused.
|
|
def _free_vram_for_training() -> None:
|
|
try:
|
|
from core.export import get_export_backend
|
|
exp_backend = get_export_backend()
|
|
# Tear down the export subprocess whenever an export is in flight,
|
|
# not just once a checkpoint is loaded: during the load phase
|
|
# current_checkpoint is still unset while the worker is already
|
|
# allocating GPU memory, so gate on is_export_active() too.
|
|
if exp_backend.current_checkpoint or exp_backend.is_export_active():
|
|
logger.info("Shutting down export subprocess to free GPU memory for training")
|
|
exp_backend._shutdown_subprocess()
|
|
exp_backend.current_checkpoint = None
|
|
exp_backend.is_vision = False
|
|
exp_backend.is_peft = False
|
|
except Exception as e:
|
|
logger.warning("Could not shut down export subprocess: %s", e)
|
|
|
|
try:
|
|
# A resident or in-flight diffusion (Images) pipeline also holds
|
|
# GPU memory the training run needs, and it can't be cheaply sized,
|
|
# so tear it down unconditionally like the export subprocess above
|
|
# (the chat block below fit-checks; diffusion can't). unload() is a
|
|
# no-op when nothing is loaded and also preempts an in-flight load;
|
|
# release the arbiter so it doesn't think the gone pipeline owns
|
|
# the GPU. Must precede the chat block, which early-returns.
|
|
from core.inference import gpu_arbiter
|
|
from core.inference.diffusion_engine_router import (
|
|
get_active_diffusion_engine,
|
|
)
|
|
|
|
# The ACTIVE engine, not the diffusers singleton: on a native
|
|
# (sd_cpp) selection the diffusers backend reports unloaded while
|
|
# the native engine still holds model state / a live generation.
|
|
diffusion = get_active_diffusion_engine()
|
|
if diffusion.is_loaded:
|
|
logger.info(
|
|
"Unloading diffusion (Images) model to free GPU memory for training"
|
|
)
|
|
diffusion.unload()
|
|
gpu_arbiter.release(gpu_arbiter.DIFFUSION)
|
|
except Exception as e:
|
|
logger.warning("Could not unload diffusion model for training: %s", e)
|
|
|
|
try:
|
|
# A resident or in-flight Video pipeline holds GPU memory the training run
|
|
# needs too, and it loads under the VIDEO arbiter owner the diffusion teardown
|
|
# above never touches. Tear it down the same way (unload is a no-op when nothing
|
|
# is loaded and preempts an in-flight load) and release VIDEO, so starting
|
|
# training while a generated-video session is resident can't OOM the run. Must
|
|
# precede the chat block, which early-returns.
|
|
from core.inference import gpu_arbiter
|
|
from core.inference.video import get_video_backend
|
|
|
|
video = get_video_backend()
|
|
if video.status().get("loaded"):
|
|
logger.info("Unloading Video model to free GPU memory for training")
|
|
video.unload()
|
|
gpu_arbiter.release(gpu_arbiter.VIDEO)
|
|
except Exception as e:
|
|
logger.warning("Could not unload video model for training: %s", e)
|
|
|
|
try:
|
|
from routes.training_vram import (
|
|
can_keep_chat_during_training,
|
|
free_chat_models_for_training,
|
|
summarize_resident_chat,
|
|
)
|
|
|
|
resident = summarize_resident_chat()
|
|
if not resident["any"]:
|
|
return
|
|
if resident.get("loading"):
|
|
# In-flight load can't be sized -> free rather than risk OOM.
|
|
freed = free_chat_models_for_training(reason = "chat model still loading")
|
|
logger.info("Freed in-flight chat load for training: %s", freed)
|
|
return
|
|
keep, info = can_keep_chat_during_training(
|
|
model_name = training_kwargs["model_name"],
|
|
hf_token = training_kwargs["hf_token"],
|
|
training_type = training_kwargs["training_type"],
|
|
load_in_4bit = training_kwargs["load_in_4bit"],
|
|
batch_size = training_kwargs["batch_size"],
|
|
max_seq_length = training_kwargs["max_seq_length"],
|
|
lora_rank = training_kwargs["lora_r"],
|
|
target_modules = training_kwargs["target_modules"],
|
|
gradient_checkpointing = training_kwargs["gradient_checkpointing"],
|
|
optimizer = training_kwargs["optim"],
|
|
gpu_ids = training_kwargs["gpu_ids"],
|
|
)
|
|
if keep:
|
|
logger.info(
|
|
"Keeping chat model(s) loaded during training "
|
|
"(free ~%s GB, needs ~%s GB): %s",
|
|
info.get("usable_gb"),
|
|
info.get("required_gb"),
|
|
resident,
|
|
)
|
|
else:
|
|
freed = free_chat_models_for_training(
|
|
reason = "insufficient VRAM to run training alongside chat",
|
|
)
|
|
logger.info("Freed chat model(s) for training: %s", freed)
|
|
except Exception as e:
|
|
logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e)
|
|
|
|
# The hook runs only once start guards pass -> VRAM freed iff training starts.
|
|
success = backend.start_training(
|
|
job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs
|
|
)
|
|
|
|
if not success:
|
|
progress_error = backend.trainer.training_progress.error
|
|
return TrainingJobResponse(
|
|
job_id = backend.current_job_id or "",
|
|
status = "error",
|
|
message = progress_error or "Failed to start training subprocess",
|
|
error = progress_error or "subprocess_start_failed",
|
|
)
|
|
|
|
return TrainingJobResponse(
|
|
job_id = job_id,
|
|
status = "queued",
|
|
message = "Training job queued and starting in subprocess",
|
|
error = None,
|
|
)
|
|
|
|
except HTTPException:
|
|
# Deliberate rejections (S3 not implemented, resume validation) must
|
|
# reach the client with their original status, not a generic 500.
|
|
raise
|
|
except ValueError as e:
|
|
logger.warning("Rejected training GPU selection: %s", e)
|
|
# Deliberate user-facing GPU-selection validation message.
|
|
validation_message = str(e)
|
|
raise HTTPException(status_code = 400, detail = validation_message)
|
|
except Exception as e:
|
|
raise log_and_http_error(
|
|
e,
|
|
500,
|
|
"Failed to start training",
|
|
event = "training.start_failed",
|
|
log = logger,
|
|
)
|
|
|
|
|
|
@router.post("/stop", response_model = TrainingStopResponse)
|
|
async def stop_training(
|
|
body: TrainingStopRequest = TrainingStopRequest(),
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""
|
|
Stop the currently running training job.
|
|
|
|
Body:
|
|
save (bool): If True (default), save the model at the current checkpoint.
|
|
"""
|
|
try:
|
|
backend = get_training_backend()
|
|
is_active = backend.is_training_active()
|
|
logger.info("Stop requested: save=%s is_active=%s", body.save, is_active)
|
|
|
|
if not is_active:
|
|
return TrainingStopResponse(
|
|
status = "idle", message = "No training job is currently running"
|
|
)
|
|
|
|
backend.stop_training(save = body.save)
|
|
|
|
return TrainingStopResponse(
|
|
status = "stopped",
|
|
message = "Stop requested. Training will stop at the next safe step.",
|
|
)
|
|
|
|
except Exception as e:
|
|
raise log_and_http_error(
|
|
e,
|
|
500,
|
|
"Failed to stop training",
|
|
event = "training.stop_failed",
|
|
log = logger,
|
|
)
|
|
|
|
|
|
@router.post("/reset")
|
|
async def reset_training(current_subject: str = Depends(get_current_subject)):
|
|
"""Reset training state so the user can return to configuration."""
|
|
try:
|
|
backend = get_training_backend()
|
|
is_active = backend.is_training_active()
|
|
|
|
if is_active:
|
|
if backend._cancel_requested:
|
|
# Cancel (save=False) requested — force-terminate to reset immediately.
|
|
logger.info("Force-terminating subprocess for immediate reset (cancel path)")
|
|
backend.force_terminate()
|
|
else:
|
|
logger.warning("Rejected reset while training active: is_active=%s", is_active)
|
|
raise HTTPException(
|
|
status_code = 409,
|
|
detail = "Training is still running. Stop training and wait for it to finish before resetting.",
|
|
)
|
|
|
|
logger.info("Reset training state: clearing runtime + metric history")
|
|
backend._should_stop = False # Clear stop flag so status returns to idle
|
|
backend.trainer._update_progress(
|
|
is_training = False,
|
|
is_completed = False,
|
|
error = None,
|
|
status_message = "Ready to train",
|
|
step = 0,
|
|
loss = None,
|
|
epoch = 0,
|
|
total_steps = 0,
|
|
)
|
|
backend.loss_history = []
|
|
backend.lr_history = []
|
|
backend.step_history = []
|
|
backend.grad_norm_history = []
|
|
backend.grad_norm_step_history = []
|
|
return {"status": "ok"}
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise log_and_http_error(
|
|
e,
|
|
500,
|
|
"Failed to reset training",
|
|
event = "training.reset_failed",
|
|
log = logger,
|
|
)
|
|
|
|
|
|
@router.get("/status")
|
|
async def get_training_status(current_subject: str = Depends(get_current_subject)):
|
|
"""
|
|
Get the current training status.
|
|
"""
|
|
try:
|
|
backend = get_training_backend()
|
|
job_id: str = getattr(backend, "current_job_id", "") or ""
|
|
|
|
is_active = backend.is_training_active()
|
|
|
|
try:
|
|
progress = backend.trainer.get_training_progress()
|
|
except Exception:
|
|
progress = None
|
|
|
|
status_message = (
|
|
getattr(progress, "status_message", None) if progress else None
|
|
) or "Ready to train"
|
|
error_message = getattr(progress, "error", None) if progress else None
|
|
|
|
trainer_stopped = getattr(backend, "_should_stop", False)
|
|
|
|
# Derive high-level phase
|
|
if error_message:
|
|
phase = "error"
|
|
elif is_active:
|
|
msg_lower = status_message.lower()
|
|
if "loading" in msg_lower or "importing" in msg_lower:
|
|
phase = "loading_model"
|
|
elif any(k in msg_lower for k in ["preparing", "initializing", "configuring"]):
|
|
phase = "configuring"
|
|
else:
|
|
phase = "training"
|
|
elif trainer_stopped:
|
|
phase = "stopped"
|
|
elif progress and getattr(progress, "is_completed", False):
|
|
phase = "completed"
|
|
else:
|
|
phase = "idle"
|
|
|
|
details = None
|
|
if progress:
|
|
details = {
|
|
"epoch": getattr(progress, "epoch", 0),
|
|
"step": getattr(progress, "step", 0),
|
|
"total_steps": getattr(progress, "total_steps", 0),
|
|
"loss": getattr(progress, "loss", None),
|
|
"learning_rate": getattr(progress, "learning_rate", None),
|
|
}
|
|
output_dir = getattr(backend, "_output_dir", None)
|
|
if output_dir:
|
|
details["output_dir"] = output_dir
|
|
|
|
# Metric history for chart recovery after SSE reconnection.
|
|
metric_history = None
|
|
if backend.step_history:
|
|
metric_history = {
|
|
"steps": list(backend.step_history),
|
|
"loss": list(backend.loss_history),
|
|
"lr": list(backend.lr_history),
|
|
"grad_norm": list(getattr(backend, "grad_norm_history", [])),
|
|
"grad_norm_steps": list(getattr(backend, "grad_norm_step_history", [])),
|
|
"eval_loss": list(backend.eval_loss_history),
|
|
"eval_steps": list(backend.eval_step_history),
|
|
}
|
|
|
|
return TrainingStatus(
|
|
job_id = job_id,
|
|
phase = phase,
|
|
is_training_running = is_active,
|
|
eval_enabled = backend.eval_enabled,
|
|
message = status_message,
|
|
error = error_message,
|
|
details = details,
|
|
metric_history = metric_history,
|
|
)
|
|
|
|
except Exception as e:
|
|
raise log_and_http_error(
|
|
e,
|
|
500,
|
|
"Failed to get training status",
|
|
event = "training.status_failed",
|
|
log = logger,
|
|
)
|
|
|
|
|
|
@router.get("/metrics", response_model = TrainingMetricsResponse)
|
|
async def get_training_metrics(current_subject: str = Depends(get_current_subject)):
|
|
"""
|
|
Get training metrics (loss, learning rate, steps).
|
|
"""
|
|
try:
|
|
backend = get_training_backend()
|
|
|
|
loss_history = backend.loss_history
|
|
lr_history = backend.lr_history
|
|
step_history = backend.step_history
|
|
grad_norm_history = getattr(backend, "grad_norm_history", [])
|
|
grad_norm_step_history = getattr(backend, "grad_norm_step_history", [])
|
|
|
|
current_loss = loss_history[-1] if loss_history else None
|
|
current_lr = lr_history[-1] if lr_history else None
|
|
current_step = step_history[-1] if step_history else None
|
|
|
|
return TrainingMetricsResponse(
|
|
loss_history = loss_history,
|
|
lr_history = lr_history,
|
|
step_history = step_history,
|
|
grad_norm_history = grad_norm_history,
|
|
grad_norm_step_history = grad_norm_step_history,
|
|
current_loss = current_loss,
|
|
current_lr = current_lr,
|
|
current_step = current_step,
|
|
)
|
|
|
|
except Exception as e:
|
|
raise log_and_http_error(
|
|
e,
|
|
500,
|
|
"Failed to get training metrics",
|
|
event = "training.metrics_failed",
|
|
log = logger,
|
|
)
|
|
|
|
|
|
@router.get("/progress")
|
|
async def stream_training_progress(
|
|
request: Request, current_subject: str = Depends(get_current_subject)
|
|
):
|
|
"""
|
|
Stream training progress via Server-Sent Events (SSE).
|
|
|
|
Real-time progress with reconnection support per the SSE spec:
|
|
- `id:` per event so the browser tracks position.
|
|
- `retry:` to control reconnection interval.
|
|
- Named `event:` types (progress, heartbeat, complete, error).
|
|
- Reads `Last-Event-ID` on reconnect to replay missed steps.
|
|
"""
|
|
# Read Last-Event-ID header for reconnection resume.
|
|
last_event_id = request.headers.get("last-event-id")
|
|
resume_from_step: Optional[int] = None
|
|
if last_event_id is not None:
|
|
try:
|
|
resume_from_step = int(last_event_id)
|
|
logger.info(f"SSE reconnect: resuming from step {resume_from_step}")
|
|
except ValueError:
|
|
logger.warning(f"Invalid Last-Event-ID: {last_event_id}")
|
|
|
|
async def event_generator():
|
|
backend = get_training_backend()
|
|
job_id: str = getattr(backend, "current_job_id", "") or ""
|
|
|
|
# ── Helpers ──────────────────────────────────────────────
|
|
def build_progress(
|
|
step: int,
|
|
loss: Optional[float],
|
|
learning_rate: Optional[float],
|
|
total_steps: int,
|
|
epoch: Optional[float] = None,
|
|
progress: Optional[Any] = None,
|
|
grad_norm_override: Optional[float] = None,
|
|
eval_loss_override: Optional[float] = None,
|
|
) -> TrainingProgress:
|
|
total = max(total_steps, 0)
|
|
if step < 0 or total == 0:
|
|
progress_percent = 0.0
|
|
else:
|
|
progress_percent = float(step) / float(total) * 100.0 if total > 0 else 0.0
|
|
|
|
# Pull values from the progress object if available.
|
|
elapsed_seconds = getattr(progress, "elapsed_seconds", None) if progress else None
|
|
eta_seconds = getattr(progress, "eta_seconds", None) if progress else None
|
|
grad_norm = grad_norm_override
|
|
if grad_norm is None and progress:
|
|
grad_norm = getattr(progress, "grad_norm", None)
|
|
num_tokens = getattr(progress, "num_tokens", None) if progress else None
|
|
eval_loss = eval_loss_override
|
|
if eval_loss is None and progress:
|
|
eval_loss = getattr(progress, "eval_loss", None)
|
|
|
|
return TrainingProgress(
|
|
job_id = job_id,
|
|
step = step,
|
|
total_steps = total,
|
|
loss = loss,
|
|
learning_rate = learning_rate,
|
|
progress_percent = progress_percent,
|
|
epoch = epoch,
|
|
elapsed_seconds = elapsed_seconds,
|
|
eta_seconds = eta_seconds,
|
|
grad_norm = grad_norm,
|
|
num_tokens = num_tokens,
|
|
eval_loss = eval_loss,
|
|
)
|
|
|
|
def format_sse(
|
|
data: str,
|
|
event: str = "progress",
|
|
event_id: Optional[int] = None,
|
|
) -> str:
|
|
"""Format a single SSE message with id/event/data fields."""
|
|
lines = []
|
|
if event_id is not None:
|
|
lines.append(f"id: {event_id}")
|
|
lines.append(f"event: {event}")
|
|
lines.append(f"data: {data}")
|
|
lines.append("") # trailing blank line
|
|
lines.append("") # double newline terminates the event
|
|
return "\n".join(lines)
|
|
|
|
# ── Retry directive ──────────────────────────────────────
|
|
# Reconnect after 3 seconds if the connection drops.
|
|
yield "retry: 3000\n\n"
|
|
|
|
# ── Replay missed steps on reconnect ─────────────────────
|
|
if resume_from_step is not None and backend.step_history:
|
|
replayed = 0
|
|
grad_norm_by_step = {
|
|
step_val: grad_val
|
|
for step_val, grad_val in zip(
|
|
getattr(backend, "grad_norm_step_history", []),
|
|
getattr(backend, "grad_norm_history", []),
|
|
)
|
|
}
|
|
for i, step_val in enumerate(backend.step_history):
|
|
if step_val > resume_from_step:
|
|
loss_val = backend.loss_history[i] if i < len(backend.loss_history) else None
|
|
lr_val = backend.lr_history[i] if i < len(backend.lr_history) else None
|
|
tp_replay = getattr(
|
|
getattr(backend, "trainer", None), "training_progress", None
|
|
)
|
|
total_replay = (
|
|
getattr(tp_replay, "total_steps", step_val) if tp_replay else step_val
|
|
)
|
|
epoch_replay = getattr(tp_replay, "epoch", None) if tp_replay else None
|
|
payload = build_progress(
|
|
step_val,
|
|
loss_val,
|
|
lr_val,
|
|
total_replay,
|
|
epoch_replay,
|
|
progress = tp_replay,
|
|
grad_norm_override = grad_norm_by_step.get(step_val),
|
|
)
|
|
yield format_sse(payload.model_dump_json(), event = "progress", event_id = step_val)
|
|
replayed += 1
|
|
if replayed:
|
|
logger.info(f"SSE reconnect: replayed {replayed} missed steps")
|
|
|
|
# ── Initial status (only on fresh connections) ───────────
|
|
if resume_from_step is None:
|
|
is_active = backend.is_training_active()
|
|
tp = getattr(getattr(backend, "trainer", None), "training_progress", None)
|
|
initial_total_steps = getattr(tp, "total_steps", 0) if tp else 0
|
|
initial_epoch = getattr(tp, "epoch", None) if tp else None
|
|
|
|
initial_progress = build_progress(
|
|
step = 0,
|
|
loss = None,
|
|
learning_rate = None,
|
|
total_steps = initial_total_steps,
|
|
epoch = initial_epoch,
|
|
progress = tp,
|
|
)
|
|
yield format_sse(initial_progress.model_dump_json(), event = "progress", event_id = 0)
|
|
|
|
# If not active, send final state and exit
|
|
if not is_active:
|
|
_live = (getattr(tp, "step", 0) or 0) if tp else 0
|
|
if backend.step_history or _live > 0:
|
|
final_step = backend.step_history[-1] if backend.step_history else 0
|
|
final_loss = backend.loss_history[-1] if backend.loss_history else None
|
|
final_lr = backend.lr_history[-1] if backend.lr_history else None
|
|
# Histories skip non-finite steps; report the live step with
|
|
# loss=None instead of the last finite pair.
|
|
if _live > final_step:
|
|
final_step = _live
|
|
final_loss = getattr(tp, "loss", None)
|
|
final_lr = getattr(tp, "learning_rate", final_lr)
|
|
final_total_steps = getattr(tp, "total_steps", final_step) if tp else final_step
|
|
final_epoch = getattr(tp, "epoch", None) if tp else None
|
|
payload = build_progress(
|
|
final_step,
|
|
final_loss,
|
|
final_lr,
|
|
final_total_steps,
|
|
final_epoch,
|
|
progress = tp,
|
|
)
|
|
yield format_sse(
|
|
payload.model_dump_json(), event = "complete", event_id = final_step
|
|
)
|
|
else:
|
|
yield format_sse(
|
|
build_progress(-1, None, None, 0, progress = tp).model_dump_json(),
|
|
event = "complete",
|
|
event_id = 0,
|
|
)
|
|
return
|
|
|
|
# ── Live polling loop ────────────────────────────────────
|
|
last_step = resume_from_step if resume_from_step is not None else -1
|
|
no_update_count = 0
|
|
# The stall timeout applies only once the run is stepping (pre-step prep
|
|
# may legitimately emit no step for a long time). On reconnect to an
|
|
# already-stepping run, seed from the resume point / history, else a worker
|
|
# that hangs after step N never times out for a client that reconnects past it.
|
|
seen_live_step = (resume_from_step is not None and resume_from_step > 0) or bool(
|
|
backend.step_history
|
|
)
|
|
|
|
while backend.is_training_active():
|
|
# Client gone: end the generator without falling through to the final
|
|
# "complete" frame, which a buffered/proxy consumer could otherwise read
|
|
# as a finished run while training is still active.
|
|
if await request.is_disconnected():
|
|
return
|
|
try:
|
|
tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None)
|
|
live_step = (getattr(tp_inner, "step", 0) or 0) if tp_inner else 0
|
|
if backend.step_history or live_step > 0:
|
|
current_step = backend.step_history[-1] if backend.step_history else 0
|
|
current_loss = backend.loss_history[-1] if backend.loss_history else None
|
|
current_lr = backend.lr_history[-1] if backend.lr_history else None
|
|
# Histories skip non-finite steps; follow the live progress
|
|
# step and report its loss (None until it recovers).
|
|
if live_step > current_step:
|
|
current_step = live_step
|
|
current_loss = getattr(tp_inner, "loss", None)
|
|
current_lr = getattr(tp_inner, "learning_rate", current_lr)
|
|
current_total_steps = (
|
|
getattr(tp_inner, "total_steps", current_step) if tp_inner else current_step
|
|
)
|
|
current_epoch = getattr(tp_inner, "epoch", None) if tp_inner else None
|
|
|
|
# Only send if the step changed.
|
|
if current_step != last_step:
|
|
progress_payload = build_progress(
|
|
current_step,
|
|
current_loss,
|
|
current_lr,
|
|
current_total_steps,
|
|
current_epoch,
|
|
progress = tp_inner,
|
|
)
|
|
yield format_sse(
|
|
progress_payload.model_dump_json(),
|
|
event = "progress",
|
|
event_id = current_step,
|
|
)
|
|
last_step = current_step
|
|
no_update_count = 0
|
|
seen_live_step = True
|
|
else:
|
|
no_update_count += 1
|
|
# Heartbeat every 10 seconds.
|
|
if no_update_count % 10 == 0:
|
|
heartbeat_payload = build_progress(
|
|
current_step,
|
|
current_loss,
|
|
current_lr,
|
|
current_total_steps,
|
|
current_epoch,
|
|
progress = tp_inner,
|
|
)
|
|
yield format_sse(
|
|
heartbeat_payload.model_dump_json(),
|
|
event = "heartbeat",
|
|
event_id = current_step,
|
|
)
|
|
else:
|
|
# No steps yet, but training is active (model loading, etc.).
|
|
no_update_count += 1
|
|
if no_update_count % 5 == 0:
|
|
# Pull total_steps + status so the frontend can show
|
|
# "Tokenizing…" etc.
|
|
tp_prep = getattr(
|
|
getattr(backend, "trainer", None),
|
|
"training_progress",
|
|
None,
|
|
)
|
|
prep_total = getattr(tp_prep, "total_steps", 0) if tp_prep else 0
|
|
preparing_payload = build_progress(
|
|
0,
|
|
None,
|
|
None,
|
|
prep_total,
|
|
progress = tp_prep,
|
|
)
|
|
yield format_sse(
|
|
preparing_payload.model_dump_json(),
|
|
event = "heartbeat",
|
|
event_id = 0,
|
|
)
|
|
|
|
# Fires only once stepping: a long pre-first-step prep phase is not
|
|
# a stall, and ending the stream there made a healthy run look frozen.
|
|
if seen_live_step and no_update_count > _PROGRESS_STALL_TIMEOUT_POLLS:
|
|
logger.warning("Progress stream timeout - no updates received")
|
|
tp_timeout = getattr(
|
|
getattr(backend, "trainer", None), "training_progress", None
|
|
)
|
|
timeout_payload = build_progress(last_step, None, None, 0, progress = tp_timeout)
|
|
yield format_sse(
|
|
timeout_payload.model_dump_json(),
|
|
event = "error",
|
|
event_id = last_step if last_step >= 0 else 0,
|
|
)
|
|
break
|
|
|
|
await asyncio.sleep(1) # Poll every second
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in progress stream: {e}", exc_info = True)
|
|
tp_error = getattr(getattr(backend, "trainer", None), "training_progress", None)
|
|
error_payload = build_progress(0, None, None, 0, progress = tp_error)
|
|
yield format_sse(
|
|
error_payload.model_dump_json(),
|
|
event = "error",
|
|
event_id = last_step if last_step >= 0 else 0,
|
|
)
|
|
break
|
|
|
|
# ── Final "complete" event ───────────────────────────────
|
|
final_step = backend.step_history[-1] if backend.step_history else last_step
|
|
final_loss = backend.loss_history[-1] if backend.loss_history else None
|
|
final_lr = backend.lr_history[-1] if backend.lr_history else None
|
|
final_tp = getattr(getattr(backend, "trainer", None), "training_progress", None)
|
|
# If the run ended on a non-finite stretch, report the live step with
|
|
# loss=None instead of rolling back to the last finite pair.
|
|
_final_live_step = (getattr(final_tp, "step", 0) or 0) if final_tp else 0
|
|
if _final_live_step > (final_step if final_step is not None else -1):
|
|
final_step = _final_live_step
|
|
final_loss = getattr(final_tp, "loss", None)
|
|
final_lr = getattr(final_tp, "learning_rate", final_lr)
|
|
final_total_steps = getattr(final_tp, "total_steps", final_step) if final_tp else final_step
|
|
final_epoch = getattr(final_tp, "epoch", None) if final_tp else None
|
|
final_payload = build_progress(
|
|
final_step,
|
|
final_loss,
|
|
final_lr,
|
|
final_total_steps,
|
|
final_epoch,
|
|
progress = final_tp,
|
|
)
|
|
yield format_sse(
|
|
final_payload.model_dump_json(),
|
|
event = "complete",
|
|
event_id = final_step if final_step >= 0 else 0,
|
|
)
|
|
|
|
return StreamingResponse(
|
|
event_generator(),
|
|
media_type = "text/event-stream",
|
|
headers = {
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"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_engine_router import get_active_diffusion_engine
|
|
|
|
# The ACTIVE engine, not the diffusers singleton: on a native (sd_cpp)
|
|
# selection the diffusers backend reports unloaded while the resident
|
|
# sd-server still holds the GPU, so unloading only the singleton is a no-op.
|
|
# Mirrors the LLM training start path.
|
|
diffusion = get_active_diffusion_engine()
|
|
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:
|
|
# A resident Video pipeline loads under the VIDEO arbiter owner, which the Images
|
|
# teardown above does not free; unload it too (no-op when nothing is loaded) and release
|
|
# VIDEO so a generated-video session left resident can't OOM the diffusion trainer.
|
|
from core.inference import gpu_arbiter
|
|
from core.inference.video import get_video_backend
|
|
|
|
video = get_video_backend()
|
|
if video.status().get("loaded"):
|
|
logger.info("Unloading resident Video pipeline to free GPU memory for training")
|
|
video.unload() # no-op when nothing is loaded; also preempts an in-flight load
|
|
gpu_arbiter.release(gpu_arbiter.VIDEO)
|
|
except Exception as e: # noqa: BLE001
|
|
logger.warning("Could not unload Video 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)
|
|
|
|
|
|
def _preflight_gated_base(base_model: str, hf_token: Optional[str]) -> None:
|
|
"""HEAD a remote base repo's model_index.json with the caller's token; raise HTTP 400 on
|
|
401/403 (gated / unauthorized) with an actionable message. Best-effort: a local path,
|
|
a non-repo string, or a network hiccup passes through so the trainer can surface any real
|
|
load error itself. Runs before GPU teardown so a doomed start never evicts a loaded model."""
|
|
import urllib.error
|
|
import urllib.request
|
|
|
|
repo = (base_model or "").strip()
|
|
# Only remote 'org/name' repos are gated; skip local paths and single-file names.
|
|
if (
|
|
not repo
|
|
or repo.count("/") != 1
|
|
or repo.startswith((".", "/", "~"))
|
|
or repo.endswith(".gguf")
|
|
):
|
|
return
|
|
url = f"https://huggingface.co/{repo}/resolve/main/model_index.json"
|
|
headers = {"Authorization": f"Bearer {hf_token}"} if hf_token else {}
|
|
req = urllib.request.Request(url, method = "HEAD", headers = headers)
|
|
try:
|
|
urllib.request.urlopen(req, timeout = 5)
|
|
except urllib.error.HTTPError as e:
|
|
if e.code in (401, 403):
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = (
|
|
f"Access to '{repo}' is gated or unauthorized. Accept the model's license "
|
|
f"on its Hugging Face page and add your HF token in Studio settings, then "
|
|
f"try again."
|
|
),
|
|
)
|
|
# 404 (e.g. a repo without a root model_index.json) and other codes are not an
|
|
# access problem -- let the trainer surface any genuine load error.
|
|
except Exception: # noqa: BLE001 -- network/DNS hiccup must not block a start
|
|
return
|
|
|
|
|
|
@router.post("/diffusion/start", response_model = DiffusionTrainingStartResponse)
|
|
async def start_diffusion_training(
|
|
body: DiffusionTrainingStartRequest,
|
|
current_subject: str = Depends(get_current_subject),
|
|
via_api_key: bool = Depends(authenticated_via_api_key),
|
|
):
|
|
"""Start an SDXL LoRA training job from an image + caption dataset."""
|
|
from core.training.diffusion_training_service import get_diffusion_training_service
|
|
|
|
# When Studio is driven as an inference API (API-key auth), refuse to start training
|
|
# while a request is in flight: _free_gpu_for_diffusion_training() below unloads the
|
|
# chat backends to reclaim VRAM, which would kill the stream. Mirrors start_training so
|
|
# a diffusion start cannot silently drop an active API inference request.
|
|
if via_api_key is True:
|
|
from core.inference.llama_keepwarm import other_inference_request_count
|
|
if other_inference_request_count(current_request_counted = False) > 0:
|
|
raise HTTPException(
|
|
status_code = 409,
|
|
detail = (
|
|
"Cannot start diffusion (Images) training over the API while an inference "
|
|
"request is in progress. Wait for it to finish, or start training from the "
|
|
"Studio UI."
|
|
),
|
|
)
|
|
|
|
# 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:
|
|
normalized_cfg = _config_from_dict(config).normalized()
|
|
except ValueError as e:
|
|
raise HTTPException(status_code = 400, detail = str(e))
|
|
|
|
# Preflight the requested DiT precision BEFORE freeing GPU residents: the DiT trainer's own
|
|
# checks (a bf16-capable GPU is required; an explicit int8 needs a functional torchao) fire
|
|
# only in the child, AFTER _free_gpu_for_diffusion_training() already evicted the user's
|
|
# chat/Images model. Fail fast (400) so a pre-Ampere GPU (T4 / V100 / RTX 20xx) or a
|
|
# stub-torchao host never tears down resident models for a run that cannot start.
|
|
from core.training.diffusion_train_common import training_precision_preflight_error
|
|
|
|
_precision_reason = training_precision_preflight_error(
|
|
normalized_cfg.resolved_family, normalized_cfg.base_precision
|
|
)
|
|
if _precision_reason:
|
|
raise HTTPException(status_code = 400, detail = _precision_reason)
|
|
|
|
# Run the trainers' trust gate here too (both assert the same predicate before
|
|
# from_pretrained), so an untrusted/typoed base 400s BEFORE freeing GPU residents
|
|
# instead of tearing down the user's chat/Images model and failing in the child.
|
|
from core.training.diffusion_train_common import _assert_trusted_base_model
|
|
|
|
try:
|
|
_assert_trusted_base_model(config.get("base_model", ""))
|
|
except ValueError as e:
|
|
raise HTTPException(status_code = 400, detail = str(e))
|
|
|
|
# Preflight access to a gated base repo with the user's token BEFORE freeing GPU
|
|
# residents, so a missing/insufficient token fails fast (400) without tearing down the
|
|
# user's loaded chat/Images model, and never surfaces as a confusing mid-load 401.
|
|
# Offloaded to a worker thread: it does a blocking urlopen HEAD (up to a 5s timeout) to
|
|
# Hugging Face, which would otherwise stall the event loop and every concurrent
|
|
# status/progress/cancel request, as the filesystem preflight just below already does.
|
|
await asyncio.to_thread(
|
|
_preflight_gated_base, config.get("base_model", ""), config.get("hf_token")
|
|
)
|
|
|
|
# Preflight the dataset too: a missing/empty/uncaptionable data_dir otherwise
|
|
# fails inside the spawned trainer AFTER the user's chat/Images model was
|
|
# evicted. Same discovery the trainer runs, so the two cannot disagree.
|
|
from core.training import diffusion_train_common as _dtc
|
|
|
|
try:
|
|
await asyncio.to_thread(
|
|
_dtc.discover_image_caption_pairs,
|
|
config["data_dir"],
|
|
instance_prompt = config.get("instance_prompt") or None,
|
|
caption_column = config.get("caption_column") or "text",
|
|
# Decode-probe every image now (cheap PIL header check) so a corrupt / zero-byte
|
|
# upload is rejected with a 400 BEFORE _free_gpu_for_diffusion_training() tears down
|
|
# the user's resident models, instead of crashing the spawned trainer post-eviction.
|
|
verify_images = True,
|
|
)
|
|
except (FileNotFoundError, ValueError) as e:
|
|
raise HTTPException(status_code = 400, detail = str(e))
|
|
|
|
service = get_diffusion_training_service()
|
|
# Reserve the training slot BEFORE freeing residents: is_active() otherwise flips true only
|
|
# at service.start(), after the free below, so a concurrent /images/load or /video/load would
|
|
# pass its training guard during the free-then-spawn window, acquire the GPU, and double-
|
|
# allocate VRAM against the trainer. reserve() is a compare-and-set: a second overlapping
|
|
# /diffusion/start raises RuntimeError (-> 409) here, before it frees anything, so two starts
|
|
# never both tear down residents and race to start(). unreserve() runs in the finally ONLY
|
|
# when THIS request acquired the reservation, so a rejected second request never clears the
|
|
# first request's claim.
|
|
reserved = False
|
|
try:
|
|
service.reserve()
|
|
reserved = True
|
|
# Free resident GPU workloads (export / Images pipeline / chat) before the trainer
|
|
# loads its own pipeline. Offload the blocking teardown (engine unload waits on the
|
|
# generation locks; the export subprocess join can take seconds) to a worker thread so
|
|
# the event loop stays free for concurrent status/progress/cancel requests, as the
|
|
# inference routes do for their blocking load/unload calls.
|
|
await asyncio.to_thread(_free_gpu_for_diffusion_training)
|
|
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 (or a start is already reserved).
|
|
raise HTTPException(status_code = 409, detail = str(e))
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise log_and_http_error(
|
|
e,
|
|
500,
|
|
"Failed to start diffusion training",
|
|
event = "diffusion_training.start_failed",
|
|
log = logger,
|
|
)
|
|
finally:
|
|
# On success the now-live proc keeps is_active() true; on any failure this clears the
|
|
# reservation so training is not left permanently "active". Only the request that actually
|
|
# reserved clears it, so a rejected overlapping start does not drop the winner's claim.
|
|
if reserved:
|
|
service.unreserve()
|
|
return DiffusionTrainingStartResponse(job_id = job_id, status = "running")
|
|
|
|
|
|
@router.post("/diffusion/stop")
|
|
async def stop_diffusion_training(
|
|
body: Optional[DiffusionTrainingStopRequest] = None,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""Request a clean stop of the running diffusion training job. The optional body's
|
|
``save`` mirrors the LLM /stop: true (default, also for an empty POST) exports the
|
|
partial adapter, false cancels without saving one."""
|
|
from core.training.diffusion_training_service import get_diffusion_training_service
|
|
|
|
save = body.save if body is not None else True
|
|
stopped = get_diffusion_training_service().stop(save = save)
|
|
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
|
|
|
|
snap = get_diffusion_training_service().status()
|
|
# Fold the service's flat history arrays into the nested metric_history the UI charts.
|
|
metric_history = DiffusionMetricHistory(
|
|
steps = snap.pop("metric_steps", []),
|
|
loss = snap.pop("metric_loss", []),
|
|
lr = snap.pop("metric_lr", []),
|
|
grad_norm = snap.pop("metric_grad_norm", []),
|
|
)
|
|
return DiffusionTrainingStatusResponse(**snap, metric_history = metric_history)
|
|
|
|
|
|
@router.get("/diffusion/runs", response_model = DiffusionTrainingRunsResponse)
|
|
async def list_diffusion_training_runs(
|
|
limit: int = 20, current_subject: str = Depends(get_current_subject)
|
|
):
|
|
"""Previous diffusion training runs (terminal), newest first, from the persisted
|
|
per-run records. Summaries only; fetch one run for its config + metric logs."""
|
|
from core.training.diffusion_training_service import list_diffusion_runs
|
|
|
|
summaries: list[DiffusionTrainingRunSummary] = []
|
|
for r in list_diffusion_runs(limit = limit):
|
|
# list_diffusion_runs already skips non-dict / missing-id records, but a record with
|
|
# a wrong-typed field (e.g. a non-numeric avg_loss) would still raise here; catch it
|
|
# per record so one bad file never breaks the whole Previous runs panel.
|
|
try:
|
|
summaries.append(DiffusionTrainingRunSummary(**r))
|
|
except ValidationError:
|
|
continue
|
|
return DiffusionTrainingRunsResponse(runs = summaries)
|
|
|
|
|
|
@router.get("/diffusion/runs/{job_id}", response_model = DiffusionTrainingRunDetail)
|
|
async def get_diffusion_training_run(
|
|
job_id: str, current_subject: str = Depends(get_current_subject)
|
|
):
|
|
"""One persisted diffusion run's full record: summary + scrubbed start config + the
|
|
step/loss/grad-norm logs (for re-plotting a past run's charts)."""
|
|
from core.training.diffusion_training_service import get_diffusion_run
|
|
|
|
rec = get_diffusion_run(job_id)
|
|
# A valid-JSON file that is not an object (a truncated / hand-edited [] record) would make
|
|
# DiffusionTrainingRunDetail(**rec) raise TypeError -- not the ValidationError caught below
|
|
# -- and 500 the endpoint. Treat any non-dict record as absent, matching the list route's
|
|
# shape check.
|
|
if not isinstance(rec, dict):
|
|
raise HTTPException(status_code = 404, detail = "No such training run.")
|
|
try:
|
|
return DiffusionTrainingRunDetail(**rec)
|
|
except ValidationError:
|
|
# A malformed on-disk record (hand-edited / older shape) should read as absent
|
|
# rather than 500 the endpoint, mirroring how the list route skips bad records.
|
|
raise HTTPException(status_code = 404, detail = "No such training run.")
|
|
|
|
|
|
# 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 _resolve_dataset_caption(
|
|
folder: Path, image_path: Path, meta_captions: dict[str, str]
|
|
) -> Optional[str]:
|
|
"""Resolve an image's caption using the same sidecar > metadata precedence the trainer
|
|
applies in ``discover_image_caption_pairs``. A per-image .txt/.caption sidecar wins and
|
|
is stripped, so an empty (tombstone) sidecar shadows metadata and yields "" -- the
|
|
trainer then skips that image (``if caption:``), so it must not count as captioned."""
|
|
caption: Optional[str] = None
|
|
for ext in (".txt", ".caption"):
|
|
sidecar = image_path.with_suffix(ext)
|
|
if sidecar.is_file():
|
|
try:
|
|
caption = sidecar.read_text(encoding = "utf-8").strip()
|
|
except OSError:
|
|
caption = None
|
|
break
|
|
if caption is None:
|
|
try:
|
|
rel = image_path.relative_to(folder).as_posix()
|
|
except ValueError:
|
|
rel = None
|
|
caption = meta_captions.get(image_path.name) or (
|
|
meta_captions.get(rel) if rel is not None else None
|
|
)
|
|
return caption
|
|
|
|
|
|
def _diffusion_dataset_summary(folder: Path) -> DiffusionDatasetSummary:
|
|
# Count an image as captioned only when it resolves to a NON-EMPTY caption via the same
|
|
# sidecar > metadata precedence the trainer uses -- an empty tombstone sidecar shadows a
|
|
# metadata row and makes the trainer skip the image, so counting it here would over-report
|
|
# caption_count and mislabel an effectively-uncaptioned dataset as captioned.
|
|
meta_captions = _load_metadata_captions(folder)
|
|
images = captions = 0
|
|
for f in folder.iterdir():
|
|
if not f.is_file() or f.suffix.lower() not in _DIFFUSION_DATASET_IMAGE_EXTS:
|
|
continue
|
|
images += 1
|
|
if _resolve_dataset_caption(folder, f, meta_captions):
|
|
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:
|
|
# Skip hidden dirs: they are never user datasets, and an in-progress example
|
|
# import stages into a dot-prefixed sibling that must not surface as a dataset.
|
|
children = sorted(
|
|
p for p in root.iterdir() if p.is_dir() and not p.name.startswith(".")
|
|
)
|
|
except OSError:
|
|
children = []
|
|
for child in children:
|
|
try:
|
|
summary = _diffusion_dataset_summary(child)
|
|
except OSError:
|
|
continue
|
|
if summary.image_count > 0:
|
|
found.append(summary)
|
|
from core.training.diffusion_train_common import family_train_infos
|
|
|
|
families = [DiffusionTrainableFamily(**info) for info in family_train_infos()]
|
|
return DiffusionTrainingInfoResponse(
|
|
datasets_root = str(root),
|
|
outputs_root = str(outputs_root()),
|
|
datasets = found,
|
|
families = families,
|
|
)
|
|
|
|
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."""
|
|
import os
|
|
import tempfile
|
|
|
|
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
|
|
# Validate every filename up front so a valid image ahead of a bad one is not left
|
|
# written on disk when the 400 fires -- make the upload all-or-nothing.
|
|
names: list[str] = []
|
|
for f in files:
|
|
# Normalise to a safe basename. Path.name does not split on a backslash on POSIX, so a
|
|
# Windows client that sends a backslash path in the multipart filename would otherwise be
|
|
# stored verbatim; fold backslashes to forward slashes first so the true basename is
|
|
# taken for both separators. The read/caption/delete endpoints run the stored name through
|
|
# _safe_dataset_image_path (rejects "\\" / ".." / path chars), so a name that still holds
|
|
# ".." here would list an image the labeling grid can never preview, caption, or delete --
|
|
# reject it now instead of persisting an unmanageable orphan.
|
|
filename = Path((f.filename or "").replace("\\", "/")).name.strip().replace("\x00", "")
|
|
ext = Path(filename).suffix.lower()
|
|
if not filename or ".." in filename or ext not in allowed:
|
|
exts = ", ".join(sorted(allowed))
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = f"Unsupported file '{f.filename}'. Allowed: {exts}",
|
|
)
|
|
# Reject a second IMAGE that shares this one's stem but differs by extension (sample.png
|
|
# vs sample.jpg): both resolve to the same <stem>.txt caption sidecar (the kohya/diffusers
|
|
# convention the reader, editor, and delete paths all use), so keeping both would silently
|
|
# share -- and corrupt -- one caption during training. Check both files already on disk
|
|
# (uploads accumulate) and earlier images validated in THIS batch (nothing is on disk yet
|
|
# in this up-front pass). Re-uploading the exact same name (same stem AND extension) stays
|
|
# an overwrite; caption/text files are exempt (sample.txt for sample.png is intended).
|
|
if ext in _DIFFUSION_DATASET_IMAGE_EXTS:
|
|
stem = Path(filename).stem
|
|
# Compare stems (and the same-name guard) case-insensitively: on Windows/macOS
|
|
# (case-insensitive filesystems) two images whose stems differ only by case
|
|
# (sample.png vs Sample.jpg) resolve to the SAME <stem>.txt caption sidecar, so a
|
|
# case-sensitive check would let both through and silently share -- and corrupt --
|
|
# one caption. Casefolding the name guard too keeps a same-name case variant
|
|
# (sample.png vs Sample.png, one file / an overwrite on those filesystems) exempt.
|
|
stem_cf = stem.casefold()
|
|
fname_cf = filename.casefold()
|
|
clash = next(
|
|
(
|
|
p.name
|
|
for p in folder.iterdir()
|
|
if p.is_file()
|
|
and p.name.casefold() != fname_cf
|
|
and p.suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS
|
|
and p.stem.casefold() == stem_cf
|
|
),
|
|
None,
|
|
)
|
|
if clash is None:
|
|
clash = next(
|
|
(
|
|
n
|
|
for n in names
|
|
if n.casefold() != fname_cf
|
|
and Path(n).suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS
|
|
and Path(n).stem.casefold() == stem_cf
|
|
),
|
|
None,
|
|
)
|
|
if clash is not None:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = (
|
|
f"Duplicate image name '{stem}'. '{clash}' is already in this "
|
|
f"dataset; two images sharing a name would share one '{stem}.txt' "
|
|
f"caption. Rename one before uploading."
|
|
),
|
|
)
|
|
names.append(filename)
|
|
# Stage each file to a temp name and only move it into place once the whole batch is
|
|
# written, so a mid-batch failure (size limit, disk error, disconnect) leaves the
|
|
# dataset untouched -- including any pre-existing file that shares a name, which a
|
|
# direct write would have truncated (repeat uploads into the same name accumulate).
|
|
staged: list[tuple[Path, Path]] = [] # (temp, final)
|
|
committed = False
|
|
try:
|
|
for f, filename in zip(files, names):
|
|
dest = folder / filename
|
|
# A filename-independent temp name so a long (but valid, <= NAME_MAX) filename
|
|
# can't overflow NAME_MAX once the staging suffix is added.
|
|
tmp = folder / f".upload-{_uuid.uuid4().hex}.part"
|
|
staged.append((tmp, dest))
|
|
with open(tmp, "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)
|
|
uploaded += 1
|
|
for tmp, dest in staged:
|
|
tmp.replace(dest) # atomic on the same filesystem
|
|
committed = True
|
|
finally:
|
|
if not committed:
|
|
for tmp, _ in staged:
|
|
try:
|
|
tmp.unlink(missing_ok = True)
|
|
except OSError:
|
|
pass
|
|
|
|
summary = _diffusion_dataset_summary(folder)
|
|
return DiffusionDatasetUploadResponse(
|
|
name = cleaned,
|
|
path = str(folder),
|
|
image_count = summary.image_count,
|
|
caption_count = summary.caption_count,
|
|
uploaded = uploaded,
|
|
)
|
|
|
|
|
|
# ── Dataset labeling (per-image caption editing) + one-click example imports ──
|
|
# Thumbnails live in a hidden subdir so they never appear in dataset listings or the
|
|
# trainer's own image discovery (both scan only top-level files).
|
|
_THUMBS_DIRNAME = ".thumbs"
|
|
_MAX_CAPTION_CHARS = 2000
|
|
|
|
|
|
def _resolve_dataset_folder(name: str, *, must_exist: bool = True) -> Path:
|
|
"""Validate ``name`` (single component, no traversal) and resolve it under the Studio
|
|
datasets root. 404 when a read target is missing."""
|
|
from utils.paths import datasets_root
|
|
|
|
cleaned = _clean_diffusion_dataset_name(name)
|
|
folder = datasets_root() / cleaned
|
|
if must_exist and not folder.is_dir():
|
|
raise HTTPException(status_code = 404, detail = f"Dataset '{cleaned}' not found.")
|
|
return folder
|
|
|
|
|
|
def _safe_dataset_image_path(folder: Path, filename: str) -> Path:
|
|
"""Resolve ``filename`` to an image path strictly inside ``folder``. Rejects any path
|
|
separators / traversal / null bytes and non-image extensions."""
|
|
raw = filename or ""
|
|
if "/" in raw or "\\" in raw or ".." in raw or "\x00" in raw or raw != Path(raw).name:
|
|
raise HTTPException(status_code = 400, detail = "Invalid image filename.")
|
|
if Path(raw).suffix.lower() not in _DIFFUSION_DATASET_IMAGE_EXTS:
|
|
exts = ", ".join(sorted(_DIFFUSION_DATASET_IMAGE_EXTS))
|
|
raise HTTPException(status_code = 400, detail = f"Not an image file. Allowed: {exts}")
|
|
path = folder / raw
|
|
# Defense in depth: the real path must stay under the dataset folder.
|
|
try:
|
|
path.resolve().relative_to(folder.resolve())
|
|
except ValueError:
|
|
raise HTTPException(status_code = 400, detail = "Invalid image filename.")
|
|
return path
|
|
|
|
|
|
def _load_metadata_captions(folder: Path) -> dict[str, str]:
|
|
"""Read metadata.jsonl / captions.jsonl into {file_name: caption}, mirroring the
|
|
trainer's discovery (keys file_name/image/file; caption in the ``text`` column)."""
|
|
import json
|
|
|
|
out: dict[str, str] = {}
|
|
for meta_name in ("metadata.jsonl", "captions.jsonl"):
|
|
meta_path = folder / meta_name
|
|
if not meta_path.is_file():
|
|
continue
|
|
try:
|
|
lines = meta_path.read_text(encoding = "utf-8").splitlines()
|
|
except OSError:
|
|
continue
|
|
for line in lines:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
row = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
key = row.get("file_name") or row.get("image") or row.get("file")
|
|
if key and "text" in row:
|
|
out[str(key)] = str(row["text"])
|
|
return out
|
|
|
|
|
|
def _image_record(
|
|
folder: Path, image_path: Path, meta_captions: dict[str, str]
|
|
) -> DiffusionDatasetImageRecord:
|
|
"""Build one image record, resolving its caption with sidecar > metadata precedence
|
|
(the same order the trainer uses). A per-image .txt / .caption sidecar wins because
|
|
it is the user's explicit edit from the labeling grid, which must override a
|
|
metadata.jsonl / captions.jsonl row for the image."""
|
|
caption: Optional[str] = None
|
|
source = "none"
|
|
for ext in (".txt", ".caption"):
|
|
sidecar = image_path.with_suffix(ext)
|
|
if sidecar.is_file():
|
|
try:
|
|
caption = sidecar.read_text(encoding = "utf-8").strip()
|
|
source = "sidecar"
|
|
except OSError:
|
|
caption = None
|
|
break
|
|
if caption is None:
|
|
# Basename first, then the relative path as written in the jsonl (as_posix so a
|
|
# Windows backslash path still matches forward-slash keys) -- the same lookup
|
|
# order discover_image_caption_pairs uses.
|
|
meta = meta_captions.get(image_path.name)
|
|
if meta is None:
|
|
try:
|
|
meta = meta_captions.get(image_path.relative_to(folder).as_posix())
|
|
except ValueError:
|
|
meta = None
|
|
if meta is not None:
|
|
caption = meta
|
|
source = "metadata"
|
|
try:
|
|
size_bytes = image_path.stat().st_size
|
|
except OSError:
|
|
size_bytes = 0
|
|
width = height = 0
|
|
try:
|
|
from PIL import Image
|
|
with Image.open(image_path) as im:
|
|
width, height = im.size
|
|
except Exception: # noqa: BLE001 -- an unreadable image still lists (0x0) rather than 500
|
|
pass
|
|
return DiffusionDatasetImageRecord(
|
|
filename = image_path.name,
|
|
caption = caption,
|
|
caption_source = source, # type: ignore[arg-type]
|
|
width = width,
|
|
height = height,
|
|
size_bytes = size_bytes,
|
|
)
|
|
|
|
|
|
@router.get("/diffusion/dataset/{name}/images", response_model = DiffusionDatasetImagesResponse)
|
|
async def list_diffusion_dataset_images(
|
|
name: str, current_subject: str = Depends(get_current_subject)
|
|
):
|
|
"""List every image in a dataset folder with its resolved caption (including
|
|
uncaptioned images), for the labeling grid."""
|
|
folder = _resolve_dataset_folder(name)
|
|
|
|
def scan() -> DiffusionDatasetImagesResponse:
|
|
meta = _load_metadata_captions(folder)
|
|
records: list[DiffusionDatasetImageRecord] = []
|
|
for p in sorted(folder.iterdir()):
|
|
if p.is_file() and p.suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS:
|
|
records.append(_image_record(folder, p, meta))
|
|
return DiffusionDatasetImagesResponse(name = folder.name, path = str(folder), images = records)
|
|
|
|
return await asyncio.to_thread(scan)
|
|
|
|
|
|
@router.get("/diffusion/dataset/{name}/image/{filename}")
|
|
async def get_diffusion_dataset_image(
|
|
name: str,
|
|
filename: str,
|
|
thumb: Optional[int] = None,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""Serve a dataset image. ``?thumb=<px>`` returns a cached downscaled JPEG (regenerated
|
|
when the source is newer), used by the labeling grid to stay light."""
|
|
from fastapi.responses import FileResponse
|
|
|
|
folder = _resolve_dataset_folder(name)
|
|
image_path = _safe_dataset_image_path(folder, filename)
|
|
if not image_path.is_file():
|
|
raise HTTPException(status_code = 404, detail = "Image not found.")
|
|
if not thumb:
|
|
return FileResponse(str(image_path))
|
|
|
|
size = max(32, min(1024, int(thumb)))
|
|
|
|
def make_thumb() -> Path:
|
|
from PIL import Image
|
|
|
|
thumbs_dir = folder / _THUMBS_DIRNAME
|
|
thumbs_dir.mkdir(exist_ok = True)
|
|
# Key on the full filename (stem + extension), not the stem: two images that
|
|
# share a stem but differ by extension (sample.png / sample.jpg) would otherwise
|
|
# collide on one cache file, and an mtime-newer cache built for the first would
|
|
# be served for the second, showing the wrong image in the labeling grid.
|
|
thumb_path = thumbs_dir / f"{image_path.name}_{size}.jpg"
|
|
src_mtime = image_path.stat().st_mtime
|
|
if thumb_path.is_file() and thumb_path.stat().st_mtime >= src_mtime:
|
|
return thumb_path
|
|
with Image.open(image_path) as im:
|
|
im = im.convert("RGB")
|
|
im.thumbnail((size, size), Image.LANCZOS)
|
|
im.save(thumb_path, format = "JPEG", quality = 85)
|
|
return thumb_path
|
|
|
|
try:
|
|
thumb_path = await asyncio.to_thread(make_thumb)
|
|
except Exception as e: # noqa: BLE001 -- fall back to the original on any decode failure
|
|
logger.warning("Thumbnail generation failed for %s: %s", image_path, e)
|
|
return FileResponse(str(image_path))
|
|
return FileResponse(str(thumb_path), media_type = "image/jpeg")
|
|
|
|
|
|
@router.put(
|
|
"/diffusion/dataset/{name}/caption/{filename}",
|
|
response_model = DiffusionDatasetImageRecord,
|
|
)
|
|
async def set_diffusion_dataset_caption(
|
|
name: str,
|
|
filename: str,
|
|
body: DiffusionCaptionUpdateRequest,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""Write (or, when blank, clear) an image's ``.txt`` caption sidecar. Returns the
|
|
updated image record."""
|
|
folder = _resolve_dataset_folder(name)
|
|
image_path = _safe_dataset_image_path(folder, filename)
|
|
if not image_path.is_file():
|
|
raise HTTPException(status_code = 404, detail = "Image not found.")
|
|
caption = (body.caption or "").strip()
|
|
if len(caption) > _MAX_CAPTION_CHARS:
|
|
raise HTTPException(
|
|
status_code = 400,
|
|
detail = f"Caption too long (max {_MAX_CAPTION_CHARS} characters).",
|
|
)
|
|
|
|
def write() -> DiffusionDatasetImageRecord:
|
|
sidecar = image_path.with_suffix(".txt")
|
|
if caption:
|
|
sidecar.write_text(caption, encoding = "utf-8")
|
|
image_path.with_suffix(".caption").unlink(missing_ok = True)
|
|
return _image_record(folder, image_path, _load_metadata_captions(folder))
|
|
# Blank must actually clear. Unlinking alone would resurface this image's
|
|
# metadata.jsonl / captions.jsonl caption (the fallback source), so when one
|
|
# exists write an EMPTY sidecar instead: both the record reader and the
|
|
# trainer's discovery treat an existing sidecar as authoritative even when
|
|
# empty, which makes it a tombstone. No metadata caption -> plain cleanup.
|
|
meta = _load_metadata_captions(folder)
|
|
try:
|
|
rel = image_path.relative_to(folder).as_posix()
|
|
except ValueError:
|
|
rel = image_path.name
|
|
if image_path.name in meta or rel in meta:
|
|
sidecar.write_text("", encoding = "utf-8")
|
|
else:
|
|
sidecar.unlink(missing_ok = True)
|
|
image_path.with_suffix(".caption").unlink(missing_ok = True)
|
|
return _image_record(folder, image_path, meta)
|
|
|
|
return await asyncio.to_thread(write)
|
|
|
|
|
|
@router.delete("/diffusion/dataset/{name}/image/{filename}")
|
|
async def delete_diffusion_dataset_image(
|
|
name: str,
|
|
filename: str,
|
|
current_subject: str = Depends(get_current_subject),
|
|
):
|
|
"""Remove an image, its caption sidecars, and any cached thumbnails."""
|
|
folder = _resolve_dataset_folder(name)
|
|
image_path = _safe_dataset_image_path(folder, filename)
|
|
if not image_path.is_file():
|
|
raise HTTPException(status_code = 404, detail = "Image not found.")
|
|
|
|
def remove() -> dict:
|
|
image_path.unlink(missing_ok = True)
|
|
for ext in (".txt", ".caption"):
|
|
image_path.with_suffix(ext).unlink(missing_ok = True)
|
|
thumbs_dir = folder / _THUMBS_DIRNAME
|
|
if thumbs_dir.is_dir():
|
|
# Thumbs are keyed on the full filename (stem + extension), so match that
|
|
# here too; a stem-only glob would leave this image's thumbs behind and
|
|
# could delete a same-stem sibling's (sample.png vs sample.jpg).
|
|
for t in thumbs_dir.glob(f"{image_path.name}_*.jpg"):
|
|
t.unlink(missing_ok = True)
|
|
return {"deleted": image_path.name}
|
|
|
|
return await asyncio.to_thread(remove)
|
|
|
|
|
|
# Curated, license-labelled example datasets for one-click import. ``loader`` picks the
|
|
# materialization strategy: "hf_dataset" streams rows from datasets.load_dataset (image +
|
|
# optional caption column); "imagefolder_jsonl" snapshot-downloads a dataset repo whose
|
|
# captions live in a *.jsonl (file_name/text) rather than a standard metadata.jsonl.
|
|
_DATASET_EXAMPLES: list[dict] = [
|
|
{
|
|
"id": "dreambooth-dog",
|
|
"label": "Dog (DreamBooth subject)",
|
|
"repo": "diffusers/dog-example",
|
|
"description": (
|
|
"5 photos of one dog. The classic DreamBooth subject set: teach the model a "
|
|
"specific subject, then summon it with the trigger prompt."
|
|
),
|
|
"license": "Released by Google for DreamBooth research/demos",
|
|
"image_cap": 10,
|
|
"suggested_trigger": "a photo of sks dog",
|
|
"loader": "hf_dataset",
|
|
"caption_column": None,
|
|
"no_checks": False,
|
|
},
|
|
{
|
|
"id": "tuxemon",
|
|
"label": "Tuxemon (captioned style set)",
|
|
"repo": "linoyts/Tuxemon",
|
|
"description": (
|
|
"Captioned cartoon monster art. A good style set: each image ships a caption, "
|
|
"so the adapter learns the look without a trigger word."
|
|
),
|
|
"license": "cc-by-sa-3.0",
|
|
"image_cap": 60,
|
|
"suggested_trigger": None,
|
|
"loader": "hf_dataset",
|
|
"caption_column": "prompt",
|
|
"no_checks": True,
|
|
},
|
|
{
|
|
"id": "tarot-1920",
|
|
"label": "1920 Tarot (public domain style set)",
|
|
"repo": "multimodalart/1920-raider-waite-tarot-public-domain",
|
|
"description": (
|
|
"Public-domain 1920 Raider-Waite tarot art with captions. A permissive style "
|
|
"set for demoing captioned LoRA training."
|
|
),
|
|
"license": "public domain",
|
|
"image_cap": 60,
|
|
"suggested_trigger": None,
|
|
"loader": "imagefolder_jsonl",
|
|
"caption_column": "text",
|
|
"no_checks": True,
|
|
},
|
|
{
|
|
"id": "smithsonian-butterflies",
|
|
"label": "Smithsonian Butterflies",
|
|
"repo": "huggan/smithsonian_butterflies_subset",
|
|
"description": (
|
|
"100 butterfly specimen photos. The classic diffusers-docs training set. No "
|
|
"captions, so pair it with the trigger prompt to teach a butterfly subject."
|
|
),
|
|
"license": "CC0 (Smithsonian Open Access)",
|
|
"image_cap": 100,
|
|
# The metadata columns are species names / boilerplate alt-text, not text-to-image
|
|
# captions, so train it as a subject set with the trigger prompt instead.
|
|
"suggested_trigger": "a photo of a sks butterfly",
|
|
"loader": "hf_dataset",
|
|
"caption_column": None,
|
|
"no_checks": False,
|
|
},
|
|
{
|
|
"id": "pixel-nouns",
|
|
"label": "Nouns (pixel avatars)",
|
|
"repo": "m1guelpf/nouns",
|
|
"description": (
|
|
"100 captioned Nouns pixel-art avatars. A captioned style set: each image ships "
|
|
"a caption, so the adapter learns the pixel look without a trigger word."
|
|
),
|
|
"license": "cc0-1.0",
|
|
"image_cap": 100,
|
|
"suggested_trigger": None,
|
|
"loader": "hf_dataset",
|
|
"caption_column": "text",
|
|
"no_checks": False,
|
|
},
|
|
]
|
|
|
|
|
|
def _example_by_id(example_id: str) -> dict:
|
|
for entry in _DATASET_EXAMPLES:
|
|
if entry["id"] == example_id:
|
|
return entry
|
|
raise HTTPException(status_code = 404, detail = f"Unknown example dataset '{example_id}'.")
|
|
|
|
|
|
@router.get("/diffusion/dataset-examples", response_model = DiffusionDatasetExamplesResponse)
|
|
async def list_diffusion_dataset_examples(current_subject: str = Depends(get_current_subject)):
|
|
"""List the curated example datasets available for one-click import."""
|
|
return DiffusionDatasetExamplesResponse(
|
|
examples = [
|
|
DiffusionDatasetExample(
|
|
id = e["id"],
|
|
label = e["label"],
|
|
repo = e["repo"],
|
|
description = e["description"],
|
|
license = e["license"],
|
|
image_cap = e["image_cap"],
|
|
suggested_trigger = e["suggested_trigger"],
|
|
)
|
|
for e in _DATASET_EXAMPLES
|
|
]
|
|
)
|
|
|
|
|
|
def _detect_image_column(features) -> Optional[str]:
|
|
"""Return the first datasets Image-feature column name, else None."""
|
|
try:
|
|
from datasets import Image as HFImage
|
|
except Exception: # noqa: BLE001
|
|
HFImage = None # type: ignore[assignment]
|
|
for col, feat in features.items():
|
|
if HFImage is not None and isinstance(feat, HFImage):
|
|
return col
|
|
if type(feat).__name__ == "Image":
|
|
return col
|
|
return None
|
|
|
|
|
|
def _detect_caption_column(entry: dict, columns: list[str]) -> Optional[str]:
|
|
"""Pick the caption column: the entry's declared one if present, else a common name."""
|
|
declared = entry.get("caption_column")
|
|
if declared and declared in columns:
|
|
return declared
|
|
for cand in ("text", "prompt", "caption", "captions"):
|
|
if cand in columns:
|
|
return cand
|
|
return None
|
|
|
|
|
|
def _materialize_hf_dataset(entry: dict, dest: Path, cap: int) -> int:
|
|
"""Stream rows from datasets.load_dataset into ``dest`` as numbered images + optional
|
|
.txt sidecars. Returns the number of images written."""
|
|
from datasets import load_dataset
|
|
|
|
kwargs = {"split": "train"}
|
|
if entry.get("no_checks"):
|
|
kwargs["verification_mode"] = "no_checks"
|
|
ds = load_dataset(entry["repo"], **kwargs)
|
|
image_col = _detect_image_column(ds.features)
|
|
if image_col is None:
|
|
raise HTTPException(
|
|
status_code = 502,
|
|
detail = f"'{entry['repo']}' has no image column to import.",
|
|
)
|
|
caption_col = _detect_caption_column(entry, list(ds.features.keys()))
|
|
written = 0
|
|
for row in ds:
|
|
if written >= cap:
|
|
break
|
|
img = row[image_col]
|
|
if img is None:
|
|
continue
|
|
img = img.convert("RGB")
|
|
stem = f"img_{written:04d}"
|
|
img.save(dest / f"{stem}.png", format = "PNG")
|
|
if caption_col:
|
|
cap_text = row.get(caption_col)
|
|
if cap_text:
|
|
(dest / f"{stem}.txt").write_text(str(cap_text).strip(), encoding = "utf-8")
|
|
written += 1
|
|
return written
|
|
|
|
|
|
def _materialize_imagefolder_jsonl(entry: dict, dest: Path, cap: int) -> int:
|
|
"""Snapshot-download a dataset repo whose captions live in *.jsonl (file_name/text),
|
|
then copy referenced images + write .txt sidecars. Returns images written."""
|
|
import json
|
|
import shutil
|
|
|
|
from huggingface_hub import snapshot_download
|
|
|
|
caption_col = entry.get("caption_column") or "text"
|
|
snap = Path(
|
|
snapshot_download(
|
|
entry["repo"],
|
|
repo_type = "dataset",
|
|
allow_patterns = [
|
|
"*.jsonl",
|
|
"*.jpg",
|
|
"*.jpeg",
|
|
"*.png",
|
|
"*.webp",
|
|
"*.bmp",
|
|
"**/*.jpg",
|
|
"**/*.jpeg",
|
|
"**/*.png",
|
|
"**/*.webp",
|
|
"**/*.bmp",
|
|
],
|
|
)
|
|
)
|
|
# Map basename -> caption from every jsonl carrying file_name + caption column.
|
|
captions: dict[str, str] = {}
|
|
for jf in sorted(snap.rglob("*.jsonl")):
|
|
for line in jf.read_text(encoding = "utf-8").splitlines():
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
row = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
fn = row.get("file_name") or row.get("image") or row.get("file")
|
|
if fn and caption_col in row:
|
|
# First writer wins over sorted manifests, for deterministic results.
|
|
captions.setdefault(Path(str(fn)).name, str(row[caption_col]))
|
|
# Copy images (those with a caption first, so a cap keeps captioned pairs).
|
|
images = sorted(
|
|
p
|
|
for p in snap.rglob("*")
|
|
if p.is_file() and p.suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS
|
|
)
|
|
images.sort(key = lambda p: (p.name not in captions, p.name))
|
|
written = 0
|
|
for src in images:
|
|
if written >= cap:
|
|
break
|
|
stem = f"img_{written:04d}"
|
|
shutil.copyfile(src, dest / f"{stem}{src.suffix.lower()}")
|
|
cap_text = captions.get(src.name)
|
|
if cap_text:
|
|
(dest / f"{stem}.txt").write_text(cap_text.strip(), encoding = "utf-8")
|
|
written += 1
|
|
return written
|
|
|
|
|
|
@router.post("/diffusion/dataset/import-example", response_model = DiffusionDatasetImportResponse)
|
|
async def import_diffusion_dataset_example(
|
|
body: DiffusionDatasetImportRequest, current_subject: str = Depends(get_current_subject)
|
|
):
|
|
"""Materialize a curated example dataset into a Studio dataset folder (images + .txt
|
|
captions), ready to train. Idempotent: a folder that already holds images is returned
|
|
as-is rather than re-downloaded."""
|
|
entry = _example_by_id(body.id)
|
|
folder = _resolve_dataset_folder(body.name or entry["id"], must_exist = False)
|
|
|
|
def do_import() -> DiffusionDatasetImportResponse:
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
|
|
folder.mkdir(parents = True, exist_ok = True)
|
|
existing = _diffusion_dataset_summary(folder)
|
|
imported = 0
|
|
if existing.image_count == 0:
|
|
cap = int(entry["image_cap"])
|
|
# Materialize into a private staging dir and promote into the dataset folder only
|
|
# after the whole import succeeds. A materialize that fails partway (a transient
|
|
# fetch/copy error after writing some images) then leaves only the staging dir,
|
|
# never a half-filled dataset -- otherwise the image_count>0 idempotency check
|
|
# above would treat that partial result as complete on the next retry (imported=0)
|
|
# and strand the user with a truncated dataset (there is no dataset-delete flow).
|
|
# Staged as a hidden sibling on the same filesystem so promotion is an atomic rename.
|
|
staging = Path(tempfile.mkdtemp(dir = folder.parent, prefix = f".{folder.name}.import-"))
|
|
try:
|
|
try:
|
|
if entry["loader"] == "imagefolder_jsonl":
|
|
imported = _materialize_imagefolder_jsonl(entry, staging, cap)
|
|
else:
|
|
imported = _materialize_hf_dataset(entry, staging, cap)
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e: # noqa: BLE001 -- surface a readable fetch/parse failure
|
|
raise HTTPException(
|
|
status_code = 502,
|
|
detail = f"Could not import '{entry['repo']}': {e}",
|
|
)
|
|
if imported == 0:
|
|
raise HTTPException(
|
|
status_code = 502,
|
|
detail = f"No images found in '{entry['repo']}'.",
|
|
)
|
|
# Promote the fully-materialized staging dir as a UNIT. A per-file move loop is
|
|
# not atomic: a hard process death (SIGKILL / OOM / power loss) between two moves
|
|
# would leave the folder with SOME images, and the image_count>0 idempotency check
|
|
# above would then accept that truncated dataset as complete on the next retry. The
|
|
# folder was created empty on this path (it only runs when it holds no images), so
|
|
# a single same-filesystem directory rename is atomic. If the folder holds
|
|
# unrelated non-image files (rmdir refuses), fall back to a per-file move rather
|
|
# than abort -- the common fresh-import path stays atomic.
|
|
try:
|
|
os.rmdir(folder)
|
|
except OSError:
|
|
for p in staging.iterdir():
|
|
shutil.move(str(p), str(folder / p.name))
|
|
else:
|
|
os.replace(str(staging), str(folder))
|
|
finally:
|
|
shutil.rmtree(staging, ignore_errors = True)
|
|
summary = _diffusion_dataset_summary(folder)
|
|
return DiffusionDatasetImportResponse(
|
|
name = folder.name,
|
|
path = str(folder),
|
|
image_count = summary.image_count,
|
|
caption_count = summary.caption_count,
|
|
imported = imported,
|
|
license = entry["license"],
|
|
source_repo = entry["repo"],
|
|
)
|
|
|
|
return await asyncio.to_thread(do_import)
|