From eee191d30e9495b96d4ea3d827189efd91513ac8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 27 Jul 2026 12:01:44 +0000 Subject: [PATCH] Tighten diffusion comments (second pass) Collapse the remaining multi-line comment blocks in the video page, training routes and service, sd.cpp server and installer, memory and speed planners, and the shared request models. Comments only, no code or behaviour changes. --- .../core/inference/diffusion_memory.py | 46 +-- .../backend/core/inference/diffusion_speed.py | 79 ++-- .../core/inference/diffusion_te_prequant.py | 45 +-- .../backend/core/inference/sd_cpp_server.py | 39 +- .../backend/core/inference/video_families.py | 70 ++-- .../core/training/diffusion_lora_trainer.py | 43 +-- .../core/training/diffusion_train_extras.py | 16 +- .../training/diffusion_training_service.py | 52 +-- studio/backend/models/inference.py | 71 ++-- studio/backend/routes/training.py | 324 +++++----------- .../src/features/video/video-page.tsx | 348 ++++++------------ studio/install_sd_cpp_prebuilt.py | 70 +--- 12 files changed, 355 insertions(+), 848 deletions(-) diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index 40f052d9e3..e167a7e8e9 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -33,20 +33,16 @@ MEMORY_MODES = ( MEMORY_MODE_LOW_VRAM, ) -# ── offload policies (what the loader does) ────────────────────────── # none -- all weights resident (fastest; fits only with room). # model -- enable_model_cpu_offload(): one top-level module on the GPU at a time. -# group -- apply_group_offloading() on the transformer: stream a few blocks at a time with a -# prefetch stream (lowest practical VRAM for the dominant module). -# sequential -- enable_sequential_cpu_offload(): submodule-level (broken for GGUF on diffusers -# 0.38, kept as an explicit escape hatch). +# group -- apply_group_offloading() on the transformer: stream a few blocks at a time with a prefetch stream (lowest practical VRAM for the dominant module). +# sequential -- enable_sequential_cpu_offload(): submodule-level (broken for GGUF on diffusers 0.38, kept as an explicit escape hatch). OFFLOAD_NONE = "none" OFFLOAD_MODEL = "model" OFFLOAD_GROUP = "group" OFFLOAD_SEQUENTIAL = "sequential" -# Transformer blocks resident per group under group offloading: fewer = lower VRAM, more -# host-to-device traffic. +# Transformer blocks resident per group under group offloading: fewer = lower VRAM, more host-to-device traffic. DEFAULT_GROUP_BLOCKS = 1 DEFAULT_IMAGE_WIDTH = 1024 @@ -195,8 +191,7 @@ def _cuda_memory(backend: str) -> tuple[Optional[int], Optional[int], str]: free, total = torch.cuda.mem_get_info() kind = "discrete_vram" try: - # Query the CURRENT device (mem_get_info reports it); hardcoding 0 would inspect the wrong GPU - # and misclassify discrete vs unified. + # Query the CURRENT device (mem_get_info reports it); hardcoding 0 would inspect the wrong GPU and misclassify discrete vs unified. props = torch.cuda.get_device_properties(torch.cuda.current_device()) if bool(getattr(props, "integrated", False) or getattr(props, "is_integrated", False)): kind = "unified_memory" # e.g. Jetson / integrated SoC @@ -403,13 +398,11 @@ def plan_diffusion_memory( } def _group_fits() -> bool: - # Group offload only helps if the resident companions fit; a too-big text encoder needs - # whole-module offload. + # Group offload only helps if the resident companions fit; a too-big text encoder needs whole-module offload. return group_floor is not None and budget is not None and group_floor <= budget if not can_offload or device_memory.is_unified: - # MPS / CPU can't stream to a separate device; on unified memory offload just shuffles bytes - # within the same pool. + # MPS / CPU cannot stream to a separate device; on unified memory offload just shuffles bytes within the same pool. policy = OFFLOAD_NONE if device_memory.is_unified: reasons.append("unified/system memory: CPU offload frees no device memory") @@ -442,8 +435,7 @@ def plan_diffusion_memory( policy = OFFLOAD_MODEL reasons.append("companions exceed budget; whole-module offload of every component") - # The legacy cpu_offload flag applies only when no memory_mode was supplied, so an explicit - # `fast` request stays resident even with the old flag on. + # The legacy cpu_offload flag applies only when no memory_mode was supplied, so an explicit `fast` request stays resident even with the old flag on. if ( explicit_offload and normalize_memory_mode(requested_mode) is None @@ -454,10 +446,8 @@ def plan_diffusion_memory( policy = OFFLOAD_MODEL reasons.append("explicit cpu_offload overrides resident placement") - # VAE savers cap the high-res decode spike. Slicing (one image at a time) is EXACT, so enable it - # on any offload tier / non-discrete backend. Tiling (spatial chunks) is only bit-identical for - # a single tile (<=1MP), so restrict it to the lowest tiers or no spare device pool. Group - # offload keeps the VAE resident for an exact full-image decode; on a roomy GPU both stay off. + # VAE savers cap the high-res decode spike. Slicing (one image at a time) is EXACT, so enable it on any offload tier / non-discrete backend. Tiling (spatial chunks) is only bit-identical for a single tile (<=1MP), so restrict it to the lowest tiers or no spare device pool. + # Group offload keeps the VAE resident for an exact full-image decode; on a roomy GPU both stay off. any_offload = policy != OFFLOAD_NONE or device_memory.backend in ("mps", "cpu") tile = policy in (OFFLOAD_MODEL, OFFLOAD_SEQUENTIAL) or device_memory.backend in ("mps", "cpu") return MemoryPlan( @@ -494,8 +484,7 @@ def apply_memory_plan( _enable_vae_saver(pipe, "enable_vae_slicing", "enable_slicing", logger) def _fallback_to_model_offload() -> None: - # The GROUP plan set vae_tiling=False (VAE stays resident). Dropping to whole-module offload is - # the low-VRAM case where the decode spike can OOM, so turn tiling on now. + # The GROUP plan set vae_tiling=False (VAE stays resident). Dropping to whole-module offload is the low-VRAM case where the decode spike can OOM, so turn tiling on now. nonlocal tiling_engaged pipe.enable_model_cpu_offload(device = device) if not tiling_engaged: @@ -554,8 +543,7 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool: import torch from diffusers.hooks import apply_group_offloading - # A dual-DiT pipeline (Ideogram 4) carries a second denoiser as large as the first; leaving it - # resident defeats this tier. Stream every DiT, keep only smaller companions. + # A dual-DiT pipeline (Ideogram 4) carries a second denoiser as large as the first; leaving it resident defeats this tier. Stream every DiT, keep only smaller companions. streamed: dict[str, Any] = {"transformer": transformer} for extra in ("transformer_2", "unconditional_transformer"): module = getattr(pipe, extra, None) @@ -571,18 +559,14 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool: "num_blocks_per_group": DEFAULT_GROUP_BLOCKS, "use_stream": use_stream, } - # On the CUDA stream path, overlap each block's H2D copy with compute: non_blocking issues it - # async, record_stream defers the free until the copy's stream is done. Lossless, and gated on - # the signature so older diffusers still works. + # On the CUDA stream path, overlap each block H2D copy with compute: non_blocking issues it async, record_stream defers the free until the copy stream is done. Lossless, and gated on the signature so older diffusers still works. if use_stream: _params = inspect.signature(apply_group_offloading).parameters if "non_blocking" in _params: gkwargs["non_blocking"] = True if "record_stream" in _params: gkwargs["record_stream"] = True - # Place the smaller components resident BEFORE attaching the transformer's group-offload hooks: - # if a companion .to() OOMs we return False with NO hooks installed, so the caller's whole-module - # fallback works (diffusers REJECTS enable_model_cpu_offload once group hooks exist). + # Place the smaller components resident BEFORE attaching the transformer group-offload hooks: if a companion .to() OOMs we return False with NO hooks installed, so the caller whole-module fallback works (diffusers REJECTS enable_model_cpu_offload once group hooks exist). for name, comp in getattr(pipe, "components", {}).items(): if name in streamed: continue @@ -594,9 +578,7 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool: return True except Exception as exc: # noqa: BLE001 — fall back to whole-module offload if installed: - # An earlier streamed module already has hooks but a later one failed: the pipe is in a PARTIAL - # group-offload state that enable_model_cpu_offload rejects, so propagate the real failure (e.g. - # the OOM) instead of a misleading hook error. The "no hooks installed" cases fall back cleanly. + # An earlier streamed module already has hooks but a later one failed: the pipe is in a PARTIAL group-offload state that enable_model_cpu_offload rejects, so propagate the real failure (e.g. the OOM) instead of a misleading hook error. The "no hooks installed" cases fall back cleanly. if logger is not None: logger.warning( "diffusion.memory: group offload failed after installing hooks on %d " diff --git a/studio/backend/core/inference/diffusion_speed.py b/studio/backend/core/inference/diffusion_speed.py index bb0714a56b..678b4328c9 100644 --- a/studio/backend/core/inference/diffusion_speed.py +++ b/studio/backend/core/inference/diffusion_speed.py @@ -186,8 +186,7 @@ def apply_speed_optims( "compiled_vae_decode": False, } mode = normalize_speed_mode(speed_mode) - # TF32 (max) and cudnn.benchmark (any non-off CUDA load) are process-global; the caller - # snapshots/restores them so a later `off` load never inherits them. + # TF32 (max) and cudnn.benchmark (any non-off CUDA load) are process-global; the caller snapshots/restores them so a later `off` load never inherits them. if mode == SPEED_OFF: return applied @@ -197,32 +196,25 @@ def apply_speed_optims( # Lossless: a channels-last VAE speeds up its convs with no numeric change. applied["channels_last"] = _vae_channels_last(pipe, logger) - # Near-lossless: cuDNN autotunes the fixed-shape VAE convs (CUDA only). It may pick a different - # conv algorithm, so it is a "default"-tier (not bit-identical) win. + # Near-lossless: cuDNN autotunes the fixed-shape VAE convs (CUDA only). It may pick a different conv algorithm, so it is a "default"-tier (not bit-identical) win. if on_cuda: applied["cudnn_benchmark"] = _enable_cudnn_benchmark(logger) - # Consumer-only: fp16 GEMMs accumulate in fp16 (~2x on GeForce-class parts; datacenter parts - # gain nothing). bf16 loads measured bit-identical with the flag on (36/36 same-seed cases), so - # on the neutral tiers it engages only when compute dtype is NOT fp16. fp16 pipelines showed - # same-seed drift (mean 2-5%), so fp16 compute gets it only under ``max``. Guarded by - # _FP16_ACCUM_DENY and the UNSLOTH_DISABLE_FP16_ACCUM kill switch. + # Consumer-only: fp16 GEMMs accumulate in fp16 (~2x on GeForce-class parts; datacenter parts gain nothing). bf16 loads measured bit-identical with the flag on (36/36 same-seed cases), so on the neutral tiers it engages only when compute dtype is NOT fp16. + # fp16 pipelines showed same-seed drift (mean 2-5%), so fp16 compute gets it only under ``max``. Guarded by _FP16_ACCUM_DENY and the UNSLOTH_DISABLE_FP16_ACCUM kill switch. if on_cuda: applied["fp16_accum"] = _enable_fp16_accumulation( family, logger, dtype = getattr(target, "dtype", None), speed_mode = mode ) # --- the compile lever, per tier --- - # default = LIGHT: GGUF compiles ONLY the dequant op chain (cheap, VRAM-free, - # resolution-invariant); dense has no dequant, so falls back to the regional block compile. - # max = FULL: regional max-autotune compile of the repeated block (subsumes the dequant fusion). - # eager = no compile. + # default = LIGHT: GGUF compiles ONLY the dequant op chain (cheap, VRAM-free, resolution-invariant); dense has no dequant, so falls back to the regional block compile. + # max = FULL: regional max-autotune compile of the repeated block (subsumes the dequant fusion). eager = no compile. if mode == SPEED_DEFAULT: if is_gguf and on_cuda and family_allows_compile: applied["compiled_dequant"] = gguf_compile.install_compiled_dequant(logger) elif compile_eligible(target, is_gguf = is_gguf, family = family): - # A U-Net (SDXL) fuses QKV BEFORE its whole-module compile: 36.3 vs 39.3 ms/step (LPIPS 0.033). - # DiTs were neutral under the regional compile, so they keep the fuse on the max tier only. + # A U-Net (SDXL) fuses QKV BEFORE its whole-module compile: 36.3 vs 39.3 ms/step (LPIPS 0.033). DiTs were neutral under the regional compile, so they keep the fuse on the max tier only. if _denoiser_unet(pipe) is not None: applied["fused_qkv"] = _fuse_qkv(pipe, logger) applied["compiled"] = _compile_repeated_blocks( @@ -241,9 +233,7 @@ def apply_speed_optims( offload_active = offload_active, ) - # A compiled U-Net family also compiles the VAE decode (a real share at SDXL's step rate: - # 4.98 to 4.25 s over 4 images, LPIPS unchanged). DiT families skip it. dynamic=True keeps it - # resolution-robust; fullgraph=False tolerates offload hooks. + # A compiled U-Net family also compiles the VAE decode (a real share at SDXL step rate: 4.98 to 4.25 s over 4 images, LPIPS unchanged). DiT families skip it. dynamic=True keeps it resolution-robust; fullgraph=False tolerates offload hooks. if applied["compiled"] and _denoiser_unet(pipe) is not None: applied["compiled_vae_decode"] = _compile_vae_decode(pipe, logger) @@ -269,12 +259,9 @@ def _vae_channels_last(pipe: Any, logger: Any) -> bool: return False -# U-Net denoisers ship no ``_repeated_blocks`` (heterogeneous block mix), so the regional compile -# can't reach them; these classes get a WHOLE-module STATIC ``torch.compile`` instead. On SDXL -# (B200, 30 steps / 1024px): 26.9 ms/step vs the 45.9 ms/step bit-exact reference, 1.61x -# end-to-end at LPIPS 0.034, while dynamic=True compiles 5x slower for less win. Static shapes -# mean a recompile per new (height, width, batch); the Mega-cache bundle carries each across -# restarts. +# U-Net denoisers ship no ``_repeated_blocks`` (heterogeneous block mix), so the regional compile cannot reach them; these classes get a WHOLE-module STATIC ``torch.compile`` instead. +# On SDXL (B200, 30 steps / 1024px): 26.9 ms/step vs the 45.9 ms/step bit-exact reference, 1.61x end-to-end at LPIPS 0.034, while dynamic=True compiles 5x slower for less win. +# Static shapes mean a recompile per new (height, width, batch); the Mega-cache bundle carries each across restarts. _UNET_WHOLE_COMPILE: frozenset[str] = frozenset({"UNet2DConditionModel"}) @@ -325,14 +312,9 @@ def _compile_repeated_blocks( unet = _denoiser_unet(pipe) if not dits else None if not dits and unet is None: return False - # default: dynamic=True -- fast cold start, no recompile on resolution change. max: - # mode="max-autotune-no-cudagraphs" + dynamic=False -- Triton autotuning for a few % more, at a - # longer compile and a recompile per resolution. CUDA-graph modes are NOT used: they crash on - # the regional block (its static output buffer is overwritten across steps). - # - # fullgraph drops to False under a step cache OR offloading: both insert an - # ``@torch.compiler.disable``d function, i.e. a graph break fullgraph=True rejects. The break is - # cheap; the rest still compiles. + # default: dynamic=True -- fast cold start, no recompile on resolution change. max: mode="max-autotune-no-cudagraphs" + dynamic=False -- Triton autotuning for a few % more, at a longer compile and a recompile per resolution. + # CUDA-graph modes are NOT used: they crash on the regional block (its static output buffer is overwritten across steps). + # fullgraph drops to False under a step cache OR offloading: both insert an ``@torch.compiler.disable``d function, i.e. a graph break fullgraph=True rejects. The break is cheap; the rest still compiles. kwargs: dict[str, Any] = { "fullgraph": not (cache_active or offload_active), "dynamic": not max_autotune, @@ -342,20 +324,15 @@ def _compile_repeated_blocks( try: import torch - # Heterogeneous-block DiTs (e.g. Z-Image) compile ~one graph per distinct block shape, and - # Z-Image needs ~11, above dynamo's default recompile_limit of 8. Past the limit a resident load - # hard-errors under fullgraph, so raise it to 64 (diffusers' documented regional-compile fix). + # Heterogeneous-block DiTs (e.g. Z-Image) compile ~one graph per distinct block shape, and Z-Image needs ~11, above dynamo default recompile_limit of 8. Past the limit a resident load hard-errors under fullgraph, so raise it to 64 (diffusers documented regional-compile fix). # NOT force_parameter_static_shapes=False: no variant-count win and ~6x slower. dynamo_cfg = getattr(getattr(torch, "_dynamo", None), "config", None) if dynamo_cfg is not None: for _limit_attr in ("recompile_limit", "cache_size_limit"): # name varies by torch ver if hasattr(dynamo_cfg, _limit_attr): setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64)) - # Match eager's intermediate rounding in inductor's fused pointwise kernels: they keep chains in - # fp32 where eager materialises bf16 between ops, a per-forward delta a multi-step denoise - # amplifies. Measured LPIPS vs eager: Qwen-Image 0.019 to 0.006, FLUX.1-dev 0.046 to 0.029 (+2% - # step), FLUX.2-klein 0.018 to 0.017, HunyuanVideo-1.5-720p 0.221 to 0.052, all at ~zero cost. - # Process-global, so snapshot_backend_flags restores it on unload. + # Match eager intermediate rounding in inductor fused pointwise kernels: they keep chains in fp32 where eager materialises bf16 between ops, a per-forward delta a multi-step denoise amplifies. + # Measured LPIPS vs eager: Qwen-Image 0.019 to 0.006, FLUX.1-dev 0.046 to 0.029 (+2% step), FLUX.2-klein 0.018 to 0.017, HunyuanVideo-1.5-720p 0.221 to 0.052, all at ~zero cost. Process-global, so snapshot_backend_flags restores it on unload. inductor_cfg = _inductor_config() if inductor_cfg is not None and hasattr(inductor_cfg, "emulate_precision_casts"): inductor_cfg.emulate_precision_casts = True @@ -363,10 +340,8 @@ def _compile_repeated_blocks( _warn(logger, "compile_repeated_blocks", exc) return False if unet is not None: - # Whole-module static compile for the U-Net classes above. fullgraph mirrors the regional - # decision (in practice only offload lowers it); dynamic is ALWAYS False, so each new - # (height, width, batch) pays its own compile. ``Module.compile`` keeps the module identity, so - # unload/status/LoRA see the same object. + # Whole-module static compile for the U-Net classes above. fullgraph mirrors the regional decision (in practice only offload lowers it); dynamic is ALWAYS False, so each new (height, width, batch) pays its own compile. + # ``Module.compile`` keeps the module identity, so unload/status/LoRA see the same object. unet_kwargs: dict[str, Any] = {"fullgraph": kwargs["fullgraph"], "dynamic": False} if max_autotune: unet_kwargs["mode"] = "max-autotune-no-cudagraphs" @@ -376,8 +351,7 @@ def _compile_repeated_blocks( except Exception as exc: # noqa: BLE001 — optimisation only _warn(logger, "unet whole-module compile", exc) return False - # Compile every denoiser DiT (dual-DiT families run both); a per-DiT failure degrades only that - # one to eager. + # Compile every denoiser DiT (dual-DiT families run both); a per-DiT failure degrades only that one to eager. engaged = False for transformer in dits: try: @@ -386,10 +360,8 @@ def _compile_repeated_blocks( except Exception as exc: # noqa: BLE001 — optimisation only _warn(logger, "compile_repeated_blocks", exc) continue - # A step cache engaged BEFORE this compile has already wrapped each block's forward in a - # @torch.compiler.disable'd hook, so the compute branch would run eager and forfeit the regional - # compile. Re-point the hooks' inner forward at compiled wrappers (no-op without cache hooks); - # the toggle path is armed by apply_step_cache. + # A step cache engaged BEFORE this compile has already wrapped each block forward in a @torch.compiler.disable'd hook, so the compute branch would run eager and forfeit the regional compile. + # Re-point the hooks inner forward at compiled wrappers (no-op without cache hooks); the toggle path is armed by apply_step_cache. try: from .diffusion_cache import _compile_hooked_block_inners _compile_hooked_block_inners(transformer, logger) @@ -436,9 +408,7 @@ def _enable_tf32(logger: Any) -> bool: return False -# Families the overflow harness found to produce non-finite activations / NEW black frames under -# fp16 accumulation. Empty by measurement: no overflow across all six families (bf16 -# bit-identical, fp16 finite; the fp16 same-seed drift is why fp16 compute is gated to ``max``). +# Families the overflow harness found to produce non-finite activations / NEW black frames under fp16 accumulation. Empty by measurement: no overflow across all six families (bf16 bit-identical, fp16 finite; the fp16 same-seed drift is why fp16 compute is gated to ``max``). _FP16_ACCUM_DENY: frozenset[str] = frozenset() @@ -486,8 +456,7 @@ def _enable_fp16_accumulation( def _fuse_qkv(pipe: Any, logger: Any) -> bool: - # Prefer the pipe-level fuse (covers every component); else fuse each denoiser DiT so a dual-DiT - # family fuses BOTH experts. + # Prefer the pipe-level fuse (covers every component); else fuse each denoiser DiT so a dual-DiT family fuses BOTH experts. fn = getattr(pipe, "fuse_qkv_projections", None) if callable(fn): try: diff --git a/studio/backend/core/inference/diffusion_te_prequant.py b/studio/backend/core/inference/diffusion_te_prequant.py index 4b9c8b6bd1..3c7af35511 100644 --- a/studio/backend/core/inference/diffusion_te_prequant.py +++ b/studio/backend/core/inference/diffusion_te_prequant.py @@ -41,14 +41,10 @@ TE_PREQUANT_FORMAT = "unsloth_prequant_text_encoder_state_dict_v1" # The one scheme hosted in v1 (see module docstring). TE_PREQUANT_SCHEMES = ("fp8",) -# Components the pipeline-assembly injection covers (text_encoder_4 is family-assembled -# separately, see diffusion_hidream.py). +# Components the pipeline-assembly injection covers (text_encoder_4 is family-assembled separately, see diffusion_hidream.py). TE_PREQUANT_COMPONENTS = ("text_encoder", "text_encoder_2", "text_encoder_3") -# Bases whose text-encoder weights are VERIFIED byte-identical, so one hosted artifact serves -# all of them (every shard's LFS sha256 compared across repos on 2026-07-18). The checkpoint -# validator accepts a base_model_id from the same group; everything else keeps the strict -# refusal. Ids are lowercased. +# Bases whose text-encoder weights are VERIFIED byte-identical, so one hosted artifact serves all of them (every shard LFS sha256 compared across repos on 2026-07-18). The checkpoint validator accepts a base_model_id from the same group; everything else keeps the strict refusal. Ids are lowercased. _TE_EQUIVALENT_BASES: tuple[frozenset[str], ...] = ( # Qwen2.5-VL-7B text encoder: 4 shards, 16,584,414,544 bytes, identical sha256 set. frozenset( @@ -57,9 +53,7 @@ _TE_EQUIVALENT_BASES: tuple[frozenset[str], ...] = ( "hunyuanvideo-community/hunyuanimage-2.1-diffusers", } ), - # T5-XXL (text_encoder_2): 2 shards, 9,524,648,584 bytes, identical sha256 set across every - # FLUX.1 release; HiDream-I1 ships the same bytes as text_encoder_3 (cross-component mapping is - # not wired yet, this entry documents the identity). + # T5-XXL (text_encoder_2): 2 shards, 9,524,648,584 bytes, identical sha256 set across every FLUX.1 release; HiDream-I1 ships the same bytes as text_encoder_3 (cross-component mapping is not wired yet, this entry documents the identity). frozenset( { "black-forest-labs/flux.1-schnell", @@ -168,8 +162,7 @@ def te_prequant_sources( if mode != TE_QUANT_FP8: return {} family = getattr(fam, "name", None) - # The per-family TE deny table ships on the video branch's precision module; the image branch has - # no denials. Resolve lazily so one module serves both. + # The per-family TE deny table ships on the video branch precision module; the image branch has no denials. Resolve lazily so one module serves both. denied = getattr(precision, "_te_family_denied", None) if callable(denied) and denied(family, mode): return {} @@ -185,9 +178,7 @@ def te_prequant_sources( return {} -# Weight files a dense encoder folder holds. Everything else in the folder (config.json, the -# shard index, tokenizer JSON) is kept when the pre-cast checkpoint replaces the weights: the -# pre-cast loader still meta-inits the encoder from the base repo's component config. +# Weight files a dense encoder folder holds. Everything else in the folder (config.json, the shard index, tokenizer JSON) is kept when the pre-cast checkpoint replaces the weights: the pre-cast loader still meta-inits the encoder from the base repo component config. _TE_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pth", ".pt", ".msgpack", ".h5") @@ -243,9 +234,7 @@ def load_prequant_text_encoder( import torch - # The layerwise-fp8 state dict is plain tensors, so weights_only=True suffices: no pickle code - # runs even for a local-path artifact. A future torchao-subclass scheme needs a format bump AND - # weights_only=False behind the same allowlist as the DiT module. + # The layerwise-fp8 state dict is plain tensors, so weights_only=True suffices: no pickle code runs even for a local-path artifact. A future torchao-subclass scheme needs a format bump AND weights_only=False behind the same allowlist as the DiT module. ckpt = torch.load(path, weights_only = True, map_location = "cpu") if not _validate_checkpoint(ckpt, scheme, component, base, logger): return None @@ -267,9 +256,7 @@ def load_prequant_text_encoder( if subfolder: config_kwargs["subfolder"] = subfolder config = transformers.AutoConfig.from_pretrained(base, **config_kwargs) - # Krea-2 ships transformers-5.x configs whose rope lives under rope_parameters; the runtime - # component loader remaps it for a 4.x runtime, and the meta-init here must match or the rebuilt - # encoder forwards with a broken rope. No-op for every other family. + # Krea-2 ships transformers-5.x configs whose rope lives under rope_parameters; the runtime component loader remaps it for a 4.x runtime, and the meta-init here must match or the rebuilt encoder forwards with a broken rope. No-op for every other family. from .diffusion_krea2 import remap_rope_parameters remap_rope_parameters(getattr(config, "text_config", config)) @@ -279,26 +266,19 @@ def load_prequant_text_encoder( with init_empty_weights(): encoder = encoder_cls(config) - # assign=True swaps in the loaded tensors rather than copying into meta; strict=True since the - # saved dict is the full state dict of the same class. + # assign=True swaps in the loaded tensors rather than copying into meta; strict=True since the saved dict is the full state dict of the same class. encoder.load_state_dict(state_dict, strict = True, assign = True) if _has_meta_tensors(encoder): - # Non-persistent buffers (built in __init__, absent from the state dict) stay on meta. Rebuild on - # CPU so they hold real values, then re-assign the cast weights. + # Non-persistent buffers (built in __init__, absent from the state dict) stay on meta. Rebuild on CPU so they hold real values, then re-assign the cast weights. encoder = encoder_cls(config) encoder.load_state_dict(state_dict, strict = True, assign = True) - # assign=True swaps in SEPARATE tensors for tied weights (the saved dict carries a copy per key), - # untying e.g. Qwen3's lm_head from embed_tokens. An untied head defeats _cast_fp8's - # tied-projection skip below, breaking bit-identity and duplicating the embedding. Re-tie to the - # builder-identical structure; a no-op for untied configs. + # assign=True swaps in SEPARATE tensors for tied weights (the saved dict carries a copy per key), untying e.g. Qwen3's lm_head from embed_tokens. An untied head defeats _cast_fp8's tied-projection skip below, breaking bit-identity and duplicating the embedding. Re-tie to the builder-identical structure; a no-op for untied configs. tie = getattr(encoder, "tie_weights", None) if callable(tie): tie() encoder.eval() - # Install the SAME upcast hooks the runtime cast applies. The weight cast inside is idempotent, - # so this only arms the per-layer upcast; without it the fp8 storage weights would meet bf16 - # activations at the first forward. A hook failure means the encoder cannot run. + # Install the SAME upcast hooks the runtime cast applies. The weight cast inside is idempotent, so this only arms the per-layer upcast; without it the fp8 storage weights would meet bf16 activations at the first forward. A hook failure means the encoder cannot run. from .diffusion_precision import _cast_fp8 class _Target: @@ -401,8 +381,7 @@ def _validate_checkpoint(ckpt: Any, scheme: str, component: str, base: str, logg return False ckpt_base = meta.get("base_model_id") if base: - # Keys matching a different base can load strict=True and encode prompts with the wrong weights. - # The builder always records base_model_id; refuse one that omits it. + # Keys matching a different base can load strict=True and encode prompts with the wrong weights. The builder always records base_model_id; refuse one that omits it. if not ckpt_base: _warn( logger, diff --git a/studio/backend/core/inference/sd_cpp_server.py b/studio/backend/core/inference/sd_cpp_server.py index beb763f66f..3d6bdc48d2 100644 --- a/studio/backend/core/inference/sd_cpp_server.py +++ b/studio/backend/core/inference/sd_cpp_server.py @@ -65,8 +65,7 @@ _TRANSPORT_ERRORS = ( httpx.WriteError, ) -# Readiness probe: the port binds only after the model loads, so any 200 means ready. Use trivial -# /v1/models, not /sdcpp/v1/capabilities (which can block enumerating metadata). +# Readiness probe: the port binds only after the model loads, so any 200 means ready. Use trivial /v1/models, not /sdcpp/v1/capabilities (which can block enumerating metadata). _READY_PATH = "/v1/models" # Native async sdcpp API. _IMG_GEN_PATH = "/sdcpp/v1/img_gen" @@ -110,8 +109,7 @@ def _diagnostic_tail( return "\n".join(chosen)[:limit] -# Grace for the best-effort native cancel to show in job status before abandoning the poll; -# without the cap a lost cancel would hold the generate lock until the job ends. +# Grace for the best-effort native cancel to show in job status before abandoning the poll; without the cap a lost cancel would hold the generate lock until the job ends. _CANCEL_GRACE_S = 5.0 @@ -210,8 +208,7 @@ class SdCppServer: wins), which is how the CPU-backend restart pins the graph off the GPU. """ with self._lifecycle_lock: - # A stop()/unload that raced in before start() took the lock already set _abort and closed the - # client; honor it rather than leak a spawned model process. + # A stop()/unload that raced in before start() took the lock already set _abort and closed the client; honor it rather than leak a spawned model process. if self._stopped or self._abort.is_set(): raise SdCppCancelled("sd-server start was cancelled before launch.") self._abort.clear() @@ -240,8 +237,7 @@ class SdCppServer: self._spawn_error: Optional[Exception] = None spawned = threading.Event() - # Spawn INSIDE the long-lived drain thread: child_popen_kwargs() sets PR_SET_PDEATHSIG, bound to - # the creating thread on Linux, so the creator must outlive the child. + # Spawn INSIDE the long-lived drain thread: child_popen_kwargs() sets PR_SET_PDEATHSIG, bound to the creating thread on Linux, so the creator must outlive the child. def _own_process() -> None: try: proc = subprocess.Popen( @@ -298,8 +294,7 @@ class SdCppServer: deadline = time.monotonic() + timeout url = f"{self.base_url}{_READY_PATH}" while time.monotonic() < deadline: - # A concurrent stop() sets _abort so this wait bails without holding the model load hostage for - # the full startup_timeout. + # A concurrent stop() sets _abort so this wait bails without holding the model load hostage for the full startup_timeout. if self._abort.is_set(): logger.info("sd-server startup aborted before ready") return False @@ -383,8 +378,7 @@ class SdCppServer: def stop(self) -> None: """Terminate the server (SIGTERM -> SIGKILL), join the drain, and release the HTTP client + atexit handler. Idempotent.""" - # Signal abort BEFORE contending for the lock so a start() readiness wait (which holds the lock - # up to startup_timeout) bails immediately instead of blocking stop(). + # Signal abort BEFORE contending for the lock so a start() readiness wait (which holds the lock up to startup_timeout) bails immediately instead of blocking stop(). self._abort.set() self._stopped = True with self._lifecycle_lock: @@ -451,8 +445,7 @@ class SdCppServer: Raises ``RuntimeError`` on submit/poll failures (including the server dying), with the log tail attached. """ - # Already stopped with the cancel event set: report cancellation (route 409), not a generic - # "server died" 500. + # Already stopped with the cancel event set: report cancellation (route 409), not a generic "server died" 500. if self._stopped or not self.is_alive(): if cancel_event is not None and cancel_event.is_set(): raise SdCppCancelled("sd-server generation was cancelled.") @@ -497,15 +490,9 @@ class SdCppServer: self.cancel(job_id) cancel_sent_at = time.monotonic() elif time.monotonic() - cancel_sent_at > _CANCEL_GRACE_S: - # Cancel not reflected within the grace window, so the job is still running and - # sd-server will not interrupt it (cancel_generating=false). Stop the process - # before reporting the cancellation, exactly as the deadline branch below does: - # abandoning the poll alone frees the generate lock while the native job keeps a - # core (or the GPU) busy to completion and holds the server's job slot, so the - # next request queues behind work nobody is waiting for. The caller does stop the - # server on unload, but a SUPERSEDING LOAD only does so after its multi-gigabyte - # download, and a load that then fails never reaches that point at all. Stopping - # here is safe: the backend reloads on the next generate. + # Cancel not reflected within the grace window, so the job is still running and sd-server will not interrupt it (cancel_generating=false). Stop the process before reporting the cancellation, exactly as the deadline branch below does: + # abandoning the poll alone frees the generate lock while the native job keeps a core (or the GPU) busy to completion and holds the server job slot, so the next request queues behind work nobody is waiting for. + # The caller does stop the server on unload, but a SUPERSEDING LOAD only does so after its multi-gigabyte download, and a load that then fails never reaches that point at all. Stopping here is safe: the backend reloads on the next generate. self.stop() raise SdCppCancelled("sd-server generation was cancelled.") if not self.is_alive(): @@ -514,8 +501,7 @@ class SdCppServer: raise SdCppCancelled("sd-server generation was cancelled.") raise RuntimeError(self._died_message("img_gen poll", None)) if time.monotonic() > deadline: - # sd-server won't interrupt an in-flight job (cancel_generating=false), so cancel + stop to free - # the slot; the backend reloads on the next generate. + # sd-server will not interrupt an in-flight job (cancel_generating=false), so cancel + stop to free the slot; the backend reloads on the next generate. self.cancel(job_id) self.stop() raise RuntimeError(f"sd-server generation timed out after {total_timeout}s") @@ -525,8 +511,7 @@ class SdCppServer: time.sleep(poll_interval) continue except RuntimeError as exc: - # A concurrent stop() closes the shared client, giving a plain RuntimeError ("client has been - # closed") rather than a transport error; map a cancel to 409. + # A concurrent stop() closes the shared client, giving a plain RuntimeError ("client has been closed") rather than a transport error; map a cancel to 409. if cancel_event is not None and cancel_event.is_set(): raise SdCppCancelled("sd-server generation was cancelled.") from exc raise diff --git a/studio/backend/core/inference/video_families.py b/studio/backend/core/inference/video_families.py index 7c6e3f02d4..498fba851b 100644 --- a/studio/backend/core/inference/video_families.py +++ b/studio/backend/core/inference/video_families.py @@ -39,19 +39,15 @@ class VideoFamily: denoiser_attr: str = "transformer" # Extra lowercased substrings (besides ``name``) that map a repo id here. aliases: tuple[str, ...] = field(default_factory = tuple) - # True when the pipeline returns synchronized audio (LTX-2): export muxes the track and size - # estimates count the audio VAE + vocoder. + # True when the pipeline returns synchronized audio (LTX-2): export muxes the track and size estimates count the audio VAE + vocoder. has_audio: bool = False - # Wan2.2-A14B dual-expert MoE: a second DiT (transformer_2) handles the low-noise steps with its - # own guidance kwarg. None/False for single-DiT. + # Wan2.2-A14B dual-expert MoE: a second DiT (transformer_2) handles the low-noise steps with its own guidance kwarg. None/False for single-DiT. transformer2_class: Optional[str] = None is_moe: bool = False cfg2_kwarg: Optional[str] = None - # HunyuanVideo-1.5 guidance: __call__ takes NO guidance kwarg; CFG lives on a ``guider`` - # component whose guidance_scale is set per request. When True, generate() writes pipe.guider. + # HunyuanVideo-1.5 guidance: __call__ takes NO guidance kwarg; CFG lives on a ``guider`` component whose guidance_scale is set per request. When True, generate() writes pipe.guider. guidance_via_guider: bool = False - # Generation defaults + shape. ``frame_step`` is the temporal compression: a valid frame count is - # k*frame_step + 1, so requests are snapped BEFORE latents are allocated. + # Generation defaults + shape. ``frame_step`` is the temporal compression: a valid frame count is k*frame_step + 1, so requests are snapped BEFORE latents are allocated. default_steps: int = 40 default_guidance: float = 4.0 default_num_frames: int = 121 @@ -61,27 +57,23 @@ class VideoFamily: resolution_multiple: int = 32 # (width, height) UI presets, landscape first; the first is the default. resolution_presets: tuple[tuple[int, int], ...] = ((768, 512),) - # Component bf16-RESIDENT sizes in decimal GB (denoiser(s), text encoder, VAE + audio - # companions): what sits on device after the dtype cast, not the download size. + # Component bf16-RESIDENT sizes in decimal GB (denoiser(s), text encoder, VAE + audio companions): what sits on device after the dtype cast, not the download size. bf16_components_gb: Optional[tuple[float, float, float]] = None # True when the DiT compiles cleanly with regional torch.compile (declares _repeated_blocks). supports_torch_compile: bool = True # Video DiTs are bf16-native, so fp16 promotes to float32; defaults True. fp16_incompatible: bool = True - # Wan's VAE decodes in float32 (loading it bf16 causes banding / black frames), so when True the - # loader pins it back to fp32. Its bf16_components_gb term is already the fp32 size. + # Wan VAE decodes in float32 (loading it bf16 causes banding / black frames), so when True the loader pins it back to fp32. Its bf16_components_gb term is already the fp32 size. vae_force_fp32: bool = False # Curated GGUF repo for the picker (the DiT as single-file GGUF quants). gguf_repo: Optional[str] = None - # Hosted PRE-CAST text-encoder checkpoints as (scheme, component, repo_id) triples; same - # semantics as DiffusionFamily.te_prequant_repos. + # Hosted PRE-CAST text-encoder checkpoints as (scheme, component, repo_id) triples; same semantics as DiffusionFamily.te_prequant_repos. te_prequant_repos: tuple[tuple[str, str, str], ...] = field(default_factory = tuple) _FAMILIES: tuple[VideoFamily, ...] = ( - # LTX-2 (diffusers >= 0.39): ~19B single-stream video DiT generating synchronized audio + video - # in one pass. The Gemma3-12B text encoder is stored fp32 on the hub (~49 GB download, ~24 GB - # resident as bf16). Base repo carries the dev config (40 steps, CFG 4); distilled runs few-step. + # LTX-2 (diffusers >= 0.39): ~19B single-stream video DiT generating synchronized audio + video in one pass. The Gemma3-12B text encoder is stored fp32 on the hub (~49 GB download, ~24 GB resident as bf16). + # Base repo carries the dev config (40 steps, CFG 4); distilled runs few-step. VideoFamily( name = "ltx-2", pipeline_class = "LTX2Pipeline", @@ -97,17 +89,13 @@ _FAMILIES: tuple[VideoFamily, ...] = ( resolution_multiple = 32, # 768x512 native default; 1216x704 the card's quality target; 704x1216 vertical. resolution_presets = ((768, 512), (1216, 704), (704, 1216), (512, 768)), - # transformer 37.8 bf16; Gemma3-12B TE ~24.4 bf16 RESIDENT (the hub stores it fp32, ~49 GB - # download, but the pipeline loads torch_dtype=bf16); VAE 2.4 + connectors 2.9 + audio 0.2. The - # old 50.4 figure double-counted the fp32 store and pushed auto toward offload. + # transformer 37.8 bf16; Gemma3-12B TE ~24.4 bf16 RESIDENT (the hub stores it fp32, ~49 GB download, but the pipeline loads torch_dtype=bf16); VAE 2.4 + connectors 2.9 + audio 0.2. The old 50.4 figure double-counted the fp32 store and pushed auto toward offload. bf16_components_gb = (37.8, 24.4, 5.5), gguf_repo = "unsloth/LTX-2.3-GGUF", # Pre-cast Gemma3-12B TE (fp32 ~49 GB on the hub, pre-cast ~13.2 GB): the biggest download win. te_prequant_repos = (("fp8", "text_encoder", "unsloth/LTX-2-FP8"),), ), - # Wan2.2-TI2V-5B (diffusers >= 0.35, verified on 0.39): ~5B single-stream video DiT (UMT5 text - # encoder). No audio, no second expert, so single-DiT. Wan VAE temporal compression 4 gives valid - # frame counts 4k+1. Pipeline defaults 50 steps / CFG 5; UI presets target 720p at 24 fps. + # Wan2.2-TI2V-5B (diffusers >= 0.35, verified on 0.39): ~5B single-stream video DiT (UMT5 text encoder). No audio, no second expert, so single-DiT. Wan VAE temporal compression 4 gives valid frame counts 4k+1. Pipeline defaults 50 steps / CFG 5; UI presets target 720p at 24 fps. VideoFamily( name = "wan2.2-ti2v-5b", pipeline_class = "WanPipeline", @@ -123,21 +111,16 @@ _FAMILIES: tuple[VideoFamily, ...] = ( default_fps = 24, # Wan VAE temporal factor 4, so valid counts are 4k+1. frame_step = 4, - # TI2V-5B VAE is 16x spatial + patch 2, so WanPipeline floors H/W to 32; snap to 32 so the - # recorded size matches the rendered clip. + # TI2V-5B VAE is 16x spatial + patch 2, so WanPipeline floors H/W to 32; snap to 32 so the recorded size matches the rendered clip. resolution_multiple = 32, # 720p-class presets (all /32); first is the default the loader plans against. resolution_presets = ((1280, 704), (704, 1280), (960, 960), (832, 480)), - # bf16-RESIDENT. transformer + VAE ship FP32 on disk (index 20.0 GB = 5B x 4), so bf16 - # transformer ~10.0; UMT5 TE ships bf16 (11.4); VAE runs fp32 (2.8). + # bf16-RESIDENT. transformer + VAE ship FP32 on disk (index 20.0 GB = 5B x 4), so bf16 transformer ~10.0; UMT5 TE ships bf16 (11.4); VAE runs fp32 (2.8). bf16_components_gb = (10.0, 11.4, 2.8), vae_force_fp32 = True, gguf_repo = "QuantStack/Wan2.2-TI2V-5B-GGUF", ), - # Wan2.2-T2V-A14B (diffusers >= 0.35, verified on 0.39): the dual-expert MoE. Both transformers - # are WanTransformer3DModel with boundary_ratio 0.875; the pipeline routes high-noise steps - # through transformer (guidance_scale) and low-noise through transformer_2 (guidance_scale_2, - # accepted only when boundary_ratio is set), so cfg2_kwarg is threaded ONLY here. + # Wan2.2-T2V-A14B (diffusers >= 0.35, verified on 0.39): the dual-expert MoE. Both transformers are WanTransformer3DModel with boundary_ratio 0.875; the pipeline routes high-noise steps through transformer (guidance_scale) and low-noise through transformer_2 (guidance_scale_2, accepted only when boundary_ratio is set), so cfg2_kwarg is threaded ONLY here. VideoFamily( name = "wan2.2-t2v-a14b", pipeline_class = "WanPipeline", @@ -145,8 +128,7 @@ _FAMILIES: tuple[VideoFamily, ...] = ( base_repo = "Wan-AI/Wan2.2-T2V-A14B-Diffusers", aliases = ("wan2.2-14b", "wan-t2v", "wan2.2-t2v", "wan-t2v-a14b", "wan-a14b"), has_audio = False, - # is_moe drives the dual-DiT optimisation layers (speed/attention/cache/quant on BOTH); - # cfg2_kwarg names the pipeline kwarg for transformer_2's guidance. + # is_moe drives the dual-DiT optimisation layers (speed/attention/cache/quant on BOTH); cfg2_kwarg names the pipeline kwarg for transformer_2's guidance. transformer2_class = "WanTransformer3DModel", is_moe = True, cfg2_kwarg = "guidance_scale_2", @@ -157,20 +139,14 @@ _FAMILIES: tuple[VideoFamily, ...] = ( default_fps = 16, # A14B runs at 16 fps (vs TI2V-5B's 24) frame_step = 4, resolution_multiple = 16, - # 480p + 720p presets (landscape + vertical). A14B's VAE is 8x so multiple 16 renders 720 exactly - # (unlike TI2V-5B's 16x VAE, which floors 720 to 704). + # 480p + 720p presets (landscape + vertical). A14B's VAE is 8x so multiple 16 renders 720 exactly (unlike TI2V-5B's 16x VAE, which floors 720 to 704). resolution_presets = ((1280, 720), (832, 480), (480, 832), (720, 1280)), - # bf16-RESIDENT. Each expert ships FP32 (index 57.15 GB = 14.3B x 4), so ~28.6 bf16 each and - # ~57.2 for BOTH (the headline before offload), NOT the 114.3 fp32 sum. UMT5 TE bf16 (11.4); - # VAE fp32 (0.5). + # bf16-RESIDENT. Each expert ships FP32 (index 57.15 GB = 14.3B x 4), so ~28.6 bf16 each and ~57.2 for BOTH (the headline before offload), NOT the 114.3 fp32 sum. UMT5 TE bf16 (11.4); VAE fp32 (0.5). bf16_components_gb = (57.2, 11.4, 0.5), vae_force_fp32 = True, # No gguf_repo: community GGUFs split the experts, and a single-file load covers only one. ), - # HunyuanVideo-1.5 (diffusers >= 0.39): 8.3B DiT, Qwen2.5-VL text encoder + ByT5 glyph encoder. - # Three quirks: (1) __call__ has NO guidance kwarg, CFG lives on the ``guider``; (2) NO - # callback_on_step_end (generate() wraps scheduler.step); (3) tencent's repo has no - # model_index.json, so only the community Diffusers repacks load. + # HunyuanVideo-1.5 (diffusers >= 0.39): 8.3B DiT, Qwen2.5-VL text encoder + ByT5 glyph encoder. Three quirks: (1) __call__ has NO guidance kwarg, CFG lives on the ``guider``; (2) NO callback_on_step_end (generate() wraps scheduler.step); (3) tencent's repo has no model_index.json, so only the community Diffusers repacks load. VideoFamily( name = "hunyuanvideo-1.5", pipeline_class = "HunyuanVideo15Pipeline", @@ -193,9 +169,7 @@ _FAMILIES: tuple[VideoFamily, ...] = ( # DiT fp32 on disk (32.0 to 16.6 bf16); VAE (4.7 to 2.4); Qwen2.5-VL TE bf16 14.0 + ByT5 0.8. bf16_components_gb = (16.6, 14.8, 2.4), ), - # The 720p t2v repack: same architecture/quirks/footprint as the 480p entry, only the trained - # resolution differs. Own family so a 720p load defaults to 720p sizes. Its full-path alias - # out-lengths (and outranks) the generic "hunyuanvideo-1.5" token for this repo only. + # The 720p t2v repack: same architecture/quirks/footprint as the 480p entry, only the trained resolution differs. Own family so a 720p load defaults to 720p sizes. Its full-path alias out-lengths (and outranks) the generic "hunyuanvideo-1.5" token for this repo only. VideoFamily( name = "hunyuanvideo-1.5-720p", pipeline_class = "HunyuanVideo15Pipeline", @@ -273,8 +247,7 @@ def snap_video_size(fam: VideoFamily, width: int, height: int) -> tuple[int, int return snap(width), snap(height) -# Default (steps, guidance) per checkpoint variant, matched by substring (picked id then base -# repo), most specific first: distilled LTX-2.3 runs few-step CFG-off, the dev base wants 40/4. +# Default (steps, guidance) per checkpoint variant, matched by substring (picked id then base repo), most specific first: distilled LTX-2.3 runs few-step CFG-off, the dev base wants 40/4. _VIDEO_GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = ( ("distilled", 8, 1.0), ("ltx", 40, 4.0), @@ -296,8 +269,7 @@ def default_video_generation_params( for identifier in identifiers: needle = (identifier or "").lower() for key, steps, guidance in _VIDEO_GENERATION_DEFAULTS: - # Match the key as a name segment: reject a preceding ASCII letter so "swan-video" or - # "taiwan-clips" doesn't false-match "wan". Trailing chars stay free. + # Match the key as a name segment: reject a preceding ASCII letter so "swan-video" or "taiwan-clips" does not false-match "wan". Trailing chars stay free. if re.search(r"(? tuple[list[int], list[str], list[str]]: - # Draw the full configured batch, not min(batch, n): the sampler refills across cycles so a - # dataset smaller than train_batch_size still yields exactly that many indices. + # Draw the full configured batch, not min(batch, n): the sampler refills across cycles so a dataset smaller than train_batch_size still yields exactly that many indices. idx = index_sampler.next_batch(cfg.train_batch_size) chosen = [pairs[i] for i in idx] return idx, [c[0] for c in chosen], [c[1] for c in chosen] @@ -552,8 +539,7 @@ def run_diffusion_lora_training( step_loss += float(loss.detach()) / cfg.gradient_accumulation_steps micro += 1 - # max_grad_norm at or below 0 disables clipping (Studio sends 0.0); passing 0.0 to - # clip_grad_norm_ would zero every gradient. + # max_grad_norm at or below 0 disables clipping (Studio sends 0.0); passing 0.0 to clip_grad_norm_ would zero every gradient. grad_norm = None if cfg.max_grad_norm and cfg.max_grad_norm > 0: # Returned value is the total PRE-clip norm, reported to the UI chart. @@ -670,8 +656,7 @@ def run_diffusion_training_process(*, event_queue: Any, stop_queue: Any, config: return got if saw else False try: - # normalized() resolves + validates the family; dispatch through the registry so a DiT family - # runs its own trainer while SDXL keeps this loop. + # normalized() resolves + validates the family; dispatch through the registry so a DiT family runs its own trainer while SDXL keeps this loop. cfg = _config_from_dict(config).normalized() trainer = get_trainer(cfg.resolved_family) trainer(cfg, on_event = on_event, should_stop = should_stop) diff --git a/studio/backend/core/training/diffusion_train_extras.py b/studio/backend/core/training/diffusion_train_extras.py index 0a42004aee..6a971e1c41 100644 --- a/studio/backend/core/training/diffusion_train_extras.py +++ b/studio/backend/core/training/diffusion_train_extras.py @@ -42,10 +42,8 @@ from typing import Any, Iterable, Optional # ── LoRA EMA ────────────────────────────────────────────────────────────────── -# Warmup horizon for the EMA decay ramp: effective decay is -# min(decay, (1 + updates) / (WARMUP_OFFSET + updates)), the standard inverse ramp. With the -# offset at 10, step 1 averages aggressively (~0.18) and the ramp reaches 0.99 after ~1000 -# updates, so a 300-step run still ends with a shadow that absorbed most of the trajectory. +# Warmup horizon for the EMA decay ramp: effective decay is min(decay, (1 + updates) / (WARMUP_OFFSET + updates)), the standard inverse ramp. +# With the offset at 10, step 1 averages aggressively (~0.18) and the ramp reaches 0.99 after ~1000 updates, so a 300-step run still ends with a shadow that absorbed most of the trajectory. _EMA_WARMUP_OFFSET = 10.0 @@ -225,8 +223,7 @@ def source_revision(ref: Any) -> str: roots += [ e.path for e in it - # vae too: cached latents come from it, so an in-place VAE swap must - # invalidate them just like an encoder change. + # vae too: cached latents come from it, so an in-place VAE swap must invalidate them just like an encoder change. if e.is_dir() and e.name.startswith(("text_encoder", "tokenizer", "vae")) ] for root in roots: @@ -341,13 +338,10 @@ class PersistentConditioningCache: # ── aspect-ratio bucketing ──────────────────────────────────────────────────── -# Pixel-dimension divisor for bucket shapes. The DiT families divide by 8 in the VAE and 2 again -# in latent patching, and regional torch.compile prefers few distinct shapes, so buckets snap to -# multiples of 64 pixels. +# Pixel-dimension divisor for bucket shapes. The DiT families divide by 8 in the VAE and 2 again in latent patching, and regional torch.compile prefers few distinct shapes, so buckets snap to multiples of 64 pixels. BUCKET_DIVISOR = 64 -# Widest aspect ratio a bucket may take; anything more extreme clamps to it (matching the common -# practice of capping panoramas). +# Widest aspect ratio a bucket may take; anything more extreme clamps to it (matching the common practice of capping panoramas). MAX_BUCKET_RATIO = 2.0 diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 4d78f9841d..20bbdc7df7 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -57,16 +57,14 @@ def _run_diffusion_child(*, event_queue: Any, stop_queue: Any, config: dict) -> def _default_target(*, event_queue: Any, stop_queue: Any, config: dict) -> None: - # First thing in the child (before torch): self-bind to parent death and scrub the native path - # secret, like the other workers (multiprocessing children can't be given a preexec_fn). + # First thing in the child (before torch): self-bind to parent death and scrub the native path secret, like the other workers (multiprocessing children cannot be given a preexec_fn). from utils.native_path_leases import run_without_native_path_secret run_without_native_path_secret( _run_diffusion_child, event_queue = event_queue, stop_queue = stop_queue, config = config ) -# Cap on retained metric points; over it, arrays are decimated (every other point) so a long run -# stays bounded while the loss chart keeps its shape. +# Cap on retained metric points; over it, arrays are decimated (every other point) so a long run stays bounded while the loss chart keeps its shape. _METRIC_CAP = 4000 @@ -89,9 +87,7 @@ def _llm_training_active() -> bool: # ── persisted run history ────────────────────────────────────────────────────── -# Every terminal run is recorded as one JSON file (summary + scrubbed config + metric logs) so -# the Train tab can show history. JSON, not the LLM sqlite, so diffusion runs stay off the LLM -# Runs page. +# Every terminal run is recorded as one JSON file (summary + scrubbed config + metric logs) so the Train tab can show history. JSON, not the LLM sqlite, so diffusion runs stay off the LLM Runs page. def _runs_dir() -> Path: from utils.paths.storage_roots import studio_root @@ -113,8 +109,7 @@ def list_diffusion_runs(limit: int = 20) -> list[dict]: rec = json.loads(p.read_text(encoding = "utf-8")) except Exception: # noqa: BLE001 -- a corrupt record never breaks the listing continue - # Skip a wrong-shape record (not a dict, or missing string job_id/status) so one bad file can't - # blow up the route's DiffusionTrainingRunSummary(**r) or the whole panel. + # Skip a wrong-shape record (not a dict, or missing string job_id/status) so one bad file cannot blow up the route DiffusionTrainingRunSummary(**r) or the whole panel. if not isinstance(rec, dict): continue if not (isinstance(rec.get("job_id"), str) and isinstance(rec.get("status"), str)): @@ -194,8 +189,7 @@ def _append_metric( floss = _finite_or_none(loss) if floss is None: # non-numeric or non-finite: skip, keep the curve JSON-safe return - # lr / grad_norm may be None or non-finite; non-finite is nulled (not dropped) to stay - # index-aligned. + # lr / grad_norm may be None or non-finite; non-finite is nulled (not dropped) to stay index-aligned. flr = _finite_or_none(lr) fgn = _finite_or_none(grad_norm) steps = state["metric_steps"] @@ -231,16 +225,11 @@ class DiffusionTrainingService: self._ctx = ctx if ctx is not None else _CTX self._target = target if target is not None else _default_target self._lock = threading.Lock() - # Set by reserve() while a start is in flight (before the route frees GPU models) so the load - # guards refuse a concurrent load during the free-then-spawn window. Cleared by unreserve(). + # Set by reserve() while a start is in flight (before the route frees GPU models) so the load guards refuse a concurrent load during the free-then-spawn window. Cleared by unreserve(). self._reserved = False - # Dataset mutations in flight (caption edit, upload commit, image delete, example import). - # A start refuses while any is open and a mutation refuses once a start is reserved, both - # decided under _lock, so neither can slip through the other's check-then-act window. + # Dataset mutations in flight (caption edit, upload commit, image delete, example import). A start refuses while any is open and a mutation refuses once a start is reserved, both decided under _lock, so neither can slip through the other check-then-act window. self._dataset_mutations = 0 - # GPU load admissions in flight (an image/video/chat load between its training guard and - # the moment it registers with the arbiter). Same two-sided rule as the dataset mutations: - # a start refuses while one is open, and an admission refuses once a start is reserved. + # GPU load admissions in flight (an image/video/chat load between its training guard and the moment it registers with the arbiter). Same two-sided rule as the dataset mutations: a start refuses while one is open, and an admission refuses once a start is reserved. self._gpu_admissions = 0 self._proc: Any = None self._stop_queue: Any = None @@ -277,19 +266,14 @@ class DiffusionTrainingService: "then start the run." ) if self._gpu_admissions: - # A load already passed its training guard and is about to take the GPU. Reserving - # now would free residents it has not registered yet, so the trainer and a - # brand-new pipeline would allocate together. Refusing is safe: the admission is - # held only across the load's registration, not the load itself. + # A load already passed its training guard and is about to take the GPU. Reserving now would free residents it has not registered yet, so the trainer and a brand-new pipeline would allocate together. + # Refusing is safe: the admission is held only across the load registration, not the load itself. raise RuntimeError( "A model is being loaded onto the GPU right now. Wait for that to finish, " "then start the run." ) - # The LLM trainer under the SAME lock, not just at the route's earlier check: that check - # and this reservation are separated by several network-bound preflights, so an LLM start - # could spawn in between and both trainers would allocate on one GPU. The LLM route holds - # gpu_load_admission() across its own spawn, so between the two either this raises or - # that one does -- never neither. + # The LLM trainer under the SAME lock, not just at the route earlier check: that check and this reservation are separated by several network-bound preflights, so an LLM start could spawn in between and both trainers would allocate on one GPU. + # The LLM route holds gpu_load_admission() across its own spawn, so between the two either this raises or that one does -- never neither. if _llm_training_active(): raise RuntimeError( "An LLM training job is already running. " @@ -366,8 +350,7 @@ class DiffusionTrainingService: _config_from_dict(config).normalized() - # Join a finished job's pump OUTSIDE the lock: its final state writes take this lock, so joining - # under it would stall the start and let the stale pump overwrite the new state. + # Join a finished job pump OUTSIDE the lock: its final state writes take this lock, so joining under it would stall the start and let the stale pump overwrite the new state. with self._lock: if self._proc is not None and self._proc.is_alive(): raise RuntimeError("A diffusion training job is already running.") @@ -549,8 +532,7 @@ class DiffusionTrainingService: elif etype == "model_load_completed": s.update(in_model_load = False, message = "Training...") elif etype == "preparing": - # A long precompute phase (e.g. VAE latent cache) before the first step; surfaced so the UI shows - # progress instead of a silent "Loading base model..." stall. + # A long precompute phase (e.g. VAE latent cache) before the first step; surfaced so the UI shows progress instead of a silent "Loading base model..." stall. done, total = ev.get("done"), ev.get("total") stage = str(ev.get("stage", "prepare")).replace("_", " ") s.update( @@ -566,8 +548,7 @@ class DiffusionTrainingService: # Non-fatal trainer notes; keep training state, surface the text. s["message"] = str(ev.get("message", "warning")) elif etype == "progress": - # Null any non-finite float so the JSON stays strict-parseable; a missing key keeps the last - # value, a present-but-non-finite one becomes None. + # Null any non-finite float so the JSON stays strict-parseable; a missing key keeps the last value, a present-but-non-finite one becomes None. loss = _finite_or_none(ev["loss"]) if "loss" in ev else s["loss"] avg_loss = _finite_or_none(ev["avg_loss"]) if "avg_loss" in ev else s["avg_loss"] learning_rate = ( @@ -602,8 +583,7 @@ class DiffusionTrainingService: ev.get("grad_norm"), ) elif etype == "complete": - # Reset in_model_load: a stop during model load emits complete with no preceding - # model_load_completed, which would otherwise leave a stale loading indicator. + # Reset in_model_load: a stop during model load emits complete with no preceding model_load_completed, which would otherwise leave a stale loading indicator. s.update( active = False, in_model_load = False, diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index fe7aaa120e..19b1e821eb 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -2318,8 +2318,7 @@ class DiffusionLoadRequest(BaseModel): @field_validator("attention_backend", mode = "before") @classmethod def _normalize_attention_backend(cls, value): - # The dispatcher accepts case/whitespace variants ("CuDNN", " sage "), but the Literal above is - # validated before any normaliser runs, so fold a string to its canonical form here. + # The dispatcher accepts case/whitespace variants ("CuDNN", " sage "), but the Literal above is validated before any normaliser runs, so fold a string to its canonical form here. return value.strip().lower() if isinstance(value, str) else value @@ -2376,8 +2375,7 @@ class ControlNetSpec(BaseModel): @model_validator(mode = "after") def _check_guidance_range(self) -> "ControlNetSpec": - # An inverted range means "act over no steps"; reject it as a clean 422 instead of letting the - # diffusers pipeline 500 deep in the denoise. + # An inverted range means "act over no steps"; reject it as a clean 422 instead of letting the diffusers pipeline 500 deep in the denoise. if self.guidance_start > self.guidance_end: raise ValueError("guidance_start must be <= guidance_end") return self @@ -2396,17 +2394,14 @@ class DiffusionGenerateRequest(BaseModel): ) steps: int = Field(9, ge = 1, le = 100, description = "Number of denoising steps") guidance: float = Field(0.0, ge = 0.0, le = 20.0, description = "Classifier-free guidance scale") - # le = 2**53-1: seeds round-trip through JSON gallery recipes, where JavaScript rounds integers - # above Number.MAX_SAFE_INTEGER and a restored recipe would generate a different image. + # le = 2**53-1: seeds round-trip through JSON gallery recipes, where JavaScript rounds integers above Number.MAX_SAFE_INTEGER and a restored recipe would generate a different image. seed: Optional[int] = Field( None, ge = 0, le = 2**53 - 1, description = "Seed for reproducibility (random if omitted)" ) batch_size: int = Field( 1, ge = 1, le = 32, description = "Images generated in one forward pass (VRAM-heavy)" ) - # Batched multi-image generation (diffusers engine): a prompt list renders one image per prompt - # in a single batched forward (txt2img only); a seed list renders one image per seed. Each image - # carries its OWN generator seed, so any batch member replays alone. + # Batched multi-image generation (diffusers engine): a prompt list renders one image per prompt in a single batched forward (txt2img only); a seed list renders one image per seed. Each image carries its OWN generator seed, so any batch member replays alone. prompts: Optional[list[str]] = Field( None, min_length = 1, @@ -2452,10 +2447,8 @@ class DiffusionGenerateRequest(BaseModel): ) return self - # Image-conditioned workflows (base64 or data-URL): init_image alone runs img2img, init_image + - # mask_image runs inpaint. Both require a family with the matching pipeline. Cap each base64 - # string so one request can't buffer a multi-GB payload (decoded dimensions are bounded - # separately); ~32 MiB fits a full 4096px image yet rejects abuse. + # Image-conditioned workflows (base64 or data-URL): init_image alone runs img2img, init_image + mask_image runs inpaint. Both require a family with the matching pipeline. + # Cap each base64 string so one request cannot buffer a multi-GB payload (decoded dimensions are bounded separately); ~32 MiB fits a full 4096px image yet rejects abuse. init_image: Optional[str] = Field( None, max_length = 32 * 1024 * 1024, @@ -2469,9 +2462,7 @@ class DiffusionGenerateRequest(BaseModel): ) strength: Optional[float] = Field( None, - # EXCLUSIVE lower bound: strength 0 does not "keep the source". Every diffusers img2img/inpaint - # pipeline derives its step count from it, so 0 leaves zero denoising steps: FLUX/Qwen/Z-Image - # raise "the number of pipeline steps is 0", and SDXL img2img crashes on empty latents (a 500). + # EXCLUSIVE lower bound: strength 0 does not "keep the source". Every diffusers img2img/inpaint pipeline derives its step count from it, so 0 leaves zero denoising steps: FLUX/Qwen/Z-Image raise "the number of pipeline steps is 0", and SDXL img2img crashes on empty latents (a 500). gt = 0.0, le = 1.0, description = "img2img/inpaint denoise strength: low values stay close to the " @@ -2508,9 +2499,7 @@ class DiffusionGenerateRequest(BaseModel): @field_validator("loras") @classmethod def _unique_lora_ids(cls, value: Optional[list[LoraSpec]]) -> Optional[list[LoraSpec]]: - # Both apply paths break alias collisions by suffixing the adapter name/file, so a repeated id - # would load the SAME adapter several times and stack its effect past the per-adapter weight - # bound. The UI already blocks duplicates; reject them for API clients too. + # Both apply paths break alias collisions by suffixing the adapter name/file, so a repeated id would load the SAME adapter several times and stack its effect past the per-adapter weight bound. The UI already blocks duplicates; reject them for API clients too. if value: seen: set[str] = set() for spec in value: @@ -2524,8 +2513,7 @@ class DiffusionGenerateRequest(BaseModel): @field_validator("reference_images") @classmethod def _bounded_reference_items(cls, value: Optional[list[str]]) -> Optional[list[str]]: - # Each reference is a base64 image; bound its length like init_image/mask_image so several - # references can't buffer a multi-GB payload. + # Each reference is a base64 image; bound its length like init_image/mask_image so several references cannot buffer a multi-GB payload. if value is not None: for item in value: if len(item) > 32 * 1024 * 1024: @@ -2535,17 +2523,14 @@ class DiffusionGenerateRequest(BaseModel): @field_validator("width", "height") @classmethod def _multiple_of_16(cls, value: int) -> int: - # Z-Image requires dimensions divisible by 16 (8x VAE downsample + 2x patch). Non-multiples crash - # deep in the pipeline, so reject them here for a clean 422. + # Z-Image requires dimensions divisible by 16 (8x VAE downsample + 2x patch). Non-multiples crash deep in the pipeline, so reject them here for a clean 422. if value % 16 != 0: raise ValueError("must be a multiple of 16") return value @model_validator(mode = "after") def _batch_seeds_json_safe(self) -> "DiffusionGenerateRequest": - # A batch derives per-image seeds as seed .. seed+batch_size-1. The base seed is capped at 2**53-1 - # to round-trip through the JSON recipe, but a derived top-of-batch seed near the cap can exceed - # it, where the frontend rounds it and a restored recipe replays a different image. + # A batch derives per-image seeds as seed .. seed+batch_size-1. The base seed is capped at 2**53-1 to round-trip through the JSON recipe, but a derived top-of-batch seed near the cap can exceed it, where the frontend rounds it and a restored recipe replays a different image. if self.seed is not None and self.seed + self.batch_size - 1 > 2**53 - 1: raise ValueError( "seed + batch_size - 1 must not exceed 2**53 - 1 so every per-image seed " @@ -2587,9 +2572,7 @@ class GalleryImage(BaseModel): controlnet: Optional[str] = Field( None, description = "ControlNet applied, formatted as 'id:control_type:strength'" ) - # Conditioned-workflow settings. The images themselves are NOT persisted (user uploads with - # their own lifetime), so these say what ran and let the client tell the user which inputs it - # needs back instead of silently restoring a conditioned image as a plain Create. + # Conditioned-workflow settings. The images themselves are NOT persisted (user uploads with their own lifetime), so these say what ran and let the client tell the user which inputs it needs back instead of silently restoring a conditioned image as a plain Create. workflow: Optional[str] = Field( None, description = "Workflow that produced it: txt2img, img2img, inpaint, upscale, edit, " @@ -2748,10 +2731,7 @@ class DiffusionStatusResponse(BaseModel): "picker's enabled state). Diffusers only, for families with a ControlNet pipeline; False " "for the native engine, GGUF-via-diffusers, and torchao fp8/int8 dense.", ) - # Additive: per-Advanced-control provenance {control: {value, source, reason}}. Present only on - # backends that record it; null when nothing is loaded. The frontend renders an "Auto: X" badge - # next to each control whose source == "auto". Declared explicitly so pydantic's extra='ignore' - # doesn't drop it. + # Additive: per-Advanced-control provenance {control: {value, source, reason}}. Present only on backends that record it; null when nothing is loaded. The frontend renders an "Auto: X" badge next to each control whose source == "auto". Declared explicitly so pydantic extra='ignore' does not drop it. resolved: Optional[Dict[str, DiffusionResolvedControl]] = Field( None, description = "Per-control resolved value + provenance (source auto|explicit + reason), " @@ -2787,10 +2767,8 @@ class DiffusionInferenceInfoResponse(BaseModel): # ── OpenAI-compatible images API (POST /v1/images/generations) ── -# -# Shapes mirror OpenAI's CreateImageRequest / ImagesResponse so off-the-shelf clients work -# unchanged. The loaded image GGUF stands in for the model; GPT-image-only knobs (quality, style, -# background, output_format, ...) are accepted and ignored, like dall-e-2. The size string is +# Shapes mirror OpenAI's CreateImageRequest / ImagesResponse so off-the-shelf clients work unchanged. The loaded image GGUF stands in for the model; GPT-image-only knobs (quality, style, background, output_format, ...) are accepted and ignored, like dall-e-2. +# The size string is parsed and `stream` rejected in the route; everything Pydantic can check declaratively is here. # parsed and `stream` rejected in the route; everything Pydantic can check declaratively is here. @@ -2813,8 +2791,7 @@ class ImageGenerationRequest(BaseModel): "url", description = "Return each image as a URL or a base64-encoded PNG." ) user: Optional[str] = Field(None, description = "End-user identifier (accepted, unused).") - # gpt-image-only; declared so we can reject it clearly instead of returning JSON to a client that - # asked for an SSE stream. + # gpt-image-only; declared so we can reject it clearly instead of returning JSON to a client that asked for an SSE stream. stream: Optional[bool] = Field( None, description = "Streaming image generation is not supported; omit or set false." ) @@ -2822,8 +2799,7 @@ class ImageGenerationRequest(BaseModel): @field_validator("n", "size", "response_format", mode = "before") @classmethod def _null_means_default(cls, value, info): - # OpenAI marks these nullable WITH a default, so an explicit null means "use the default": - # coalesce it instead of 400-ing a spec-valid body. + # OpenAI marks these nullable WITH a default, so an explicit null means "use the default": coalesce it instead of 400-ing a spec-valid body. if value is None: return cls.model_fields[info.field_name].default return value @@ -2955,8 +2931,7 @@ class VideoLoadRequest(BaseModel): @field_validator("attention_backend", mode = "before") @classmethod def _normalize_attention_backend(cls, value): - # The dispatcher accepts case/whitespace variants ("CuDNN", " sage "), but the Literal above is - # validated before any normaliser runs, so fold a string to its canonical form here. + # The dispatcher accepts case/whitespace variants ("CuDNN", " sage "), but the Literal above is validated before any normaliser runs, so fold a string to its canonical form here. return value.strip().lower() if isinstance(value, str) else value @@ -2967,8 +2942,7 @@ class VideoGenerateRequest(BaseModel): negative_prompt: Optional[str] = Field( None, description = "What to avoid (if the model supports it)" ) - # Width/height/num_frames/fps default per loaded family (the backend snaps them to its required - # multiples/lattice), so they are optional here. + # Width/height/num_frames/fps default per loaded family (the backend snaps them to its required multiples/lattice), so they are optional here. width: Optional[int] = Field( None, ge = 32, le = 2048, description = "Frame width in pixels (family multiple)" ) @@ -2999,8 +2973,7 @@ class VideoGenerateRequest(BaseModel): "pipeline default it to the main guidance. Ignored by single-DiT families (their pipeline " "signature has no second guidance kwarg).", ) - # le = 2**53-1: seeds round-trip through JSON gallery recipes, where JavaScript rounds integers - # above Number.MAX_SAFE_INTEGER and a restored recipe would generate a different clip. + # le = 2**53-1: seeds round-trip through JSON gallery recipes, where JavaScript rounds integers above Number.MAX_SAFE_INTEGER and a restored recipe would generate a different clip. seed: Optional[int] = Field( None, ge = 0, le = 2**53 - 1, description = "Seed for reproducibility (random if omitted)" ) @@ -3151,9 +3124,7 @@ class VideoStatusResponse(BaseModel): defaults: Optional[VideoGenerationDefaults] = Field( None, description = "Per-family generation defaults + shape constraints; null when unloaded" ) - # Additive: per-Advanced-control provenance {control: {value, source, reason}}. Same shape as the - # diffusion status; null when nothing is loaded. The frontend renders an "Auto: X" badge next to - # each control whose source == "auto". + # Additive: per-Advanced-control provenance {control: {value, source, reason}}. Same shape as the diffusion status; null when nothing is loaded. The frontend renders an "Auto: X" badge next to each control whose source == "auto". resolved: Optional[Dict[str, DiffusionResolvedControl]] = Field( None, description = "Per-control resolved value + provenance (source auto|explicit + reason), " diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 6a950c00ff..0fc858ea4f 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -226,9 +226,7 @@ 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 - # refuse while one is active or the two trainers contend for VRAM. 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. Symmetric with the check in start_diffusion_training. if _diffusion_training_active(): return TrainingJobResponse( job_id = "", @@ -470,18 +468,14 @@ async def start_training( logger.warning("Could not shut down export subprocess: %s", e) try: - # 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() no-ops 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. + # A resident or in-flight Images pipeline also holds GPU memory the run needs and cannot be cheaply sized, so tear it down unconditionally like the export subprocess above (the chat block below fit-checks; diffusion cannot). + # unload() no-ops when nothing is loaded and preempts an in-flight load; release the arbiter so it does not 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( @@ -493,9 +487,7 @@ 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 run needs too, and loads under the - # VIDEO arbiter owner the diffusion teardown above never touches. Tear it down the same way and - # release VIDEO, so a resident video session can't OOM the run. Must precede the chat block. + # 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 and release VIDEO, so a resident video session cannot OOM the run. Must precede the chat block. from core.inference import gpu_arbiter from core.inference.video import get_video_backend @@ -538,17 +530,9 @@ async def start_training( from utils.transformers_version import SidecarSwapInProgress try: - # 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 freeze every concurrent - # status/cancel/UI request. Overlapping starts are serialized by the backend's own guard. - # - # The diffusion admission is held ACROSS the spawn so the cross-trainer decision is - # atomic. The _diffusion_training_active() check above is separated from this point by - # dataset validation and memory coordination, and the diffusion route likewise checks - # this backend well before it reserves, so two near-simultaneous starts of different - # types could both pass their checks and train on the same GPU. Entering this context - # re-tests the diffusion state under the service's own lock, and while it is held + # Offloaded to a worker thread: the hook 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 freeze every concurrent status/cancel/UI request. Overlapping starts are serialized by the backend own guard. + # The diffusion admission is held ACROSS the spawn so the cross-trainer decision is atomic: the _diffusion_training_active() check above is separated from this point by dataset validation and memory coordination, and the diffusion route likewise checks this backend well before it reserves, so two near-simultaneous starts of different types could both pass their checks and train on the same GPU. + # Entering this context re-tests the diffusion state under the service own lock, and while it is held reserve() refuses -- so exactly one of the two wins. # reserve() refuses -- so exactly one of the two wins. with _diffusion_gpu_admission(): success = await asyncio.to_thread( @@ -1162,9 +1146,7 @@ async def stream_training_progress( # ── Diffusion (SDXL) LoRA training ──────────────────────────────────────────── -# 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). +# 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). def _diffusion_training_active() -> bool: @@ -1280,9 +1262,7 @@ 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") @@ -1292,9 +1272,7 @@ 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 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. + # A resident Video pipeline loads under the VIDEO arbiter owner the Images teardown above does not free; unload it too (no-op when nothing is loaded) and release VIDEO so a resident video session cannot OOM the diffusion trainer. from core.inference import gpu_arbiter from core.inference.video import get_video_backend @@ -1307,8 +1285,7 @@ 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 (like the LLM path does for an in-flight load) rather than risk an OOM. + # The SDXL trainer footprint cannot 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") @@ -1349,8 +1326,7 @@ def _preflight_gated_base(base_model: str, hf_token: Optional[str]) -> None: f"try again." ), ) - # 404 (e.g. a repo without a root model_index.json) and other codes are not an access problem; - # let the trainer surface any genuine load error. + # 404 (e.g. a repo without a root model_index.json) and other codes are not an access problem; let the trainer surface any genuine load error. except Exception: # noqa: BLE001 -- network/DNS hiccup must not block a start return @@ -1373,9 +1349,7 @@ def _resolve_diffusion_data_dir(raw: str) -> Path: # A single component that is not "..", so joining under datasets_root() cannot escape it. if not p.is_absolute() and len(p.parts) == 1 and p.parts[0] != "..": direct = datasets_root() / value - # Route a bare name through the same protected resolver the CRUD routes use, so a name to - # external-directory symlink is rejected here too (is_dir() follows the link). A broken symlink - # is included so it is rejected, not passed to resolve_dataset_path. + # Route a bare name through the same protected resolver the CRUD routes use, so a name to external-directory symlink is rejected here too (is_dir() follows the link). A broken symlink is included so it is rejected, not passed to resolve_dataset_path. if direct.is_dir() or direct.is_symlink(): return _resolve_dataset_folder(value) return resolve_dataset_path(raw) @@ -1390,9 +1364,7 @@ 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 - # 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. + # 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. if via_api_key is True: from core.inference.llama_keepwarm import other_inference_request_count if ( @@ -1408,8 +1380,7 @@ 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( @@ -1424,26 +1395,20 @@ 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 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 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 config["data_dir"] = str(_resolve_diffusion_data_dir(config["data_dir"])) config["output_dir"] = str(resolve_output_dir(config["output_dir"])) - # The persistent conditioning cache is another directory the TRAINER writes to, so it gets - # the same containment as output_dir rather than the trainer's cwd. Blank/None means the - # in-memory cache (the trainer's own "off"), so it must not resolve to the outputs root. + # The persistent conditioning cache is another directory the TRAINER writes to, so it gets the same containment as output_dir rather than the trainer cwd. Blank/None means the in-memory cache (the trainer own "off"), so it must not resolve to the outputs root. cond_cache = str(config.get("cond_cache_dir") or "").strip() config["cond_cache_dir"] = str(resolve_output_dir(cond_cache)) if cond_cache else None except ValueError as e: raise HTTPException(status_code = 400, detail = str(e)) - # 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. + # Validate the config BEFORE freeing resident GPU workloads, so a start then refused (bad numbers, non-SDXL base) never tears down the user chat/Images model. service.start() re-runs this cheaply before spawn. from core.training.diffusion_lora_trainer import _config_from_dict try: @@ -1451,11 +1416,8 @@ async def start_diffusion_training( except ValueError as e: raise HTTPException(status_code = 400, detail = str(e)) - # Only the DiT trainer reads cond_cache_dir. The SDXL trainer builds a per-process in-memory - # latent cache and never touches the persistent store, so accepting the option there promised - # cross-run reuse that never happened and silently re-encoded the dataset every run. Refuse it - # instead of ignoring it. Checked against the RESOLVED family, not the request field, so a - # request that omits model_family and lets an SDXL base be detected is caught too. + # Only the DiT trainer reads cond_cache_dir. The SDXL trainer builds a per-process in-memory latent cache and never touches the persistent store, so accepting the option there promised cross-run reuse that never happened and silently re-encoded the dataset every run. + # Refuse it instead of ignoring it. Checked against the RESOLVED family, not the request field, so a request that omits model_family and lets an SDXL base be detected is caught too. if cond_cache and normalized_cfg.resolved_family == "sdxl": raise HTTPException( status_code = 400, @@ -1467,10 +1429,8 @@ async def start_diffusion_training( ), ) - # 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 or stub-torchao host never tears down residents for a run that cannot start. + # Preflight the requested DiT precision BEFORE freeing GPU residents: the trainer 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 model. + # Fail fast (400) so a pre-Ampere GPU 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( @@ -1479,8 +1439,7 @@ async def start_diffusion_training( if _precision_reason: raise HTTPException(status_code = 400, detail = _precision_reason) - # Run the trainers' trust gate here too (both assert the same predicate before from_pretrained), - # so an untrusted/typoed base 400s BEFORE freeing GPU residents rather than failing in the child. + # 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 rather than failing in the child. from core.training.diffusion_train_common import _assert_trusted_base_model try: @@ -1488,10 +1447,8 @@ 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 model and never - # surfaces as a confusing mid-load 401. Offloaded to a worker thread: it does a blocking urlopen - # HEAD (5s timeout) that would otherwise stall the event loop. + # Preflight access to a gated base repo with the user token BEFORE freeing GPU residents, so a missing/insufficient token fails fast (400) without tearing down the user model and never surfaces as a confusing mid-load 401. + # Offloaded to a worker thread: it does a blocking urlopen HEAD (5s timeout) that would otherwise stall the event loop. await asyncio.to_thread( _preflight_gated_base, config.get("base_model", ""), config.get("hf_token") ) @@ -1499,42 +1456,31 @@ async def start_diffusion_training( from core.training import diffusion_train_common as _dtc service = get_diffusion_training_service() - # Reserve the training slot BEFORE the dataset preflight (not just before freeing residents): - # is_active() otherwise flips true only at service.start(), so during this scan -- which - # decode-probes every image and can take a while -- a concurrent upload/caption/delete would pass - # _require_diffusion_dataset_mutable() and mutate the dataset the trainer is about to read, and a - # concurrent /images/load or /video/load would double-allocate VRAM. reserve() is a - # compare-and-set, so a second overlapping start 409s before touching anything; unreserve() runs - # in the finally ONLY when THIS request reserved. + # Reserve the training slot BEFORE the dataset preflight (not just before freeing residents): is_active() otherwise flips true only at service.start(), so during this scan -- which decode-probes every image and can take a while -- a concurrent upload/caption/delete would pass _require_diffusion_dataset_mutable() and mutate the dataset the trainer is about to read, and a concurrent /images/load or /video/load would double-allocate VRAM. + # reserve() is a compare-and-set, so a second overlapping start 409s before touching anything; unreserve() runs in the finally ONLY when THIS request reserved. reserved = False try: service.reserve() reserved = True - # Preflight the dataset: 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. + # Preflight the dataset: a missing/empty/uncaptionable data_dir otherwise fails inside the spawned trainer AFTER the user model was evicted. Same discovery the trainer runs, so the two cannot disagree. try: await asyncio.to_thread( _dtc.discover_image_caption_pairs, config["data_dir"], instance_prompt = config.get("instance_prompt") or None, caption_column = config.get("caption_column") or "text", - # Decode-probe every image now (cheap PIL header check) so a corrupt/zero-byte upload 400s BEFORE - # _free_gpu_for_diffusion_training() tears down the user's models. + # 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 models. verify_images = True, ) except (FileNotFoundError, ValueError) as e: raise HTTPException(status_code = 400, detail = str(e)) - # Free resident GPU workloads (export / Images pipeline / chat) before the trainer loads its own - # 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 responsive. + # 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 responsive. await asyncio.to_thread(_free_gpu_for_diffusion_training) job_id = service.start(config) except ValueError as e: raise HTTPException(status_code = 400, detail = str(e)) except RuntimeError as e: - # A job is already running (or a start is already reserved), or a dataset mutation is open - # (DatasetMutationInFlight) -- the same interlock from the other side, so also a 409. + # A job is already running (or a start is already reserved), or a dataset mutation is open (DatasetMutationInFlight) -- the same interlock from the other side, so also a 409. raise HTTPException(status_code = 409, detail = str(e)) except HTTPException: raise @@ -1547,8 +1493,7 @@ async def start_diffusion_training( log = logger, ) finally: - # 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. + # On success the now-live proc keeps is_active() true; on failure this clears the reservation so training is not left permanently "active". Only the request that reserved clears it. if reserved: service.unreserve() return DiffusionTrainingStartResponse(job_id = job_id, status = "running") @@ -1595,8 +1540,7 @@ 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 wrong-typed field 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 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: @@ -1613,21 +1557,17 @@ 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) makes - # DiffusionTrainingRunDetail(**rec) raise TypeError -- not the ValidationError caught below -- and - # 500 the endpoint. Treat any non-dict record as absent, like the list route. + # 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, 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) reads as absent rather than 500 the - # endpoint, like 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.") -# Extensions accepted into an image-training dataset folder: images the trainer reads, plus its -# caption sources (per-image sidecars and metadata/captions jsonl). +# Extensions accepted into an image-training dataset folder: images the trainer reads, plus its caption sources (per-image sidecars and metadata/captions jsonl). _DIFFUSION_DATASET_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"} _DIFFUSION_DATASET_TEXT_EXTS = {".txt", ".caption", ".jsonl"} @@ -1648,10 +1588,8 @@ def _resolve_dataset_caption( try: caption = sidecar.read_text(encoding = "utf-8").strip() except (OSError, UnicodeError): - # Unreadable / invalid UTF-8 sidecar: the EMPTY TOMBSTONE, not "no sidecar", which - # is what the trainer does with it. Uploads accept raw sidecar bytes, so reading it - # as absent let the grid and the dataset summary show a metadata caption that the - # run would silently replace with the instance prompt (or skip the image over). + # Unreadable / invalid UTF-8 sidecar: the EMPTY TOMBSTONE, not "no sidecar", which is what the trainer does with it. + # Uploads accept raw sidecar bytes, so reading it as absent let the grid and the dataset summary show a metadata caption that the run would silently replace with the instance prompt (or skip the image over). caption = "" break if not sidecar_present: @@ -1700,9 +1638,7 @@ def _import_response( 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 - # over metadata precedence the trainer uses: an empty tombstone sidecar shadows a metadata row and - # makes the trainer skip the image, so counting it would mislabel an uncaptioned dataset. + # Count an image as captioned only when it resolves to a NON-EMPTY caption via the same sidecar over metadata precedence the trainer uses: an empty tombstone sidecar shadows a metadata row and makes the trainer skip the image, so counting it would mislabel an uncaptioned dataset. meta_captions = _load_metadata_captions(folder) images = captions = 0 for f in folder.iterdir(): @@ -1728,13 +1664,11 @@ async def diffusion_training_info(current_subject: str = Depends(get_current_sub root = datasets_root() found: list[DiffusionDatasetSummary] = [] try: - # 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. + # 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() - # Skip symlinked dirs: the CRUD resolver rejects them, so discovery must not advertise one as - # selectable (the read/caption/delete routes would refuse it). + # Skip symlinked dirs: the CRUD resolver rejects them, so discovery must not advertise one as selectable (the read/caption/delete routes would refuse it). if p.is_dir() and not p.is_symlink() and not p.name.startswith(".") ) except OSError: @@ -1762,9 +1696,7 @@ async def diffusion_training_info(current_subject: str = Depends(get_current_sub _DATASET_NAME_RE = None # compiled lazily; module keeps its import block torch-free -# Reserved in EVERY directory on Windows, with or without an extension (NUL.txt is NUL). The -# superscript COM/LPT digits are recognised as digits by Win32 and are reserved too. -# https://learn.microsoft.com/en-us/windows/win32/fileio/naming-a-file +# Reserved in EVERY directory on Windows, with or without an extension (NUL.txt is NUL). The superscript COM/LPT digits are recognised as digits by Win32 and are reserved too. https://learn.microsoft.com/en-us/windows/win32/fileio/naming-a-file _WINDOWS_RESERVED_NAMES = frozenset( {"con", "prn", "aux", "nul"} | {f"com{d}" for d in "123456789¹²³"} @@ -1832,17 +1764,12 @@ async def upload_diffusion_dataset( _require_diffusion_dataset_mutable() cleaned = _clean_diffusion_dataset_name(name) - # Run the same symlink + root-containment check as the read/caption/delete endpoints before any - # write, so a name to external-directory symlink can't make the staged upload write outside root. + # Run the same symlink + root-containment check as the read/caption/delete endpoints before any write, so a name to external-directory symlink cannot make the staged upload write outside root. folder = _resolve_dataset_folder(name, must_exist = False) folder.mkdir(parents = True, exist_ok = True) - # Serialize against a concurrent import into the SAME folder. The training interlock counts - # mutations rather than excluding them, and only imports took this lock, so an upload could add - # files while an import was materializing: the import's atomic promotion (os.rmdir + rename) - # then failed on the now-non-empty folder and fell back to a per-file move, silently merging - # the curated set with the uploaded one, and a failure partway through that move left a mixed - # dataset the image_count > 0 idempotency check accepts as complete. The duplicate-stem - # validation below reads the folder too, so it has to be inside the lock as well. + # Serialize against a concurrent import into the SAME folder. The training interlock counts mutations rather than excluding them, and only imports took this lock, so an upload could add files while an import was materializing: + # the import atomic promotion (os.rmdir + rename) then failed on the now-non-empty folder and fell back to a per-file move, silently merging the curated set with the uploaded one, and a failure partway through that move left a mixed dataset the image_count > 0 idempotency check accepts as complete. + # The duplicate-stem validation below reads the folder too, so it has to be inside the lock as well. _lock = _dataset_import_lock(folder) if not _lock.acquire(blocking = False): raise HTTPException( @@ -1857,15 +1784,11 @@ 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 isn't left on disk when the - # 400 fires; the upload is all-or-nothing. + # Validate every filename up front so a valid image ahead of a bad one is not left on disk when the 400 fires; the upload is all-or-nothing. names: list[str] = [] for f in files: - # 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 first so the true basename is taken for both separators. The read/caption/delete - # endpoints run the stored name through _safe_dataset_image_path, so a name still holding ".." - # here would list an image the grid can never preview, caption, or delete. + # Normalise to a safe basename. Path.name does not split on a backslash on POSIX, so a Windows client sending a backslash path in the multipart filename would be stored verbatim; fold backslashes first so the true basename is taken for both separators. + # The read/caption/delete endpoints run the stored name through _safe_dataset_image_path, so a name still holding ".." here would list an image the grid can never preview, caption, or delete. 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: @@ -1874,11 +1797,8 @@ 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 from different folders, or an API - # client repeating a part). The same-name exemption below is for SEPARATE repeat uploads, a - # deliberate overwrite; inside one batch the two parts are distinct files staged to the same - # destination on EVERY filesystem, so the later replace would silently discard the earlier one. - # Exact match only: a case VARIANT pair stays exempt per the stem guard. + # 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; + # inside one batch the two parts are distinct files staged to the same destination on EVERY filesystem, so the later replace would silently discard the earlier one. Exact match only: a case VARIANT pair stays exempt per the stem guard. fname_cf = filename.casefold() if filename in names: raise HTTPException( @@ -1889,18 +1809,12 @@ async def upload_diffusion_dataset( "uploading." ), ) - # 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 and earlier images in THIS batch. Re-uploading the exact same name stays - # an overwrite; caption/text files are exempt. + # 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 and earlier images in THIS batch. Re-uploading the exact same name stays an overwrite; caption/text files are exempt. if ext in _DIFFUSION_DATASET_IMAGE_EXTS: stem = Path(filename).stem - # Compare stems (and the same-name guard) case-insensitively: on case-insensitive filesystems two - # images whose stems differ only by case resolve to the SAME .txt sidecar, so a - # case-sensitive check would let both corrupt one caption. A same-name case variant is exempt ONLY - # when its stem also differs in case (one file on case-insensitive filesystems, separate sidecars - # on Linux). An EXTENSION-case variant (cat.PNG vs cat.png) has equal stems, so it is rejected. + # Compare stems (and the same-name guard) case-insensitively: on case-insensitive filesystems two images whose stems differ only by case resolve to the SAME .txt sidecar, so a case-sensitive check would let both corrupt one caption. + # A same-name case variant is exempt ONLY when its stem also differs in case (one file on case-insensitive filesystems, separate sidecars on Linux). An EXTENSION-case variant (cat.PNG vs cat.png) has equal stems, so it is rejected. stem_cf = stem.casefold() def _shares_sidecar(other_name: str) -> bool: @@ -1911,8 +1825,7 @@ async def upload_diffusion_dataset( or other.stem.casefold() != stem_cf ): 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). + # A casefold-equal full name is exempt unless the stems match EXACTLY (extension-case variants collide on one sidecar on case-sensitive filesystems). return other.stem == stem or other_name.casefold() != fname_cf clash = next( @@ -1931,16 +1844,13 @@ async def upload_diffusion_dataset( ), ) names.append(filename) - # 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. + # 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: for f, filename in zip(files, names): dest = folder / filename - # A filename-independent temp name so a long (but valid) filename can't overflow NAME_MAX once the - # staging suffix is added. + # A filename-independent temp name so a long (but valid) filename cannot overflow NAME_MAX once the staging suffix is added. tmp = folder / f".upload-{_uuid.uuid4().hex}.part" staged.append((tmp, dest)) with open(tmp, "wb") as out: @@ -1956,21 +1866,15 @@ async def upload_diffusion_dataset( ), ) out.write(chunk) - # Reject a decompression bomb before commit: a small compressible PNG can pass the byte limit yet - # decode to huge pixels and OOM the trainer's latent cache, so bound each image's dimensions from - # the header (mirrors diffusion._decode_b64_image). + # Reject a decompression bomb before commit: a small compressible PNG can pass the byte limit yet decode to huge pixels and OOM the trainer latent cache, so bound each image dimensions from the header (mirrors diffusion._decode_b64_image). if Path(filename).suffix.lower() in _DIFFUSION_DATASET_IMAGE_EXTS: _validate_uploaded_training_image(tmp, filename) uploaded += 1 - # Re-check the interlock immediately before the commit: the entry guard only saw the pre-upload - # state, so a /diffusion/start could have reserved the training slot while we were streaming. - # Committing now would move images/captions underneath the trainer; a 409 here leaves the staged - # temps to the finally below. + # Re-check the interlock immediately before the commit: the entry guard only saw the pre-upload state, so a /diffusion/start could have reserved the training slot while we were streaming. + # Committing now would move images/captions underneath the trainer; a 409 here leaves the staged temps to the finally below. _require_diffusion_dataset_mutable() - # Commit every staged file as one transaction. A plain replace loop is not atomic across files: a - # mid-loop failure leaves earlier destinations already overwritten while the request errors. Back - # up each pre-existing destination first, then on any failure drop the versions this request - # installed and restore every displaced original. + # Commit every staged file as one transaction. A plain replace loop is not atomic across files: a mid-loop failure leaves earlier destinations already overwritten while the request errors. + # Back up each pre-existing destination first, then on any failure drop the versions this request installed and restore every displaced original. backups: list[tuple[Path, Optional[Path]]] = [] # (dest, backup path or None) installed: list[Path] = [] try: @@ -2025,8 +1929,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 -# image discovery (both scan only top-level files). +# Thumbnails live in a hidden subdir so they never appear in dataset listings or the trainer image discovery (both scan only top-level files). _THUMBS_DIRNAME = ".thumbs" _MAX_CAPTION_CHARS = 2000 @@ -2039,9 +1942,7 @@ def _resolve_dataset_folder(name: str, *, must_exist: bool = True) -> Path: cleaned = _clean_diffusion_dataset_name(name) root = datasets_root().resolve() folder = root / cleaned - # Reject a symlinked dataset directory and prove the resolved folder stays under root: - # _safe_dataset_image_path only checks each image path, so a folder symlinked to an external - # directory would let read / caption / delete operate on files outside Studio. + # Reject a symlinked dataset directory and prove the resolved folder stays under root: _safe_dataset_image_path only checks each image path, so a folder symlinked to an external directory would let read / caption / delete operate on files outside Studio. if folder.is_symlink(): raise HTTPException( status_code = 400, @@ -2059,8 +1960,7 @@ def _resolve_dataset_folder(name: str, *, must_exist: bool = True) -> Path: return folder -# Per-side dimension bound for uploaded training images, matching diffusion._decode_b64_image's -# 4096px inference guard, so a compressible PNG can't smuggle huge pixels past the byte limit. +# Per-side dimension bound for uploaded training images, matching diffusion._decode_b64_image 4096px inference guard, so a compressible PNG cannot smuggle huge pixels past the byte limit. _MAX_TRAINING_IMAGE_SIDE = 4096 @@ -2076,9 +1976,7 @@ def _validate_uploaded_training_image(path: Path, original_name: str) -> None: with Image.open(path) as image: width, height = image.size except Image.DecompressionBombError: - # Past Pillow's own hard limit (~179 MP) Image.open() raises before .size can be read. That error - # derives straight from Exception (not OSError/ValueError), so letting it escape 500s the upload; - # it is exactly the oversized image this guard rejects. + # Past Pillow own hard limit (~179 MP) Image.open() raises before .size can be read. That error derives straight from Exception (not OSError/ValueError), so letting it escape 500s the upload; it is exactly the oversized image this guard rejects. raise HTTPException( status_code = 400, detail = ( @@ -2126,8 +2024,7 @@ def _load_metadata_captions(folder: Path) -> dict[str, str]: meta_path = folder / meta_name if not meta_path.is_file(): continue - # Tolerate a bad upload (invalid UTF-8, or a line of non-object JSON): skip the record so the - # info / labeling / caption / summary endpoints don't 500. + # Tolerate a bad upload (invalid UTF-8, or a line of non-object JSON): skip the record so the info / labeling / caption / summary endpoints do not 500. try: lines = meta_path.read_text(encoding = "utf-8").splitlines() except (OSError, UnicodeError): @@ -2168,17 +2065,13 @@ def _image_record( caption = sidecar.read_text(encoding = "utf-8").strip() source = "sidecar" except (OSError, UnicodeError): - # Unreadable / invalid UTF-8 sidecar (uploads store text sidecars as raw bytes): - # UnicodeDecodeError is a ValueError, not an OSError, so an OSError-only guard let it 500 the - # whole labeling grid. The trainer treats ANY existing sidecar as the empty tombstone and - # never reads metadata for that image, so showing a metadata caption here would display a - # label the run silently replaces with the instance prompt. + # Unreadable / invalid UTF-8 sidecar (uploads store text sidecars as raw bytes): UnicodeDecodeError is a ValueError, not an OSError, so an OSError-only guard let it 500 the whole labeling grid. + # The trainer treats ANY existing sidecar as the empty tombstone and never reads metadata for that image, so showing a metadata caption here would display a label the run silently replaces with the instance prompt. caption = None source = "sidecar" break if caption is None and not sidecar_present: - # 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. + # 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 order. meta = meta_captions.get(image_path.name) if meta is None: try: @@ -2253,9 +2146,7 @@ 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 sharing a stem but - # differing by extension would otherwise collide on one cache file, and an mtime-newer cache for - # the first would be served for the second. + # Key on the full filename (stem + extension), not the stem: two images sharing a stem but differing by extension 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: @@ -2305,10 +2196,8 @@ 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, 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. - # With no metadata caption it is a plain cleanup. + # Blank must actually clear. Unlinking alone would resurface this image metadata.jsonl / captions.jsonl caption, so when one exists write an EMPTY sidecar instead: + # both the reader and the trainer discovery treat an existing sidecar as authoritative even when empty, a tombstone. With no metadata caption it is a plain cleanup. meta = _load_metadata_captions(folder) try: rel = image_path.relative_to(folder).as_posix() @@ -2342,9 +2231,7 @@ async def delete_diffusion_dataset_image( import glob as _glob image_path.unlink(missing_ok = True) - # Sidecars are keyed on the STEM, so cat.jpg and cat.png share cat.txt: both the trainer's - # pair discovery and the labeling grid resolve either image to it. Deleting it with one of - # them would silently strip the survivor's caption and change what the next run trains on. + # Sidecars are keyed on the STEM, so cat.jpg and cat.png share cat.txt: both the trainer pair discovery and the labeling grid resolve either image to it. Deleting it with one of them would silently strip the survivor caption and change what the next run trains on. # New collisions are refused at upload, but hand-made and legacy folders still have them. stem_still_used = any( p.is_file() @@ -2358,9 +2245,7 @@ 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 strand this image's thumbs or delete a same-stem sibling's. Escape the name: a raw - # glob metacharacter would match siblings' thumbs while leaving its own behind. + # Thumbs are keyed on the full filename (stem + extension), so match that here too; a stem-only glob would strand this image thumbs or delete a same-stem sibling. Escape the name: a raw glob metacharacter would match siblings thumbs while leaving its own behind. for t in thumbs_dir.glob(f"{_glob.escape(image_path.name)}_*.jpg"): t.unlink(missing_ok = True) return {"deleted": image_path.name} @@ -2368,10 +2253,8 @@ async def delete_diffusion_dataset_image( return await asyncio.to_thread(remove) -# Curated, license-labelled example datasets for one-click import. ``loader`` picks the -# materialization strategy: "hf_dataset" streams rows from datasets.load_dataset (image + optional -# caption column); "imagefolder_jsonl" snapshot-downloads a dataset repo whose captions live in a -# *.jsonl (file_name/text) not a standard metadata.jsonl. +# 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) not a standard metadata.jsonl. _DATASET_EXAMPLES: list[dict] = [ { "id": "dreambooth-dog", @@ -2416,8 +2299,7 @@ _DATASET_EXAMPLES: list[dict] = [ "description": "100 butterfly photos. No captions, so use the trigger prompt.", "license": "CC0", "image_cap": 100, - # The metadata columns are species names / boilerplate alt-text, not 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, @@ -2510,19 +2392,15 @@ def _materialize_hf_dataset(entry: dict, dest: Path, cap: int) -> int: kwargs = {"split": "train"} if entry.get("no_checks"): kwargs["verification_mode"] = "no_checks" - # Stream rather than prepare the whole split: the loop keeps at most `cap` rows (10-100) while - # these curated repos run to 49,859 rows / 328 MB (m1guelpf/nouns) and 1,000 rows / 237 MB - # (huggan/smithsonian_butterflies_subset), all of which a prepared load downloads and converts - # before the first row is read. A repo that cannot stream (loading script, no listed data files) - # falls back to the prepared load so the one-click import still works. + # Stream rather than prepare the whole split: the loop keeps at most `cap` rows (10-100) while these curated repos run to 49,859 rows / 328 MB (m1guelpf/nouns) and 1,000 rows / 237 MB (huggan/smithsonian_butterflies_subset), all of which a prepared load downloads and converts before the first row is read. + # A repo that cannot stream (loading script, no listed data files) falls back to the prepared load so the one-click import still works. try: ds = load_dataset(entry["repo"], streaming = True, **kwargs) features = ds.features except Exception: # noqa: BLE001 -- not streamable; the prepared load is the fallback ds = load_dataset(entry["repo"], **kwargs) features = ds.features - # Streaming can hand back a dataset whose features are only known once a row is read, so the - # columns are resolved from the first row in that case. + # Streaming can hand back a dataset whose features are only known once a row is read, so the columns are resolved from the first row in that case. image_col = _detect_image_column(features) if features else None if image_col is None and features: raise HTTPException( @@ -2638,11 +2516,8 @@ async def import_diffusion_dataset_example( folder.mkdir(parents = True, exist_ok = True) if _diffusion_dataset_summary(folder).image_count > 0: return _import_response(entry, folder, imported = 0) - # One import at a time per dataset folder. The training interlock COUNTS mutations rather - # than excluding them, so two imports of different examples into the same empty name both - # passed the emptiness check; the loser then merged its files into the winner's folder, - # overwriting same-numbered images and leaving a dataset whose images and captions came - # from two sources. Refusing the second is honest: the first is already filling that name. + # One import at a time per dataset folder. The training interlock COUNTS mutations rather than excluding them, so two imports of different examples into the same empty name both passed the emptiness check; the loser then merged its files into the winner folder, overwriting same-numbered images and leaving a dataset whose images and captions came from two sources. + # Refusing the second is honest: the first is already filling that name. lock = _dataset_import_lock(folder) if not lock.acquire(blocking = False): raise HTTPException( @@ -2663,17 +2538,12 @@ async def import_diffusion_dataset_example( import tempfile imported = 0 - # Re-read under the lock: a winner may have promoted its staging dir while this request - # was checking, so the folder may no longer be empty. Returning it as-is matches the - # idempotent path rather than mixing two imports. + # Re-read under the lock: a winner may have promoted its staging dir while this request was checking, so the folder may no longer be empty. Returning it as-is matches the idempotent path rather than mixing two imports. existing = _diffusion_dataset_summary(folder) 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 partial materialize 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 and strand a truncated dataset (there is no dataset-delete flow). Staged as a - # hidden same-filesystem sibling so promotion is an atomic rename. + # Materialize into a private staging dir and promote into the dataset folder only after the whole import succeeds. A partial materialize 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 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: @@ -2693,20 +2563,13 @@ async def import_diffusion_dataset_example( status_code = 502, detail = f"No images found in '{entry['repo']}'.", ) - # Promote the fully-materialized staging dir as a UNIT: a same-filesystem rename is - # atomic, so a hard process death leaves either the old folder or the finished - # import, never a half-filled one that the image_count>0 check above would accept as - # complete on retry. rmdir needs an empty target, and an image-empty folder can still - # hold files (a .thumbs cache, or a metadata.jsonl / captions from an earlier - # upload), so fold those INTO the staging dir first and keep one atomic promotion. - # Moving them one by one into a live folder instead -- the old fallback -- gave up - # exactly the atomicity this whole staging dance exists for. + # Promote the fully-materialized staging dir as a UNIT: a same-filesystem rename is atomic, so a hard process death leaves either the old folder or the finished import, never a half-filled one that the image_count>0 check above would accept as complete on retry. + # rmdir needs an empty target, and an image-empty folder can still hold files (a .thumbs cache, or a metadata.jsonl / captions from an earlier upload), so fold those INTO the staging dir first and keep one atomic promotion. + # Moving them one by one into a live folder instead -- the old fallback -- gave up exactly the atomicity this whole staging dance exists for. for p in sorted(folder.iterdir()): dest = staging / p.name if dest.exists(): - # Same name in both: the import's own file wins, exactly as the previous - # per-file move did by overwriting it. Drop the old one so the folder can - # still be emptied for the rename. + # Same name in both: the import own file wins, exactly as the previous per-file move did by overwriting it. Drop the old one so the folder can still be emptied for the rename. if p.is_dir(): shutil.rmtree(p, ignore_errors = True) else: @@ -2716,8 +2579,7 @@ async def import_diffusion_dataset_example( try: os.rmdir(folder) except OSError as e: - # Something landed in the folder in the meantime. Fail with the dataset - # untouched rather than promoting it piecemeal. + # Something landed in the folder in the meantime. Fail with the dataset untouched rather than promoting it piecemeal. raise HTTPException( status_code = 409, detail = ( diff --git a/studio/frontend/src/features/video/video-page.tsx b/studio/frontend/src/features/video/video-page.tsx index b1c92fb62c..8880352e69 100644 --- a/studio/frontend/src/features/video/video-page.tsx +++ b/studio/frontend/src/features/video/video-page.tsx @@ -82,27 +82,21 @@ import { unloadVideoModel, } from "./api"; -// Curated models come from the shared catalog: one canonical group per model with its -// artifacts as data (the HunyuanVideo group carries both the 480p and 720p repacks), and the -// load kind per artifact via loadSpecFor (replacing the old PIPELINE_MODELS table). The picker -// renders groups with a format second level -- which also surfaces LTX-2.3 in Recommended (its +// Curated models come from the shared catalog: one canonical group per model with its artifacts as data (the HunyuanVideo group carries both the 480p and 720p repacks), and the load kind per artifact via loadSpecFor (replacing the old PIPELINE_MODELS table). +// The picker renders groups with a format second level -- which also surfaces LTX-2.3 in Recommended (its HF pipeline_tag is image-to-video, so the live text-to-video listing missed it). // HF pipeline_tag is image-to-video, so the live text-to-video listing missed it). const VIDEO_MODELS: ModelOption[] = catalogToModelOptions(VIDEO_CATALOG); -// Per-model generation defaults (steps + guidance), matched by repo-id substring, most -// specific first. The distilled model wants very few steps and no guidance; the full base -// model wants more steps and real CFG. +// Per-model generation defaults (steps + guidance), matched by repo-id substring, most specific first. The distilled model wants very few steps and no guidance; the full base model wants more steps and real CFG. const DEFAULT_GEN = { steps: 8, guidance: 1 }; const MODEL_DEFAULTS: Array<{ match: string; steps: number; guidance: number }> = [ // "distilled" before the generic "ltx": the distilled model runs at 8 steps, guidance 1. { match: "distilled", steps: 8, guidance: 1 }, { match: "ltx", steps: 40, guidance: 4 }, - // Wan2.2 pipelines default to 50 steps at CFG 5.0 (WanPipeline defaults, verified in - // diffusers 0.39). The backend supplies the fps per family (24 for TI2V-5B, 16 for A14B). + // Wan2.2 pipelines default to 50 steps at CFG 5.0 (WanPipeline defaults, verified in diffusers 0.39). The backend supplies the fps per family (24 for TI2V-5B, 16 for A14B). { match: "wan", steps: 50, guidance: 5 }, - // HunyuanVideo-1.5 runs 50 steps; guidance 6 matches the guider the repo ships - // (the backend writes it onto the guider component, there is no pipeline kwarg). + // HunyuanVideo-1.5 runs 50 steps; guidance 6 matches the guider the repo ships (the backend writes it onto the guider component, there is no pipeline kwarg). { match: "hunyuanvideo", steps: 50, guidance: 6 }, ]; @@ -111,42 +105,32 @@ function defaultsFor(repoId: string): { steps: number; guidance: number } { return MODEL_DEFAULTS.find((d) => id.includes(d.match)) ?? DEFAULT_GEN; } -// Resolution presets offered before a model is loaded (default first). Once loaded, the -// backend's status.defaults.resolution_presets replaces these. +// Resolution presets offered before a model is loaded (default first). Once loaded, the backend status.defaults.resolution_presets replaces these. const FALLBACK_RESOLUTION_PRESETS: Array<[number, number]> = [ [768, 512], [1216, 704], [704, 1216], ]; -// Fallbacks used to build the duration presets before a model is loaded, so the duration -// select is populated and valid on first paint. +// Fallbacks used to build the duration presets before a model is loaded, so the duration select is populated and valid on first paint. const FALLBACK_FRAME_STEP = 8; const FALLBACK_FPS = 24; // Module cache of the backend-persisted gallery, so a tab switch re-renders instantly. -// The srcById entries are short-lived signed links, not object URLs: nothing is pinned in the -// webview, so they survive a remount and need no budget, eviction or revoke. The clip's bytes -// are streamed by the media element itself, which fetches ranges as it plays and drops what it -// no longer needs -- the whole point of not blob-ing a file that runs to hundreds of MB. +// The srcById entries are short-lived signed links, not object URLs: nothing is pinned in the webview, so they survive a remount and need no budget, eviction or revoke. The clip bytes are streamed by the media element itself, which fetches ranges as it plays -- the whole point of not blob-ing a file that runs to hundreds of MB. const galleryCache: { videos: GalleryVideo[]; hasMore: boolean; selectedId: string | null; quant: string | null; - // id -> the signed link and when it was minted. The link is short-lived (the backend expires it, - // and its signing secret is per-process, so a server restart invalidates every outstanding one), - // while this cache deliberately survives navigation -- so an entry has to be re-mintable rather - // than final, or playback, seeking and Save would 401 until a full page reload. + // id -> the signed link and when it was minted. The link is short-lived (the backend expires it, and its signing secret is per-process, so a server restart invalidates every outstanding one), + // while this cache deliberately survives navigation -- so an entry has to be re-mintable rather than final, or playback, seeking and Save would 401 until a full page reload. srcById: Map; - // Ids re-minted once after a media error already, so a clip that is broken for any other reason - // cannot spin in a mint/error loop. + // Ids re-minted once after a media error already, so a clip that is broken for any other reason cannot spin in a mint/error loop. refreshed: Set; // Ids with a mint in flight, so concurrent ensureSrc calls don't double-request. inflight: Set; - // Ids deleted while their link was still being minted, so a reply that lands after the delete - // isn't cached for a card that no longer exists. Clear-all bumps the epoch instead of listing - // every id. + // Ids deleted while their link was still being minted, so a reply that lands after the delete is not cached for a card that no longer exists. Clear-all bumps the epoch instead of listing every id. deleted: Set; epoch: number; } = { @@ -161,8 +145,7 @@ const galleryCache: { epoch: 0, }; -// Re-mint a cached link once it is this old. Comfortably inside the backend's own expiry, so a -// long-lived tab keeps working without waiting for a 401 to tell it the link died. +// Re-mint a cached link once it is this old. Comfortably inside the backend own expiry, so a long-lived tab keeps working without waiting for a 401 to tell it the link died. const VIDEO_LINK_REFRESH_MS = 6 * 60 * 60 * 1000; // Videos loaded per infinite-scroll page. @@ -188,9 +171,7 @@ function saveLink(href: string, filename: string) { link.click(); } -// MP4 saves the original file straight from its signed link (same-origin, so the download -// attribute is honoured); WebM / GIF are transcoded by the backend on demand (501 with a -// readable reason when the codec is absent). +// MP4 saves the original file straight from its signed link (same-origin, so the download attribute is honoured); WebM / GIF are transcoded by the backend on demand (501 with a readable reason when the codec is absent). async function downloadVideo( src: string, video: GalleryVideo, @@ -220,21 +201,17 @@ function clipMeta(video: GalleryVideo): string { return `${secs} · ${video.width}×${video.height}`; } -// Bar label for an in-flight generation: the phase ("Denoising step X/Y" during denoise, -// "Encoding video…" during export) plus an ETA once known. +// Bar label for an in-flight generation: the phase ("Denoising step X/Y" during denoise, "Encoding video..." during export) plus an ETA once known. function genStepLabel(p: VideoGenerateProgress): string { if (p.phase === "export") return "Encoding video…"; - // Text encoding and the first-step warmup run inside the pipeline before the first - // scheduler tick, so step 0 means "working, not denoising yet" -- up to a minute at - // 720p. Label that phase honestly instead of sitting on "Denoising step 0/N". + // Text encoding and the first-step warmup run inside the pipeline before the first scheduler tick, so step 0 means "working, not denoising yet" -- up to a minute at 720p. Label that phase honestly instead of sitting on "Denoising step 0/N". if (p.step === 0) return "Preparing (text encoding + warmup)…"; const base = p.total > 0 ? `Denoising step ${p.step}/${p.total}` : "Denoising…"; const eta = p.eta_seconds != null ? formatEta(p.eta_seconds) : ""; return eta ? `${base} · ~${eta}` : base; } -// The chat tab's model-load toast styling, reused verbatim so the video load toast is -// visually identical (persistent, progress bar, same chrome). +// The chat tab model-load toast styling, reused verbatim so the video load toast is visually identical (persistent, progress bar, same chrome). const LOAD_TOAST_CLASSNAMES = { toast: "chat-model-load-toast items-center gap-2.5", content: "gap-0.5 flex-1 min-w-0", @@ -242,8 +219,7 @@ const LOAD_TOAST_CLASSNAMES = { description: "mt-0 w-full", } as const; -// The download total for a video load can only be estimated from a companion base repo, so -// the toast shows a byte count rather than a hard percentage until the total is known. +// The download total for a video load can only be estimated from a companion base repo, so the toast shows a byte count rather than a hard percentage until the total is known. function loadFraction(p: VideoLoadProgress): number | null { if (!p.expected_bytes || p.expected_bytes <= 0) return null; return Math.min(1, p.downloaded_bytes / p.expected_bytes); @@ -275,8 +251,7 @@ function loadToastDescription(p: VideoLoadProgress) { ); } -// Toast args mirroring chat's: persistent, closeable, content in `description`. Pass `id` -// to update the existing toast in place instead of stacking a new one. +// Toast args mirroring chat: persistent, closeable, content in `description`. Pass `id` to update the existing toast in place instead of stacking a new one. function loadToastArgs(p: VideoLoadProgress, id?: string | number) { return { ...(id != null ? { id } : {}), @@ -346,9 +321,7 @@ function Field({ ); } -// The engaged value of a resolved Advanced control, formatted for its "Auto: X" badge. -// Short scheme/mode tokens go uppercase (FBCACHE); the attention backend the backend reports -// as `_native_cudnn` shows as cuDNN. +// The engaged value of a resolved Advanced control, formatted for its "Auto: X" badge. Short scheme/mode tokens go uppercase (FBCACHE); the attention backend the backend reports as `_native_cudnn` shows as cuDNN. function formatResolvedValue(value: string | boolean | null): string { if (value === null || value === "") return "Off"; if (typeof value === "boolean") return value ? "On" : "Off"; @@ -356,10 +329,8 @@ function formatResolvedValue(value: string | boolean | null): string { return value.toUpperCase(); } -// The "Auto: X" badge for one Advanced control: rendered only when the backend resolved that -// control itself (source === "auto"); an explicit user choice renders nothing. The reason is -// surfaced as a hover tooltip. Muted pill matching the panel's other chips. Reuses the same -// markup as the images page's ResolvedBadge. +// The "Auto: X" badge for one Advanced control: rendered only when the backend resolved that control itself (source === "auto"); an explicit user choice renders nothing. +// The reason is surfaced as a hover tooltip. Muted pill matching the panel other chips, same markup as the images page ResolvedBadge. function ResolvedBadge({ status, controlKey, @@ -444,8 +415,7 @@ function RecipePopover({ onRestore: (video: GalleryVideo) => void; active: boolean; }) { - // Controlled + force-closed off-tab: PopoverContent portals to body, so the hidden/inert - // page wrapper can't contain it when the page is kept mounted. + // Controlled + force-closed off-tab: PopoverContent portals to body, so the hidden/inert page wrapper cannot contain it when the page is kept mounted. const [open, setOpen] = useState(false); useEffect(() => { if (!active) setOpen(false); @@ -528,11 +498,9 @@ export function VideoPage({ active = true }: { active?: boolean }) { const [resolutionIdx, setResolutionIdx] = useState(0); // The chosen frame count (must lie on the family's temporal lattice: k*frame_step+1). const [numFrames, setNumFrames] = useState(FALLBACK_FRAME_STEP * 3 + 1); - // Advanced options live in a right-docked panel (like Chat's settings panel). Closed by - // default; a single fixed toggle in the top bar opens/closes it. + // Advanced options live in a right-docked panel (like Chat settings panel). Closed by default; a single fixed toggle in the top bar opens/closes it. const [advancedOpen, setAdvancedOpen] = useState(false); - // Advanced (load-time) options. "auto"/"off" map to the backend defaults (sent through on - // load). They apply when a model loads; a "Reapply" button reloads with the new values. + // Advanced (load-time) options. "auto"/"off" map to the backend defaults (sent through on load). They apply when a model loads; a "Reapply" button reloads with the new values. const [memoryMode, setMemoryMode] = useState<"auto" | "fast" | "balanced" | "low_vram">("auto"); const [speedMode, setSpeedMode] = useState<"auto" | "off" | "eager" | "default" | "max">("auto"); const [attentionBackend, setAttentionBackend] = useState< @@ -542,44 +510,34 @@ export function VideoPage({ active = true }: { active?: boolean }) { const [transformerQuant, setTransformerQuant] = useState< "auto" | "none" | "fp8" | "int8" | "nvfp4" | "mxfp8" >("auto"); - // The last load descriptor, so "Reapply" can reload the same model with new advanced - // options without the user re-picking it from the dropdown. + // The last load descriptor, so "Reapply" can reload the same model with new advanced options without the user re-picking it from the dropdown. const lastLoad = useRef<{ repoId: string; kind: "gguf" | "single_file" | "pipeline"; filename?: string } | null>( null, ); - // Whether this session holds a reapply descriptor (set only by our own loads). On a - // mount/refresh with a model already resident, status.loaded is true but lastLoad is null, so - // Reapply would do nothing -- hide the button rather than offer a dead control. + // Whether this session holds a reapply descriptor (set only by our own loads). On a mount/refresh with a model already resident, status.loaded is true but lastLoad is null, so Reapply would do nothing -- hide the button rather than offer a dead control. const [canReapply, setCanReapply] = useState(false); const [busy, setBusy] = useState(null); // Live per-step progress (phase / step / total + ETA) polled during generation. const [genStep, setGenStep] = useState(null); const genPollTimer = useRef | null>(null); - // visibilitychange handler active while a generation poll runs: background tabs clamp - // setInterval to >=1s (and can suspend it outright after ~5 min), so returning to the - // tab fires one immediate poll instead of waiting for a throttled tick. + // visibilitychange handler active while a generation poll runs: background tabs clamp setInterval to >=1s (and can suspend it outright after ~5 min), so returning to the tab fires one immediate poll instead of waiting for a throttled tick. const genVisibilityListener = useRef<(() => void) | null>(null); const [status, setStatus] = useState(null); - // Controlled so the body-portaled overlays force-close when this page is mounted but - // off-tab (a hidden/inert parent can't contain a body portal): the model selector. + // Controlled so the body-portaled overlays force-close when this page is mounted but off-tab (a hidden/inert parent cannot contain a body portal): the model selector. const [selectorOpen, setSelectorOpen] = useState(false); // Records come from the backend (durable); srcById maps each id to its object URL. const [videos, setVideos] = useState(() => galleryCache.videos); const [hasMore, setHasMore] = useState(() => galleryCache.hasMore); const [selectedId, setSelectedId] = useState(() => galleryCache.selectedId); - // Autoplay replays per selected clip (3 total plays, then pause). Reset on - // every selection change so a new generation or pick gets its own 3 plays. + // Autoplay replays per selected clip (3 total plays, then pause). Reset on every selection change so a new generation or pick gets its own 3 plays. const playCountRef = useRef(0); useEffect(() => { playCountRef.current = 0; }, [selectedId]); - // Pause the preview when this page stops being the visible one. The keep-alive layout only - // hides it, and display:none does not pause a media element, so a clip the user unmuted - // would keep decoding and playing its audio over whatever page they opened next. + // Pause the preview when this page stops being the visible one. The keep-alive layout only hides it, and display:none does not pause a media element, so a clip the user unmuted would keep decoding and playing its audio over whatever page they opened next. const previewRef = useRef(null); - // The media element's own handlers fire while the page is hidden, so they read `active` - // through a ref rather than closing over a stale render's value. + // The media element own handlers fire while the page is hidden, so they read `active` through a ref rather than closing over a stale render value. const activeRef = useRef(active); useEffect(() => { activeRef.current = active; @@ -590,8 +548,7 @@ export function VideoPage({ active = true }: { active?: boolean }) { ); // Guards a "load more" so a fast scroll can't fire several at once. const loadingMore = useRef(false); - // False once the page truly unmounts (app close / chat-only eject). The page stays mounted - // across tab switches, so a switch does NOT flip this. + // False once the page truly unmounts (app close / chat-only eject). The page stays mounted across tab switches, so a switch does NOT flip this. const isMounted = useRef(true); const pollTimer = useRef | null>(null); // The persistent load toast's id, so each poll updates it in place (chat-style). @@ -601,9 +558,7 @@ export function VideoPage({ active = true }: { active?: boolean }) { // The quant to restore if the current optimistic swap fails. const quantRevert = useRef<{ prev: string | null } | null>(null); // The Reapply target (and its canReapply flag) to restore if the optimistic swap fails. - // handleLoad overwrites lastLoad.current with the pending pick at load start; if the load then - // fails AFTER starting (error/eviction during download) the previous model stays resident, so - // the poll rolls lastLoad back rather than leave Reapply pointing at the failed pick. + // handleLoad overwrites lastLoad.current with the pending pick at load start; if the load then fails AFTER starting (error/eviction during download) the previous model stays resident, so the poll rolls lastLoad back rather than leave Reapply pointing at the failed pick. const lastLoadRevert = useRef<{ prev: typeof lastLoad.current; canReapply: boolean } | null>(null); const dismissLoadToast = useCallback(() => { @@ -625,8 +580,7 @@ export function VideoPage({ active = true }: { active?: boolean }) { ); const selectedSrc = selected ? srcById[selected.id] : undefined; - // The resolution presets + temporal lattice for the currently loaded family, or the - // fallbacks before anything is loaded. + // The resolution presets + temporal lattice for the currently loaded family, or the fallbacks before anything is loaded. const resolutionPresets = useMemo>(() => { const presets = status?.defaults?.resolution_presets; if (presets && presets.length > 0) { @@ -638,8 +592,7 @@ export function VideoPage({ active = true }: { active?: boolean }) { const frameStep = status?.defaults?.frame_step ?? FALLBACK_FRAME_STEP; const fps = status?.defaults?.fps ?? FALLBACK_FPS; - // Duration presets: valid frame counts (k*frame_step+1) closest to ~1s/2s/3s/5s at the - // current fps. Deduped so two targets that snap to the same count don't repeat. + // Duration presets: valid frame counts (k*frame_step+1) closest to ~1s/2s/3s/5s at the current fps. Deduped so two targets that snap to the same count do not repeat. const durationOptions = useMemo>(() => { const targets = [1, 2, 3, 5]; const seen = new Set(); @@ -666,9 +619,7 @@ export function VideoPage({ active = true }: { active?: boolean }) { const familyChanged = loadedFamily !== prevFamilyRef.current; prevFamilyRef.current = loadedFamily; setNumFrames((cur) => { - // A newly loaded family brings its own default clip length (121 frames for - // LTX-2); without this the pre-load fallback (25 frames, still on the new - // lattice) silently sticks and every default run is a ~1s clip. + // A newly loaded family brings its own default clip length (121 frames for LTX-2); without this the pre-load fallback (25 frames, still on the new lattice) silently sticks and every default run is a ~1s clip. if (familyChanged && loadedFamily && familyDefaultFrames) { const best = durationOptions.reduce((a, b) => Math.abs(b.frames - familyDefaultFrames) < Math.abs(a.frames - familyDefaultFrames) @@ -683,13 +634,8 @@ export function VideoPage({ active = true }: { active?: boolean }) { }); }, [durationOptions, loadedFamily, familyDefaultFrames]); - // Seed steps/guidance from the loaded model's backend defaults. On mount with a model already - // loaded (browser refresh, or a load from another client) only refreshStatus runs -- - // handleModelSelect never fires -- so the controls otherwise stick at the pre-load DEFAULT_GEN - // (8/1) and a base checkpoint wanting 40/4 generates a degraded clip. Key on the repo id so it - // fires once per newly-loaded model (distilled vs base of the same family differ); a later user - // edit isn't clobbered (the key changes only on model change), and a gallery restore (same - // repo) is left untouched. + // Seed steps/guidance from the loaded model backend defaults. On mount with a model already loaded (browser refresh, or a load from another client) only refreshStatus runs -- handleModelSelect never fires -- so the controls otherwise stick at the pre-load DEFAULT_GEN (8/1) and a base checkpoint wanting 40/4 generates a degraded clip. + // Key on the repo id so it fires once per newly-loaded model (distilled vs base of the same family differ); a later user edit is not clobbered (the key changes only on model change), and a gallery restore (same repo) is left untouched. const loadedModelKey = status?.loaded ? status.repo_id : null; const defaultSteps = status?.defaults?.steps; const defaultGuidance = status?.defaults?.guidance; @@ -703,10 +649,8 @@ export function VideoPage({ active = true }: { active?: boolean }) { } }, [loadedModelKey, defaultSteps, defaultGuidance]); - // Mint (once) a playable link for a record's MP4; cached across remounts. Unlike the images - // gallery this does NOT download the file: the link goes straight into the