diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 382266db52..3471b67df1 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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]: diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 6cf049ff29..cd45378110 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -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): diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 7a29e75ab2..1ef804afc2 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -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 .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 .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 .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 .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: