ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Tighten comments in the image stack routes

Browse source
This commit is contained in:
Daniel Han 2026-07-12 12:00:44 +00:00
commit 3a6a007038
3 changed files with 281 additions and 336 deletions
Download patch file Download diff file Expand all files Collapse all files

183
studio/backend/routes/inference.py
View file

@ -3756,8 +3756,8 @@ def _guard_chat_load_against_training(
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.
# fit-checked here, so refuse the chat load 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,
@ -3913,9 +3913,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
# Reclaim the GPU for chat (evicting a resident Images/Video pipeline) only once the
# load is known viable: the already-loaded fast paths below re-assert CHAT ownership
# themselves, and the real handoff is deferred past identifier / gpu_ids / training-memory
# validation so a doomed chat load (bad id, unsupported gpu_ids on GGUF, or a training
# 409) can't evict a working image/video model and then error. Mirrors the image/video
# loaders, which validate before acquire_for.
# validation so a doomed load (bad id, unsupported gpu_ids on GGUF, training 409) can't
# evict a working image/video model and then error. Mirrors the image/video loaders,
# which validate before acquire_for.
from core.inference.gpu_arbiter import acquire_for, CHAT
# ── Already-loaded check: skip reload if the exact model is active ──
@ -3951,9 +3951,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
_gguf_audio = getattr(llama_backend, "_audio_type", None)
_gguf_is_audio = getattr(llama_backend, "_is_audio", False)
# The requested GGUF chat model is already resident: assert CHAT ownership (a
# no-op when it already holds it) so a drifted arbiter owner is corrected. This
# is a guaranteed-success path, not a doomed load, so evicting here is correct.
# Requested GGUF chat model already resident: assert CHAT ownership (no-op when
# held) to correct a drifted arbiter owner. Guaranteed-success path, so evicting
# here is correct.
await asyncio.to_thread(acquire_for, CHAT)
return LoadResponse(
status = "already_loaded",
@ -4007,9 +4007,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
_sf_flags = _detect_safetensors_features(backend, _chat_template)
_sf_supports_reasoning = _sf_flags["supports_reasoning"]
_sf_reasoning_style = _sf_flags["reasoning_style"]
# The requested chat model is already resident: assert CHAT ownership (no-op when
# it already holds it) to correct a drifted arbiter owner. Guaranteed-success
# path, not a doomed load, so evicting here is correct.
# Requested chat model already resident: assert CHAT ownership (no-op when held)
# to correct a drifted arbiter owner. Guaranteed-success path, so evicting here
# is correct.
await asyncio.to_thread(acquire_for, CHAT)
return LoadResponse(
status = "already_loaded",
@ -4094,11 +4094,10 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1),
)
# The load is now known viable (valid identifier, gpu_ids ok, fits alongside any active
# Load now known viable (valid identifier, gpu_ids ok, fits alongside any active
# training): reclaim the GPU for chat, evicting a resident Images/Video pipeline. Doing
# this only here -- not before the validation above -- is what keeps a doomed chat load
# from evicting a working image/video model and then erroring. No-op when chat already
# owns the GPU.
# this only here -- not before the validation above -- keeps a doomed load from evicting
# a working image/video model and then erroring. No-op when chat already owns the GPU.
await asyncio.to_thread(acquire_for, CHAT)
# ── GGUF path: load via llama-server ──────────────────────
@ -14236,10 +14235,9 @@ async def _openai_passthrough_non_streaming_upstream(
# ──────────────────────────────────────────────────────────────────────────
# Diffusion (local text-to-image)
#
# Studio-only routes (studio_router is not mounted under /v1). The diffusion
# backend runs in-process and is synchronous, so the blocking load/generate/
# unload calls are offloaded with asyncio.to_thread to keep the event loop free.
# This is the single error boundary: backend methods raise, we map to HTTP here.
# Studio-only routes (studio_router is not mounted under /v1). The diffusion backend
# runs in-process and synchronously, so blocking load/generate/unload calls are
# offloaded with asyncio.to_thread. Single error boundary: backend raises, we map to HTTP.
# ──────────────────────────────────────────────────────────────────────────
@ -14265,9 +14263,9 @@ def _guard_diffusion_load_against_training() -> None:
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.
# 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, or the pipeline contends with the trainer
# for VRAM. Symmetric with the diffusion-start interlock.
if not llm_active and not _diffusion_training_active():
return
raise HTTPException(
@ -14303,19 +14301,18 @@ async def load_diffusion_model(
# engine selection, and the load all agree. A bad explicit kind raises here -> 400.
kind = resolve_model_kind(request.gguf_filename, request.model_kind)
# A local On-Device pick can be a bare single-file .safetensors directory (no
# model_index.json): the scanner advertises it as a text-to-image model, but the local
# picker starts it as a pipeline with no filename, so a pipeline load would 400 on the
# missing model_index.json and the advertised model is unusable. If the directory holds
# exactly one checkpoint, reinterpret the pick as a single_file load of it (the only
# loadable shape for that dir), so validation, engine selection, and the load all agree.
# model_index.json): the scanner advertises it as text-to-image, but the picker starts
# it as a pipeline with no filename, so a pipeline load would 400 on the missing
# model_index.json. If the directory holds exactly one checkpoint, reinterpret the pick
# as a single_file load of it (its only loadable shape), so all three paths agree.
if kind == "pipeline" and not request.gguf_filename:
sole = await asyncio.to_thread(resolve_local_single_file, request.model_path)
if sole is not None:
request.gguf_filename = sole
kind = resolve_model_kind(sole)
# Validate cheaply BEFORE touching the GPU: an unloadable pick (bad family,
# missing local GGUF, a non-unsloth non-GGUF repo) must not evict a working chat
# model and then 400. The validated family also drives engine selection below.
# Validate cheaply BEFORE touching the GPU: an unloadable pick (bad family, missing
# local GGUF, non-unsloth non-GGUF repo) must not evict a working chat model and then
# 400. The validated family also drives engine selection below.
fam = await asyncio.to_thread(
backend.validate_load_request,
request.model_path,
@ -14329,18 +14326,16 @@ async def load_diffusion_model(
# same via _guard_chat_load_against_training; this is its image sibling.
_guard_diffusion_load_against_training()
# Pick the engine for this host (diffusers on GPU, native sd.cpp with no GPU),
# installing the sd-cli binary if needed -- all BEFORE evicting chat, so a
# native fallback never strands a half-loaded state. Non-GGUF kinds force diffusers.
# installing the sd-cli binary if needed -- all BEFORE evicting chat, so a native
# fallback never strands a half-loaded state. Non-GGUF kinds force diffusers.
engine = await asyncio.to_thread(
select_and_activate_engine, fam, hf_token = request.hf_token, model_kind = kind
)
# Take the GPU from the chat backend only when this load will actually use it,
# which is exactly the resolved device being non-CPU. diffusers on an accelerator
# and a force-native sd.cpp load on CUDA/XPU/MPS both resolve to that device; a
# native sd.cpp load on a pure-CPU host does not. Crucially, a CPU-only host with
# no usable sd-cli falls back to diffusers ON CPU -- that also never touches GPU
# memory, so keying off the engine name (not the device) would wrongly evict a
# resident chat model for a load that cannot use the GPU. Gate on the device.
# Take the GPU from chat only when this load will actually use it, i.e. the resolved
# device is non-CPU. diffusers on an accelerator and a force-native sd.cpp load on
# CUDA/XPU/MPS both resolve to a device; a native sd.cpp load on a pure-CPU host, and a
# CPU-only host falling back to diffusers ON CPU, do not. So gate on the device, not the
# engine name -- else we'd evict a resident chat model for a load that can't use the GPU.
device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device)
needs_gpu = device != "cpu"
if needs_gpu:
@ -14350,10 +14345,9 @@ async def load_diffusion_model(
else:
# A CPU-only native load never touches the GPU, so it neither acquires nor is
# tracked by the arbiter. But switching here FROM a previous diffusers/GPU load
# (select_and_activate_engine unloaded it above) leaves DIFFUSION still marked
# as the arbiter owner; a later chat acquire would then "evict" this CPU model
# for no reason. Release that stale ownership -- release() is owner-guarded, so
# it is a no-op when diffusion never owned the GPU.
# leaves DIFFUSION still marked as arbiter owner, so a later chat acquire would
# "evict" this CPU model for no reason. Release that stale ownership -- release()
# is owner-guarded, so it's a no-op when diffusion never owned the GPU.
await asyncio.to_thread(release, DIFFUSION)
status_dict = await asyncio.to_thread(
engine.begin_load,
@ -14429,13 +14423,11 @@ async def generate_diffusion_image(
# doesn't support) — a 400 with the reason, not a generic 500.
raise HTTPException(status_code = 400, detail = str(exc))
except RuntimeError as exc:
# Only "no model loaded" / user-cancelled are client-state (409); both engines
# raise these two EXACT messages. The native sd.cpp engine also raises
# RuntimeError for execution failures (nonzero exit, timeout, missing output)
# whose text can embed the raw sd-cli tail (local paths / argv) -- those are
# server errors (500) returned as a fixed literal, never echoed. Match the
# sentinels exactly, not as a substring, so an sd-cli failure that merely
# contains "cancelled" can't misroute to 409 and leak that output.
# Only "no model loaded" / user-cancelled are client-state (409); both engines raise
# these two EXACT messages. The native sd.cpp engine also raises RuntimeError for
# execution failures whose text can embed the raw sd-cli tail (local paths / argv) --
# those are 500s returned as a fixed literal, never echoed. Match the sentinels exactly
# (not as substrings) so an sd-cli failure containing "cancelled" can't misroute to 409.
msg = str(exc)
if msg in (DIFFUSION_NOT_LOADED_MSG, DIFFUSION_CANCELLED_MSG):
raise HTTPException(status_code = 409, detail = msg)
@ -14445,9 +14437,9 @@ async def generate_diffusion_image(
logger.error("diffusion.generate_failed: %s", exc, exc_info = True)
raise HTTPException(status_code = 500, detail = "Image generation failed.")
# Persist each image with its full recipe embedded. The diffusers batch shares
# one seed (drawn sequentially from one generator); the native sd.cpp batch uses a
# distinct seed per image and returns them in ``seeds`` so each is reproducible.
# Persist each image with its full recipe embedded. The diffusers batch shares one seed
# (drawn sequentially from one generator); the native sd.cpp batch uses a distinct seed
# per image, returned in ``seeds`` so each is reproducible.
created_at = time.time()
per_image_seeds = result.get("seeds")
@ -14465,21 +14457,20 @@ async def generate_diffusion_image(
{
"prompt": request.prompt,
"negative_prompt": request.negative_prompt,
# Persist the ACTUAL output size, not the request sliders: Transform/
# Inpaint/Edit derive it from the uploaded image, Extend grows the
# canvas, and Upscale resizes it, so request.width/height would record
# (and later restore) the wrong dimensions for those workflows. For
# plain txt2img the image size equals the sliders anyway.
# Persist the ACTUAL output size, not the request sliders:
# Transform/Inpaint/Edit derive it from the uploaded image, Extend grows
# the canvas, Upscale resizes it, so request.width/height would record the
# wrong dims. For plain txt2img the size equals the sliders anyway.
"width": getattr(image, "width", None) or request.width,
"height": getattr(image, "height", None) or request.height,
"steps": request.steps,
"guidance": request.guidance,
"seed": seed,
# Position within the batch: shared timestamp, so the export
# filename needs this to stay unique.
# Position within the batch (shared timestamp), so the export filename
# stays unique.
"batch_index": index,
# The batch shares one seed, so reproducing image batch_index>0
# needs the original batch_size: persist it so restore can replay.
# The batch shares one seed, so reproducing a batch_index>0 image needs
# the original batch_size: persist it so restore can replay.
"batch_size": request.batch_size,
"model": result.get("repo_id"),
"loras": (
@ -14488,9 +14479,8 @@ async def generate_diffusion_image(
"controlnet": (
f"{request.controlnet.id}:{request.controlnet.control_type}:"
f"{request.controlnet.strength:g}"
# strength 0 is treated as disabled and skipped before loading /
# conditioning, so the image is unconditioned; don't claim a
# ControlNet was applied in the recipe/metadata.
# strength 0 is disabled and skipped before loading/conditioning,
# so don't claim a ControlNet was applied in the recipe/metadata.
if request.controlnet and request.controlnet.strength > 0
else None
),
@ -14524,11 +14514,9 @@ async def list_gallery_images(
# Fetch one extra to learn whether more remain, without a second scan.
records = await asyncio.to_thread(image_gallery.list_images, limit + 1, offset)
has_more = len(records) > limit
# Build the response per record and drop any that fail schema validation: a PNG
# whose recipe chunk has all required keys but a wrong value type (e.g. a
# hand-dropped or corrupted file) passes the presence-only read but would raise
# inside GalleryImage(**r). Skipping it keeps one bad file from 500-ing the whole
# gallery listing.
# Drop records that fail schema validation: a PNG whose recipe chunk has all keys but a
# wrong value type (hand-dropped or corrupted) passes the presence-only read yet raises
# inside GalleryImage(**r). Skipping it keeps one bad file from 500-ing the listing.
images = []
for r in records[:limit]:
try:
@ -14581,11 +14569,10 @@ async def unload_diffusion_model(current_subject: str = Depends(get_current_subj
status_dict = await asyncio.to_thread(get_active_diffusion_engine().unload)
# Drop DIFFUSION ownership only if nothing is resident AND no new load is in flight: a
# concurrent /images/load that re-acquired DIFFUSION while this (slow) unload ran must keep
# ownership, or a later chat load would see no owner, skip eviction, and OOM against the newly
# ownership, or a later chat load would see no owner, skip eviction, and OOM the newly
# resident pipeline. An in-flight load has is_loaded False for its whole download/finalize
# window, so gate on loading_repo_ids() too (both engines expose it), not just the committed
# state. release() is owner-guarded and identity-less, so an unconditional release here would
# clear the newer load's claim.
# window, so gate on loading_repo_ids() too, not just committed state. release() is
# owner-guarded and identity-less, so an unconditional release would clear the newer claim.
engine = get_active_diffusion_engine()
if not engine.loading_repo_ids() and not engine.is_loaded:
release(DIFFUSION)
@ -14623,20 +14610,17 @@ async def diffusion_generate_progress(current_subject: str = Depends(get_current
# ──────────────────────────────────────────────────────────────────────────
# OpenAI-compatible images API (POST /v1/images/generations)
#
# The inference router is mounted at both /api/inference and /v1, so this also
# answers /v1/images/generations for off-the-shelf OpenAI clients. It maps
# OpenAI's CreateImageRequest onto the in-process diffusion backend and returns
# an ImagesResponse. Studio's own Image tab uses the richer /images/generate
# route above; this is the spec-shaped surface, and the single error boundary
# mapping backend exceptions to OpenAI error envelopes (the global /v1 handler
# wraps HTTPException detail into the envelope).
# The inference router is mounted at both /api/inference and /v1, so this also answers
# /v1/images/generations for off-the-shelf OpenAI clients, mapping CreateImageRequest onto
# the in-process diffusion backend. Studio's Image tab uses the richer /images/generate
# above; this is the spec-shaped surface and the single error boundary mapping backend
# exceptions to OpenAI error envelopes (the global /v1 handler wraps HTTPException detail).
# ──────────────────────────────────────────────────────────────────────────
# Diffusion dims must land in [256, 2048] on a multiple of 16 (8x VAE downsample
# x 2x patch); the named OpenAI sizes (1024x1024, 1536x1024, 256x256, ...) all
# satisfy this. Mirrors DiffusionGenerateRequest's width/height bounds so both
# generate paths accept the same geometry.
# Diffusion dims must land in [256, 2048] on a multiple of 16 (8x VAE downsample x 2x patch);
# the named OpenAI sizes (1024x1024, 1536x1024, 256x256, ...) all satisfy this. Mirrors
# DiffusionGenerateRequest's width/height bounds so both generate paths accept the same geometry.
_IMAGE_SIZE_RE = _re.compile(r"^(\d{1,5})\s*x\s*(\d{1,5})$")
_IMAGE_DIM_MIN, _IMAGE_DIM_MAX = 256, 2048
# Sanitized 503 detail shared by the pre-check and the unload-race branch, so both
@ -14702,9 +14686,8 @@ async def openai_image_generations(
status_code = 400, detail = openai_error_body(str(exc), status = 400, param = "size")
)
# Use the active engine (diffusers OR native sd.cpp on a no-GPU host), the same
# accessor /images/generate uses, so a model loaded on the native engine isn't
# wrongly reported unloaded here.
# Use the active engine (diffusers OR native sd.cpp on a no-GPU host), the same accessor
# /images/generate uses, so a native-engine model isn't wrongly reported unloaded here.
backend = get_active_diffusion_engine()
status = backend.status()
if not status.get("loaded"):
@ -14712,9 +14695,8 @@ async def openai_image_generations(
# isn't loaded; the global handler turns this into the OpenAI envelope.
raise HTTPException(status_code = 503, detail = _NO_IMAGE_MODEL_MSG)
# An edit-only model (Qwen-Image-Edit, FLUX Kontext) needs an input image this API
# cannot supply; refuse up front with a 400 instead of letting the backend's
# ValueError surface as a sanitized 500.
# An edit-only model (Qwen-Image-Edit, FLUX Kontext) needs an input image this API can't
# supply; refuse up front with a 400 rather than let the backend ValueError become a 500.
workflows = status.get("workflows") or []
if workflows and "txt2img" not in workflows:
raise HTTPException(
@ -14741,10 +14723,9 @@ async def openai_image_generations(
batch_size = body.n,
)
except Exception as exc: # noqa: BLE001 (single boundary, sanitized envelope)
# A RuntimeError with the model now unloaded means it was evicted/unloaded
# between the readiness check above and the call (a transient race): 503.
# Every other failure (CUDA OOM, a diffusers shape/device error, both also
# RuntimeError) is a real 500, and its raw message must not reach the client.
# A RuntimeError with the model now unloaded means it was evicted between the readiness
# check and the call (a transient race): 503. Every other failure (CUDA OOM, a diffusers
# shape/device error) is a real 500 whose raw message must not reach the client.
if isinstance(exc, RuntimeError) and not backend.is_loaded:
raise HTTPException(status_code = 503, detail = _NO_IMAGE_MODEL_MSG)
logger.error("openai_images.generate_failed: %s", exc)
@ -14752,8 +14733,8 @@ async def openai_image_generations(
created = int(time.time())
want_b64 = body.response_format == "b64_json"
# Persist each image with its full recipe embedded, like /images/generate, so
# the response_format=url links resolve and the images show up in the gallery.
# Persist each image with its full recipe, like /images/generate, so response_format=url
# links resolve and the images show up in the gallery.
recipe = {
"prompt": body.prompt,
"negative_prompt": None,
@ -14762,14 +14743,14 @@ async def openai_image_generations(
"steps": steps,
"guidance": guidance,
# The batch shares one base seed, so restoring a batch_index>0 sibling needs the
# original batch_size to replay it (same as /images/generate); persist it.
# original batch_size to replay (same as /images/generate); persist it.
"batch_size": body.n,
"model": result.get("repo_id"),
"created_at": float(created),
}
# The diffusers batch shares one seed; the native sd.cpp batch uses a distinct seed
# per image (returned in ``seeds``), so record each image's own seed, like
# /images/generate, or a native batch_index>0 image shows the wrong seed.
# The diffusers batch shares one seed; the native sd.cpp batch uses a distinct seed per
# image (returned in ``seeds``), so record each image's own seed like /images/generate,
# or a native batch_index>0 image shows the wrong seed.
per_image_seeds = result.get("seeds")
def _persist() -> list[ImageGenerationData]:

119
studio/backend/routes/models.py
View file

@ -28,18 +28,15 @@ class CachedModelRepo(BaseModel):
repo_id: str
size_bytes: int
last_modified: Optional[float] = None
# "text-to-image" for cached diffusers image repos; response_model would silently
# drop the value the handler sets, letting image-only repos pass the chat picker's
# task gate.
# "text-to-image" for cached diffusers image repos; declared here or response_model
# drops it, letting image-only repos pass the chat picker's task gate.
task: Optional[str] = None
# True when the snapshot is incomplete (a cancelled/partial download left only some
# weights). The picker must not treat a partial base repo as a usable download, or an
# On Device click routes to a fresh multi-GB re-download instead of the complete GGUF.
# True when the snapshot is incomplete (cancelled/partial download): the picker must
# not treat it as usable, or an On Device click re-downloads the full GGUF.
partial: Optional[bool] = None
# True for a diffusion-tagged repo with NO top-level model_index.json: a single-file
# checkpoint that needs from_single_file + a filename. The task-scoped pickers must not
# offer it as a pipeline load (from_pretrained on it fails after the GPU handoff)
# unless the curated catalog carries its artifact.
# checkpoint needing from_single_file + a filename. Pickers must not offer it as a
# pipeline load (from_pretrained fails) unless the curated catalog carries its artifact.
single_file: Optional[bool] = None
@ -919,7 +916,7 @@ async def list_local_models(
try:
models = collect_local_models(models_root)
# Tag each model with its task so the Images picker can filter to diffusion
# (GGUF by architecture; local diffusers checkpoints by pipeline / family).
# (GGUF by architecture; local checkpoints by pipeline / family).
models = [m.model_copy(update = {"task": _local_model_task(m)}) for m in models]
return LocalModelListResponse(
@ -3153,16 +3150,14 @@ def _repo_gguf_last_modified(repo_info) -> float:
return latest
# GGUF general.architecture values that denote a diffusion (image) model;
# everything else is treated as a text model. Lets the Images picker show only
# image GGUFs in its On Device list.
# GGUF general.architecture values that denote a diffusion (image) model (everything
# else is text); lets the Images picker show only image GGUFs in its On Device list.
_DIFFUSION_GGUF_ARCHS = frozenset(
{
# ONLY the families the diffusion backend can actually assemble (see
# diffusion_families._FAMILIES). Other on-device diffusion archs (SD1/2/3,
# SDXL, PixArt, Lumina2, AuraFlow, Wan, HunyuanVideo, ...) would pass this
# Images-picker filter and then fail validate_load with a 400, so they are
# deliberately excluded until the backend supports them.
# ONLY the families the diffusion backend can assemble (see
# diffusion_families._FAMILIES). Other diffusion archs (SD1/2/3, SDXL,
# PixArt, Lumina2, AuraFlow, Wan, HunyuanVideo, ...) would pass this filter
# then 400 in validate_load, so they stay excluded until the backend supports them.
"flux", # flux.1
"flux2", # flux.2-klein
"qwen_image", # qwen-image
@ -3172,13 +3167,11 @@ _DIFFUSION_GGUF_ARCHS = frozenset(
}
)
# Known diffusion / image-video GGUF archs the backend can NOT assemble yet. These
# are the GGUF general.architecture values llama.cpp also has no architecture for,
# kept in sync with core.inference.llama_cpp.LlamaCppBackend._DIFFUSION_ARCHES
# (minus the loadable set above). Tagging them with a dedicated, non-loadable task
# keeps them OUT of the chat picker -- loading one as a chat model dies with
# "unknown model architecture" -- while also keeping them out of the Images picker
# (the task is not an IMAGE_GEN_TASK), where they would 400 in validate_load.
# Diffusion / image-video GGUF archs the backend can NOT assemble yet (llama.cpp also
# lacks an architecture for them); kept in sync with
# core.inference.llama_cpp.LlamaCppBackend._DIFFUSION_ARCHES minus the loadable set above.
# A dedicated non-loadable task keeps them out of the chat picker (they die with
# "unknown model architecture") and out of Images (not an IMAGE_GEN_TASK; would 400).
_UNSUPPORTED_DIFFUSION_GGUF_ARCHS = frozenset(
{
"sd1",
@ -3218,15 +3211,13 @@ def _arch_to_task(arch: Optional[str], name_hints: tuple[Optional[str], ...] = (
if a in _DIFFUSION_GGUF_ARCHS:
return "text-to-image"
if a in _VIDEO_GGUF_ARCHS:
# Advertise as loadable video only when a VideoFamily actually resolves. Some archs map
# straight from the arch (ltxv); others are ambiguous at the arch level -- bare "wan"
# covers both the single-DiT TI2V-5B (GGUF-loadable) and the dual-expert A14B MoE whose
# single file the loader refuses -- so when the bare arch does not resolve, fall back to
# the repo/file names like the loader's own detect_video_family does (each tried
# separately, since it matches on name segments not substrings), and surface only a
# non-MoE (loadable) match. Without a name we cannot disambiguate, so a bare-arch Wan
# GGUF (which the loader also cannot resolve) stays in the unsupported bucket rather than
# advertising a GGUF that would 400 on load.
# Advertise as loadable video only when a VideoFamily resolves. Some archs map
# straight from the arch (ltxv); bare "wan" is ambiguous -- it covers both the
# GGUF-loadable single-DiT TI2V-5B and the dual-expert A14B MoE the loader refuses --
# so when the bare arch doesn't resolve, fall back to repo/file names (each tried
# separately, matching name segments not substrings) like the loader's own
# detect_video_family, surfacing only a non-MoE match. Without a name we can't
# disambiguate, so a bare-arch Wan GGUF stays unsupported rather than 400ing on load.
from core.inference.video_families import detect_video_family
fam = detect_video_family("", override = a)
@ -3239,8 +3230,8 @@ def _arch_to_task(arch: Optional[str], name_hints: tuple[Optional[str], ...] = (
if fam is not None and not getattr(fam, "is_moe", False):
return _VIDEO_GEN_TASK
return _UNSUPPORTED_DIFFUSION_TASK
# A diffusion arch the backend can't assemble: hide it from chat (it would die
# in llama.cpp) without surfacing it in Images (it would 400 in validate_load).
# A diffusion arch the backend can't assemble: hide from chat (dies in llama.cpp)
# without surfacing in Images (would 400 in validate_load).
if a in _UNSUPPORTED_DIFFUSION_GGUF_ARCHS:
return _UNSUPPORTED_DIFFUSION_TASK
return "text-generation"
@ -3289,11 +3280,10 @@ def _local_model_task(model: "LocalModelInfo") -> Optional[str]:
pass
return None
if _local_is_diffusers(model):
# A local diffusers pipeline can be a VIDEO family (LTX / Wan / Hunyuan), not just an
# image one. Tag it text-to-video so it surfaces in the Video On-Device picker instead
# of the Images picker (where the image loader would reject it), mirroring the
# cached-repo _cached_repo_task. Gated on _local_is_diffusers, so only a real loadable
# pipeline dir (model_index.json) or a name-matched checkpoint reaches this check.
# A local diffusers pipeline can be a VIDEO family (LTX / Wan / Hunyuan), not just
# image. Tag it text-to-video so it surfaces in the Video On-Device picker instead of
# Images (which would reject it), mirroring _cached_repo_task. Gated on
# _local_is_diffusers, so only a real pipeline dir or name-matched checkpoint reaches here.
try:
from core.inference.video import _is_trusted_video_repo
from core.inference.video_families import detect_video_family
@ -3458,9 +3448,8 @@ def _cached_repo_task(repo_info) -> Optional[str]:
from core.inference.video import _is_trusted_video_repo
from core.inference.video_families import detect_video_family
# Both gates: a detected video family (so unsloth image repos don't
# match) AND the load path's own trust rule (so an untrusted video repo
# isn't advertised as loadable).
# Both gates: a detected video family (so image repos don't match) AND the
# load path's trust rule (so an untrusted video repo isn't advertised as loadable).
if detect_video_family(repo_id) is not None and _is_trusted_video_repo(repo_id):
return _VIDEO_GEN_TASK
except Exception:
@ -3511,10 +3500,9 @@ async def list_cached_models(
key = repo_id.lower()
existing = seen_lower.get(key)
is_partial = _cached_repo_partial(repo_id, Path(repo_info.repo_path))
# Prefer the most COMPLETE snapshot, then the largest. The picker drops partial
# rows, so a partial copy in one cache root must not shadow a smaller COMPLETE
# copy in another (that would make a usable model vanish from On Device).
# Completeness wins outright; size only breaks ties among equal completeness.
# Prefer the most COMPLETE snapshot, then largest. The picker drops partial
# rows, so a partial copy in one cache root must not shadow a smaller complete
# copy in another (size only breaks ties among equal completeness).
if existing is None or (not is_partial, total_size) > (
not bool(existing.get("partial")),
existing["size_bytes"],
@ -3527,8 +3515,8 @@ async def list_cached_models(
if is_partial:
row["partial"] = True
# Flag diffusion repos with no pipeline index: loadable only via
# from_single_file with a checkpoint filename, so the pickers must
# not offer them as pipeline loads unless the catalog carries them.
# from_single_file, so pickers must not offer them as pipeline
# loads unless the catalog carries them.
if row["task"] is not None and not _repo_has_pipeline_index(repo_info):
row["single_file"] = True
# Keep the newest timestamp across duplicate caches;
@ -3610,13 +3598,12 @@ async def delete_cached_model(
except Exception:
pass
# Also refuse if the diffusion (Images) backend has this repo loaded; its
# delete guard is otherwise chat-only, so its GGUF could be removed from
# under a live pipeline. Repo-level match, like the chat guards above.
# Also refuse if the Images backend has this repo loaded (guards above are
# chat-only), or its GGUF could be removed from under a live pipeline.
try:
# The ACTIVE engine (diffusers or native sd_cpp): on a native selection the
# diffusers singleton reports unloaded while sd-cli still generates from the
# cached GGUF, so checking it alone would let the files be deleted mid-use.
# cached GGUF, so checking it alone would let files be deleted mid-use.
from core.inference.diffusion_engine_router import get_active_diffusion_engine
engine = get_active_diffusion_engine()
@ -3628,11 +3615,11 @@ async def delete_cached_model(
status_code = 400,
detail = "Unload the model before deleting",
)
# The native sd.cpp one-shot engine re-reads its companion VAE / text-encoder files
# from the HF cache on every generation, so deleting a companion repo (e.g.
# comfyanonymous/flux_text_encoders) while a native GGUF is loaded would brick the
# next generation. status().repo_id only covers the main GGUF, so also refuse the
# committed companion repos the loaded engine reads from disk.
# The native sd.cpp engine re-reads companion VAE / text-encoder files from the HF
# cache every generation, so deleting a companion repo (e.g.
# comfyanonymous/flux_text_encoders) while a native GGUF is loaded bricks the next
# generation. status().repo_id covers only the main GGUF, so also refuse the
# committed companion repos the engine reads from disk.
for lid in getattr(engine, "loaded_repo_ids", tuple)():
if _loaded_id_matches_repo(str(lid).lower(), repo_id):
raise HTTPException(
@ -3640,8 +3627,8 @@ async def delete_cached_model(
detail = "Unload the model before deleting",
)
# Also refuse while a background image load is DOWNLOADING this repo (or its
# companion base): status().loaded is still False in that window, but deleting
# would remove blobs from under the in-flight download/assembly.
# companion base): status().loaded is still False then, but deleting would
# remove blobs from under the in-flight download/assembly.
loading_ids = getattr(engine, "loading_repo_ids", tuple)()
for lid in loading_ids:
lid = str(lid).lower()
@ -3656,9 +3643,9 @@ async def delete_cached_model(
pass
# And refuse if the Video backend has this repo loaded or is downloading it: cached non-GGUF
# video repos now surface in the Video On-Device picker with the normal delete action, but the
# guards above only cover chat + the Images engine, so without this a loaded/loading Wan / LTX /
# Hunyuan pipeline could have its HF snapshot removed from under it. Mirror the Images guard.
# video repos now surface in the Video picker with a delete action, but the guards above cover
# only chat + Images, so without this a loaded/loading Wan / LTX / Hunyuan pipeline could lose
# its HF snapshot from under it. Mirror the Images guard.
try:
from core.inference.video import get_video_backend
@ -3672,8 +3659,8 @@ async def delete_cached_model(
detail = "Unload the model before deleting",
)
# Also refuse while a background VIDEO load is DOWNLOADING this repo (or its companion
# base): status().loaded is still False in that window, but deleting would remove blobs
# from under the in-flight download/assembly -- same as the Images guard above.
# base): status().loaded is still False then, but deleting would remove blobs from under
# the in-flight download/assembly -- same as the Images guard above.
for lid in getattr(video_backend, "loading_repo_ids", tuple)():
lid = str(lid).lower()
if _loaded_id_matches_repo(lid, repo_id):

315
studio/backend/routes/training.py
View file

@ -212,9 +212,9 @@ 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.
# A diffusion (SDXL) LoRA job runs in its own subprocess on the same GPU, so an LLM
# start must refuse while one is active, or 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 = "",
@ -436,21 +436,20 @@ async def start_training(
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.
# A resident or in-flight Images pipeline also holds GPU memory the run needs
# and 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 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.
# 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(
@ -462,12 +461,11 @@ async def start_training(
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.
# A resident or in-flight Video pipeline holds GPU memory the run needs too, and
# loads under the VIDEO arbiter owner the diffusion teardown above never touches.
# Tear it down the same way (unload no-ops when nothing is loaded, preempts an
# in-flight load) and release VIDEO, so a resident video session 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
@ -524,12 +522,10 @@ async def start_training(
logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e)
# The hook runs only once start guards pass -> VRAM freed iff training starts.
# Offloaded to a worker thread: the hook's diffusion/video unload() waits on the
# engines' generation locks until an in-flight denoise step reaches its cancel
# callback (and the export subprocess teardown can take seconds), which would
# otherwise block the event loop and freeze every concurrent status/cancel/UI
# request -- the same reason start_diffusion_training runs
# _free_gpu_for_diffusion_training via asyncio.to_thread. Overlapping starts are
# Offloaded to a worker thread: the hook's diffusion/video unload() waits on the engines'
# generation locks until an in-flight denoise step hits its cancel callback (and the
# export subprocess teardown can take seconds), which would otherwise block the event
# loop and freeze every concurrent status/cancel/UI request. Overlapping starts are
# serialized by the backend's own start-in-progress guard.
success = await asyncio.to_thread(
backend.start_training,
@ -1125,10 +1121,10 @@ async def stream_training_progress(
# ── 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).
# A separate, lightweight job path from the LLM 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:
@ -1164,10 +1160,9 @@ def _free_gpu_for_diffusion_training() -> None:
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.
# 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")
@ -1177,9 +1172,9 @@ def _free_gpu_for_diffusion_training() -> None:
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.
# A resident Video pipeline loads under the VIDEO arbiter owner the Images teardown
# above doesn't free; unload it too (no-op when nothing is loaded) and release VIDEO
# so a resident video session can't OOM the diffusion trainer.
from core.inference import gpu_arbiter
from core.inference.video import get_video_backend
@ -1192,9 +1187,9 @@ def _free_gpu_for_diffusion_training() -> None:
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.
# The SDXL trainer's footprint can't be cheaply sized against a resident chat model,
# so free chat unconditionally (like the LLM path does for an in-flight 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")
@ -1273,10 +1268,9 @@ async def start_diffusion_training(
"""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.
# Under API-key auth, refuse to start training while a request is in flight:
# _free_gpu_for_diffusion_training() below unloads the chat backends, killing the stream.
# Mirrors start_training so a diffusion start can't silently drop an active API request.
if via_api_key is True:
from core.inference.llama_keepwarm import other_inference_request_count
if (
@ -1292,8 +1286,8 @@ async def start_diffusion_training(
),
)
# 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.
# 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(
@ -1308,9 +1302,9 @@ async def start_diffusion_training(
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.
# 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_output_dir
@ -1319,9 +1313,9 @@ async def start_diffusion_training(
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.
# Validate the config BEFORE freeing resident GPU workloads, so a start then refused (bad
# numbers, non-SDXL base) never tears down the user's chat/Images model. service.start()
# re-runs this cheaply before spawn.
from core.training.diffusion_lora_trainer import _config_from_dict
try:
@ -1329,11 +1323,11 @@ async def start_diffusion_training(
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.
# Preflight the requested DiT precision BEFORE freeing GPU residents: the trainer's own
# checks (bf16-capable GPU required; explicit int8 needs a functional torchao) fire only in
# the child, AFTER _free_gpu_for_diffusion_training() evicted the user's model. Fail fast
# (400) so a pre-Ampere GPU (T4 / V100 / RTX 20xx) or stub-torchao host never tears down
# residents for a run that cannot start.
from core.training.diffusion_train_common import training_precision_preflight_error
_precision_reason = training_precision_preflight_error(
@ -1343,8 +1337,8 @@ async def start_diffusion_training(
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_pretrained), so an untrusted/typoed base 400s BEFORE freeing GPU residents rather
# than tearing down the user's model and failing in the child.
from core.training.diffusion_train_common import _assert_trusted_base_model
try:
@ -1352,19 +1346,18 @@ async def start_diffusion_training(
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.
# 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 model, and
# never surfaces as a confusing mid-load 401. Offloaded to a worker thread: it does a blocking
# urlopen HEAD (5s timeout) to HF, which would otherwise stall the event loop and every
# concurrent status/progress/cancel request (as the filesystem preflight below also 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.
# Preflight the dataset too: a missing/empty/uncaptionable data_dir otherwise fails inside
# the spawned trainer AFTER the user's model was evicted. Same discovery the trainer runs,
# so the two cannot disagree.
from core.training import diffusion_train_common as _dtc
try:
@ -1373,32 +1366,29 @@ async def start_diffusion_training(
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.
# Decode-probe every image now (cheap PIL header check) so a corrupt/zero-byte upload
# 400s BEFORE _free_gpu_for_diffusion_training() tears down the user's models, rather
# than 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.
# Reserve the training slot BEFORE freeing residents: is_active() otherwise flips true only at
# service.start(), after the free, so a concurrent /images/load or /video/load would pass its
# training guard during the free-then-spawn window and double-allocate VRAM. reserve() is a
# compare-and-set: a second overlapping /diffusion/start raises RuntimeError (-> 409) before
# freeing anything, so two starts never both tear down residents. unreserve() runs in the
# finally ONLY when THIS request reserved, so a rejected second request can't clear the 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.
# Free resident GPU workloads (export / Images pipeline / chat) before the trainer loads
# its own pipeline. Offload the blocking teardown (engine unload waits on generation
# locks; export subprocess join can take seconds) to a worker thread so the event loop
# stays free for concurrent status/progress/cancel requests.
await asyncio.to_thread(_free_gpu_for_diffusion_training)
job_id = service.start(config)
except ValueError as e:
@ -1417,9 +1407,9 @@ async def start_diffusion_training(
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.
# On success the now-live proc keeps is_active() true; on failure this clears the
# reservation so training isn't left permanently "active". Only the request that reserved
# clears it, so a rejected overlapping start doesn't drop the winner's claim.
if reserved:
service.unreserve()
return DiffusionTrainingStartResponse(job_id = job_id, status = "running")
@ -1466,9 +1456,9 @@ async def list_diffusion_training_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.
# list_diffusion_runs already skips non-dict / missing-id records, but 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:
@ -1485,17 +1475,16 @@ async def get_diffusion_training_run(
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
# A valid-JSON file that is not an object (a truncated / hand-edited [] record) makes
# 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.
# -- and 500 the endpoint. Treat any non-dict record as absent, like the list route.
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.
# A malformed on-disk record (hand-edited / older shape) reads as absent rather than
# 500 the endpoint, like the list route skips bad records.
raise HTTPException(status_code = 404, detail = "No such training run.")
@ -1535,8 +1524,8 @@ def _resolve_dataset_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.
# metadata row and makes the trainer skip the image, so counting it would over-report
# caption_count and mislabel an uncaptioned dataset as captioned.
meta_captions = _load_metadata_captions(folder)
images = captions = 0
for f in folder.iterdir():
@ -1562,8 +1551,8 @@ async def diffusion_training_info(current_subject: str = Depends(get_current_sub
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.
# Skip hidden dirs: 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(".")
)
@ -1635,17 +1624,17 @@ async def upload_diffusion_dataset(
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.
# Validate every filename up front so a valid image ahead of a bad one isn't left 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.
# Normalise to a safe basename. Path.name doesn't split on a backslash on POSIX, so a
# Windows client sending a backslash path in the multipart filename would 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 still holding
# ".." here would list an image the 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:
@ -1654,15 +1643,13 @@ async def upload_diffusion_dataset(
status_code = 400,
detail = f"Unsupported file '{f.filename}'. Allowed: {exts}",
)
# Reject an EXACT duplicate name within THIS batch (two cat.png dragged from
# different folders, or an API client repeating a part). The same-name exemption
# below exists for SEPARATE repeat uploads, where re-sending a name is a
# deliberate overwrite of the file on disk; inside one batch the two parts are
# distinct files staged to the same destination on EVERY filesystem, so the later
# tmp.replace(dest) in the commit loop would silently discard the earlier one
# while `uploaded` still counts both. Exact match only: a case VARIANT pair
# (pic.png vs Pic.png) stays exempt like the stem guard documents -- one file /
# an overwrite on case-insensitive filesystems, two files on Linux.
# Reject an EXACT duplicate name within THIS batch (two cat.png from different folders,
# or an API client repeating a part). The same-name exemption below is for SEPARATE
# repeat uploads, a deliberate overwrite of the file on disk; inside one batch the two
# parts are distinct files staged to the same destination on EVERY filesystem, so the
# later tmp.replace(dest) would silently discard the earlier one while `uploaded` counts
# both. Exact match only: a case VARIANT pair (pic.png vs Pic.png) stays exempt per the
# stem guard -- one file / overwrite on case-insensitive filesystems, two on Linux.
fname_cf = filename.casefold()
if filename in names:
raise HTTPException(
@ -1673,25 +1660,22 @@ async def upload_diffusion_dataset(
"uploading."
),
)
# 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).
# Reject a second IMAGE sharing this stem but differing by extension (sample.png vs
# sample.jpg): both resolve to the same <stem>.txt sidecar (the kohya/diffusers
# convention the reader, editor, and delete paths use), so keeping both would silently
# share -- and corrupt -- one caption. Check files already on disk (uploads accumulate)
# and earlier images in THIS batch. Re-uploading the exact same name (stem AND extension)
# stays an overwrite; caption/text files are exempt (sample.txt for sample.png is fine).
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. A same-name case variant is exempt ONLY when its stem also differs
# in case (sample.png vs Sample.png): one file / an overwrite on case-insensitive
# filesystems, and on Linux the two files write SEPARATE sidecars (sample.txt vs
# Sample.txt). An EXTENSION-case variant (cat.PNG vs cat.png) has exactly equal
# stems, so on Linux both files land and both resolve to ONE cat.txt -- the very
# collision this guard exists for -- and is rejected like any other stem clash.
# Compare stems (and the same-name guard) case-insensitively: on case-insensitive
# filesystems (Windows/macOS) two images whose stems differ only by case (sample.png
# vs Sample.jpg) resolve to the SAME <stem>.txt sidecar, so a case-sensitive check
# would let both share -- and corrupt -- one caption. A same-name case variant is
# exempt ONLY when its stem also differs in case (sample.png vs Sample.png): one file /
# overwrite on case-insensitive filesystems, SEPARATE sidecars on Linux. An
# EXTENSION-case variant (cat.PNG vs cat.png) has equal stems, so on Linux both land
# and resolve to ONE cat.txt -- the collision this guard exists for -- and is rejected.
stem_cf = stem.casefold()
def _shares_sidecar(other_name: str) -> bool:
@ -1703,8 +1687,7 @@ async def upload_diffusion_dataset(
):
return False
# A casefold-equal full name is exempt unless the stems match EXACTLY
# (extension-case variants collide on one sidecar on case-sensitive
# filesystems).
# (extension-case variants collide on one sidecar on case-sensitive FS).
return other.stem == stem or other_name.casefold() != fname_cf
clash = next(
@ -1723,10 +1706,9 @@ async def upload_diffusion_dataset(
),
)
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).
# Stage each file to a temp name and move it into place only once the whole batch is written,
# so a mid-batch failure (size limit, disk error, disconnect) leaves the dataset untouched --
# including any pre-existing same-name file a direct write would have truncated.
staged: list[tuple[Path, Path]] = [] # (temp, final)
committed = False
try:
@ -1773,7 +1755,7 @@ async def upload_diffusion_dataset(
# ── 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).
# trainer's image discovery (both scan only top-level files).
_THUMBS_DIRNAME = ".thumbs"
_MAX_CAPTION_CHARS = 2000
@ -1855,9 +1837,8 @@ def _image_record(
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.
# Basename first, then the relative path as written in the jsonl (as_posix so a Windows
# backslash path still matches forward-slash keys) -- discover_image_caption_pairs's order.
meta = meta_captions.get(image_path.name)
if meta is None:
try:
@ -1932,10 +1913,9 @@ async def get_diffusion_dataset_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.
# Key on the full filename (stem + extension), not the stem: two images sharing a stem
# but differing by extension (sample.png / sample.jpg) would otherwise collide on one
# cache file, and an mtime-newer cache for the first would be served for the second.
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:
@ -1983,11 +1963,10 @@ async def set_diffusion_dataset_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.
# Blank must actually clear. Unlinking alone would resurface this image's metadata.jsonl
# / captions.jsonl caption (the fallback), so when one exists write an EMPTY sidecar
# instead: both the reader and the trainer's discovery treat an existing sidecar as
# authoritative even when empty, a tombstone. No metadata caption -> plain cleanup.
meta = _load_metadata_captions(folder)
try:
rel = image_path.relative_to(folder).as_posix()
@ -2021,9 +2000,8 @@ async def delete_diffusion_dataset_image(
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).
# Thumbs are keyed on the full filename (stem + extension), so match that here too;
# a stem-only glob would strand this image's thumbs or delete a same-stem sibling's.
for t in thumbs_dir.glob(f"{image_path.name}_*.jpg"):
t.unlink(missing_ok = True)
return {"deleted": image_path.name}
@ -2034,7 +2012,7 @@ async def delete_diffusion_dataset_image(
# 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.
# captions live in a *.jsonl (file_name/text) not a standard metadata.jsonl.
_DATASET_EXAMPLES: list[dict] = [
{
"id": "dreambooth-dog",
@ -2091,8 +2069,8 @@ _DATASET_EXAMPLES: list[dict] = [
),
"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.
# The metadata columns are species names / boilerplate alt-text, not 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,
@ -2285,12 +2263,12 @@ async def import_diffusion_dataset_example(
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.
# after the whole import succeeds. A partial materialize (a transient fetch/copy
# error after 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 as complete on retry (imported=0) and strand a truncated dataset (there is
# no dataset-delete flow). Staged as a hidden same-filesystem sibling so promotion is
# an atomic rename.
staging = Path(tempfile.mkdtemp(dir = folder.parent, prefix = f".{folder.name}.import-"))
try:
try:
@ -2311,11 +2289,10 @@ async def import_diffusion_dataset_example(
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
# not atomic: a hard process death (SIGKILL / OOM / power loss) mid-loop would
# leave SOME images, which the image_count>0 idempotency check above would accept
# as complete on retry. The folder was created empty here (runs only when it holds
# no images), so a single same-filesystem 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: