From 9fde4b9991b52b7ae24481f34f02f4d291db0d25 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Jul 2026 11:40:05 +0000 Subject: [PATCH] Tighten comments across the remaining image stack files --- studio/backend/core/inference/diffusion.py | 1145 ++++++----------- .../backend/core/inference/sd_cpp_backend.py | 227 ++-- .../backend/core/inference/sd_cpp_server.py | 94 +- studio/backend/core/inference/video.py | 503 +++----- .../backend/core/inference/video_families.py | 210 +-- studio/backend/routes/video.py | 43 +- 6 files changed, 722 insertions(+), 1500 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 444a80c042..2473c5defa 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -114,17 +114,10 @@ from .diffusion_transformer_quant import ( logger = get_logger(__name__) -# A load resolves to exactly one of these "kinds", which decide how the transformer -# (and the rest of the pipeline) is built: -# "gguf" -- a single-file GGUF transformer dequantised on-device via -# GGUFQuantizationConfig; the VAE / text encoders / scheduler come -# from the companion base diffusers repo. The original behaviour. -# "single_file" -- a single-file *.safetensors transformer loaded with from_single_file -# WITHOUT the GGUF dequant config (e.g. an fp8 checkpoint); companions -# still come from the base repo. -# "pipeline" -- a full diffusers repo loaded with pipeline_cls.from_pretrained(repo_id), -# which pulls every component (transformer included) and re-applies any -# embedded quantization_config (e.g. a bnb-4bit pipeline) automatically. +# A load resolves to one "kind", deciding how the transformer + pipeline is built: +# "gguf" -- single-file GGUF transformer dequantised on-device; companions from base repo. +# "single_file" -- single-file *.safetensors transformer (e.g. fp8), no GGUF dequant; companions from base. +# "pipeline" -- full diffusers repo via from_pretrained, re-applying any embedded quant config. _MODEL_KINDS = frozenset({"gguf", "single_file", "pipeline"}) @@ -205,18 +198,13 @@ def _decode_b64_image(data: str, *, mode: str = "RGB") -> Any: blob = base64.b64decode(raw, validate = False) except (binascii.Error, ValueError) as exc: raise ValueError(f"Invalid base64 image data: {exc}") from exc - # Bound the decoded size. Every image-conditioned workflow (img2img / inpaint / upscale / - # reference / edit) decodes through here, so this single guard protects init, mask, and - # each reference image uniformly. PIL only WARNS in its 89-178MP "decompression bomb" soft - # zone and still loads (~0.5 GB RGB each, times up to 4 with multi-reference); cap the side - # well below that. 4096px covers txt2img's 2048 max, upscales, and normal outpaint canvases; - # anything larger is rejected with a clear 400 instead of risking an OOM. + # Bound the decoded size (this is the single decode path for every image-conditioned + # workflow). 4096px covers txt2img's 2048 max, upscales, and outpaint canvases; larger 400s. max_side = 4096 try: img = Image.open(io.BytesIO(blob)) - # Read the declared dimensions from the header (Image.open is lazy) and reject an - # over-limit image BEFORE img.load() decompresses its pixels, so a crafted - # small-payload/huge-dimension file can't spike memory before the guard runs. + # Reject an over-limit image from the header BEFORE img.load() decompresses pixels, so a + # crafted small-payload/huge-dimension file can't spike memory first. w, h = img.size if w > max_side or h > max_side: raise ValueError(f"Image is too large ({w}x{h}); maximum is {max_side}px per side.") @@ -282,27 +270,16 @@ def _compile_shape_dims(workflow: str, init_pil: Any, width: int, height: int) - return int(iw), int(ih) -# A small allowlist of well-known official base repos that may load as a full -# (non-GGUF) pipeline even though they are not under ``unsloth/``. These are -# safetensors-only checkpoints from their original publisher (no pickle, no remote -# code) that some architectures require: SDXL ships only as a full pipeline and has -# no unsloth-hosted GGUF, so without this its curated catalog entry could not load. -# Exact-match, lowercased, so it cannot be widened by a typo-squat. Extend -# deliberately, and never add a repo that carries pickled weights or remote code. -# The SDXL refiner is intentionally NOT here: it is an img2img-only refiner pipeline -# (StableDiffusionXLImg2ImgPipeline), but this backend loads every ``sdxl`` repo as the -# base txt2img StableDiffusionXLPipeline and advertises txt2img, so allowlisting the -# refiner would surface the wrong workflow and call it without its required input image. +# Official base repos that may load as a full (non-GGUF) pipeline despite not being under +# unsloth/. Safetensors-only, no pickle/remote code; exact-match lowercased (typo-squat safe). +# Extend deliberately; never add pickled weights or remote code. The SDXL refiner is +# intentionally NOT here (img2img-only; this backend loads every sdxl repo as base txt2img). _TRUSTED_NON_GGUF_REPOS = frozenset( { "stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo", - # Official vendor, safetensors-only base repos allowlisted as LoRA TRAINING bases - # (diffusion training loads the full pipeline from these) and as the official - # BF16 artifact behind each catalog group (model-catalog.ts). Same rule as above: - # no pickled weights, no remote code, exact-match lowercased. FLUX.1-dev/schnell/ - # Kontext are gated on the Hub (need the user's token); the Qwen and Z-Image repos - # are open. All verified as diffusers model_index pipelines. + # Vendor safetensors-only bases: LoRA TRAINING bases + the BF16 artifact behind each + # catalog group. FLUX.1 repos are Hub-gated (need the user's token); Qwen/Z-Image are open. "black-forest-labs/flux.1-dev", "black-forest-labs/flux.1-schnell", "black-forest-labs/flux.1-kontext-dev", @@ -310,16 +287,12 @@ _TRUSTED_NON_GGUF_REPOS = frozenset( "qwen/qwen-image", "qwen/qwen-image-2512", "qwen/qwen-image-edit-2511", - # Krea 2: official vendor repos, safetensors-only, no remote code. Loaded - # per-component via core/inference/diffusion_krea2.py (no GGUF variant yet). - # Turbo is the inference model; Raw is the undistilled base Krea recommends - # training LoRAs on (train on Raw, run adapters on Turbo). + # Krea 2: assembled per-component (diffusion_krea2.py). Turbo = inference; Raw = the + # undistilled base to train LoRAs on (train on Raw, run adapters on Turbo). "krea/krea-2-turbo", "krea/krea-2-raw", - # Ideogram 4: official vendor repos, safetensors-only diffusers pipelines, no - # remote code. The vendor ships no bf16 checkpoint: -fp8 stores the two DiTs - # as raw float8 (highest precision available, the family base); the two nf4 - # repos are identical bnb-4bit exports (both listed so either id loads). + # Ideogram 4: no bf16 ships. -fp8 stores the two DiTs as raw float8 (the family base); + # the two nf4 repos are identical bnb-4bit exports (both listed so either id loads). "ideogram-ai/ideogram-4-fp8", "ideogram-ai/ideogram-4-nf4", "ideogram-ai/ideogram-4-nf4-diffusers", @@ -388,63 +361,48 @@ class _LoadState: device: str dtype: str cpu_offload: bool - # The resolved memory profile (Phase 2A). Appended with defaults so older - # positional constructions (and the back-compat status shape) keep working. + # Resolved memory profile; defaulted so older positional constructions keep working. offload_policy: str = OFFLOAD_NONE vae_tiling: bool = False memory_mode: str = "auto" - # The resolved load kind: "gguf" | "single_file" | "pipeline". Surfaced in status so the - # UI can gate GGUF-only controls (the dense transformer_quant fast path only engages on - # the gguf kind; on single_file/pipeline it is a silent no-op). + # Resolved load kind ("gguf"|"single_file"|"pipeline"); surfaced so the UI can gate + # GGUF-only controls (the dense transformer_quant fast path engages only on gguf). kind: str = "gguf" - # The opt-in speed profile (Phase 3). + # The opt-in speed profile. speed_mode: str = SPEED_OFF speed_optims: tuple = () - # Process-wide torch backend flags (TF32 / cudnn.benchmark) captured before the - # speed layer mutated them, restored on unload so a later `off` load is not - # contaminated by this one's globals. None when nothing was changed. + # Process-wide torch backend flags (TF32 / cudnn.benchmark) captured before the speed + # layer mutated them; restored on unload so a later `off` load isn't contaminated. backend_flags_before: Optional[dict] = None - # Text-encoder quantisation actually engaged: "fp8" | "nvfp4" | None (Phase 2B/2C). + # Text-encoder quant engaged: "fp8" | "nvfp4" | None. text_encoder_quant: Optional[str] = None - # Transformer quant actually engaged on the opt-in dense fast path: "int8" | "fp8" - # | "nvfp4" | "mxfp8" | None. None means the default GGUF transformer was loaded. + # Transformer quant engaged on the dense fast path ("int8"|"fp8"|"nvfp4"|"mxfp8") or None (GGUF loaded). transformer_quant: Optional[str] = None - # Attention backend engaged via the diffusers dispatcher (e.g. "_native_cudnn"), or - # None for the default SDPA. Set before compile; orthogonal to the weight quant. + # Attention backend engaged via the diffusers dispatcher, or None for default SDPA. attention_backend: Optional[str] = None - # The caller's ORIGINAL attention request (None / "auto" left it to the backend, else - # an explicit alias like "native" / "sage" / "flash"). Carried so the deferred-speed - # engagement re-runs the SAME selection the load-time path did, instead of forcing the - # auto cuDNN upgrade -- otherwise an explicitly pinned backend (e.g. "native" to avoid - # cuDNN) is silently discarded when the 3rd generation engages the deferred profile. + # The caller's ORIGINAL attention request, carried so the deferred-speed engagement + # re-runs the SAME selection instead of forcing the auto cuDNN upgrade over an explicit pin. attention_request: Optional[str] = None # Step cache engaged ("fbcache") or None. Opt-in, for many-step models. transformer_cache: Optional[str] = None - # True when the cache decision was AUTO on a cache-capable transformer: generate() - # then re-checks the actual step count and toggles FBCache across FBCACHE_MIN_STEPS. - # An explicit request (off / fbcache) is never toggled. + # AUTO on a cache-capable transformer: generate() re-checks the step count and toggles + # FBCache across FBCACHE_MIN_STEPS. An explicit request (off / fbcache) is never toggled. cache_auto: bool = False # Inputs the generation-time toggle re-applies (quantised threshold + override). cache_quant_active: bool = False cache_threshold: Optional[float] = None - # Shared eager monkey-patches (diffusion_eager_patches) installed for this load (any - # non-off speed tier). Uninstalled on unload so a later `off` load is bit-identical. + # Shared eager patches installed for this load (any non-off tier); uninstalled on unload. eager_patched: bool = False - # Deferred speed auto (dense models with speed_mode unset): the load stays fully - # eager/bit-identical, and generate() engages the `default` compile profile at the - # start of the 3rd generation this session -- repeated use is established by then, - # so the one-time compile warmup amortises. Cleared once engaged (or failed). + # Deferred speed auto: the load stays eager/bit-identical; generate() engages the `default` + # compile profile at the 3rd generation this session. Cleared once engaged (or failed). speed_deferred: bool = False # Successful generations on this load; drives the deferred engagement above. generation_count: int = 0 - # Pre-warmed torch.compile cache context (diffusion_compile_cache.CacheContext) when a - # compiled tier ran, else None. Carries the per-key inductor dir + bundle for save/restore. + # Pre-warmed torch.compile cache context when a compiled tier ran, else None. compile_cache_ctx: Any = None # Token kept so LoRA adapters selected at generate time can be fetched from the Hub. hf_token: Optional[str] = None - # Per-control provenance from the auto-policy: {control: {value, source, reason}}. - # source is "auto" when the backend decided (request left unset / "auto") and - # "explicit" when the caller pinned the value. Surfaced via status for the UI badges. + # Per-control provenance {control: {value, source, reason}} (source auto/explicit), for status badges. resolved: Optional[dict] = None @@ -496,42 +454,28 @@ class DiffusionBackend: """Holds at most one loaded diffusers pipeline. All mutations are serialised.""" def __init__(self) -> None: - # _lock serialises the small state mutations (the load swap, _loading, - # _load_token, _gen). status() / load_progress() / generate_progress() - # read those references WITHOUT it, so polling never blocks a slow load. + # _lock serialises the small state mutations; status()/load_progress()/ + # generate_progress() read lock-free so polling never blocks a slow load. self._lock = threading.Lock() - # _generate_lock serialises generations and is the ONLY lock the denoise - # holds, so a long generation never blocks status()/unload()/a new load. + # _generate_lock serialises generations and is the ONLY lock the denoise holds. self._generate_lock = threading.Lock() self._state: Optional[_LoadState] = None self._loading: Optional[_LoadingState] = None - # Bumped on every begin_load and unload so a worker whose load was - # superseded (a new load) or cancelled (unload, incl. an arbiter eviction) - # neither commits its pipeline nor stamps progress onto the current load. + # Bumped on every begin_load/unload so a superseded/cancelled worker neither + # commits its pipeline nor stamps progress onto the current load. self._load_token = 0 - # Set by unload() to abort an in-flight download (which runs without the - # lock, like the chat backend), so an eviction/unload can preempt a slow - # load instead of blocking on the lock for the whole download. + # Set by unload() to abort an in-flight (lock-free) download so an eviction preempts it. self._cancel_event = threading.Event() - # The cancel Event of the generation currently in flight (or None). Set - # under _lock by unload() / a superseding load to abort that specific - # denoise (its step callback flips pipe._interrupt). Per-generation rather - # than one shared flag the next generate would clear, so a cancel can't be - # lost to a racing generate nor leak onto the wrong one. + # Cancel Event of the in-flight generation (or None), set under _lock to abort THAT + # denoise. Per-generation so a cancel can't be lost to a racing generate nor leak. self._active_generate_cancel: Optional[threading.Event] = None - # The callback mutates _gen and generate_progress() reads it, both lock-free, - # so per-step progress polling stays live during a generation. + # Written by the callback, read lock-free by generate_progress(). self._gen: Optional[_GenState] = None - # Cache of image-conditioned workflow pipelines (img2img / inpaint) built via - # Pipeline.from_pipe around the loaded text-to-image pipe. They share its already - # resident modules (no extra VRAM, no reload), so we build each once per load and - # reuse it. Keyed by pipeline class name; cleared on unload with the base pipe. + # Image-conditioned workflow pipes (img2img/inpaint) built via from_pipe around the + # loaded pipe (shared modules, no extra VRAM). Keyed by class name; cleared on unload. self._aux_pipes: dict[str, Any] = {} - # Cache of loaded ControlNet models (id -> module) and the ControlNet workflow - # pipelines built around them ((pipeline_class, cn_id) -> pipe). ControlNet models - # are a small extra module loaded via from_pretrained; the pipeline is assembled via - # Pipeline.from_pipe(base, controlnet=model), reusing the resident base modules (no - # reload). Both are cleared on unload with the base pipe. + # Loaded ControlNet models (id -> module) and their from_pipe pipelines + # ((class, cn_id) -> pipe), reusing resident base modules; cleared on unload. self._cn_models: dict[str, Any] = {} self._cn_pipes: dict[tuple[str, str], Any] = {} @@ -589,33 +533,24 @@ class DiffusionBackend: mode = normalize_transformer_quant(raw) if mode is None: return False - # An explicit Speed="off" (bit-exact) load suppresses the auto-dtype default in - # load_pipeline and stays GGUF-as-is, so the dense path never runs -- don't widen the - # prefetch for it either. + # An explicit Speed="off" load stays GGUF-as-is (dense path never runs); don't widen the prefetch. speed = kwargs.get("speed_mode") if speed is not None and str(speed).strip().lower() == SPEED_OFF: return False try: - # A definite-offload memory policy forces load_pipeline onto offload regardless of the - # dense candidate's smaller footprint, so its re-plan never flips to OFFLOAD_NONE and - # the dense build never runs. balanced -> OFFLOAD_GROUP and low_vram -> OFFLOAD_MODEL are - # set unconditionally in plan_diffusion_memory; the legacy cpu_offload flag forces - # OFFLOAD_MODEL when no memory_mode overrides it. In those cases the GGUF path runs - # offloaded and never touches the base transformer/ shards, so widening the prefetch only - # wastes a multi-GB download -- and a disk-full on that begin_load pull has NO GGUF - # fallback (unlike the in-load_pipeline dense failure). Mirror those offload gates here. + # A definite-offload policy forces load_pipeline onto offload, so the dense build + # never runs and never touches the base transformer/ shards. Widening the prefetch + # would only waste a multi-GB pull -- and a disk-full here has NO GGUF fallback + # (unlike an in-load_pipeline dense failure). Mirror those offload gates. mm = normalize_memory_mode(kwargs.get("memory_mode")) if mm in (MEMORY_MODE_BALANCED, MEMORY_MODE_LOW_VRAM): return False if mm is None and kwargs.get("cpu_offload"): return False target = self._resolve_device_target(fam) - # Only widen the prefetch when the loader would actually take the dense path: resolve - # the SAME dense-quant candidate load_pipeline re-plans against, which also checks the - # cache volume has room for the extra bf16 transformer/ shards. When disk (or scheme / - # support / a prequant checkpoint) rules the dense build out, do NOT eagerly pull those - # shards -- otherwise the widened prefetch fills the disk and hard-fails the load in a - # spot unload/cancel cannot preempt, instead of the disk guard falling back to the GGUF. + # Only widen when the loader would actually take the dense path: resolve the SAME + # candidate load_pipeline re-plans against (which also checks the cache has disk room). + # When disk/scheme/support/a prequant rule it out, don't eagerly pull the shards. candidate = resolve_dense_quant_candidate( fam = fam, target = target, @@ -624,12 +559,8 @@ class DiffusionBackend: prequant_path = kwargs.get("transformer_prequant_path"), logger = None, ) - # A prequant candidate loads from the small pre-quantized checkpoint (+ config / - # companions), NOT the base repo's full dense transformer/ shards, so widening the - # prefetch to pull those shards both defeats the prequant download savings and can - # hard-fail the load: the widened pull runs in begin_load, where a disk-full has no - # GGUF fallback (unlike the in-load_pipeline dense failure). Only widen for a real - # dense build. + # A prequant candidate loads a small checkpoint, not the dense transformer/ shards, + # so widening for it defeats the savings and can disk-full the fallback-less begin_load pull. return candidate is not None and not candidate.prequant except Exception: # noqa: BLE001 — widening the prefetch is best-effort only return False @@ -701,77 +632,56 @@ class DiffusionBackend: f"pass family_override with that family name. (Video models and image models " f"whose diffusers transformer has no single-file loader are not supported.)" ) - # A GGUF load builds a transformer-only file via the generic GGUF branch - # (UNet2DConditionModel.from_single_file(subfolder="transformer", GGUFQuantizationConfig)). # Families whose single file IS the whole pipeline (SDXL) have no transformer-only - # GGUF path, so reject GGUF here -- before the route evicts the current model and - # the background load fails deep in from_single_file. + # GGUF path; reject GGUF here, before the route evicts the current model. if kind == "gguf" and fam.single_file_is_pipeline: raise ValueError( f"'{fam.name}' checkpoints are whole-pipeline single files and have no GGUF " f"transformer variant; load the .safetensors pipeline instead of a GGUF." ) - # A family that assembles MULTIPLE denoisers per-component (Ideogram 4's dual - # DiTs) has no transformer-only single-file or GGUF path: those kinds build one - # transformer and would assemble a pipeline missing its second DiT (or fail deep - # in from_pretrained). Reject them here -- before the route evicts the current - # model -- so only a full pipeline load reaches the per-component loader. + # A multi-denoiser family (Ideogram 4's dual DiTs) has no transformer-only path; + # a single-file/GGUF load would miss its second DiT. Reject here, before eviction. if kind in ("gguf", "single_file") and fam.pipeline_only: raise ValueError( f"'{fam.name}' loads only as a full diffusers pipeline (it assembles " f"multiple transformers), not from a single-file or GGUF checkpoint; " f"select the pipeline repo." ) - # Non-GGUF loads (a single-file safetensors transformer, or a full pipeline) - # are gated to the unsloth org or a local path -- they fetch + deserialise - # weights, so an arbitrary remote repo is rejected here, before any work. + # Non-GGUF loads fetch + deserialise weights, so gate to unsloth/ or a local path. if kind != "gguf" and not _is_trusted_diffusion_repo(repo_id): raise ValueError( f"Non-GGUF diffusion loads are restricted to unsloth/* repos (or a local " f"path); got '{repo_id}'. Pass a gguf_filename to load a GGUF instead." ) - # A companion base repo also loads via from_pretrained (its diffusers pipeline is - # assembled around the GGUF/single-file transformer), so it must clear the same trust - # bar as a non-GGUF repo_id -- otherwise a trusted GGUF model_path could smuggle in an - # arbitrary remote base that gets downloaded and deserialised. Gate it here (before the - # route evicts the resident model), mirroring the video loader's base_repo check. + # A companion base repo also loads via from_pretrained, so it must clear the same + # trust bar (else a GGUF pick could smuggle in an arbitrary remote base). Gate here. if base_repo and base_repo.strip() and not _is_trusted_diffusion_repo(base_repo): raise ValueError( f"base_repo is restricted to unsloth/* repos (or a local path); got " f"'{base_repo}'." ) - # An existing LOCAL base_repo is loaded as a full pipeline (from_pretrained(base) / - # config=base), which needs a model_index.json. Any existing path passes the trust check - # above, so reject a non-pipeline local base here -- before the route evicts the resident - # model -- rather than deep in the background load. Mirrors the repo_id check below. + # A local base_repo loads as a full pipeline (needs model_index.json); reject a + # non-pipeline local base here, before eviction. _assert_local_base_is_pipeline(base_repo) - # Reject a bad LOCAL pick now (the same checks the load would hit later), so - # the route never evicts a working chat model for a request that can't load. - # A path-shaped repo_id (absolute / ~ / ./ / ..) is meant to be on disk, so a - # missing one is an error here; a bare "org/name" id is a remote HF repo and - # is left for the background load to resolve. + # Reject a bad LOCAL pick now so the route never evicts chat for an unloadable request. + # A path-shaped repo_id is meant to be on disk; a bare "org/name" is a remote HF repo. local_root = Path(repo_id).expanduser() - # POSIX path-shaped, a "."/".." prefix (covers ./ ../ and their Windows .\ ..\ - # forms), a Windows separator anywhere (never present in a bare "org/name" HF - # id), or an absolute path on this OS. + # Path-shaped: "."/".." prefix, a backslash (never in "org/name"), or an absolute path. path_shaped = ( repo_id.startswith(("/", "\\", "~", ".")) or "\\" in repo_id or local_root.is_absolute() ) if kind in ("gguf", "single_file"): if not gguf_filename: raise ValueError(f"a single-file checkpoint name is required for a '{kind}' load.") - # Fail a kind/extension mismatch here (before the route evicts chat and grabs the - # GPU), instead of deep in the background from_single_file: a "gguf" load needs a - # .gguf file, and a "single_file" load must not be handed a .gguf. + # Fail a kind/extension mismatch here, before the handoff: gguf needs .gguf, + # single_file must not be handed a .gguf. is_gguf_name = gguf_filename.lower().endswith(".gguf") if kind == "gguf" and not is_gguf_name: raise ValueError("a 'gguf' load requires a .gguf checkpoint name.") if kind == "single_file" and is_gguf_name: raise ValueError("a .gguf checkpoint needs model_kind 'gguf', not 'single_file'.") - # A single-file load must name an actual checkpoint: an arbitrary repo file - # (README.md, config.json) would pass preflight, evict the chat model, and - # only fail in the background from_single_file -- the eviction this - # validation exists to prevent. + # A single-file load must name an actual .safetensors checkpoint (else it evicts + # chat and only fails in the background from_single_file). if kind == "single_file" and not gguf_filename.lower().endswith(".safetensors"): raise ValueError( f"'{gguf_filename}' is not a loadable single-file checkpoint " @@ -794,11 +704,8 @@ class DiffusionBackend: elif path_shaped: raise FileNotFoundError(f"Local model path does not exist: {repo_id}") elif repo_id.upper().endswith("-GGUF"): - # A remote "*-GGUF" id is a single-file GGUF repo, not a full diffusers - # pipeline: loading it as a pipeline passes the trusted-repo check, evicts - # chat, then fails in the background when from_pretrained finds no - # model_index.json. Reject the certain case here (no network round-trip) - # so the bad pick fails before the GPU handoff, as the route expects. + # A remote "*-GGUF" id is a GGUF repo, not a pipeline; loading it as a pipeline + # would evict chat then fail on the missing model_index.json. Reject here. raise ValueError( f"'{repo_id}' is a single-file GGUF repo; load it with model_kind 'gguf' " f"and a .gguf filename, not as a full pipeline." @@ -828,13 +735,10 @@ class DiffusionBackend: model_kind: Optional[str] = None, ) -> dict[str, Any]: """Validate, then run the (slow) load on a daemon thread. Returns at once.""" - # A blank token (the Studio default when none is configured) must mean - # "anonymous", not an explicit empty credential the Hub rejects with 401. + # A blank token must mean "anonymous", not an empty credential the Hub 401s. hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None - # base_repo is gated at the /images/load route's pre-eviction validate_load_request - # (the client entry point); the re-validation here is a redundant cheap-fail guard for - # the resolved repo/family, so it does not re-gate base_repo (which internal callers pass - # through already-validated). + # base_repo is already gated at the route's pre-eviction validate; this re-validation + # is a cheap-fail guard for the resolved repo/family and does not re-gate base_repo. fam = self.validate_load_request( repo_id, gguf_filename = gguf_filename, @@ -848,11 +752,9 @@ class DiffusionBackend: raise RuntimeError("A diffusion load is already in progress.") self._load_token += 1 token = self._load_token - # Best-effort download preemption only; the token (not this event) is - # the real guard that a superseded worker can't commit its pipeline. + # Best-effort download preemption; the token is the real commit guard. self._cancel_event.clear() - # Seed with the family fallback; the worker resolves the real base - # (a network lookup) and updates this, so begin_load never blocks. + # Seed with the family fallback; the worker resolves the real base and updates this. self._loading = _LoadingState(repo_id = repo_id, base_repo = fam.base_repo) threading.Thread( @@ -883,16 +785,14 @@ class DiffusionBackend: def _run_load(self, **kwargs: Any) -> None: token = kwargs.get("_load_token") try: - # Resolve the base repo and estimate sizes on this thread (both network - # calls) so begin_load returns instantly; the bar shows raw bytes until - # the total lands. This is the only writer of _loading's fields here. + # Resolve the base repo and estimate sizes on this thread (both network calls) so + # begin_load returns instantly. Only writer of _loading's fields here. fam = detect_family_for_pick( kwargs["repo_id"], kwargs.get("gguf_filename"), kwargs.get("family_override") ) kind = resolve_model_kind(kwargs.get("gguf_filename"), kwargs.get("model_kind")) if kind == "pipeline": - # The full pipeline IS the repo: from_pretrained pulls every component - # (transformer included) from it, so the base repo is the repo itself. + # The full pipeline IS the repo, so the base repo is the repo itself. base = kwargs["repo_id"] else: base = _resolve_base_repo( @@ -906,23 +806,17 @@ class DiffusionBackend: kwargs.get("hf_token"), kind = kind, single_file_is_pipeline = bool(fam and fam.single_file_is_pipeline), - # The dense transformer-quant path downloads the base repo's - # transformer/ shards via from_pretrained(subfolder="transformer") - # INSIDE the locked finalize phase, where unload/cancellation cannot - # preempt the multi-GB pull. When that path can actually run, pull the - # shards here in the preemptible prefetch instead. (Pipeline loads - # already include transformer/ via their own filter.) + # The dense-quant path otherwise pulls the base transformer/ shards inside the + # locked finalize (unpreemptable); when it can run, pull them in the prefetch here. include_transformer = kind == "gguf" and self._dense_quant_prefetch_needed(fam, kwargs), ) with self._lock: - # Stamp progress only if this load is still current; a superseding - # load (or unload) has its own token and its own _LoadingState. + # Stamp progress only if this load is still current (a superseder has its own token). if self._load_token == token and self._loading is not None: self._loading.base_repo = base self._loading.expected_bytes = expected - # Download outside the lock so unload()/an eviction can preempt the - # multi-GB pull; load_pipeline below then assembles from the cache. + # Download outside the lock so unload/an eviction can preempt the pull. kwargs["_base_local_dir"] = self._prefetch_files( kwargs["repo_id"], kwargs.get("gguf_filename"), @@ -932,27 +826,21 @@ class DiffusionBackend: ) self.load_pipeline(**kwargs) with self._lock: - # Only clear the marker if this load is still the current one; a - # newer begin_load (or an unload) has its own token. + # Only clear the marker if this load is still current (a superseder has its own token). if self._load_token == token: self._loading = None except Exception as exc: # noqa: BLE001 — surfaced to the client via load_progress - # A cancelled/superseded load raised below; don't log it as a failure - # or stamp its error onto whatever load is current now. + # A cancelled/superseded load raised below; don't log/stamp it onto the current load. if self._load_token != token: return logger.error("diffusion.load_failed: %s", exc) - # Free the debris of a failed construction (e.g. a load-time OOM): _state was - # never committed, and the next load's _unload_locked early-returns on a None - # state, so nothing else releases the reserved VRAM. Guarded: a sticky CUDA - # error makes synchronize() raise, which would skip stamping the REAL error - # below and leave the client polling forever. + # Free the debris of a failed construction (uncommitted _state, so nothing else + # reclaims the VRAM). Guarded so a sticky CUDA error can't skip stamping the real error. try: clear_gpu_cache() except Exception: # noqa: BLE001 pass - # Redact native paths: this error is surfaced verbatim via the - # load-progress poll, and Studio can run as a shared server. + # Redact native paths: this error is surfaced verbatim and Studio can be shared. from utils.native_path_leases import redact_native_paths with self._lock: @@ -967,15 +855,13 @@ class DiffusionBackend: if loading is None: return _progress("ready" if self._state is not None else None) - # Sum the checkpoint repo + companion base cache. For a full-pipeline load the - # base IS the repo, so count it once (else the bar double-counts to "finalizing"). + # Sum checkpoint + companion base cache; for a full-pipeline load base IS the repo, + # so count it once (else the bar double-counts). downloaded = self._cache_bytes(loading.repo_id) if loading.base_repo and loading.base_repo != loading.repo_id: downloaded += self._cache_bytes(loading.base_repo) expected = loading.expected_bytes - # Downloads done but pipeline still dequantising / moving to GPU. The cache - # scan can slightly exceed the estimate (extra cached quants, blob padding), - # so clamp the reported bytes/fraction so the bar never overshoots 100%. + # Downloads done, still finalizing. The cache scan can exceed the estimate, so clamp to 100%. if expected > 0 and downloaded >= expected * 0.999: return _progress("finalizing", min(downloaded, expected), expected, 1.0) fraction = min(downloaded / expected, 1.0) if expected > 0 else 0.0 @@ -1022,8 +908,7 @@ class DiffusionBackend: if kind == "pipeline": info = api.model_info(repo_id, files_metadata = True, token = hf_token) picked = [s for s in info.siblings if _pipeline_file_downloaded(s.rfilename)] - # diffusers prefers safetensors per component: drop a .bin whose - # directory also carries a picked .safetensors weight. + # diffusers prefers safetensors: drop a .bin whose dir also has a picked .safetensors. st_dirs = { s.rfilename.rsplit("/", 1)[0] for s in picked @@ -1035,15 +920,12 @@ class DiffusionBackend: base_files.append(s.rfilename) total += s.size or 0 return total, base_files - # Skip the Hub size lookup for a LOCAL gguf path: model_info(repo_id) would - # raise on a filesystem path and (caught below) skip the base-repo lookup too, - # so the companion VAE/text-encoder files would never be prefetched and would - # instead download synchronously under the load lock. + # Skip the Hub size lookup for a LOCAL gguf path: model_info would raise on a + # filesystem path and (caught below) skip the base lookup, forcing a synchronous companion pull. if gguf_filename and not Path(repo_id).expanduser().exists(): info = api.model_info(repo_id, files_metadata = True, token = hf_token) total += sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename) - # A whole-pipeline single file (SDXL) needs only the base repo's config/tokenizer, - # not its (unused, multi-GB) weight files. + # A whole-pipeline single file (SDXL) needs only the base's config/tokenizer, not its weights. if kind == "single_file" and single_file_is_pipeline: base_filter = _base_config_file_downloaded else: @@ -1146,9 +1028,7 @@ class DiffusionBackend: numel *= dim total += numel return total - except Exception: # noqa: BLE001 — best-effort estimate; a corrupt/crafted shard - # (bad header length, non-dict header, odd shape) must degrade to 0 so the caller - # gates on the plain plan, never crash the load. + except Exception: # noqa: BLE001 — corrupt/crafted shard degrades to 0, never crashes the load return 0 @staticmethod @@ -1195,21 +1075,16 @@ class DiffusionBackend: _load_token: Optional[int] = None, _base_local_dir: Optional[str] = None, ) -> dict[str, Any]: - # A blank / whitespace-only token must degrade to anonymous access, not be passed - # as an explicit credential (from_single_file / from_pretrained / the Hub client - # can error on a malformed token instead of falling back). Normalize once here so - # every load branch and the size estimate below use a real token or None. + # A blank/whitespace token must degrade to anonymous, not be passed as a credential + # the Hub client can error on. Normalize once for every branch below. hf_token = hf_token.strip() if isinstance(hf_token, str) else hf_token hf_token = hf_token or None - # Validate first (cheap, no torch/diffusers) so a direct call with a bad - # family fails with ValueError even in a no-diffusers runtime. Sanitize the - # token here too (direct callers bypass begin_load): a blank string must - # load anonymously, not 401 as an explicit empty credential. + # Validate first (cheap, no torch/diffusers) so a bad family fails even in a no-diffusers + # runtime. Re-sanitize the token (direct callers bypass begin_load). hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None - # base_repo is gated at the route before eviction (validate_load_request there); this - # direct-load re-validation only cheap-fails the resolved repo/family, so it does not - # re-gate an already-validated base_repo. + # base_repo is gated at the route before eviction; this re-validation cheap-fails the + # resolved repo/family and does not re-gate an already-validated base_repo. fam = self.validate_load_request( repo_id, gguf_filename = gguf_filename, @@ -1217,17 +1092,14 @@ class DiffusionBackend: model_kind = model_kind, ) kind = resolve_model_kind(gguf_filename, model_kind) - # Validate every mode string that can raise NOW, before this load evicts the - # previous pipeline below: their first in-line uses all sit past _unload_locked, - # where a bad request would cost the user their working model. Validate-only for - # transformer_quant: the raw value keeps the unset/auto vs explicit-off tri-state. + # Validate every mode string that can raise NOW, before this load evicts the previous + # pipeline. Validate-only for transformer_quant (keep the unset/auto vs explicit-off tri-state). normalize_transformer_quant(transformer_quant) normalize_speed_mode(speed_mode) normalize_attention_backend(attention_backend) normalize_transformer_cache(transformer_cache) normalize_te_quant(text_encoder_quant) - # For a full pipeline the repo itself supplies every component, so it is its - # own base; the single-file kinds resolve the companion base diffusers repo. + # A full pipeline is its own base; single-file kinds resolve the companion base repo. base = ( repo_id if kind == "pipeline" else _resolve_base_repo(repo_id, base_repo, fam, hf_token) ) @@ -1236,14 +1108,10 @@ class DiffusionBackend: import diffusers - # Pre-install the optional attention kernel BEFORE taking the load locks. The - # wheel-only pip install can run up to 600s, and doing it under _lock / - # _generate_lock (as the in-lock apply_attention_backend otherwise would) blocks - # unload() and cancellation for that whole window. Only an explicit backend pulls - # a package -- auto resolves to cuDNN / native, which ship with torch -- and an - # explicit backend's resolution ignores the speed tier, so it can run here without - # effective_speed. Best-effort: the authoritative resolve + set still happens under - # the lock, where the now-satisfied install call is a fast no-op. + # Pre-install the optional attention kernel BEFORE the load locks: the wheel-only pip + # install can run up to 600s, and doing it under the lock would block unload/cancel that + # whole window. Only an explicit backend pulls a package (auto uses cuDNN/native from + # torch). Best-effort; the authoritative resolve + set under the lock is then a no-op. try: preinstall_backend = select_attention_backend( target, attention_backend, speed_active = True @@ -1253,33 +1121,25 @@ class DiffusionBackend: except Exception: # noqa: BLE001 — the locked path re-resolves and validates pass - # Signal an in-flight denoise to abort, then take _generate_lock to WAIT for - # it to actually exit before allocating the replacement: a load is about to - # claim VRAM, so unlike unload() it must not overlap a still-live pipeline. - # The cancel makes that wait ~one step (or the rest of the denoise for a - # pipeline that ignores the step callback). + # Signal an in-flight denoise to abort, then take _generate_lock to WAIT for it to exit + # before allocating the replacement (a load claims VRAM, so it must not overlap a live pipe). with self._lock: - # Bail BEFORE signalling any cancel if this load was already superseded (an - # unload/eviction or a newer load bumped the token while we were resolving / - # downloading). Otherwise a stale worker would abort an unrelated, still-live - # generation from the CURRENT model and only then discover it has nothing to do. + # Bail BEFORE signalling a cancel if this load was already superseded, else a stale + # worker would abort an unrelated live generation from the CURRENT model. if _load_token is not None and _load_token != self._load_token: raise RuntimeError("Diffusion load was cancelled.") if self._active_generate_cancel is not None: self._active_generate_cancel.set() with self._generate_lock: with self._lock: - # Re-check under the generate lock: a newer load/unload may have superseded - # this one while we waited for the in-flight denoise to exit. + # Re-check: a newer load/unload may have superseded this one while we waited. if _load_token is not None and _load_token != self._load_token: raise RuntimeError("Diffusion load was cancelled.") - # Free the old pipeline before allocating the new one so two - # checkpoints never sit in VRAM at once. + # Free the old pipeline before allocating the new one (never two in VRAM). self._unload_locked() - # The single-file kinds resolve a checkpoint path (GGUF or safetensors); - # the pipeline kind has none (from_pretrained pulls the repo directly). + # Single-file kinds resolve a checkpoint path; the pipeline kind has none. single_file_path = ( self._resolve_gguf_path(repo_id, gguf_filename, hf_token) if kind in ("gguf", "single_file") @@ -1288,9 +1148,8 @@ class DiffusionBackend: transformer_cls = getattr(diffusers, fam.transformer_class) pipeline_cls = getattr(diffusers, fam.pipeline_class) - # Decide placement up front (the weights are still on CPU, so free VRAM is - # the real budget). This plan budgets the GGUF file and places the plain - # load; the dense-quant fast path is preflighted separately below. + # Decide placement up front (weights still on CPU, so free VRAM is the budget). + # Budgets the GGUF file; the dense-quant fast path is preflighted separately below. plan = self._plan_memory( target, single_file_path, @@ -1302,42 +1161,29 @@ class DiffusionBackend: repo_id = repo_id, ) - # Dtype tri-state: an UNSET request (or "auto") hands the decision to - # the hardware ladder -- on a dense-capable GPU the quantised build - # (int8 minimum, fp8 on data-center silicon) beats running the GGUF - # as-is, so auto is the DEFAULT. An explicit "none"/"off" pins - # GGUF-as-is and an explicit scheme pins that scheme. The overwritten - # "auto" still records source=auto in the resolved provenance. + # Dtype tri-state: unset/"auto" -> hardware ladder picks a quantised build (int8 + # min, fp8 on datacenter silicon) over GGUF-as-is; "none"/"off" pins GGUF-as-is; + # an explicit scheme pins it. An overwritten "auto" still records source=auto. if transformer_quant is None or str(transformer_quant).strip().lower() in ( "", "auto", ): - # An explicit Speed="off" (bit-exact) load must stay GGUF-as-is: promoting the - # unset dtype to auto-quant here would engage int8/fp8 + compile and silently - # break the user's bit-exact request (an auto DEFAULT overriding an EXPLICIT - # control). Suppress the auto default when speed was explicitly pinned off; - # otherwise auto (the dense-capable default) applies. + # An explicit Speed="off" load must stay GGUF-as-is: auto-quant would engage + # int8/fp8 + compile and break the bit-exact request. "off" -> None (GGUF-as-is). speed_off = ( speed_mode is not None and str(speed_mode).strip().lower() == SPEED_OFF ) - # "off" normalizes to None (no dense quant), keeping the GGUF-as-is path. transformer_quant = "off" if speed_off else TQ_AUTO # Default-on fast path: load the DENSE bf16 transformer and torchao-quantise it - # (int8 / fp8 / fp4 tensor cores), which beats GGUF's bf16-rate per-matmul - # dequant on both speed and quality, at the cost of a higher-memory dense - # load. Gated on CUDA + bf16 + a resident fit; ANY failure (unsupported arch - # / scheme, OOM, partial quant) falls back to the GGUF build below. Only the - # GGUF kind offers it: it materialises the dense bf16 transformer from the - # base repo, which the safetensors kinds (a single-file or already-quantized - # pipeline) do not have. + # (int8/fp8/fp4 tensor cores), which beats GGUF's per-matmul dequant on speed AND + # quality at a higher-memory dense load. CUDA + bf16 + resident fit; ANY failure + # falls back to the GGUF build. GGUF kind only (it has the dense bf16 to materialise). pipe = None transformer_quant_engaged = None quant_plan = None - # The GGUF-size `plan` can mis-budget the dense-quant fast path two ways, so - # preflight the real footprint BEFORE evicting the current pipeline. Both - # branches need the base repo + a resolved scheme, so gate on the dense-path - # preconditions first. + # The GGUF-size `plan` can mis-budget the fast path two ways, so preflight the real + # footprint BEFORE eviction; both branches need the base repo + a resolved scheme. dense_declined = False if ( kind == "gguf" @@ -1345,12 +1191,9 @@ class DiffusionBackend: and dense_transformer_supported(target) ): if plan.offload_policy != OFFLOAD_NONE: - # The GGUF-size plan picked offload, but the dense-quant artifact has - # a DIFFERENT footprint: int8/fp8 weights are ~half the bf16 bytes, and - # a pre-quantized checkpoint never materialises dense bf16 at all. Ask - # the auto-policy for the candidate's estimate and re-plan against it: - # a resident quantised build beats an offloaded GGUF on speed AND - # quality, so it must be attempted before settling for offload. + # The GGUF plan picked offload, but the quantised artifact is smaller + # (int8/fp8 ~half bf16; a prequant never materialises dense). Re-plan + # against the candidate's estimate -- a resident quant build beats an offloaded GGUF. candidate = resolve_dense_quant_candidate( fam = fam, target = target, @@ -1372,35 +1215,26 @@ class DiffusionBackend: transformer_resident_override_mib = ( candidate.transient_transformer_mib ), - # The dense path prefetches the base transformer/ shards into the - # cache _companion_cache_bytes reads; pass the auto-policy's own - # companion estimate so the re-plan does not double-count them. + # Pass the auto-policy's companion estimate so the prefetched base + # transformer/ shards in the cache aren't double-counted. companion_override_mib = candidate.companions_mib, ) if replanned.offload_policy == OFFLOAD_NONE: quant_plan = replanned else: - # The GGUF fits resident, but this path first materialises the base - # repo's dense bf16 transformer -- bigger than the quantised GGUF -- so - # re-check the fit against THAT. A card that fits the GGUF but not the - # dense transformer must skip the fast path up front, not evict the - # current pipeline then OOM in finalization. A prequant checkpoint loads - # a small quantised file (no dense bf16), so skip the re-check there - # (mirrors the prefetch guard). _dense_transformer_resident_bytes reads - # the on-disk shard headers, so it also covers families the size table - # (resolve_dense_quant_candidate) does not list; it returns 0 when the - # shards are absent, in which case the fast path keeps today's behaviour. + # The GGUF fits resident, but this path first materialises the base's dense + # bf16 transformer (bigger), so re-check the fit against THAT -- a card that + # fits the GGUF but not the dense must skip the fast path up front, not OOM + # after eviction. A prequant loads a small file (no dense), so skip the + # re-check there. _dense_transformer_resident_bytes returns 0 if shards are absent. scheme = select_transformer_quant_scheme( target, transformer_quant, # normalized above family = getattr(fam, "name", None), ) - # usable_prequant_source (not resolve_): a request-supplied local - # path that is missing or outside the allowlist must NOT count as a - # prequant source here, or it would skip the dense-fit re-check and - # _load_dense_quant_pipeline's refusal would fall back to - # materialising the dense bf16 transformer AFTER the eviction -- - # exactly the OOM this re-check exists to prevent. + # usable_prequant_source (not resolve_): a missing/non-allowlisted local + # path must NOT count as prequant here, or it skips the dense-fit re-check + # and OOMs materialising the dense transformer after eviction. prequant = ( usable_prequant_source( fam, scheme, path_override = transformer_prequant_path @@ -1453,69 +1287,54 @@ class DiffusionBackend: ) pipe = None transformer_quant_engaged = None - # Drop the exception (and its traceback) BEFORE clearing the cache: - # exc.__traceback__ keeps _load_dense_quant_pipeline's frame -- and - # thus its partially-built dense bf16 transformer/pipe -- alive, so - # clear_gpu_cache() could not otherwise reclaim that VRAM before the - # GGUF build (the OOM-fallback path this cleanup exists for). + # Drop the exception BEFORE clearing the cache: its traceback keeps the + # partially-built dense transformer/pipe alive, blocking VRAM reclaim. del exc - # Guarded: after an OOM/sticky CUDA error synchronize() can - # raise, and this fallback path must still reach the GGUF build. + # Guarded: a sticky CUDA error can raise; the fallback must reach the GGUF build. try: clear_gpu_cache() except Exception: # noqa: BLE001 pass if transformer_quant_engaged is not None and quant_plan is not None: - # The re-planned resident placement is the one the engaged dense build - # actually uses; the GGUF-size plan stays in force for the fallback. + # The engaged dense build uses the re-planned placement; the GGUF-size plan stays for fallback. plan = quant_plan if pipe is None: if kind == "pipeline": - # Full diffusers repo: from_pretrained pulls every component - # (transformer + VAE + text encoders + scheduler) from the repo - # and re-applies any embedded quantization_config (e.g. bnb-4bit), - # so a pre-quantized pipeline reloads quantized with no extra config. + # Full diffusers repo: from_pretrained pulls every component and re-applies + # any embedded quantization_config (e.g. bnb-4bit). if fam.name == KREA2_FAMILY_NAME: - # The krea repo ships transformers-5.x style configs the 4.x - # line cannot parse; assemble the pipeline per-component - # (see diffusion_krea2.py for the exact compat story). + # krea ships transformers-5.x configs the 4.x line can't parse; assemble + # per-component (see diffusion_krea2.py). pipe = load_krea2_pipeline(repo_id, dtype, hf_token = hf_token) elif fam.name == IDEOGRAM4_FAMILY_NAME: - # The ideogram repos ship the same transformers-5.x style Qwen - # text stack as krea (rope under rope_parameters, a slow-only - # tokenizer pin without its vocab files), so this family is - # assembled per-component too (see diffusion_ideogram4.py). + # ideogram ships the same transformers-5.x Qwen stack as krea; assemble + # per-component too (see diffusion_ideogram4.py). pipe = load_ideogram4_pipeline(repo_id, dtype, hf_token = hf_token) else: pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype} if hf_token: pipe_kwargs["token"] = hf_token - # The prefetched snapshot dir keeps from_pretrained off the - # hub: its own snapshot sweep re-downloads files the scoped - # prefetch skipped (packaged root singles, 24 GB per FLUX.1). + # The prefetched snapshot dir keeps from_pretrained off the hub (its + # sweep re-pulls files the scoped prefetch skipped: 24 GB per FLUX.1). pipe = pipeline_cls.from_pretrained( _base_local_dir or repo_id, **pipe_kwargs ) elif kind == "single_file" and fam.single_file_is_pipeline: - # A single-file SDXL-style checkpoint is the WHOLE pipeline - # (U-Net + VAE + both text encoders), not a transformer-only file, - # so load it through the pipeline class. ``config`` points at the - # base repo so diffusers builds the correct structure/scheduler - # around the single-file weights instead of guessing from the file. + # A single-file SDXL-style checkpoint is the WHOLE pipeline, so load it + # through the pipeline class; ``config`` points at the base repo so diffusers + # builds the correct structure around the single-file weights. sf_pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "config": base} if hf_token: sf_pipe_kwargs["token"] = hf_token pipe = pipeline_cls.from_single_file(single_file_path, **sf_pipe_kwargs) else: - # Single-file transformer; the VAE / text-encoder / scheduler come - # from the base diffusers repo (the single file is transformer-only). + # Transformer-only single file; VAE/text-encoder/scheduler come from the base repo. sf_kwargs: dict[str, Any] = { "torch_dtype": dtype, "config": base, "subfolder": "transformer", - # Forward the token: the config is fetched from the (possibly - # gated) base repo before from_pretrained can authenticate. + # Config is fetched from the (possibly gated) base before auth. "token": hf_token, } if kind == "gguf": @@ -1523,8 +1342,7 @@ class DiffusionBackend: sf_kwargs["quantization_config"] = diffusers.GGUFQuantizationConfig( compute_dtype = dtype ) - # A safetensors single-file (e.g. fp8) carries its own dtype, so no - # GGUF dequant config is passed. + # A safetensors single-file (fp8) carries its own dtype: no GGUF dequant config. transformer = transformer_cls.from_single_file( single_file_path, **sf_kwargs ) @@ -1541,46 +1359,31 @@ class DiffusionBackend: _base_local_dir or base, **pipe_kwargs ) - # Resolve the effective speed mode: GGUF models default to the - # near-lossless `default` profile (compile is ~2.2x and sits below - # the quant noise floor), dense models stay bit-identical `off`. An - # explicit speed_mode (incl. "off") is honored verbatim. + # Effective speed: GGUF defaults to near-lossless `default` (compile ~2.2x, below + # the quant noise floor); dense stays bit-identical `off`. Explicit is honored. effective_speed = resolve_speed_mode(speed_mode, is_gguf = kind == "gguf") - # A torchao-quantized dense transformer runs its matmuls through the - # regional torch.compile; UNcompiled (eager) it is ~30x slower and would - # lose to the GGUF fallback. A dense model otherwise resolves to `off`, so - # force at least `default` (regional compile) whenever the quant engaged, - # or the opt-in "fast" path silently commits an eager, pathologically slow - # pipeline. + # A torchao-quantized dense transformer must be compiled (eager is ~30x slower and + # loses to GGUF), so force at least `default` when quant engaged. if transformer_quant_engaged is not None and effective_speed == SPEED_OFF: logger.info( "diffusion.transformer_quant: forcing speed_mode=default " "(quantized transformer must be compiled; eager is ~30x slower)" ) effective_speed = SPEED_DEFAULT - # Deferred speed auto for dense models: the load stays eager (a one-off - # image should not pay the 25-60s compile warmup, and eager is the - # bit-identical reference), but a user starting their 3rd image in one - # session has revealed repeated use -- generate() then engages the - # `default` profile, where the warmup starts paying back. Only when - # the request left speed unset, nothing forced a compiled tier, and - # this device/family could actually compile. + # Deferred speed auto for dense models: stay eager (a one-off image shouldn't pay + # the 25-60s compile), but generate() engages `default` on the 3rd image, where + # repeated use amortises it. Only when speed was unset, nothing forced compile, and this device can compile. speed_deferred = ( speed_mode is None and effective_speed == SPEED_OFF and transformer_quant_engaged is None and compile_eligible(target, is_gguf = False, family = fam) ) - # Opt-in speed optims run BEFORE placement (channels_last / compile - # must precede CPU offload). Snapshot the process-wide backend flags - # first so unload can restore them: TF32 / cudnn.benchmark are global, - # and a later `off` load must not inherit this load's settings. + # Speed optims run BEFORE placement (channels_last/compile precede offload). + # Snapshot the global backend flags (TF32/cudnn.benchmark) first for unload restore. backend_flags_before = snapshot_backend_flags() - # Pick the attention kernel BEFORE compile (compile traces attention). auto - # upgrades to cuDNN fused attention on NVIDIA when a speed profile is active - # (~1.18x, near-lossless); an explicit backend is honored, falling back to - # the diffusers default if its kernel is unavailable. Orthogonal to the - # weight quant -- it speeds the QK/PV matmuls torchao does not touch. + # Pick the attention kernel BEFORE compile. auto upgrades to cuDNN fused attention + # on NVIDIA when a speed profile is active (~1.18x); explicit is honored. attention_engaged = apply_attention_backend( pipe, select_attention_backend( @@ -1588,14 +1391,10 @@ class DiffusionBackend: ), logger = logger, ) - # Step caching (First-Block-Cache), also before compile. For many-step - # models it reuses the transformer tail across steps (~1.4x on Flux at - # LPIPS ~0.08). When engaged, compile must drop fullgraph (the cache's - # per-step decision is a graph break), so pass it through. - # Tri-state request: unset / "auto" lets the step-count policy decide - # (engage when this model's DEFAULT schedule reaches FBCACHE_MIN_STEPS, - # then re-check against the actual step count on every generation); - # explicit "off" / "fbcache" are pinned and never toggled. + # Step caching (First-Block-Cache), also before compile: reuses the transformer + # tail across steps (~1.4x on Flux at LPIPS ~0.08); when engaged, compile drops + # fullgraph (graph break). Tri-state: unset/"auto" -> step-count policy decides + # (engage when the DEFAULT schedule reaches FBCACHE_MIN_STEPS); "off"/"fbcache" pinned. cache_request = normalize_transformer_cache(transformer_cache) cache_auto = transformer_cache is None or cache_request == TC_AUTO cache_quant_active = transformer_quant_engaged is not None or bool(gguf_filename) @@ -1609,15 +1408,12 @@ class DiffusionBackend: pipe, mode = cache_request, threshold = transformer_cache_threshold, - # GGUF transformers are quantized too (the default Studio path), so the - # cache needs the higher quantized threshold to still trigger -- not just - # the dense-quant fast path. + # GGUF transformers are quantized too, so the cache needs the higher threshold. quant_active = cache_quant_active, logger = logger, ) - # An auto decision can flip at generation time, but only on a transformer - # that supports caching at all; a non-CacheMixin transformer (e.g. - # Z-Image) can never engage, so compile keeps fullgraph there. + # An auto decision can flip at generation time, but only on a cache-capable + # transformer (a non-CacheMixin one keeps fullgraph). cache_may_toggle = cache_auto and callable( getattr(getattr(pipe, "transformer", None), "enable_cache", None) ) @@ -1636,33 +1432,18 @@ class DiffusionBackend: ) else: cache_reason = "requested" - # Install the shared compile-safe eager patches (fused RMSNorm / - # AdaLayerNorm) for any active speed tier. They are class-level, idempotent - # and math-equivalent (FMA / fused -> neutral under compile, equal-or-more - # accurate), so they help eager AND compiled runs. The bit-identical `off` - # reference path must run with them UNINSTALLED, so uninstall there. - # - # Everything from here to the _LoadState commit mutates PROCESS-WIDE state - # (class patches, TORCHINDUCTOR_CACHE_DIR, backend flags). _unload_locked only - # reverses it via _state, so a failure BEFORE the commit would leak it (and - # break the next `off` load's bit-identity). Guard the whole block: on any - # pre-commit failure, restore everything; on success the commit transfers - # ownership to _state and _unload_locked takes over. - # The GGUF-specific speed lever (compiled dequant) applies only when the - # GGUF transformer was ACTUALLY loaded. On the dense torchao-quant - # fast path (fp8 / int8 / fp4) `gguf_filename` is still set as the fallback, - # but `pipe.transformer` is dense (no GGUFLinear), and those schemes need the - # REGIONAL block compile (dynamic quant is ~30x slower eager), not the GGUF - # dequant compile -- so treat the transformer as non-GGUF here. The - # safetensors kinds (single_file / pipeline) likewise have no GGUFLinear. + # Everything from here to the _LoadState commit mutates PROCESS-WIDE state (class + # patches, TORCHINDUCTOR_CACHE_DIR, backend flags). _unload_locked reverses it via + # _state, so a pre-commit failure would leak it; the try/finally below restores on failure. + # gguf_transformer: the GGUF-specific compiled dequant applies only when the GGUF + # was actually loaded. On the dense fast path gguf_filename is still set (fallback) + # but pipe.transformer is dense (needs REGIONAL block compile), so treat it non-GGUF. gguf_transformer = kind == "gguf" and transformer_quant_engaged is None eager_patched = False compile_ctx = None state_committed = False - # Lazy import: these patch modules import torch at module level, so - # importing them here (not at module load) keeps diffusion.py torch-free - # to import, letting get_diffusion_backend() run on a torchless native install. + # Lazy import (these modules import torch) keeps diffusion.py torch-free to import. from .diffusion_eager_patches import ( install_compile_safe_patches, uninstall_patches, @@ -1675,39 +1456,32 @@ class DiffusionBackend: try: if effective_speed != SPEED_OFF: install_compile_safe_patches() - # Per-arch compile-safe fusions (qwen _modulate / z-image residual - # addcmul, etc.). Also neutral under compile, so on for every active - # tier; tracked by the same eager_patched flag for uninstall. + # Per-arch compile-safe fusions (qwen _modulate / z-image residual, etc.); + # neutral under compile, tracked by the same eager_patched flag. install_arch_patches() eager_patched = True else: uninstall_patches() uninstall_arch_patches() - # Pre-warmed torch.compile cache (Mega-cache): when a compiled tier will - # run, point inductor at a per-fingerprint dir and load a matching bundle - # BEFORE the first compiled forward, so the one-time 25-58s compile can be - # paid once (by us / a first run) and reused. A miss is silent -> local - # compile, exactly as today. + # Pre-warmed torch.compile cache: point inductor at a per-fingerprint dir and + # load a matching bundle before the first compiled forward, so the 25-58s + # compile is paid once and reused. A miss is silent -> local compile. if effective_speed in (SPEED_DEFAULT, SPEED_MAX) and compile_eligible( target, is_gguf = gguf_transformer, family = fam ): compile_ctx = compile_cache.begin( family = fam.name, - # U-Net families (SDXL) carry the denoiser as pipe.unet; the - # fingerprint needs the module actually compiled. + # U-Net families (SDXL) carry the denoiser as pipe.unet. transformer = getattr(pipe, "transformer", None) or getattr(pipe, "unet", None), dtype = getattr(target, "dtype", None), quant = transformer_quant_engaged, attention_backend = attention_engaged, compile_kwargs = { - # Mirrors apply_speed_optims' fullgraph decision: an active - # step cache OR a planned offload graph-breaks, so the cached - # bundle must be keyed on the same fullgraph setting. An auto - # cache that could still engage mid-session also drops - # fullgraph: enabling FBCache under a fullgraph-compiled - # transformer would crash the first cached generation. + # Mirrors apply_speed_optims' fullgraph decision: an active or + # still-toggleable step cache OR a planned offload graph-breaks, + # so the cached bundle must key on the same fullgraph setting. "fullgraph": cache_engaged is None and not cache_may_toggle and plan.offload_policy == OFFLOAD_NONE, @@ -1726,24 +1500,20 @@ class DiffusionBackend: family = fam, speed_mode = effective_speed, cache_active = cache_engaged is not None or cache_may_toggle, - # The planned offload policy: group/model/sequential offload installs - # compiler-disabled onload hooks, so compile must drop fullgraph. + # Offload installs compiler-disabled onload hooks, so compile drops fullgraph. offload_active = plan.offload_policy != OFFLOAD_NONE, logger = logger, ) if transformer_quant_engaged is not None and not speed_applied.get("compiled"): - # Promotion above could not engage compile (e.g. the family is not - # compile-friendly, or compile_repeated_blocks failed): the quantized - # transformer is now running eager, which is far slower than the GGUF - # path it replaced. Surface it loudly rather than hiding the regression. + # Compile couldn't engage: the quantized transformer runs eager, far slower + # than the GGUF it replaced. Surface it loudly. logger.warning( "diffusion.transformer_quant: %s engaged but the transformer is NOT " "compiled; eager torchao quant is ~30x slower than GGUF here", transformer_quant_engaged, ) - # Quantise the dense companion text encoder(s) (opt-in fp8 / fp8_dynamic / - # int8 / nvfp4), also before placement so the offload hooks move the smaller - # weights. int8 needs a per-family keep-bf16 schedule, so pass the family. + # Quantise the dense companion text encoder(s) (opt-in), before placement so + # offload moves the smaller weights. Family drives int8's keep-bf16 schedule. te_quant = quantize_text_encoders( pipe, target, @@ -1753,19 +1523,14 @@ class DiffusionBackend: logger = logger, ) - # Apply the placement planned above (from MEASURED free device memory vs - # the model's estimated resident size). apply_memory_plan returns the - # (policy, tiling) ACTUALLY engaged (it may fall back to whole-module - # offload, and tiling is a no-op on a pipeline with no tiling control), so - # status stays honest. The dense fast path already placed the pipe - # resident; for the `none` policy this is an idempotent re-placement. + # Apply the planned placement; apply_memory_plan returns the (policy, tiling) + # ACTUALLY engaged so status stays honest. Idempotent for the `none` policy. effective_policy, effective_tiling = apply_memory_plan( pipe, plan, device = device, logger = logger ) - # Per-control provenance for status: what engaged and who decided it - # (the caller, or this backend's auto resolution). cpu_offload=False is - # the unset default, so only True counts as an explicit request. + # Per-control provenance for status. cpu_offload=False is the unset default, + # so only True is an explicit request. resolved = build_resolved_record( { "speed_mode": ( @@ -1784,9 +1549,7 @@ class DiffusionBackend: "transformer_quant": ( transformer_quant, transformer_quant_engaged or "off", - # The None reason must match the load kind: only a GGUF - # load has a GGUF transformer; a dense pipeline / - # single-file load simply keeps its dense weights. + # The None reason matches the load kind (GGUF loaded vs dense kept). ( "not engaged (GGUF transformer loaded)" if kind == "gguf" @@ -1857,19 +1620,15 @@ class DiffusionBackend: ) state_committed = True finally: - # Pre-commit failure: nothing owns the process-wide mutations yet, so - # roll them back here (symmetric with _unload_locked). + # Pre-commit failure: roll back the process-wide mutations (symmetric with _unload_locked). if not state_committed: restore_backend_flags(backend_flags_before) compile_cache.restore(compile_ctx) - # apply_speed_optims may have installed the compiled GGUF dequant - # before a later step failed; uninstall is idempotent. - gguf_compile.uninstall_all() + gguf_compile.uninstall_all() # idempotent if eager_patched: uninstall_patches() uninstall_arch_patches() - # Also free the half-built pipe's VRAM: the failed load never - # commits _state, so nothing else reclaims it until the next unload. + # Free the half-built pipe's VRAM (uncommitted _state -> nothing else reclaims it). clear_gpu_cache() logger.info( @@ -1917,11 +1676,9 @@ class DiffusionBackend: # 1. Pre-quantized checkpoint, when one is configured for the resolved scheme. scheme = select_transformer_quant_scheme(target, mode, family = getattr(fam, "name", None)) if scheme is None: - # Bail BEFORE the (multi-GB) dense download: an explicit unsupported scheme - # (e.g. fp8 on Ampere, nvfp4 off Blackwell) would otherwise materialise the - # dense transformer and move the pipe to CUDA only to fail at quantize below -- - # a long finalization under the load lock after the old model was already - # evicted. load_pipeline catches this and builds the GGUF pipeline instead. + # Bail BEFORE the multi-GB dense download: an unsupported scheme (fp8 on Ampere, + # nvfp4 off Blackwell) would otherwise materialise the transformer only to fail at + # quantize, after eviction. load_pipeline catches this and builds the GGUF pipeline. raise RuntimeError("transformer quant unsupported for this device/scheme") if fam is not None: source = resolve_prequant_source(fam, scheme, path_override = prequant_path) @@ -1934,12 +1691,10 @@ class DiffusionBackend: dtype = dtype, hf_token = hf_token, scheme = scheme, - # Reject a checkpoint built with a different Linear filter than the - # dense path uses, so the prequant and runtime-quant models match. + # Reject a checkpoint with a different Linear filter so prequant matches runtime-quant. min_features = DEFAULT_MIN_LINEAR_FEATURES, - # Only enforced when the caller forces fp8 fast-accum: a checkpoint that - # baked the other choice would ignore the request, so fall to the dense - # path (which applies it) instead of silently using the baked kernels. + # Only enforced when the caller forces fp8 fast-accum; a checkpoint that baked + # the other choice falls to the dense path instead of using the baked kernels. fast_accum = fast_accum, logger = logger, ) @@ -2020,10 +1775,8 @@ class DiffusionBackend: top of transformer_resident_override_mib (a double-count of the transformer).""" device_memory = snapshot_device_memory(target) if kind == "pipeline": - # The whole repo (transformer + companions) is one cached download; the - # cached bytes are the resident estimate (bnb-4bit / fp8 stay compressed). - # A LOCAL pipeline path isn't in the HF blob cache, so sum its on-disk weights - # (transformer included) instead of folding to zero and skipping offload. + # The whole repo is one cached download; cached bytes are the resident estimate + # (bnb-4bit/fp8 stay compressed). A LOCAL path isn't cached, so sum its on-disk weights. local_repo = Path(repo_id).expanduser() if repo_id else None if local_repo is not None and local_repo.is_dir(): cached = self._local_dir_weight_bytes(local_repo, exclude_transformer = False) @@ -2031,14 +1784,9 @@ class DiffusionBackend: cached = self._cache_bytes(repo_id) if repo_id else 0 cached_mib = int(cached // (1024 * 1024)) if cached else None model_dense_mib = estimate_safetensors_dense_mib(cached_mib) - # A repo can store weights in a NARROWER dtype than they occupy after the - # loader's torch_dtype cast: ideogram-4's base repo ships its two DiTs as - # raw float8, so the cached bytes undershoot the bf16-resident footprint - # by ~2x and auto planning would pick a resident placement that OOMs. - # When the family size table knows the bf16-resident total for THIS repo - # (the family base -- prequant repos like the bnb-4bit exports have - # different ids and really do stay compressed), plan against the larger - # of the two estimates. + # A repo can store weights NARROWER than the loaded dtype: ideogram-4's base ships its + # two DiTs as raw float8, so cached bytes undershoot the bf16 footprint ~2x and auto + # would OOM. When the size table knows the bf16 total for THIS repo, plan against the larger. is_narrow_base = bool(repo_id) and repo_id.strip().lower() == fam.base_repo.lower() if ( not is_narrow_base @@ -2046,20 +1794,15 @@ class DiffusionBackend: and local_repo is not None and local_repo.is_dir() ): - # A LOCAL directory mirror of the fp8 base never string-matches base_repo, - # so detect the fp8 layout from its transformer shard headers and reserve - # the bf16 footprint too (a local nf4 mirror has no fp8 scales and stays - # compressed). Header-only read, so this stays cheap and network-free. + # A local fp8 mirror never string-matches base_repo, so detect fp8 from the shard + # headers and reserve the bf16 footprint (a local nf4 mirror stays compressed). is_narrow_base = ideogram4_repo_is_fp8(repo_id) if is_narrow_base: table = family_bf16_components_gb(fam, fam.base_repo) if table is not None: - # family_bf16_components_gb is a network-free constant, so reserve the bf16 - # footprint even when the cache-derived estimate is absent (empty blob cache, - # or a best-effort download probe that swallowed a transient HF error and - # returned nothing). Otherwise model_dense_mib stays None and the planner - # reads "size unknown -> stay resident", so the ~54 GB fp8 pipeline plans a - # resident placement and OOMs a card that offload would have fit. + # Reserve the bf16 footprint from this network-free constant even when the + # cache estimate is absent; else model_dense_mib stays None ("size unknown -> + # resident") and the ~54 GB fp8 pipeline OOMs a card offload would have fit. table_mib = int(sum(table) * (1000.0**3) / (1024.0 * 1024.0)) model_dense_mib = ( table_mib if model_dense_mib is None else max(model_dense_mib, table_mib) @@ -2067,15 +1810,12 @@ class DiffusionBackend: companion_mib = None else: if transformer_resident_override_mib is not None: - # Planning for a different artifact than the file on disk (the dense - # transformer-quant candidate): the auto-policy's estimate replaces the - # file-size derivation; companions below stay measured from the cache. + # Planning for the dense-quant candidate (not the file on disk): the auto-policy's + # estimate replaces the file-size derivation; companions stay measured from cache. transformer_resident = transformer_resident_override_mib elif kind == "single_file": - # An fp8 transformer checkpoint loads via from_single_file with a bf16 - # compute dtype and no quantization_config, so diffusers upcasts it to - # bf16 (~2x resident); detect it from the basename. Excludes the - # single-file-is-pipeline (SDXL) case, which is already a bf16 pipeline. + # An fp8 checkpoint upcasts to bf16 on load (~2x resident); detect from the + # basename. Excludes the SDXL (single_file_is_pipeline) case (already bf16). fp8_upcast = not getattr(fam, "single_file_is_pipeline", False) and ( "fp8" in Path(single_file_path).name.lower() if single_file_path else False ) @@ -2084,18 +1824,12 @@ class DiffusionBackend: ) else: transformer_resident = estimate_gguf_resident_mib(file_size_mib(single_file_path)) - # The companion components (VAE + text encoders) load near their on-disk - # size; sum whatever the prefetch placed in the base-repo cache, or -- for a - # LOCAL diffusers base -- the on-disk component weights (the blob cache is - # empty for a local path, which would otherwise fold multi-GB companions to 0 - # and let auto planning pick a resident placement that OOMs). + # Companions (VAE + text encoders) load near on-disk size; sum the base-repo cache, + # or a LOCAL base's on-disk weights (the blob cache is empty for a local path). if companion_override_mib is not None: - # Re-planning the dense transformer-quant candidate: the dense path - # prefetches the base repo's transformer/ shards into the SAME blob cache - # _companion_cache_bytes sums, so reading it here would count the - # transformer AGAIN on top of transformer_resident_override_mib and make - # the resident quant plan look far too large. Use the auto-policy's own - # companion (text-encoder + VAE) estimate for this artifact instead. + # Re-planning the dense candidate: the prefetched transformer/ shards land in the + # SAME cache _companion_cache_bytes sums, so use the auto-policy's companion estimate + # instead of double-counting the transformer. companion_mib = companion_override_mib else: companion = self._companion_cache_bytes(base) @@ -2103,10 +1837,8 @@ class DiffusionBackend: model_dense_mib = None if transformer_resident is not None: model_dense_mib = transformer_resident + (companion_mib or 0) - # Feed the variant hint (single-file basename + base/repo) next to the family name - # so estimate_image_runtime_mib sees distilled markers ("turbo"/"schnell") that - # detect_family normalizes out of fam.name -- distilled models need ~15% less - # activation headroom, and over-reserving can force needless offload / tiling. + # Feed the variant hint (basename + repo) so estimate_image_runtime_mib sees distilled + # markers ("turbo"/"schnell") normalized out of fam.name (distilled needs ~15% less headroom). variant_hint = " ".join( p for p in ( @@ -2142,20 +1874,14 @@ class DiffusionBackend: return cached import diffusers - # torch_dtype=None is load-bearing: diffusers' from_pipe defaults torch_dtype to - # torch.float32 and then runs new_pipeline.to(dtype=float32) over EVERY component. - # That recast (a) needlessly upcasts the reused bf16 modules and (b) hard-crashes - # on the dense-quant fast path -- a torchao-quantized + torch.compiled transformer - # has tensor-subclass Linear weights that torch.nn.Module._apply cannot swap_tensors - # ("Couldn't swap Linear.weight"). Passing None makes from_pipe skip the cast and - # reuse the resident modules AT THEIR LOADED dtype, which is the whole point of - # from_pipe (component reuse, no reload, no extra VRAM). + # torch_dtype=None is load-bearing: from_pipe otherwise recasts EVERY component to fp32, + # which upcasts the reused bf16 modules and hard-crashes the dense-quant path (a + # torchao+compiled transformer's tensor-subclass weights can't swap_tensors). None reuses + # the resident modules at their loaded dtype (the point of from_pipe). pipe = getattr(diffusers, class_name).from_pipe(state.pipe, torch_dtype = None) - # Only publish to the shared aux cache if THIS load is still current. from_pipe runs - # under _generate_lock but NOT _lock, so an unload()/superseding load can clear - # _aux_pipes and null _state while it builds; caching unconditionally would re-insert - # a wrapper over now-stale modules that a later same-workflow load would reuse (or - # keep the old VRAM pinned). This generation still uses the returned pipe. + # Publish to the shared aux cache only if THIS load is still current: from_pipe runs under + # _generate_lock but NOT _lock, so an unload can null _state while it builds; caching then + # would hand a wrapper over stale modules to a later load. with self._lock: if self._state is state: self._aux_pipes[class_name] = pipe @@ -2179,48 +1905,37 @@ class DiffusionBackend: if cn_model is None: if cancel.is_set(): raise RuntimeError(DIFFUSION_CANCELLED_MSG) - # resolve_controlnet accepts a bare owner/name repo without the non-GGUF base - # trust gate, and from_pretrained below downloads and deserializes it. A - # malicious pickle .bin would execute on load, so run the same Hub malware - # preflight the chat/export loaders use before any remote ControlNet load. A - # local dir the user picked has no Hub scan and is exempt (fail-open there). + # resolve_controlnet accepts a bare owner/name without the base trust gate, and + # from_pretrained deserializes it (a malicious pickle would execute), so run the same + # Hub malware preflight the chat/export loaders use. A local dir is exempt (fail-open). if not getattr(resolved_cn, "is_local", False): from utils.security import evaluate_file_security _cn_fs = evaluate_file_security(resolved_cn.path, hf_token = state.hf_token or None) if _cn_fs.blocked: raise ValueError(_cn_fs.reason) - # Keep at most one ControlNet resident: evict the previous module + its - # from_pipe wrapper before loading the new one, or swapping ControlNets - # within a base-model load accumulates until OOM. + # Keep at most one ControlNet resident: evict the previous module + wrapper first, + # or swapping ControlNets within a load accumulates until OOM. if self._cn_models or self._cn_pipes: self._cn_models.clear() self._cn_pipes.clear() clear_gpu_cache() import torch - # state.dtype is the display string saved at load ("bfloat16"), NOT a - # torch.dtype; pass the real dtype so diffusers loads the ControlNet at the - # base compute dtype instead of silently defaulting to float32 (extra VRAM). + # state.dtype is the display string ("bfloat16"), not a torch.dtype; pass the real + # dtype so diffusers loads at the base compute dtype, not float32 (extra VRAM). cn_dtype = getattr(torch, str(state.dtype).replace("torch.", ""), None) cn_model = getattr(diffusers, model_cls_name).from_pretrained( resolved_cn.path, torch_dtype = cn_dtype, - # An empty / malformed token means anonymous access; the HF client can - # raise on a blank credential instead of falling back, so coerce to None. - token = state.hf_token or None, + token = state.hf_token or None, # blank -> anonymous ) if cancel.is_set(): - # An unload/eviction raced the blocking download above and may have already - # cleared the load. Bail BEFORE any device placement so we don't allocate - # several GB onto the GPU after _unload_locked() freed it (which would OOM - # or make the unload appear to free memory only to repopulate it). + # An unload raced the blocking download; bail BEFORE placement so we don't + # allocate onto a GPU _unload_locked() just freed. del cn_model raise RuntimeError(DIFFUSION_CANCELLED_MSG) - # Placement must follow the base model's offload policy. A resident base moves - # the ControlNet resident too; an offloaded (low-VRAM) base streams it through - # the device with group offloading instead of forcing the whole module onto the - # GPU, which would defeat the offload and risk an OOM. Best-effort: any failure - # falls back to the resident placement (the prior behaviour). + # Placement follows the base's offload policy: a resident base places it resident, an + # offloaded base streams it via group offloading. Best-effort; failure -> resident. if getattr(state, "offload_policy", OFFLOAD_NONE) != OFFLOAD_NONE and ( _offload_controlnet_module(cn_model, state.device, logger) ): @@ -2228,8 +1943,7 @@ class DiffusionBackend: else: cn_model = cn_model.to(state.device) if cancel.is_set(): - # An unload raced the blocking download above and already cleared the - # ControlNet caches; caching now would pin the module past the unload. + # An unload raced the download and cleared the caches; caching now would pin it. del cn_model raise RuntimeError(DIFFUSION_CANCELLED_MSG) self._cn_models[resolved_cn.id] = cn_model @@ -2240,9 +1954,8 @@ class DiffusionBackend: state.pipe, controlnet = cn_model, torch_dtype = None ) with self._lock: - # Same race as the model cache above: an unload/superseding load may - # have cleared _cn_pipes while from_pipe ran; caching now would pin a - # pipeline built around the UNLOADED base and hand it to the next load. + # Same race as the model cache: an unload may have cleared _cn_pipes while + # from_pipe ran; caching now would pin a pipeline over the unloaded base. if cancel.is_set() or self._state is not state: del pipe raise RuntimeError(DIFFUSION_CANCELLED_MSG) @@ -2264,10 +1977,8 @@ class DiffusionBackend: if denoiser is None or vae is None: return try: - # Read the dtype from the parameters (not denoiser.dtype): a plain nn.Module - # has no .dtype, and a torch.compile'd / wrapped denoiser can obscure it. Take - # the first FLOATING dtype: a GGUF-quantized transformer's leading params are - # packed uint8 storage, and nn.Module.to() rejects integer dtypes outright. + # Read the dtype from the parameters (a plain/compiled nn.Module may hide .dtype). + # Take the first FLOATING dtype (a GGUF transformer's leading params are packed uint8). target_dtype = next( (p.dtype for p in denoiser.parameters() if p.dtype.is_floating_point), None, @@ -2319,9 +2030,7 @@ class DiffusionBackend: ) resolved = diffusion_lora.resolve_specs(specs, hf_token = state.hf_token, cancel_event = cancel) - # The shared catalog scans both .safetensors and .gguf, but diffusers' - # load_lora_weights only takes safetensors; a .gguf adapter would otherwise fail - # deep in generation. Reject it here as a clean 400 before touching the pipe. + # diffusers load_lora_weights takes safetensors only; reject a .gguf adapter as a clean 400. bad = [r.id for r in resolved if r.fmt != "safetensors"] if bad: raise ValueError( @@ -2475,40 +2184,30 @@ class DiffusionBackend: negative_prompt: Optional[str] = None, width: int = 1024, height: int = 1024, - # Fallbacks for a caller that passes nothing; the route always sends the - # per-model values the UI seeds (few steps / no CFG for distilled models, - # more steps / real CFG for full ones). + # Fallbacks; the route always sends the per-model values the UI seeds. steps: int = 9, guidance: float = 0.0, seed: Optional[int] = None, batch_size: int = 1, - # Image-conditioned workflows (base64 / data-URL): an init image alone selects - # img2img; an init image + mask selects inpaint. ``strength`` is the img2img/ - # inpaint denoise strength (0 = keep source, 1 = full redraw). None = txt2img. + # Image-conditioned (base64/data-URL): init alone = img2img; init + mask = inpaint. + # ``strength`` is the denoise strength (0 = keep source, 1 = full redraw). None = txt2img. init_image: Optional[str] = None, mask_image: Optional[str] = None, strength: Optional[float] = None, - # Upscale (hires fix): a factor > 1 with an init image enlarges the input and - # re-denoises it at low strength to paint detail at the higher resolution. + # Upscale (hires fix): factor > 1 with an init image enlarges then re-denoises at low strength. upscale: Optional[float] = None, - # Reference workflow (FLUX.2): ADDITIONAL reference images beyond ``init_image``. The - # pipeline accepts a list, so multiple references can be combined (subject + style, - # character + scene). Ignored by non-reference workflows. + # Reference (FLUX.2): additional reference images beyond init_image (a list). Ignored elsewhere. reference_images: Optional[list[str]] = None, - # LoRA adapters as (id, weight) pairs; loaded onto the pipe (non-fused) and activated - # with set_adapters for this generation. None/empty = no LoRA (adapters cleared). + # LoRA (id, weight) pairs; loaded non-fused and activated for this generation. None/empty clears. loras: Optional[list[tuple[str, float]]] = None, - # ControlNet as (id, control_image_b64, control_type, strength, guidance_start, - # guidance_end); conditions the text-to-image path on a spatial control map. None = off. + # ControlNet (id, control_image_b64, control_type, strength, guidance_start, guidance_end). None = off. controlnet: Optional[tuple[str, str, str, float, float, float]] = None, ) -> dict[str, Any]: import torch from PIL import Image - # A per-generation cancel Event: unload()/a superseding load set THIS event - # (registered under _lock below) to abort just this denoise. _generate_lock - # serialises generations and is the only lock the denoise holds, so a slow - # generation never blocks status()/unload()/a new load. + # Per-generation cancel Event that unload()/a superseding load set (registered under + # _lock below) to abort just this denoise. _generate_lock is the only lock the denoise holds. cancel = threading.Event() with self._generate_lock: with self._lock: @@ -2516,41 +2215,26 @@ class DiffusionBackend: if state is None: raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) # Register under _lock so unload()/a load can signal THIS generation. - # A cancel that arrived before now either nulled _state (we raised - # above) or targets an older generation, so nothing is lost. self._active_generate_cancel = cancel try: - # Snapshot taken: the local `state` ref keeps the pipe alive even if - # unload() nulls _state mid-denoise, so the call below needs no _lock. + # The local `state` ref keeps the pipe alive even if unload() nulls _state. generator = torch.Generator(device = state.device) if seed is None: - # Draw a fresh random seed but keep it within JS's safe-integer - # range (< 2**53), so the reported seed round-trips through JSON - # and actually reproduces the image (a raw 64-bit seed would lose - # precision in the browser and the recipe couldn't be replayed). + # Keep the seed in JS's safe-integer range (< 2**53) so it round-trips + # through JSON and reproduces the image (a raw 64-bit seed loses precision). seed = generator.seed() & ((1 << 53) - 1) else: seed = int(seed) generator.manual_seed(seed) - # Deferred speed auto: by the 3rd image in one session repeated use is - # established, so engage the compile profile now -- before the LoRA / - # workflow wiring, matching the load-time ordering. Best-effort: a - # failure logs, leaves the eager pipe running, and never retries - # (the helper clears the flag first). - # - # But NOT when this generation requests a LoRA: a compiled transformer rejects - # LoRA (supports_lora is False once compiled), and _apply_loras raises before its - # unchanged-selection no-op, so engaging compile here would permanently break every - # LoRA generation on this load. Compile and LoRA are mutually exclusive; keep the - # pipe eager and let compile defer to a later LoRA-free generation. + # Deferred speed auto: engage the compile profile on the 3rd image, before the LoRA/ + # workflow wiring (load-time ordering). Best-effort; a failure stays eager and never retries. + # NOT when a LoRA is requested: a compiled transformer rejects LoRA, so compiling + # here would permanently break every LoRA generation on this load. lora_requested = any(w != 0 for (_id, w) in (loras or [])) - # Also stay eager while adapters from a PRIOR generation are still attached: this - # request may clear them (lora_requested False), but _apply_loras runs AFTER the - # engage below, so compiling here would bake the resident adapter into the graph and - # the subsequent unload_lora_weights() (swallowed on a compiled pipe) would leave it - # active forever -- silent wrong output on every later LoRA-free generation. Deferring - # lets _apply_loras clear it on the still-eager pipe; compile engages a gen later. + # Also stay eager while a PRIOR generation's adapters are attached: _apply_loras runs + # AFTER the engage below, so compiling would bake the adapter in and the swallowed + # unload_lora_weights() would leave it active forever. Defer until a later gen. loras_attached = bool(getattr(state.pipe, "_unsloth_loras", ())) if ( state.speed_deferred @@ -2566,22 +2250,18 @@ class DiffusionBackend: exc, ) - # Apply/adjust LoRA adapters on the resident pipe (non-fused) before picking - # the workflow pipe; from_pipe pipes share the transformer, so it propagates. + # Apply/adjust LoRA before picking the workflow pipe; from_pipe pipes share the transformer. self._apply_loras(state, loras, cancel) - # Select the pipeline for this workflow. txt2img uses the loaded pipe; - # img2img/inpaint reuse its resident modules via from_pipe (no reload); - # an edit model's OWN loaded pipe is already the edit pipeline. + # Select the workflow pipeline: txt2img uses the loaded pipe; img2img/inpaint reuse + # its modules via from_pipe; an edit model's own pipe is already the edit pipeline. pipe = state.pipe init_pil = mask_pil = None control_pil = None cn_scale = cn_gstart = cn_gend = cn_mode = None ref_extra: list = [] - # Validate parameter dependencies up front: mask / upscale / reference all - # need an input image, and reference conditioning needs a family that - # supports it. Without these guards an unsupported combination would be - # silently ignored and quietly fall back to txt2img / img2img. + # Validate dependencies up front: mask/upscale/reference need an input image, and + # reference needs a supporting family (else the combo silently falls back to txt2img). if init_image is None: if mask_image is not None: raise ValueError("mask_image requires an input image (init_image).") @@ -2595,9 +2275,8 @@ class DiffusionBackend: "model family." ) if getattr(state.family, "edit", False): - # Instruction editing: the loaded pipe is the edit pipeline. It always - # needs an input image; the prompt is the edit instruction. No mask, no - # from_pipe (the model has no plain text-to-image mode). + # Instruction editing: the loaded pipe IS the edit pipeline; always needs an + # input image, prompt is the instruction. No mask, no from_pipe. if init_image is None: raise ValueError( f"{state.family.name} is an image-editing model: provide an input image." @@ -2616,28 +2295,21 @@ class DiffusionBackend: init_pil = _decode_b64_image(init_image, mode = "RGB") mask_pil = _decode_b64_image(mask_image, mode = "L") elif init_image is not None and upscale is not None and upscale > 1.0: - # Upscale (hires fix): enlarge the input with Lanczos, then re-run the - # img2img pipeline on it at a low denoise strength so the transformer - # adds high-frequency detail without redrawing the content. Shares the - # img2img pipeline/modules via from_pipe (no extra VRAM, no reload). + # Upscale (hires fix): enlarge with Lanczos, then re-run img2img at low strength + # to add detail without redrawing. Shares the img2img pipeline via from_pipe. workflow = "upscale" pipe = self._workflow_pipe(state, state.family.img2img_pipeline_class, workflow) init_pil = _decode_b64_image(init_image, mode = "RGB") iw, ih = init_pil.size - # Cap the factor, THEN cap the absolute output: a large input times the - # factor (e.g. 1024 at 4x = 4096, or a big upload) would otherwise OOM the - # VAE/transformer. Bound the longest side to 2048 (txt2img's own max), - # scaling both dims to keep the aspect ratio; round to a multiple of 16 - # (VAE downsample + patch size require it for our families). + # Cap the factor, then the absolute output (longest side 2048, txt2img's max) + # to avoid an OOM-scale latent; round to a multiple of 16 (VAE downsample + patch). factor = max(1.0, min(float(upscale), 4.0)) tw_f, th_f = iw * factor, ih * factor max_side = 2048 fit = min(1.0, max_side / max(tw_f, th_f)) tw = max(16, int(round(tw_f * fit / 16.0)) * 16) th = max(16, int(round(th_f * fit / 16.0)) * 16) - # After the absolute cap, the target must still exceed the input, or - # "upscale" would shrink it (e.g. a 3000px source at 2x clamps to 2048). - # Reject rather than silently return a smaller image than uploaded. + # After the cap, the target must still exceed the input (else upscale shrinks it). if max(tw, th) <= max(iw, ih): raise ValueError( f"Upscale would not enlarge this image: its longest side " @@ -2646,22 +2318,14 @@ class DiffusionBackend: ) init_pil = init_pil.resize((tw, th), Image.LANCZOS) if strength is None: - # Hires-fix default: low enough to preserve content, high enough to - # synthesise new detail at the higher resolution. - strength = 0.35 + strength = 0.35 # hires-fix default: preserve content, add detail elif getattr(state.family, "reference", False) and init_image is not None: - # FLUX.2-style reference conditioning: the loaded pipe (Flux2KleinPipeline) - # takes the reference image directly via its `image` arg and generates a - # fresh image at the REQUESTED size, guided by both the prompt and the - # reference. No from_pipe (the loaded pipe already supports it), no strength - # (reference-conditioning, not a denoise blend), and the output size comes - # from the sliders (the pipeline resizes the reference to ~1MP itself). - # Checked AFTER inpaint/upscale so a mask/upscale request on a reference - # family (FLUX.2-klein also has an inpaint pipeline) still routes correctly. + # FLUX.2 reference conditioning: the loaded pipe takes the reference via `image` + # and generates at the REQUESTED size. No from_pipe, no strength; output size + # from the sliders. After inpaint/upscale so a mask/upscale on a reference family routes right. workflow = "reference" init_pil = _decode_b64_image(init_image, mode = "RGB") - # Additional references (FLUX.2 accepts a list): decode them so the - # conditioning combines all of them. Capped to keep VRAM bounded. + # Additional references (FLUX.2 combines a list); capped to bound VRAM. ref_extra = [ _decode_b64_image(x, mode = "RGB") for x in (reference_images or [])[:3] ] @@ -2672,15 +2336,12 @@ class DiffusionBackend: else: workflow = "txt2img" - # ControlNet conditioning (diffusers): applies to the plain text-to-image path. - # Builds the family's ControlNet pipeline around the resident modules (no reload) - # and passes a control map. v1 conditions txt2img only (not img2img/inpaint/edit). + # ControlNet (diffusers): txt2img only (not img2img/inpaint/edit). Builds the + # family's CN pipeline around resident modules and passes a control map. if controlnet is not None: from core.inference import diffusion_controlnet cn_id, cn_image_b64, cn_type, cn_strength, cn_gs, cn_ge = controlnet - # strength 0 disables ControlNet (documented on the request model, and the - # frontend slider allows it): skip the whole path so a no-op selection never - # pays the multi-GB ControlNet download / VRAM cost. + # strength 0 disables CN: skip the whole path so a no-op never pays the download/VRAM. if cn_strength in (None, 0, 0.0): controlnet = None else: @@ -2703,10 +2364,8 @@ class DiffusionBackend: "diffusers engine (needs a bf16 or bnb-4bit load of a family with a " "ControlNet pipeline; not GGUF-via-diffusers or torchao fp8/int8)." ) - # Decode + preprocess the control image FIRST so a malformed / unsupported - # image fails as a clean 400 BEFORE any ControlNet download or pipe build, - # rather than after paying that cost. Control map at the OUTPUT size so it - # aligns with the generated latents. + # Decode + preprocess the control image FIRST so a bad image 400s before + # any CN download/build. Control map at the OUTPUT size to align with latents. src = _decode_b64_image(cn_image_b64, mode = "RGB") control_pil = diffusion_controlnet.preprocess_control(src, cn_type).resize( (width, height), Image.LANCZOS @@ -2716,26 +2375,19 @@ class DiffusionBackend: cn_id, family = state.family.name ) except FileNotFoundError as exc: - # An unknown / missing ControlNet id is a bad selection -> 400, not a - # generic 500 (the route maps ValueError, not FileNotFoundError). + # An unknown CN id -> 400, not 500 (the route maps ValueError). raise ValueError(str(exc)) from exc pipe = self._controlnet_pipe(state, resolved_cn, cancel) workflow = "controlnet" cn_scale, cn_gstart, cn_gend = cn_strength, cn_gs, cn_ge - # Flux Union ControlNet selects the active mode by an integer - # ``control_mode`` (canny/depth/pose/...); map the chosen control type so - # the union model applies the right head instead of a default/wrong one. + # Flux Union CN selects its head by an integer control_mode; map the type. cn_mode = diffusion_controlnet.union_control_mode(cn_id, cn_type) - # Auto-resize odd-sized inputs to a multiple of 16 for the workflows whose - # OUTPUT size is taken from the input image (img2img / inpaint / extend / edit), - # so an upload like 186px tall no longer fails the pipeline's divisibility check. - # txt2img/reference use the validated slider size; upscale already produced a /16 - # target. The mask is matched to the snapped image so inpaint stays aligned. + # Snap odd-sized inputs to a multiple of 16 for workflows whose OUTPUT size comes + # from the input image (img2img/inpaint/edit); txt2img/reference use the slider, + # upscale already produced a /16 target. Mask is matched to the snapped image. if init_pil is not None and workflow in ("img2img", "inpaint", "edit"): - # img2img/inpaint derive the OUTPUT size from the uploaded image, so bound the - # longest side to txt2img's own 2048 ceiling first -- otherwise a normal phone - # photo (up to the 4096/side decode cap) drives an OOM-scale latent and an - # opaque 500. edit is exempt: its pipeline resizes the input to ~1MP internally. + # img2img/inpaint take output size from the upload, so bound the longest side to + # 2048 first (else a phone photo drives an OOM-scale latent). edit resizes internally. if workflow in ("img2img", "inpaint"): init_pil = _clamp_max_side(init_pil, 2048) init_pil = _snap_to_multiple(init_pil, 16) @@ -2744,57 +2396,40 @@ class DiffusionBackend: mask_pil = mask_pil.resize(init_pil.size, _PILImage.NEAREST) if init_pil is not None: # Keep the VAE encode dtype consistent with the input image. - # state.family is always a DiffusionFamily, which defines denoiser_attr. self._align_vae_dtype(pipe, state.family.denoiser_attr) - # Pipelines vary in which kwargs they accept (img2img derives size from the - # input image and may reject width/height; a distilled pipe may take no - # negative prompt or step callback), so gate every optional kwarg on the - # actual signature. + # Pipelines vary in accepted kwargs, so gate every optional one on the signature. call_params = inspect.signature(pipe.__call__).parameters kwargs: dict[str, Any] = { "prompt": prompt, "num_inference_steps": steps, - # Most pipelines take guidance via "guidance_scale"; Qwen-Image - # uses "true_cfg_scale" (its distilled guidance is off). + # Most pipelines use "guidance_scale"; Qwen-Image uses "true_cfg_scale". state.family.cfg_kwarg: guidance, "generator": generator, - # Generate the whole batch in one forward pass (VRAM-heavy). All - # share this call's seed, drawn sequentially from one generator. + # Whole batch in one forward pass; all share this call's seed. "num_images_per_prompt": batch_size, } if state.family.name == IDEOGRAM4_FAMILY_NAME: - # Ideogram 4 drives CFG through EITHER a constant guidance_scale OR - # a per-step guidance_schedule; its check_inputs rejects the call - # when both are set, and the schedule DEFAULTS to the recommended - # 45x7.0 + 3x3.0 polish taper (valid only at exactly 48 steps). At - # the family's advertised defaults, drop the constant so the - # recommended taper engages; any other request nulls the schedule - # so the constant broadcasts legally to the chosen step count. + # Ideogram 4 drives CFG via EITHER a constant guidance_scale OR a per-step + # guidance_schedule (check_inputs rejects both). At the advertised defaults drop + # the constant so the recommended 48-step taper engages; else null the schedule. if steps == 48 and abs(float(guidance) - 7.0) < 1e-6: kwargs.pop(state.family.cfg_kwarg, None) else: kwargs["guidance_schedule"] = None if init_pil is not None: - # Reference with extra images passes the whole list (FLUX.2 combines them); - # every other workflow takes the single image. + # Reference passes the whole list (FLUX.2 combines); others take the single image. kwargs["image"] = [init_pil, *ref_extra] if ref_extra else init_pil if mask_pil is not None and "mask_image" in call_params: kwargs["mask_image"] = mask_pil if strength is not None and "strength" in call_params: kwargs["strength"] = strength - # width/height. txt2img uses the requested slider size. Image-conditioned - # pipes must use the INPUT IMAGE's own size, NOT the slider: the output is - # the redrawn/extended input, and the denoise builds latents from the image, - # so a slider size that differs from the image mismatches (e.g. a 1536px - # outpaint vs a 1024 slider -> "tensor a (128) must match tensor b (192)"). - # Many img2img/inpaint pipelines drop width/height entirely; pass them only - # when accepted, derived from the image so they are always consistent. + # width/height: txt2img uses the slider; image-conditioned pipes must use the INPUT + # IMAGE's own size (a differing slider mismatches the latents). Many img2img/inpaint + # pipes drop them entirely, so pass only when accepted, derived from the image. if workflow in ("txt2img", "reference", "controlnet"): - # txt2img, FLUX.2 reference, and ControlNet all generate at the REQUESTED - # size; the reference/control image is resized to match, so it must not be - # pinned to an input image's size like img2img/inpaint/upscale are. + # These generate at the REQUESTED size (reference/control image resized to match). kwargs["width"] = width kwargs["height"] = height elif init_pil is not None: @@ -2806,9 +2441,8 @@ class DiffusionBackend: if negative_prompt and "negative_prompt" in call_params: kwargs["negative_prompt"] = negative_prompt if workflow == "controlnet" and control_pil is not None: - # The ControlNet pipeline takes the control map + its conditioning scale; - # guidance start/end bound the step range it acts over. Every kwarg is gated - # on the pipe signature so a family whose CN pipe omits one still runs. + # CN pipeline takes the control map + scale; guidance start/end bound its step + # range. Every kwarg is signature-gated so a family that omits one still runs. if "control_image" in call_params: kwargs["control_image"] = control_pil elif "image" in call_params: # some CN pipelines name it "image" @@ -2819,8 +2453,7 @@ class DiffusionBackend: kwargs["control_guidance_start"] = cn_gstart if "control_guidance_end" in call_params and cn_gend is not None: kwargs["control_guidance_end"] = cn_gend - # Union ControlNet mode index (Flux); only when the pipe accepts it and the - # selected control type maps to a known mode. + # Union CN mode index (Flux); only when accepted and the type maps to a mode. if "control_mode" in call_params and cn_mode is not None: kwargs["control_mode"] = cn_mode @@ -2835,8 +2468,7 @@ class DiffusionBackend: gen.eta_seconds = _estimate_eta( gen.total_steps, gen.step, gen.first_step_at, now ) - # Preempt a long denoise on unload/eviction or a superseding load: - # diffusers checks pipe._interrupt and stops after the current step. + # Preempt a long denoise on unload/superseding load (diffusers checks _interrupt). if cancel.is_set(): pipe._interrupt = True return callback_kwargs @@ -2844,22 +2476,13 @@ class DiffusionBackend: if "callback_on_step_end" in call_params: kwargs["callback_on_step_end"] = _on_step - # An AUTO cache decision is re-checked against the ACTUAL step count: - # a 28-step dev-style request gains FBCache even when the load's default - # schedule kept it off, and a few-step turbo request drops it (skipping - # a step there is a large quality hit). Explicit choices never toggle. + # Re-check an AUTO cache decision against the ACTUAL step count (a 28-step request + # gains FBCache, a few-step turbo drops it); explicit choices never toggle. if state.cache_auto: - # Key the policy on the EFFECTIVE denoise steps: an img2img/upscale/ - # inpaint request at strength < 1 only denoises a fraction of `steps` - # (e.g. a 28-step upscale at strength 0.35 runs ~10 steps), so passing - # the raw request would wrongly engage FBCache on exactly the short - # trajectory the policy keeps uncached. Only fold in `strength` when it - # is ACTUALLY applied to the pipe (same gate as the kwarg below), so a - # stray strength on a txt2img request never shortens the count. - # The pipe denoises `steps * strength`. When the request omits strength the - # kwarg above is NOT passed, so the pipe runs its OWN signature default (< 1 for - # every img2img/inpaint pipeline here, e.g. 0.6) -- still a short trajectory the - # policy must key on, or FBCache engages on a fraction of the advertised steps. + # Key on the EFFECTIVE denoise steps: an img2img/upscale request at strength < 1 + # only denoises a fraction of `steps`, so folding in `strength` (only when + # actually applied) keeps FBCache off the short trajectory. When strength is + # omitted the pipe's own default (< 1) still applies, so key on that too. strength_applied = effective_request_strength( strength, init_pil is not None, @@ -2875,9 +2498,8 @@ class DiffusionBackend: logger = logger, ) if toggled != state.transformer_cache: - # _LoadState is frozen (loads swap it as one unit); this is the - # one deliberate in-place update, tracking the pipe-level toggle - # that already happened so status() reports the true cache state. + # _LoadState is frozen; the one deliberate in-place update, tracking the + # pipe-level toggle so status() is truthful. object.__setattr__(state, "transformer_cache", toggled) entry = (state.resolved or {}).get("transformer_cache") if isinstance(entry, dict): @@ -2887,38 +2509,28 @@ class DiffusionBackend: + ("reaches" if toggled else "is below") + f" {FBCACHE_MIN_STEPS}" ) - # Start each generation from a clean step cache: FBCache residuals from - # a prior request on this resident pipe would otherwise be compared - # against this generation's first step (shape mismatch on a resolution/ - # batch change, or stale reuse). No-op when no cache is engaged. + # Start each generation from a clean step cache: prior FBCache residuals would + # otherwise be compared against this first step (shape mismatch / stale reuse). if state.transformer_cache: self._reset_step_cache(state.pipe) self._gen = gen try: - # inference_mode is strictly faster than the no_grad diffusers - # uses internally and numerically identical for inference. + # inference_mode is faster than no_grad and numerically identical here. with torch.inference_mode(): images = pipe(**kwargs).images finally: self._gen = None - # A cancelled denoise returns early with a partial/garbage image; - # don't hand it back to be persisted. + # A cancelled denoise returns a partial/garbage image; don't persist it. if cancel.is_set(): raise RuntimeError(DIFFUSION_CANCELLED_MSG) - # The first compiled generation just paid the compile cost; persist the - # warm torch.compile cache bundle when saving is enabled (the auto - # default / distributor mode). A STATIC compile (max tier, U-Net - # whole-module) produces new artifacts per (width, height, batch), so - # register this generation's shape first: a shape the bundle does not - # cover re-dirties the context and the save below rewrites the bundle - # with the enriched set. Idempotent + best-effort -- never fails a - # generation. + # Persist the warm torch.compile bundle after the first compiled generation. A + # STATIC compile makes new artifacts per (w,h,batch), so register this shape first + # (an uncovered shape re-dirties the context and the save rewrites the bundle). + # Idempotent + best-effort. try: - # Register the dims the forward ACTUALLY compiled with: the - # image-conditioned workflows run at the input image's size, not the - # slider's (see _compile_shape_dims), and a mis-registered slider - # shape would keep the truly-used shape out of the saved bundle. + # Register the dims the forward ACTUALLY compiled with (image-conditioned + # workflows run at the input image's size, not the slider; see _compile_shape_dims). reg_width, reg_height = _compile_shape_dims(workflow, init_pil, width, height) compile_cache.register_shape( state.compile_cache_ctx, @@ -2929,15 +2541,12 @@ class DiffusionBackend: compile_cache.save(state.compile_cache_ctx, logger = logger) except Exception: # noqa: BLE001 — cache persistence is best-effort pass - # Count the finished generation (drives the deferred speed - # engagement above); a batch of N images is one generation. + # Count the finished generation (drives deferred speed); a batch is one generation. object.__setattr__(state, "generation_count", state.generation_count + 1) - # Return the PIL images (not yet encoded): the route embeds each - # image's recipe and persists it via the gallery. + # Return the PIL images (unencoded); the route embeds recipes and persists them. return {"images": list(images), "seed": int(seed), "repo_id": state.repo_id} finally: - # Deregister so a later unload/load can't poke a finished generation - # (only if still ours — a newer generation may have replaced it). + # Deregister so a later unload/load can't poke a finished generation (if still ours). with self._lock: if self._active_generate_cancel is cancel: self._active_generate_cancel = None @@ -2962,27 +2571,19 @@ class DiffusionBackend: } def unload(self) -> dict[str, Any]: - # Abort an in-flight download so unload/an eviction returns promptly instead - # of waiting it out (the download runs without _lock and checks this event). + # Abort an in-flight (lock-free) download so unload/eviction returns promptly. self._cancel_event.set() with self._lock: - # Abort an in-flight denoise too by setting ITS cancel event, so the step - # callback stops it. The running generate keeps its own pipe reference, so - # freeing _state here can't crash it; its VRAM is reclaimed when it exits - # (within ~one step thanks to the cancel). + # Abort an in-flight denoise via ITS cancel event; the running generate keeps its + # own pipe ref, so freeing _state can't crash it (VRAM reclaimed when it exits). if self._active_generate_cancel is not None: self._active_generate_cancel.set() self._unload_locked() - # Cancel any in-flight load (its worker checks this token before - # committing) and drop the marker so the next load starts clean. + # Cancel any in-flight load (its worker checks this token) and drop the marker. self._load_token += 1 self._loading = None - # Wait for the signalled denoise to actually exit before reporting unloaded: - # callers treat this return as "VRAM is free" (the GPU arbiter hands the GPU - # to chat next; the training routes size their run against it), and the - # denoise holds its pipe until the next step callback. generate() holds - # _generate_lock for its full body, so a bare acquire is the exit barrier - # (never while holding _lock -- generate takes _lock inside _generate_lock). + # Barrier: wait for the signalled denoise to exit before reporting unloaded (callers + # treat this return as "VRAM is free"). generate() holds _generate_lock for its full body. with self._generate_lock: pass return self.status() @@ -2991,33 +2592,22 @@ class DiffusionBackend: state = self._state if state is None: return - # Restore the process-wide backend flags (TF32 / cudnn.benchmark) this load - # may have flipped, so the next `off` load is bit-identical again. + # Restore the process-wide backend flags this load flipped, so the next `off` load is + # bit-identical. compile_cache.restore + gguf_compile.uninstall_all likewise; all idempotent. restore_backend_flags(state.backend_flags_before) - # Restore TORCHINDUCTOR_CACHE_DIR and uninstall the shared eager patches, so a - # later `off` load runs the bit-identical reference path. Both are idempotent. compile_cache.restore(state.compile_cache_ctx) - # Uninstall the GGUF dequant accelerators (compiled dequant / global weight - # buffer) this load may have installed, so a later `off` load runs the stock, - # bit-identical dequant. Idempotent. gguf_compile.uninstall_all() if state.eager_patched: - # Lazy import (torch at module level) to keep diffusion.py torch-free to import. + # Lazy import to keep diffusion.py torch-free to import. from .diffusion_eager_patches import uninstall_patches from .diffusion_arch_patches import uninstall_arch_patches uninstall_patches() uninstall_arch_patches() - # NOTE: we deliberately do NOT call state.pipe.unload_lora_weights() here. unload() - # only acquires _generate_lock AFTER this teardown, so a LoRA-backed denoise - # can still be running on this same pipe for up to one more callback; mutating its - # adapter layers now would race that in-flight generation. The whole pipe is dropped - # just below (self._state = None; del state; clear_gpu_cache()), so the adapter - # tensors are freed with it -- no explicit unload is needed for memory or for a - # later load (which builds a fresh pipe). - # Drop the workflow pipes built around this load's modules so they don't pin the - # freed pipeline (they only re-wire its components, but holding the wrappers - # would keep the modules alive past unload). + # NOTE: deliberately NOT unload_lora_weights() here. unload() acquires _generate_lock only + # AFTER this teardown, so a LoRA-backed denoise may still run for one more callback; mutating + # its adapters now would race it. The whole pipe is dropped below, freeing the adapters with it. + # Drop the workflow pipes so they don't pin the freed pipeline's modules past unload. self._aux_pipes.clear() # Drop any ControlNet models + pipelines so the freed load carries no extra modules. self._cn_pipes.clear() @@ -3073,8 +2663,7 @@ class DiffusionBackend: "attention_backend": state.attention_backend, "transformer_cache": state.transformer_cache, "resolved": state.resolved, - # Image-conditioned workflows the loaded family supports, so the UI can gate - # its tabs. txt2img is always available on the diffusers engine. + # Workflows the loaded family supports, so the UI can gate its tabs. "workflows": _family_workflows(state.family), "supports_lora": diffusion_lora.supports_lora( engine = "diffusers", @@ -3097,24 +2686,20 @@ class DiffusionBackend: def _family_workflows(fam: DiffusionFamily) -> list[str]: """The workflow ids the diffusers engine can run for ``fam`` (drives UI gating).""" - # Instruction-editing families have no plain text-to-image mode: their pipeline always - # takes an input image + instruction, so they expose only the "edit" workflow. + # Instruction-editing families have no txt2img mode, so expose only "edit". if getattr(fam, "edit", False): return ["edit"] workflows = ["txt2img"] - # Reference families (FLUX.2) keep txt2img and add reference conditioning via their own - # pipeline's optional image arg (no img2img/inpaint classes needed). + # Reference families (FLUX.2) add reference conditioning via their pipeline's image arg. if getattr(fam, "reference", False): workflows.append("reference") if getattr(fam, "img2img_pipeline_class", None): - # Upscale (hires fix) runs on the img2img pipeline, so it is available exactly - # when img2img is. + # Upscale runs on the img2img pipeline, so available exactly when img2img is. workflows.append("img2img") workflows.append("upscale") if getattr(fam, "inpaint_pipeline_class", None): workflows.append("inpaint") - # Outpaint (extend) reuses the inpaint pipeline with a padded canvas + border mask, - # so it needs an inpaint pipeline that preserves the (larger) canvas size. + # Outpaint reuses the inpaint pipeline with a padded canvas, so needs one that preserves size. if getattr(fam, "inpaint_preserves_size", True): workflows.append("outpaint") return workflows diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index cf39441285..a6847ae020 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -69,23 +69,17 @@ from utils.subprocess_compat import windows_hidden_subprocess_kwargs logger = get_logger(__name__) -# A sampling-progress line like " 4/4" / "[ 12/ 28]" / "sampling: 50%|...| 14/28". -# We only trust a match whose denominator equals the requested step count, so an -# unrelated "1/100" elsewhere in the log can't move the bar. +# A sampling-progress line ("4/4", "[ 12/ 28]", "sampling: 50%|...| 14/28"). Only a match +# whose denominator equals the requested step count is trusted, so a stray "1/100" can't move the bar. _STEP_RE = re.compile(r"(\d+)\s*/\s*(\d+)") -# Serialises the one-time binary install so concurrent first-loads don't race on the -# download / extract / chmod. +# Serialises the one-time binary install so concurrent first-loads don't race. _install_lock = threading.Lock() -# sd-server accepts at most this many images per img_gen job; larger Studio batches -# (the request model allows up to 32) are split into chunks of this size, the way the -# one-shot path did them one image at a time. +# Max images per img_gen job; larger Studio batches (up to 32) are split into these chunks. _MAX_SERVER_BATCH = 8 -# Per-image wall-clock budget for a server job, so a batch gets a timeout proportional to -# its image count (matching the one-shot path, where each image had its own budget) rather -# than one fixed deadline the whole batch has to finish within. +# Per-image server-job budget, so a batch's timeout scales with image count. _SERVER_PER_IMAGE_TIMEOUT_S = 1800.0 @@ -121,12 +115,10 @@ def _server_binary_runnable(binary: str) -> bool: ) except OSError: return False # cannot exec at all (wrong arch / no execute bit / missing loader) - except Exception: # noqa: BLE001 -- timeout or anything odd: don't block on a flaky probe + except Exception: # noqa: BLE001 -- don't block on a flaky probe (timeout etc.) return True - # A negative return code is a signal death (e.g. -4 SIGILL from an incompatible - # prebuilt on an older CPU): the binary launches but immediately crashes, so treat it - # as unavailable and let the load fall back to diffusers instead of routing to a - # server that will die on startup. + # Negative return code = signal death (e.g. -4 SIGILL from an incompatible prebuilt on + # an older CPU): launches then crashes, so treat as unavailable and fall back to diffusers. return proc.returncode >= 0 and proc.returncode not in (126, 127) @@ -253,8 +245,7 @@ class _SdLoading: repo_id: str base_repo: str - # Companion asset repos (VAE / text encoders) this load fetches, so the - # delete-cached guard protects them for the whole download/finalize window. + # Companion asset repos (VAE / text encoders) so the delete-cached guard protects them. asset_repos: tuple[str, ...] = () expected_bytes: int = 0 downloaded_bytes: int = 0 @@ -301,18 +292,16 @@ class SdCppDiffusionBackend: self._lock = threading.Lock() self._generate_lock = threading.Lock() self._engine = engine # resolved lazily on first load so import stays cheap - # An engine passed in is an EXPLICIT injection (the test seam / escape hatch) and - # pins one-shot mode; an engine cached later by a runtime fallback must NOT, so a - # now-available server can still be used on the next load. + # An injected engine (test seam / escape hatch) pins one-shot mode; a fallback-cached + # engine must NOT, so a now-available server can still be used on the next load. self._engine_injected = engine is not None self._state: Optional[_SdState] = None self._loading: Optional[_SdLoading] = None self._load_token = 0 self._cancel_event = threading.Event() self._active_generate_cancel: Optional[threading.Event] = None - # The sd-server being started for an in-flight load, before it is committed to - # _state. Tracked so an unload / superseding load can stop it mid-startup instead - # of leaving it loading (and holding the generate lock) for the whole timeout. + # sd-server started for an in-flight load, before it commits to _state; tracked so an + # unload / superseding load can stop it mid-startup instead of waiting out the timeout. self._pending_server: Optional[SdCppServer] = None self._gen: Optional[_SdGen] = None @@ -343,10 +332,8 @@ class SdCppDiffusionBackend: """ if self._engine_injected and self._engine is not None: return "oneshot", None, self._resolve_engine() - # Install the sd-server build matching the resolved device backend (ROCm / Vulkan / - # CUDA), not the default CPU build: a forced/enabled native load on a GPU host must - # not silently fetch the plain-CPU server. Lazy import avoids an import cycle with - # the router, which imports this backend during engine selection. + # Install the server build matching the resolved backend (ROCm/Vulkan/CUDA), not the + # default CPU build. Lazy import avoids an import cycle with the router. from core.inference.diffusion_engine_router import _install_accelerator_for accelerator = _install_accelerator_for( @@ -375,8 +362,8 @@ class SdCppDiffusionBackend: cpu_offload: bool = False, memory_mode: Optional[str] = None, speed_mode: Optional[str] = None, - # diffusers-only knobs accepted (so the route calls both engines uniformly) - # and ignored -- sd.cpp has no torchao quant / SDPA dispatcher / fbcache. + # diffusers-only knobs accepted for a uniform call and ignored (sd.cpp has no + # torchao quant / SDPA dispatcher / fbcache). text_encoder_quant: Optional[str] = None, transformer_quant: Optional[str] = None, transformer_quant_fast_accum: Optional[bool] = None, @@ -384,22 +371,18 @@ class SdCppDiffusionBackend: attention_backend: Optional[str] = None, transformer_cache: Optional[str] = None, transformer_cache_threshold: Optional[float] = None, - # Accepted for a uniform engine interface; the native engine is GGUF-only, so a - # non-GGUF kind never routes here (the router forces diffusers for those). + # Accepted for interface parity; native is GGUF-only (router forces diffusers otherwise). model_kind: Optional[str] = None, ) -> dict[str, Any]: """Validate, then fetch assets on a daemon thread. Returns at once.""" - # An empty / whitespace token is "no token": passing "" verbatim to HfApi / - # hf_hub_download is treated as an explicit (invalid) credential and breaks the - # anonymous fallback for public repos. + # Empty/whitespace token = "no token"; "" verbatim breaks the anonymous fallback. hf_token = hf_token.strip() if hf_token and hf_token.strip() else None if not gguf_filename: raise ValueError( "gguf_filename is required: the native engine loads single-file GGUF checkpoints only." ) - # Use the filename-fallback detector the route validated with, so a local - # .gguf pick whose family keyword lives only in the basename doesn't pass - # validation and then dead-end here on a no-GPU (native-routed) host. + # Filename-fallback detector (as the route validated) so a local .gguf whose family + # keyword lives only in the basename doesn't dead-end here on a native-routed host. fam = detect_family_for_pick(repo_id, gguf_filename, family_override) if fam is None: raise ValueError( @@ -414,9 +397,8 @@ class SdCppDiffusionBackend: with self._lock: if self._loading is not None and self._loading.error is None: raise RuntimeError("A diffusion load is already in progress.") - # A superseding load must stop any in-flight generation, or the old sd-cli - # keeps running against the previous model and can still return / persist an - # image after the new load has started (matches unload()'s cancel). + # A superseding load must stop any in-flight generation, else the old run can + # still persist an image after the new load starts (matches unload()'s cancel). if self._active_generate_cancel is not None: self._active_generate_cancel.set() self._load_token += 1 @@ -465,15 +447,12 @@ class SdCppDiffusionBackend: _load_token: int, ) -> None: try: - # Resolve the backend mode (persistent sd-server preferred, one-shot sd-cli - # fallback) and binary up front so an install / missing-binary failure - # surfaces before the multi-GB asset pull. + # Resolve mode (server preferred, one-shot fallback) + binary up front so an + # install / missing-binary failure surfaces before the multi-GB asset pull. mode, server_binary, engine = self._resolve_backend() if mode == "server": - # Probe the server binary before the multi-GB asset pull: a present but - # unrunnable build (wrong arch / missing libs) would otherwise download - # everything and only then fail to start. If it cannot run, fall back to - # the one-shot engine now (when it is usable), else surface the failure. + # Probe the server binary before the pull: a present-but-unrunnable build + # would download everything then fail. Fall back to one-shot if usable. assert server_binary is not None if not _server_binary_runnable(server_binary): logger.warning( @@ -488,9 +467,8 @@ class SdCppDiffusionBackend: raise RuntimeError("sd-server binary is present but not runnable.") mode, server_binary, engine = "oneshot", None, self._resolve_engine() if mode == "oneshot": - # Probe the binary: version() returns None when the present binary cannot - # run (bad perms / missing libs), so fail now rather than commit a "ready" - # state that crashes on the first generation. + # version() is None when a present binary can't run; fail now, not on the + # first generation. assert engine is not None if engine.version() is None: raise RuntimeError("sd-cli binary is present but not runnable.") @@ -509,22 +487,17 @@ class SdCppDiffusionBackend: qwen2vl = paths.get("qwen2vl"), ) device = resolve_diffusion_device_target().device - # Honor the requested speed everywhere; offload only off-CPU (forced - # sd_cpp / MPS), since on CPU the weights are resident in RAM and the - # offload flags are no-ops. + # Honor speed everywhere; offload only off-CPU (on CPU weights are resident, + # so the flags are no-ops). offload: tuple[str, ...] = () if device != "cpu": offload = tuple(offload_flags(_memory_policy(memory_mode, cpu_offload))) native_speed = _native_speed_for(speed_mode) - # Tear down any previously-loaded model, then commit the new one. A generation - # that started during the (slow) asset download is still running against the OLD - # model: abort it and WAIT on _generate_lock for it to exit before swapping, or - # a stale run could finish afterward and persist an image from the previous - # model. For server mode we stop the old server and start (load) the new one - # HERE, under _generate_lock, so generation never races a half-loaded server and - # two resident models never coexist. _generate_lock is taken only now, not during - # the download, so the long fetch never serialises against generation. + # Tear down the old model then commit the new one under _generate_lock: abort and + # WAIT for any generation that started during the download, so a stale run can't + # persist an image afterward and two resident servers never coexist. The lock is + # taken only now (not during the fetch), so the long download never serialises generation. with self._lock: if self._load_token != _load_token: return # superseded / cancelled @@ -542,33 +515,27 @@ class SdCppDiffusionBackend: if mode == "server": assert server_binary is not None server = SdCppServer(server_binary) - # Publish the not-yet-committed server so unload() / a superseding load - # can stop it mid-startup (SdCppServer.stop aborts the readiness wait - # without waiting on the lifecycle lock), instead of it loading for the - # full startup timeout while holding the generate lock. + # Publish the uncommitted server so unload() / a superseding load can stop + # it mid-startup (stop() aborts the readiness wait) instead of waiting out + # the full startup timeout while holding the generate lock. with self._lock: self._pending_server = server try: - # Blocks until the server has loaded the model and is answering - # (its readiness check); raises with the log tail on a failed load. + # Blocks until the model is loaded and answering; raises with the log tail on failure. server.start( files, vae_format = fam.sd_cpp_vae_format, offload = list(offload), native_speed = native_speed, - # Pin the CPU backend to physical cores; sd.cpp's own - # default oversubscribes hyperthreads (see _default_threads). + # Pin to physical cores (sd.cpp's default oversubscribes; see _default_threads). threads = _default_threads(), ) except SdCppCancelled: - # Startup was aborted by an unload / superseding load: stop the - # half-started server and bail (the outer handler returns cleanly). + # Aborted by unload / superseding load: stop the half-started server and bail. server.stop() raise except Exception as start_exc: # noqa: BLE001 - # A present-but-unusable sd-server must be no worse than the - # one-shot engine: fall back to sd-cli when it is usable, else - # surface the server error. + # Fall back to one-shot sd-cli if usable, else surface the server error. logger.warning( "sd-server failed to start (%s); falling back to one-shot sd-cli.", start_exc, @@ -595,8 +562,7 @@ class SdCppDiffusionBackend: vae_format = fam.sd_cpp_vae_format, native_speed = native_speed, offload_flags = offload, - # One-shot sd-cli reads this per generation (state.threads); pin to - # physical cores for the same reason as the server (see _default_threads). + # One-shot sd-cli reads this per generation; pin to physical cores. threads = _default_threads(), sampling_method = fam.sd_cpp_sampling_method, flow_shift = fam.sd_cpp_flow_shift, @@ -606,8 +572,7 @@ class SdCppDiffusionBackend: ) with self._lock: if self._load_token != _load_token: - # Superseded / unloaded while we were loading: discard the server - # we just started so it doesn't leak (and keep _state unloaded). + # Superseded / unloaded while loading: discard the started server so it doesn't leak. if server is not None: server.stop() return @@ -619,9 +584,7 @@ class SdCppDiffusionBackend: if self._load_token != _load_token: return logger.error("sd_cpp.load_failed: %s", exc) - # Redact filesystem paths before this reaches /images/load-progress: an - # asset-fetch / local-path / cache-IO failure can embed absolute paths - # (e.g. /home//...), and the diffusers load path scrubs the same way. + # Redact filesystem paths before this reaches /images/load-progress (as diffusers does). from utils.native_path_leases import redact_native_paths with self._lock: @@ -649,8 +612,7 @@ class SdCppDiffusionBackend: from huggingface_hub import HfApi api = HfApi(token = hf_token) for repo, fn, kind in assets: - # Only the transformer can be a local path; for the others ``repo`` is - # an HF id (a same-named local dir must not skip the size estimate). + # Only the transformer can be a local path; others are always HF ids. if kind == "diffusion_model" and Path(repo).expanduser().exists(): continue try: @@ -748,24 +710,17 @@ class SdCppDiffusionBackend: guidance: float = 0.0, seed: Optional[int] = None, batch_size: int = 1, - # Accepted for a uniform engine interface. The native engine is text-to-image - # only for now (sd-cli's init-img/mask plumbing is not wired), so an image- - # conditioned request is rejected clearly rather than silently dropping the input. + # Accepted for interface parity; native is text-to-image only, so image-conditioned + # requests are rejected clearly below rather than silently dropped. init_image: Optional[str] = None, mask_image: Optional[str] = None, strength: Optional[float] = None, - # Accepted for the uniform engine interface; upscale needs an init image, so the - # init_image guard below rejects it on the native engine like img2img/inpaint. - upscale: Optional[float] = None, - # Reference workflow is GPU/diffusers-only (FLUX.2); accepted for interface parity. - reference_images: Optional[list[str]] = None, - # LoRA adapters as (id, weight) pairs; resolved up front, then applied per engine - # path: prompt tags for one-shot sd-cli, structured `lora` entries - # for the resident sd-server. None/empty = no LoRA. + upscale: Optional[float] = None, # needs an init image; rejected by the guard below + reference_images: Optional[list[str]] = None, # GPU/diffusers-only (FLUX.2) + # LoRA (id, weight) pairs; resolved up front then applied per path: prompt tags for + # one-shot sd-cli, structured `lora` for sd-server. None/empty = no LoRA. loras: Optional[list[tuple[str, float]]] = None, - # Accepted for the uniform engine interface; the guard below rejects it on the native - # engine (ControlNet is diffusers-only) like img2img/inpaint, so a direct API call with - # ControlNet set fails clearly instead of TypeError'ing on an unexpected kwarg. + # ControlNet is diffusers-only; rejected by the guard below (accepted for parity). controlnet: Optional[tuple[str, str, str, float, float, float]] = None, ) -> dict[str, Any]: import tempfile @@ -780,15 +735,11 @@ class SdCppDiffusionBackend: or reference_images or (upscale is not None and upscale > 1) ): - # upscale needs an input image, so a direct API call with upscale > 1 but no - # init_image must be rejected too rather than silently returning a plain, - # un-upscaled text-to-image result (the diffusers backend rejects the same). raise ValueError( "img2img / inpaint / reference / upscale are not yet supported on the native " "sd.cpp engine; run on a GPU (diffusers) for image-conditioned workflows." ) - # strength 0 (or None) disables ControlNet (documented on the request model, matches - # the diffusers path), so it must be a no-op here too, not a hard 400. + # strength 0/None disables ControlNet (matches diffusers), so no-op it rather than 400. if controlnet is not None and controlnet[3] in (None, 0, 0.0): controlnet = None if controlnet is not None: @@ -803,9 +754,8 @@ class SdCppDiffusionBackend: state = self._state if state is None: raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) - # A resident server can exit while idle; if a client generates without first - # polling status, drop the stale loaded state and report not-loaded so it gets - # the recoverable reload path instead of a 500 from img_gen (not running). + # A resident server can exit while idle; drop stale state and report not-loaded + # so the client gets the recoverable reload path, not a 500 from img_gen. if ( state.mode == "server" and state.server is not None @@ -820,11 +770,9 @@ class SdCppDiffusionBackend: else: seed = int(seed) cfg_scale, flux_guidance = _map_guidance(state.family, guidance) - # Resolve any selected LoRA adapters up front (downloads land in the HF - # cache; a bad id fails here as a clear 400 before we generate). Drop - # weight-0 rows BEFORE the support gate: weight 0 disables an adapter, so a - # request carrying only disabled rows stays a no-op even on a family where - # native LoRA is unsupported, rather than 400 on a dead selection. + # Resolve selected LoRAs up front (a bad id -> clear 400 before generating). + # Drop weight-0 rows BEFORE the support gate so a request of only-disabled + # rows stays a no-op even where native LoRA is unsupported. lora_resolved: list = [] active_loras = [(i, w) for (i, w) in (loras or []) if w != 0] if active_loras: @@ -874,8 +822,7 @@ class SdCppDiffusionBackend: ) if cancel.is_set(): raise RuntimeError(DIFFUSION_CANCELLED_MSG) - # ``seeds`` is the per-image seed (image i used seed+i), so the route can - # persist the real seed for every image in the batch. + # ``seeds`` is the per-image seed (image i used seed+i) for the route to persist. return { "images": images, "seed": int(seed), @@ -930,15 +877,12 @@ class SdCppDiffusionBackend: assert state.server is not None total = max(1, int(batch_size)) - # sd.cpp's image seed is signed int64; mask the base (and every derived seed) so a - # large explicit seed is not rejected / wrapped inconsistently by the server. + # sd.cpp's image seed is signed int64; mask base and derived seeds to that range. base_seed = int(seed) & ((1 << 63) - 1) images: list = [] seeds: list[int] = [] - # Stage selected LoRAs into a per-request subdir of the server's lora-model-dir so a - # previous request's adapters can't leak into this one; reference them by the path - # relative to that dir (what the server's recursive scan resolves against). The - # subdir is removed after the batch. supports_lora already gated the family upstream. + # Stage LoRAs into a per-request subdir of the server's lora-model-dir (so a prior + # request's adapters can't leak in), referenced by path relative to that dir; removed after. lora_payload: Optional[list[dict]] = None lora_stage: Optional[Path] = None if lora_resolved: @@ -979,9 +923,7 @@ class SdCppDiffusionBackend: cancel_event = cancel, total_timeout = _SERVER_PER_IMAGE_TIMEOUT_S * count, ) - # All-or-nothing per chunk, like the one-shot path: if the server returns fewer - # blobs than requested (e.g. one image in the batch failed to encode), fail - # rather than silently dropping images from the user's requested batch. + # All-or-nothing per chunk: fail rather than silently drop images from the batch. if not cancel.is_set() and len(blobs) != count: raise RuntimeError( f"sd-server returned {len(blobs)} of {count} requested images in the batch." @@ -1032,9 +974,7 @@ class SdCppDiffusionBackend: images = [] seeds: list[int] = [] with tempfile.TemporaryDirectory(prefix = "sdcpp_gen_") as tmpdir: - # Materialize selected LoRAs into a managed dir sd-cli can scan, and inject - # matching tags into the prompt (deduped against any the user - # typed). Empty -> prompt/dir unchanged. + # Materialize LoRAs into a scan dir and inject tags (deduped). Empty -> unchanged. eff_prompt = prompt lora_dir: Optional[str] = None if lora_resolved: @@ -1046,10 +986,8 @@ class SdCppDiffusionBackend: for index in range(max(1, int(batch_size))): if cancel.is_set(): raise RuntimeError(DIFFUSION_CANCELLED_MSG) - # Distinct seed per batch image, reproducible image-by-image from the base - # seed. Mask to int64, NOT 53 bits: the request model and the diffusers - # backend both accept large explicit seeds, so a tight 2**53 mask would - # truncate them and collide distinct requested seeds onto the same image. + # Distinct reproducible seed per image; mask to int64 (not 53 bits, which + # would truncate large explicit seeds and collide distinct ones). seed_i = (seed + index) & ((1 << 63) - 1) out_path = str(Path(tmpdir) / f"img_{index}.png") params = SdCppGenParams( @@ -1123,32 +1061,25 @@ class SdCppDiffusionBackend: self._state = None self._load_token += 1 self._loading = None - # A load may be mid server.start() with the server not yet committed to _state; - # grab it too so we can stop it (its startup is abortable) instead of leaving it - # loading for the full startup timeout. + # Grab a mid-start() uncommitted server too so we can stop it (startup is abortable). pending = self._pending_server self._pending_server = None - # Stop the resident server outside the lock (terminate can take a few seconds). A - # mid-flight generation had its cancel event set above, so its poll loop unwinds - # as the process goes away. + # Stop the resident server outside the lock (terminate can take seconds); a mid-flight + # generation had its cancel set above and unwinds as the process goes away. if state is not None and state.server is not None: state.server.stop() if pending is not None and pending is not (state.server if state else None): pending.stop() - # Wait for a signalled one-shot generation to actually exit before reporting - # unloaded: callers (the GPU arbiter, training cleanup) treat this return as - # "the device is free", but a one-shot sd-cli child killed by the cancel above - # unwinds under _generate_lock. A bare acquire is the exit barrier (never taken - # while holding _lock; same pattern as DiffusionBackend.unload). + # Barrier: wait for a signalled one-shot generation to exit before reporting unloaded, + # since callers treat this return as "device is free" (same pattern as DiffusionBackend.unload). with self._generate_lock: pass return self.status() def status(self) -> dict[str, Any]: state = self._state - # A resident sd-server can exit after load (OOM-killed / crashed while idle). If so, - # drop the stale loaded state so status reports not-loaded and clients reload, - # instead of every generation failing with a 500 against a dead process. + # A resident sd-server can exit after load (OOM/crash while idle); drop stale state so + # status reports not-loaded and clients reload, not a 500 per generation on a dead process. if ( state is not None and state.mode == "server" @@ -1193,10 +1124,7 @@ class SdCppDiffusionBackend: "base_repo": state.base_repo, "device": state.device, "dtype": "gguf", - # Reflect the offload flags actually passed to sd-cli, so a balanced/low_vram - # (or cpu_offload) load is verifiable from status instead of always reading - # "none". On CPU _run_load leaves offload_flags empty (the flags are no-ops), - # so this correctly stays "none" there. + # Reflect the offload flags actually passed to sd-cli (empty on CPU -> "none"). "cpu_offload": bool(state.offload_flags), "offload_policy": "active" if state.offload_flags else "none", "vae_tiling": False, @@ -1218,10 +1146,7 @@ class SdCppDiffusionBackend: "supports_controlnet": False, # "server" = resident sd-server (load once); "oneshot" = legacy per-image sd-cli. "native_mode": state.mode, - # The native engine supports plain text-to-image only (generate() rejects - # img2img / inpaint / reference / upscale), so advertise just txt2img. Without - # this the status omits workflows, the UI reads [], and it disables the Create - # tab for a loaded native model, stranding the user on an image-only tab. + # Native supports txt2img only; advertise it so the UI doesn't disable the Create tab. "workflows": ["txt2img"], } diff --git a/studio/backend/core/inference/sd_cpp_server.py b/studio/backend/core/inference/sd_cpp_server.py index d66f3cdded..dff17d8b31 100644 --- a/studio/backend/core/inference/sd_cpp_server.py +++ b/studio/backend/core/inference/sd_cpp_server.py @@ -53,8 +53,7 @@ from utils.subprocess_compat import windows_hidden_subprocess_kwargs logger = logging.getLogger(__name__) -# httpx transport errors meaning "the server is gone / connection refused" -- treated -# as "not ready yet" while polling readiness, and as a fatal "server died" mid-request. +# "server gone / connection refused": not-ready while polling, fatal mid-request. _TRANSPORT_ERRORS = ( httpx.ConnectError, httpx.ReadError, @@ -62,10 +61,8 @@ _TRANSPORT_ERRORS = ( httpx.WriteError, ) -# Readiness probe. Upstream binds the port only AFTER the model is loaded, so any 200 -# means ready. We use /v1/models (a trivial, always-fast handler) rather than -# /sdcpp/v1/capabilities: the capabilities handler can block in some builds (it enumerates -# model metadata), which would stall readiness even though the server is up. +# Readiness probe: port binds only after the model loads, so any 200 means ready. Use +# trivial /v1/models, not /sdcpp/v1/capabilities (can block enumerating metadata). _READY_PATH = "/v1/models" # Native async sdcpp API. _IMG_GEN_PATH = "/sdcpp/v1/img_gen" @@ -75,10 +72,8 @@ _TERMINAL_OK = "completed" _TERMINAL_FAIL = "failed" _TERMINAL_CANCELLED = "cancelled" -# After a cancel is requested, how long to let the server reflect it in job status before -# abandoning the poll. The native cancel is best-effort, so without this cap a server that -# ignores/loses the cancel would keep this call (and the backend's generate lock) alive -# until the job finishes naturally, blocking a superseding load from swapping the model. +# 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 @@ -95,21 +90,15 @@ class SdCppServer: self.host = host self.port: Optional[int] = None self._process: Optional[subprocess.Popen] = None - # Fixed-size, thread-safe tail buffer: the drain thread appends while lifecycle / - # request threads read it for diagnostics, so a deque(maxlen) is safer and cheaper - # than a list with manual slicing. + # Bounded tail buffer shared by the drain thread (appends) and readers (diagnostics). self._tail: deque[str] = deque(maxlen = 200) self._stdout_thread: Optional[threading.Thread] = None self._lifecycle_lock = threading.Lock() - # Set (lock-free) by stop() so a blocking start()/readiness wait can be aborted - # promptly without waiting on the lifecycle lock start() holds. + # Set lock-free by stop() so a blocking start()/readiness wait bails promptly. self._abort = threading.Event() - # Set for the duration of a generation so the continuous stdout drain can feed - # the active request's step-progress callback; cleared in img_gen's finally. + # Set during a generation so the stdout drain feeds the step-progress callback. self._step_listener: Optional[Callable[[str], None]] = None - # trust_env=False: this client only ever talks to the loopback sd-server, so it must - # not route through HTTP_PROXY/HTTPS_PROXY (a proxy without 127.0.0.1 in NO_PROXY - # would break readiness/generation). Matches the local llama-server clients. + # trust_env=False: loopback-only client must not route through HTTP(S)_PROXY. self._client = httpx.Client(timeout = 30.0, trust_env = False) self._scratch_dir: Optional[str] = None self._stopped = False @@ -156,16 +145,13 @@ class SdCppServer: a concurrent start/stop can't interleave. """ with self._lifecycle_lock: - # A stop()/unload that raced in AFTER the backend published this server as - # _pending_server but BEFORE start() took the lock has already set _abort and - # closed the httpx client. Honor that delivered stop instead of clearing the - # abort and spawning a model process the cancelled load would then leak. + # 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() port = self._find_free_port() - # An empty scratch dir for sd-server's LoRA / upscaler / embeddings scans - # (it recursively iterates them per request and errors on a missing dir). + # Empty scratch dir for sd-server's LoRA/upscaler/embeddings scans (errors if missing). self._scratch_dir = tempfile.mkdtemp(prefix = "sdcpp_dirs_") cmd = build_sd_cpp_server_command( self.binary, @@ -183,19 +169,14 @@ class SdCppServer: if env: run_env.update(env) logger.info("starting sd-server: %s", " ".join(cmd)) - # Clear in place: reassigning to [] drops the deque(maxlen=200) bound, so the - # continuous stdout drain would then grow the tail without limit for the whole - # resident-server lifetime. + # Clear in place; reassigning [] would drop the maxlen bound and grow unbounded. self._tail.clear() self._spawn_error: Optional[Exception] = None spawned = threading.Event() - # Spawn INSIDE the drain thread, which then reads stdout for the process's whole - # lifetime. child_popen_kwargs() sets PR_SET_PDEATHSIG, which on Linux is bound to - # the CREATING THREAD -- so the child must be created by a thread that outlives it, - # or a transient spawner thread ending would kill the server. The drain thread is - # exactly that long-lived owner; it dies only when the process exits or the - # interpreter goes away (the case we DO want to reap the GPU-resident server). + # 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 (a transient spawner ending would kill the server). def _own_process() -> None: try: proc = subprocess.Popen( @@ -217,8 +198,7 @@ class SdCppServer: adopt_pid(proc.pid) # so a global shutdown sweep also reaps it spawned.set() self._drain_stdout(proc) - # stdout closed == the process exited; reap it so it is not left a zombie - # until the next stop()/reload. + # stdout closed == process exited; reap it so it is not left a zombie. try: proc.wait(timeout = 5) except Exception: # noqa: BLE001 @@ -253,8 +233,8 @@ class SdCppServer: deadline = time.monotonic() + timeout url = f"{self.base_url}{_READY_PATH}" while time.monotonic() < deadline: - # A concurrent stop() (unload / superseding load) sets _abort so this wait can - # bail 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 @@ -280,7 +260,7 @@ class SdCppServer: line = raw.rstrip() if not line: continue - self._tail.append(line) # deque(maxlen) discards the oldest automatically + self._tail.append(line) logger.debug("[sd-server] %s", line) cb = self._step_listener if cb is not None: @@ -294,9 +274,8 @@ 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 lifecycle lock: a concurrent start() holds - # that lock for the whole (up to startup_timeout) readiness wait, so setting the - # event lets that wait bail immediately instead of stop() blocking behind it. + # 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: @@ -363,9 +342,8 @@ class SdCppServer: Raises ``RuntimeError`` on submit/poll failures (including the server dying), with the log tail attached. """ - # If the server was already stopped for a cancel/unload/superseding load that set - # the cancel event before this submit began, report it as a cancellation (which the - # route maps to a client-state 409) rather than 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.") @@ -410,22 +388,17 @@ class SdCppServer: self.cancel(job_id) cancel_sent_at = time.monotonic() elif time.monotonic() - cancel_sent_at > _CANCEL_GRACE_S: - # The best-effort cancel was not reflected in job status within the - # grace window; abandon the poll so the caller can stop the server - # instead of holding the generate lock until the job finishes. + # Cancel not reflected within the grace window; abandon the poll so + # the caller can stop the server instead of holding the generate lock. raise SdCppCancelled("sd-server generation was cancelled.") if not self.is_alive(): - # If we're unwinding a cancel (e.g. unload killed the server), surface a - # clean cancellation rather than a generic "server died" error. + # Unwinding a cancel (e.g. unload killed the server) -> clean cancellation. if cancel_event is not None and cancel_event.is_set(): raise SdCppCancelled("sd-server generation was cancelled.") raise RuntimeError(self._died_message("img_gen poll", None)) if time.monotonic() > deadline: - # Best-effort cancel, then tear the server down: current sd-server does - # not interrupt an already-generating job (cancel_generating=false / 409), - # so leaving it up would keep denoising the abandoned job and block later - # generations/reloads behind it. Stopping frees the slot; the backend sees - # the dead server on the next generate and takes the recoverable reload path. + # 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. self.cancel(job_id) self.stop() raise RuntimeError(f"sd-server generation timed out after {total_timeout}s") @@ -435,10 +408,8 @@ class SdCppServer: time.sleep(poll_interval) continue except RuntimeError as exc: - # A concurrent stop()/unload closes the shared httpx client; httpx then - # raises a plain RuntimeError ("client has been closed") that is NOT a - # transport error. When we are being cancelled, report it as a clean - # cancellation (route -> 409) instead of a generic 500 generation failure. + # A concurrent stop() closes the shared client -> plain RuntimeError + # ("client has been closed"), not a transport error; map cancel -> 409. if cancel_event is not None and cancel_event.is_set(): raise SdCppCancelled("sd-server generation was cancelled.") from exc raise @@ -479,8 +450,7 @@ class SdCppServer: @staticmethod def _decode_images(job: dict[str, Any]) -> list[bytes]: - # Defensive against an unexpected response shape (a misbehaving/older server): - # verify each level is the type we index before calling dict/list methods. + # Type-check each level before indexing (guards a misbehaving/older server). result = job.get("result") if isinstance(job, dict) else None images = result.get("images") if isinstance(result, dict) else None items = [it for it in images if isinstance(it, dict)] if isinstance(images, list) else [] diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 05d9825df5..6446906d54 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -95,26 +95,22 @@ from utils.hardware import clear_gpu_cache logger = get_logger(__name__) -# Load kinds, mirroring the image backend: "gguf" (single-file GGUF DiT + -# companion base repo), "single_file" (safetensors DiT, e.g. the fp8 LTX-2.3 -# checkpoints), "pipeline" (a full diffusers repo via from_pretrained). +# Load kinds (mirror the image backend): gguf (single-file GGUF DiT + base repo), +# single_file (safetensors DiT, e.g. fp8 LTX-2.3), pipeline (full diffusers repo). _MODEL_KINDS = frozenset({"gguf", "single_file", "pipeline"}) -# Official vendor base repos allowed to load as full (non-GGUF) artifacts even -# though they are not under unsloth/. Exact-match, lowercased, safetensors-only, -# no remote code -- same bar as the image backend's allowlist. +# Vendor base repos allowed to load as full (non-GGUF) artifacts despite not being +# under unsloth/. Exact-match, lowercased, safetensors-only, no remote code. _TRUSTED_NON_GGUF_VIDEO_REPOS = frozenset( { "lightricks/ltx-2", "lightricks/ltx-2.3", "lightricks/ltx-2.3-fp8", - # Wan2.2 official diffusers base repos (Wan-AI org): safetensors-only, no - # remote code, so allowed as full (pipeline-kind) loads like the LTX-2 bases. + # Wan2.2 official diffusers base repos: safetensors-only, no remote code. "wan-ai/wan2.2-ti2v-5b-diffusers", "wan-ai/wan2.2-t2v-a14b-diffusers", # HunyuanVideo-1.5 community Diffusers repacks (tencent's own repo is the - # original non-diffusers layout: config.json, no model_index.json, so it - # cannot load through HunyuanVideo15Pipeline at all). + # non-diffusers layout with no model_index.json, unloadable here). "hunyuanvideo-community/hunyuanvideo-1.5-diffusers-480p_t2v", "hunyuanvideo-community/hunyuanvideo-1.5-diffusers-720p_t2v", } @@ -158,10 +154,8 @@ def _picked_gguf_arch(repo_id: str, gguf_filename: str) -> Optional[str]: path = Path(repo_id).expanduser() / gguf_filename if not path.is_file(): - # Not a local dir: resolve a cached HUB blob from the HF cache (no network). The - # cached-gguf picker only offers already-downloaded repos, so the blob is on disk -- - # but that listing scans the active, legacy, AND default cache roots, so probe all - # three here or a GGUF cached in a non-active root would be offered yet 400 on load. + # Not a local dir: resolve a cached HUB blob (no network). Probe active, legacy, + # AND default cache roots (as the picker's listing does) or a non-active-root GGUF 400s. from huggingface_hub import try_to_load_from_cache cached = try_to_load_from_cache(repo_id, gguf_filename) @@ -231,11 +225,8 @@ def _detect_load_family( else None ) if fam is None and gguf_filename and not family_override: - # The picker admits a GGUF (local dir OR cached hub repo) by its general.architecture, but - # its path/name may carry no whole-segment family token (e.g. a renamed "model.gguf"), so - # the name-based detection above misses it. Resolve the same family the picker offered by - # reading the arch -- its string ("ltxv") is a family alias. A video arch with no backend - # family (e.g. "wan") still yields None, so an unsupported pick 400s exactly as before. + # A renamed GGUF carries no family token in its name; resolve via general.architecture + # (its string, e.g. "ltxv", is a family alias). No-backend archs still yield None -> 400. arch = _picked_gguf_arch(repo_id, gguf_filename) if arch: fam = detect_video_family(repo_id, override = arch) @@ -274,21 +265,17 @@ class _VideoLoadState: backend_flags: Optional[dict] = None attention_backend: Optional[str] = None transformer_cache: Optional[str] = None - # True when the cache decision was AUTO on a cache-capable DiT: generate() then - # re-checks the actual step count and toggles FBCache across FBCACHE_MIN_STEPS. - # An explicit request (off / fbcache) is never toggled. + # AUTO on a cache-capable DiT: generate() re-checks the step count and toggles FBCache + # across FBCACHE_MIN_STEPS. An explicit request (off / fbcache) is never toggled. cache_auto: bool = False # Inputs the generation-time toggle re-applies (quantised threshold + override). cache_quant_active: bool = False cache_threshold: Optional[float] = None - # Dense transformer quant actually engaged ("int8" | "fp8" | "nvfp4" | "mxfp8") or - # None. Mirrors the image backend's _LoadState.transformer_quant: on a pipeline-kind - # load the dense DiT(s) can be torchao-quantised in place onto the low-precision - # tensor cores; None means they run at their loaded (bf16) precision. + # Dense transformer quant engaged ("int8"|"fp8"|"nvfp4"|"mxfp8") or None (loaded bf16). + # Pipeline-kind only; torchao-quantised in place onto the low-precision tensor cores. transformer_quant: Optional[str] = None - # Text-encoder quant actually engaged ("fp8" | "fp8_dynamic" | "int8" | "nvfp4") or None. - # The companion text encoder (UMT5 / Gemma3 / Qwen2.5-VL) loads dense bf16 and is often the - # largest resident component; this shrinks it in place, mirroring the image backend. + # Text-encoder quant engaged ("fp8"|"fp8_dynamic"|"int8"|"nvfp4") or None. The companion + # encoder (UMT5/Gemma3/Qwen2.5-VL) is often the largest resident; shrunk in place. text_encoder_quant: Optional[str] = None resolved: Optional[dict] = None @@ -306,17 +293,11 @@ def _progress(phase: Optional[str], **extra: Any) -> dict[str, Any]: # ── dual-DiT (Wan2.2-A14B MoE) helpers ──────────────────────────────────────── -# -# The imported optimisation helpers (apply_speed_optims / apply_attention_backend / -# apply_step_cache) and the dense quantiser all read ``pipe.transformer`` and act on -# that ONE denoiser -- correct for every single-DiT family (LTX-2, Wan2.2-TI2V-5B). -# Wan2.2-A14B is a dual-expert MoE: ``transformer`` handles the high-noise steps and -# ``transformer_2`` the low-noise steps (pipeline_wan.py routes by boundary_ratio), so -# an optimisation applied only to ``transformer`` would leave the second expert eager / -# unquantised / on the wrong attention kernel for half the schedule. Rather than fork -# each helper, present the second DiT to them AS ``pipe.transformer`` via a thin proxy -# and call the helper a second time, so the helpers stay untouched and single-DiT loads -# are bit-identical (the proxy is only built for is_moe families). +# The optimisation helpers and the quantiser all act on ``pipe.transformer`` -- fine for +# single-DiT families. Wan2.2-A14B is a dual-expert MoE (transformer = high-noise steps, +# transformer_2 = low-noise), so an optimisation on ``transformer`` alone leaves the second +# expert unoptimised for half the schedule. Rather than fork each helper, present the second +# DiT AS ``pipe.transformer`` via a thin proxy (built only for is_moe) and call the helper again. def _transformer_names(pipe: Any, fam: VideoFamily) -> tuple[str, ...]: @@ -339,7 +320,6 @@ class _SecondDiTView: ``transformer_2``. Only ever wrapped around an MoE pipe (guarded by fam.is_moe).""" def __init__(self, pipe: Any) -> None: - # Store on the instance dict under a name __getattr__ never fires for. object.__setattr__(self, "_pipe", pipe) @property @@ -347,14 +327,12 @@ class _SecondDiTView: return self._pipe.transformer_2 def __getattr__(self, name: str) -> Any: - # Only reached for attributes not found on the instance/class (i.e. not - # ``transformer`` / ``_pipe``), so everything else delegates to the real pipe. + # Only reached for attrs not on the instance/class, so delegate to the real pipe. return getattr(object.__getattribute__(self, "_pipe"), name) def __setattr__(self, name: str, value: Any) -> None: - # Writes must land on the real pipe, or a helper's side effect (for example - # reassigning the transformer it optimised) would vanish with the view. - # ``transformer`` mirrors the read property onto the second expert. + # Writes land on the real pipe (else a helper's reassignment vanishes with the + # view); ``transformer`` mirrors onto the second expert. pipe = object.__getattribute__(self, "_pipe") setattr(pipe, "transformer_2" if name == "transformer" else name, value) @@ -382,9 +360,8 @@ class VideoBackend: self._active_generate_cancel: Optional[threading.Event] = None # Generation progress, written by the step callback / phase transitions. self._gen: dict[str, Any] = {"active": False} - # True from begin_generate() until its worker records a terminal state, so - # a second begin_generate() is refused while the first still runs (or is - # about to run: generate() only sets _gen after taking its locks). + # True from begin_generate() until its worker records a terminal state, so a second + # begin_generate() is refused while the first still runs. self._generate_job_active = False # ── validation ─────────────────────────────────────────────────────────── @@ -402,10 +379,8 @@ class VideoBackend: ) -> VideoFamily: """Cheap, network-free validation shared by the route and the load path.""" kind = resolve_video_model_kind(gguf_filename, model_kind) - # A -GGUF repo picked without a quant filename resolves to the pipeline - # kind and would only fail minutes later in from_pretrained (no - # model_index.json), AFTER the route evicted the current GPU owner. - # Reject it here, where failing is still free. + # A -GGUF repo picked without a quant filename resolves to pipeline kind and would + # only fail in from_pretrained (no model_index.json) after the route evicts the owner. if kind == "pipeline" and repo_id.strip().lower().rstrip("/").endswith("-gguf"): raise ValueError( f"'{repo_id}' is a GGUF repo: pick one of its .gguf files " @@ -423,39 +398,31 @@ class VideoBackend: f"Non-GGUF video loads are limited to unsloth/* repos, the official " f"family base repos, and local paths; '{repo_id}' is neither." ) - # The companions load with from_pretrained too, so an explicit base repo is - # held to the same bar as a non-GGUF repo id: a GGUF pick must not smuggle - # in an arbitrary remote base. + # Companions load with from_pretrained, so a base repo is held to the non-GGUF bar: + # a GGUF pick must not smuggle in an arbitrary remote base. if base_repo and (base_repo or "").strip() and not _is_trusted_video_repo(base_repo): raise ValueError( f"base_repo is limited to unsloth/* repos, the official family base " f"repos, and local paths; '{base_repo}' is neither." ) - # An existing LOCAL base_repo loads as a full pipeline (from_pretrained(base) / config=base), - # which needs a model_index.json. The pipeline-kind shape check below covers only repo_id, - # and an explicit base_repo is only meaningful for gguf/single_file kinds, so a non-pipeline - # local base would otherwise pass here and fail deep in the background load AFTER the route - # evicted the resident model. Shared helper, so image/video/training stay in sync. + # A local base_repo loads as a full pipeline (needs model_index.json); reject a + # non-pipeline local base here, before the load. Shared helper keeps image/video/training in sync. from core.inference.diffusion import _assert_local_base_is_pipeline _assert_local_base_is_pipeline(base_repo) if kind in ("gguf", "single_file") and not gguf_filename: raise ValueError("A gguf/single_file load needs the checkpoint filename.") if kind in ("gguf", "single_file") and fam.is_moe: - # A single checkpoint carries only one expert; the pipeline would then pull - # the other expert dense bf16 from the base repo, outside the memory plan. + # A single checkpoint carries one expert; the other would load dense bf16, off-plan. raise ValueError( f"'{fam.name}' is a dual-expert model: a single {kind} file covers only " f"one of its two transformers. Load the diffusers pipeline repo " f"('{fam.base_repo}') instead." ) - # A local checkpoint that cannot exist must fail HERE, before the route evicts - # a resident chat/image model for a load that dies at resolve time. + # A missing local checkpoint must fail HERE, before the route evicts a resident model. if kind in ("gguf", "single_file"): - # Fail a kind/extension mismatch before the GPU handoff instead of deep in the - # background loader: a "gguf" load needs a .gguf file, a "single_file" load must not be - # handed a .gguf and must name an actual .safetensors checkpoint. Mirrors the image - # loader's kind/extension gate in diffusion.validate_load_request. + # Fail a kind/extension mismatch before the GPU handoff: gguf needs .gguf, + # single_file needs .safetensors (mirrors the image loader's gate). is_gguf_name = (gguf_filename or "").lower().endswith(".gguf") if kind == "gguf" and not is_gguf_name: raise ValueError("a 'gguf' load requires a .gguf checkpoint name.") @@ -467,10 +434,8 @@ class VideoBackend: f"(expected a .safetensors name; use a .gguf name for a GGUF load)." ) root = Path(repo_id).expanduser() - # POSIX path-shaped, a "."/".." prefix (covers ./ ../ and Windows .\ ..\), a Windows - # separator anywhere (never in a bare "org/name" id), or an absolute path on this OS - # (covers Windows C:\ / C:/). Mirrors the image loader so a missing Windows-shaped - # local pick fails before the GPU handoff instead of being treated as a Hub repo. + # Path-shaped: "."/".." prefix, a backslash (never in "org/name"), or an absolute + # path -- so a missing Windows-shaped local pick fails before the handoff, not as a Hub repo. path_shaped = ( repo_id.startswith(("/", "\\", "~", ".")) or "\\" in repo_id or root.is_absolute() ) @@ -481,12 +446,8 @@ class VideoBackend: except Exception as exc: # noqa: BLE001 -- surface as client input error raise ValueError(str(exc)) from exc elif root.is_file(): - # The loader hands a local FILE straight to the gguf/single_file loader - # (_resolve_checkpoint_path returns the file itself, ignoring gguf_filename), - # so the file's OWN suffix must match the kind. Otherwise a .gguf picked as - # single_file (or a .safetensors picked as gguf) slips past the gguf_filename - # checks above, evicts the resident model in the route, and only then fails - # in from_single_file / the GGUF reader. Reject it here, before the handoff. + # The loader hands a local FILE straight through (ignoring gguf_filename), so + # the file's OWN suffix must match the kind; reject a mismatch before the handoff. suffix = root.suffix.lower() if kind == "gguf" and suffix != ".gguf": raise ValueError( @@ -500,27 +461,20 @@ class VideoBackend: ) elif path_shaped: raise ValueError(f"Local model path '{repo_id}' does not exist.") - # A local pipeline pick must be a real diffusers directory (model_index.json), or it - # would only fail deep in from_pretrained AFTER the route evicted the resident model. - # Mirrors the image loader's local-pipeline shape check in diffusion.validate_load_request. + # A local pipeline pick must be a diffusers directory (model_index.json), else it would + # only fail in from_pretrained after eviction (mirrors the image loader). if kind == "pipeline": root = Path(repo_id).expanduser() - # Gate on .exists() (not .is_dir()) so a local FILE picked as a pipeline is rejected - # too: a bare .safetensors file is not a diffusers directory, so from_pretrained would - # still fail in the background load after the eviction. Mirrors the image loader, which - # uses .exists() here. + # Gate on .exists() (not .is_dir()) so a local FILE picked as a pipeline is rejected too. if root.exists() and not (root.is_dir() and (root / "model_index.json").is_file()): raise ValueError( f"Local pipeline path is not a diffusers directory " f"(no model_index.json): {repo_id}" ) - # Reject a malformed transformer_quant scheme cheaply, before the GPU handoff - # (normalize_transformer_quant raises ValueError on an unknown scheme). It applies - # only on pipeline-kind loads (the dense DiT from the base repo); an ignored value - # on a gguf/single_file load is left to the loader, matching the image backend. + # Reject a malformed transformer_quant cheaply, before the handoff (applies on + # pipeline-kind loads; ignored on gguf/single_file, matching the image backend). normalize_transformer_quant(transformer_quant) - # Reject a malformed text_encoder_quant the same way (applies to any load kind: the dense - # text encoder is resident for pipeline / gguf / single_file alike). + # Reject a malformed text_encoder_quant the same way (any kind: the encoder is always dense). normalize_te_quant(text_encoder_quant) _ensure_mp4_encoder_available() return fam @@ -605,10 +559,8 @@ class VideoBackend: if self._load_token == token and self._loading is not None: self._loading.base_repo = base self._loading.expected_bytes = expected - # The GGUF/single-file checkpoint downloads outside the lock so an - # unload/eviction can preempt the multi-GB pull; the pipeline - # companions pre-download the same way (scoped file list, cancellable, - # resumes from the cache so a cancelled pull costs nothing). + # Checkpoint downloads outside the lock so an unload/eviction can preempt the + # multi-GB pull; companions pre-download the same way (scoped, cancellable, resumable). checkpoint_local: Optional[Path] = None if kwargs.get("gguf_filename") and not Path(kwargs["repo_id"]).expanduser().exists(): from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback @@ -620,19 +572,17 @@ class VideoBackend: cancel_event = self._cancel_event, ) ) - # An LTX-2.3 checkpoint replaces the base VAEs/vocoder/connectors too, so - # its base pull shrinks to scheduler + text encoder + tokenizer; the - # estimate is recomputed to match (detectable only once the checkpoint - # header is on disk, hence after the pull above). + # An LTX-2.3 checkpoint supplies the VAEs/vocoder/connectors, so the base pull + # shrinks to scheduler + text encoder + tokenizer; recompute the estimate to match + # (detectable only once the checkpoint header is on disk). ltx23 = False if fam is not None and fam.name == "ltx-2" and kind != "pipeline": from .video_ltx2 import is_ltx23_checkpoint probe = checkpoint_local if probe is None: - # Local repos: a bare file, or a directory whose child the same - # resolver load_pipeline uses picks out. Unresolvable here means - # load_pipeline will surface the real error; keep the wide pull. + # Local repos: a bare file, or a dir child via the same resolver load_pipeline + # uses. Unresolvable -> load_pipeline surfaces the real error; keep the wide pull. root = Path(kwargs["repo_id"]).expanduser() if root.is_file(): probe = root @@ -659,28 +609,23 @@ class VideoBackend: if self._load_token == token and self._loading is not None: self._loading.expected_bytes = expected base_local = self._predownload_base(base, kwargs.get("hf_token"), kind, ltx23 = ltx23) - # The 2.3 assembly pulls per component from the hub id (its snapshot here - # deliberately lacks the base VAEs), so it only gets the warmed cache; the - # generic from_pretrained paths get the complete local snapshot. + # The 2.3 assembly pulls per component from the hub id (its snapshot lacks the base + # VAEs), so it only gets the warmed cache; generic paths get the full local snapshot. kwargs["_base_local_dir"] = None if ltx23 else base_local self.load_pipeline(**kwargs) with self._lock: if self._load_token == token: self._loading = None except Exception as exc: # noqa: BLE001 -- surfaced via load_progress - # A failed or cancelled load never commits _VideoLoadState, so the - # teardown path has no snapshot to restore: roll back the process-wide - # speed globals here (token-scoped, so a superseded load cannot clobber - # the globals a newer in-flight load now owns). + # A failed/cancelled load never commits _VideoLoadState, so roll back the + # process-wide speed globals here (token-scoped, so a superseded load can't clobber a newer one's). self._rollback_precommit_globals(token) if self._load_token != token: return logger.error("video.load_failed: %s", exc) - # Free the debris of a failed construction (mirrors diffusion.py's _run_load): - # no _VideoLoadState was committed, so no later unload releases the VRAM a - # partially built pipeline (OOM in from_pretrained / quant / placement) left - # reserved in the caching allocator -- which would OOM the next load. Guarded so - # a sticky CUDA error cannot skip stamping the real error below. + # Free the debris of a failed construction (mirrors diffusion.py): no state was + # committed, so nothing else releases the VRAM a partial pipeline reserved. Guarded + # so a sticky CUDA error can't skip stamping the real error below. try: clear_gpu_cache() except Exception: # noqa: BLE001 -- cleanup is best-effort @@ -710,9 +655,8 @@ class VideoBackend: diffusion_gguf_compile.uninstall_all() - # Base-repo subfolders an LTX-2.3 assembly reads: the checkpoint (plus the GGUF - # repo's extras files) supplies the DiT, connectors, both VAEs and the vocoder, - # so only the 2.0 base's scheduler / text encoder / tokenizer are pulled. + # LTX-2.3 gets DiT/connectors/VAEs/vocoder from the checkpoint + extras, so only the + # 2.0 base's scheduler / text encoder / tokenizer are pulled. _LTX23_BASE_PREFIXES = ("scheduler/", "text_encoder/", "tokenizer/") @staticmethod @@ -738,10 +682,8 @@ class VideoBackend: files: list[tuple[str, int]] = [] for sibling in info.siblings or []: name, size = sibling.rfilename, sibling.size or 0 - # .jinja: tokenizer/chat_template.jinja ships as a standalone file in the - # LTX-2 and HunyuanVideo-1.5 repos (not embedded in tokenizer_config.json) - # and apply_chat_template needs it at generation time, so a snapshot - # without it loads fine and then crashes the first generation. + # .jinja: tokenizer/chat_template.jinja is a standalone file apply_chat_template + # needs at generation time; a snapshot without it crashes the first generation. if not name.endswith((".safetensors", ".json", ".model", ".txt", ".jinja")): continue if "/" not in name and name.endswith(".safetensors"): @@ -812,9 +754,8 @@ class VideoBackend: snapshot_root: Optional[Path] = None for name, _ in files: - # Explicit per-file check: a fully-cached file returns without ever - # consulting the event, so a warm-cache sweep would otherwise run to - # completion after an unload already cancelled this load. + # Explicit check: a cached file returns without consulting the event, so a + # warm-cache sweep would otherwise run to completion after an unload cancelled. if self._cancel_event.is_set(): raise RuntimeError(VIDEO_CANCELLED_MSG) local = Path( @@ -859,10 +800,8 @@ class VideoBackend: phase = "downloading" if expected and downloaded >= expected: phase = "finalizing" - # The cache scan counts every blob of the repo(s), including files a - # previous (or broader) pull left behind that this load never reads, so - # the raw counter can exceed the scoped estimate. Clamp: everything the - # load needs is present, which is what the bar reports. + # The cache scan counts every blob (incl. files this load never reads), so the raw + # counter can exceed the scoped estimate; clamp to what the bar reports. downloaded = expected return _progress( phase, @@ -922,59 +861,40 @@ class VideoBackend: with self._lock: if _load_token is not None and _load_token != self._load_token: raise RuntimeError("Video load was cancelled or superseded.") - # Signal only a generation from the PREVIOUS model; the token check - # above already bailed a superseded worker before this point. + # Signal a generation from the PREVIOUS model (the token check above bailed a superseded worker). if self._active_generate_cancel is not None: self._active_generate_cancel.set() - # Wait for the signalled generation to actually exit before tearing the old - # pipeline down: the denoise loop holds its own pipe reference until the - # next step callback, and freeing/reallocating under it would put two - # models in VRAM at once. generate() holds _generate_lock for its full - # body, so a bare acquire is the exit barrier (never while holding _lock). + # Barrier: wait for the signalled generation to exit before teardown, or two models + # coexist in VRAM (the denoise loop holds its pipe ref until the next callback). with self._generate_lock: pass - # The barrier wait can outlive this load: an unload or a newer load may - # have superseded it while blocked, and tearing down now would destroy - # the model that should remain current (or waste minutes building a - # pipeline nobody wants). Recheck before touching shared state. + # The barrier wait can outlive this load (a newer load / unload superseded it); recheck + # before touching shared state so we don't destroy the current model or build a dead pipe. if _load_token is not None and _load_token != self._load_token: raise RuntimeError("Video load was cancelled or superseded.") self._teardown_state() target = resolve_diffusion_device_target() device = target.device - # Video DiTs are bf16-native; fp16 overflows them, so a resolved fp16 - # promotes to float32 (the same rule as the fp16-incompatible image - # families). CPU stays float32. + # Video DiTs are bf16-native; fp16 overflows, so a resolved fp16 promotes to float32 + # (same rule as fp16-incompatible image families). CPU stays float32. dtype = target.dtype if fam.fp16_incompatible and dtype is torch.float16: dtype = torch.float32 - # The size tables below are bf16 (2-byte) figures. When the promotion - # above lands fp32 weights on an accelerator (a pre-bf16 GPU), every - # dense estimate doubles; budgeting the 2-byte figure would let auto - # pick a resident plan that OOMs inside from_pretrained. GGUF weights - # stay quantised on disk and in memory, so only dense estimates scale. + # Size tables below are bf16 (2-byte); when the promotion lands fp32 on an accelerator, + # dense estimates double, so scale them (GGUF stays quantised, so only dense scales). dtype_scale = 2.0 if device != "cpu" and dtype is torch.float32 else 1.0 - # Precision tri-state, mirroring the image backend: an UNSET request (or - # "auto") hands the decision to the hardware ladder -- on a dense-capable - # GPU the quantised DiT (int8 minimum, fp8 on data-center silicon) is - # faster at the same resident-or-better footprint. An explicit - # "none"/"off" pins dense bf16 and an explicit scheme pins that scheme. - # Only the pipeline kind can engage it (gguf/single_file checkpoints - # already carry their own precision), and the offload guard below still - # skips it when the plan moves the DiT. + # Precision tri-state (mirror image backend): unset/"auto" -> hardware ladder picks a + # quantised DiT (int8 min, fp8 on datacenter silicon); "none"/"off" pins dense bf16; an + # explicit scheme pins it. Pipeline-kind only; the offload guard below still skips it. if transformer_quant is None or str(transformer_quant).strip().lower() in ( "", "auto", ): - # An explicit Speed="off" (bit-exact) load must stay dense bf16: promoting the unset - # precision to auto-quant here would engage int8/fp8 + regional compile and silently - # break the user's bit-exact request (an auto DEFAULT overriding an EXPLICIT control), - # and the quant path below would then also force effective_speed back to default. - # Suppress the auto default when speed was explicitly pinned off, mirroring the image - # backend (diffusion.py); otherwise auto (the dense-capable default) applies. "off" - # normalizes to None (no dense quant), keeping the dense bf16 path. + # An explicit Speed="off" (bit-exact) load must stay dense bf16: auto-quant would + # engage int8/fp8 + regional compile and break the bit-exact request. Suppress the + # auto default when speed was pinned off (mirrors diffusion.py); else auto applies. speed_off = speed_mode is not None and str(speed_mode).strip().lower() == SPEED_OFF transformer_quant = "off" if speed_off else TQ_AUTO @@ -1002,10 +922,8 @@ class VideoBackend: if components is not None else None ) - # The resident check budgets ALL weights (the image backend's contract): - # the companions stay resident even when only the transformer would fit, - # so budgeting the transformer alone lets auto pick OFFLOAD_NONE and OOM - # while from_pretrained loads the text encoder / VAEs. + # Budget ALL weights (image-backend contract): companions stay resident, so + # budgeting the transformer alone lets auto pick OFFLOAD_NONE and OOM. model_dense_mib = ( transformer_mib + (companion_mib or 0) if transformer_mib is not None else None ) @@ -1022,10 +940,9 @@ class VideoBackend: companion_dense_mib = companion_mib, requested_mode = normalize_memory_mode(memory_mode), ) - # Parity with the image dense-quant path: the bf16-table plan can force offload - # a quantised DiT would not need (offload also disables quant entirely). Re-plan - # with the scheme's steady factor and keep the resident placement when it fits; - # if quantisation later fails, the load falls back to this bf16 plan. + # Parity with the image dense-quant path: the bf16-table plan can force offload a + # quantised DiT would not need. Re-plan with the scheme's steady factor and keep the + # resident placement if it fits; fall back to this bf16 plan if quant later fails. bf16_plan = plan quant_replanned = False if ( @@ -1066,24 +983,20 @@ class VideoBackend: pipeline_cls = getattr(diffusers, fam.pipeline_class) pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype} if getattr(fam, "vae_force_fp32", False): - # Wan's VAE must decode in float32, but a scalar torch_dtype casts EVERY component - # (VAE included) to the pipe dtype during load -- AutoencoderKLWan has no - # _keep_in_fp32_modules, so from_pretrained truncates its fp32 weights to bf16 and a - # later .to(float32) only widens the already-lossy values (banding / black frames). - # diffusers >= 0.39 takes a per-component dtype dict, so load the VAE at fp32 directly; - # "default" MUST be set or unlisted components fall back to fp32 (over-widening the DiT). + # Wan's VAE must decode in float32. A scalar torch_dtype truncates its fp32 weights + # to bf16 (no _keep_in_fp32_modules); a later .to(float32) only widens lossy values + # (banding / black frames). Use the per-component dtype dict; "default" MUST be set + # or unlisted components fall back to fp32 (over-widening the DiT). pipe_kwargs["torch_dtype"] = {"vae": torch.float32, "default": dtype} if hf_token: pipe_kwargs["token"] = hf_token if kind == "pipeline": - # The pre-downloaded snapshot dir keeps from_pretrained off the hub (its - # own snapshot sweep would also pull the repo's packaged root checkpoints - # and duplicate text-encoder shards); hub id when pre-download was skipped. + # The pre-downloaded snapshot dir keeps from_pretrained off the hub (its sweep would + # also pull root checkpoints + duplicate shards); hub id when pre-download was skipped. pipe = pipeline_cls.from_pretrained(_base_local_dir or repo_id, **pipe_kwargs) else: transformer_cls = getattr(diffusers, fam.transformer_class) - # checkpoint_path was already resolved (and downloaded) by the memory - # planning branch above for every non-pipeline kind. + # checkpoint_path was resolved (and downloaded) by the memory-planning branch above. sf_kwargs: dict[str, Any] = { "torch_dtype": dtype, "config": base, @@ -1097,9 +1010,8 @@ class VideoBackend: from .video_ltx2 import is_ltx23_checkpoint, load_ltx23_pipeline if fam.name == "ltx-2" and is_ltx23_checkpoint(checkpoint_path): - # 2.3 checkpoints need the full assembly: new transformer config - # flags, key renames the stock converter lacks, and the 2.3 - # connectors/VAEs/vocoder the 2.0 base repo does not carry. + # 2.3 checkpoints need the full assembly: new config flags, key renames the + # stock converter lacks, and the 2.3 connectors/VAEs/vocoder the base lacks. pipe = load_ltx23_pipeline( checkpoint_path, base_repo = base, @@ -1113,10 +1025,8 @@ class VideoBackend: _base_local_dir or base, transformer = transformer, **pipe_kwargs ) - # The per-component torch_dtype above already loads the Wan VAE at float32 (bf16_components_gb - # budgets it at that fp32 size, so the memory plan stays consistent). Belt-and-suspenders for - # any path that bypassed the dict (e.g. a passed-in vae=): re-pin an fp32-force VAE that came - # back at a lower precision. This is a no-op on the primary path (the load already fp32'd it). + # The dtype dict already loads the Wan VAE at float32. Belt-and-suspenders for any path + # that bypassed it (e.g. a passed-in vae=): re-pin an fp32-force VAE that came back lower. if getattr(fam, "vae_force_fp32", False): vae = getattr(pipe, "vae", None) if vae is not None and getattr(vae, "dtype", None) is not torch.float32: @@ -1127,19 +1037,13 @@ class VideoBackend: clear_gpu_cache() raise RuntimeError("Video load was cancelled or superseded.") - # For a dual-DiT MoE family (Wan2.2-A14B), every optimisation site below must - # cover BOTH experts: ``views`` is (pipe, _SecondDiTView(pipe)) so a helper that - # reads ``pipe.transformer`` runs once per denoiser. A single-DiT load resolves to - # (pipe,), so it behaves exactly as before. + # For a dual-DiT MoE (Wan2.2-A14B), every optimisation site below covers BOTH experts: + # ``views`` is (pipe, _SecondDiTView(pipe)); a single-DiT load resolves to (pipe,). views = _views_for(pipe, fam) - # ── dense transformer quant (opt-in, pipeline-kind only): load the dense bf16 - # DiT from the base repo and torchao-quantise it in place onto the low-precision - # tensor cores, mirroring the image backend's transformer_quant fast path. Only - # the pipeline kind materialises the dense weights (gguf/single_file already carry - # their own precision), and only on CUDA + bf16. Best-effort: any failure leaves - # the DiT dense. Quant must precede compile (dynamic quant is ~30x slower eager), - # so it runs before apply_speed_optims below -- same order as diffusion.py. + # ── dense transformer quant (opt-in, pipeline-kind only): torchao-quantise the dense + # bf16 DiT in place onto the low-precision tensor cores (image-backend fast path). CUDA + + # bf16 only; best-effort. Quant must precede compile (eager dynamic quant is ~30x slower). transformer_quant_engaged: Optional[str] = None quant_skipped_for_offload = False if ( @@ -1148,13 +1052,10 @@ class VideoBackend: and dense_transformer_supported(target) and plan.offload_policy != "none" ): - # Offload hooks move modules with Module.to(), which torchao quantized - # tensors reject (aten._has_compatible_shallow_copy_type is - # unimplemented) -- observed as a hard crash on the Wan2.2-A14B gate - # run, where the 114 GB dual DiT plans model offload. A dense DiT - # under offload beats a crashed one, so quant is skipped, surfaced in - # the resolved record, and the user can force it by pinning a - # resident memory mode. + # Offload hooks move modules with Module.to(), which torchao quantized tensors reject + # (aten._has_compatible_shallow_copy_type unimplemented) -- a hard crash on the + # Wan2.2-A14B gate run (114 GB dual DiT plans model offload). Skip quant (dense-under- + # offload beats a crash); surfaced in the resolved record, forceable via a resident mode. logger.info( "video.transformer_quant: skipped (offload policy '%s' moves the " "DiT via Module.to(), unsupported for torchao quantized tensors); " @@ -1169,10 +1070,8 @@ class VideoBackend: ): engaged = [] for view in views: - # quantize_transformer reads ``pipe.transformer`` and returns the scheme it - # engaged (or None); pass each expert's view so both DiTs are quantised with - # the same arch-chosen scheme. The family name drives the per-family deny - # table (_FAMILY_SCHEME_DENY) exactly as on the image side. + # Pass each expert's view so both DiTs quantise with the same arch-chosen scheme. + # The family name drives the per-family deny table (_FAMILY_SCHEME_DENY). scheme = quantize_transformer( view, target, @@ -1182,10 +1081,8 @@ class VideoBackend: ) if scheme is not None: engaged.append(scheme) - # Quant must engage on every DiT or none: the first expert is mutated in - # place, so a second-expert failure cannot fall back to dense (the schedule - # would run at mismatched precision with quant reported off). Fail the load - # cleanly instead; a full miss (nothing engaged) stays best-effort dense. + # All experts or none: the first is mutated in place, so a second-expert failure + # can't fall back to dense (mismatched precision). Fail cleanly; a full miss stays dense. if engaged and len(engaged) < len(views): del pipe clear_gpu_cache() @@ -1195,18 +1092,13 @@ class VideoBackend: ) if engaged: transformer_quant_engaged = engaged[0] - # The quant-sized plan is only valid when quant actually engaged; a dense - # fallback must keep the conservative bf16 placement. + # The quant-sized plan is valid only when quant engaged; a dense fallback keeps bf16 placement. if quant_replanned and transformer_quant_engaged is None: plan = bf16_plan - # ── dense text-encoder quant (opt-in): the DiT arrives quantised in a GGUF, but the - # companion encoder (Gemma3 / UMT5 / Qwen2.5-VL) loads dense bf16 from the base repo and - # is often the largest resident component. Quantise it in place, mirroring the image - # backend (diffusion.py): applied for every kind (the encoder is dense regardless of how - # the DiT was sourced) and before placement so the offload hooks move the smaller weights. - # Best-effort: quantize_text_encoders leaves any encoder it can't cast dense. int8 needs a - # per-family keep-bf16 schedule, so the family name is passed. + # ── dense text-encoder quant (opt-in): the companion encoder (Gemma3/UMT5/Qwen2.5-VL) + # loads dense bf16 and is often the largest resident. Quantise in place for every kind, + # before placement (so offload moves the smaller weights). Best-effort; family drives int8's keep-bf16 schedule. text_encoder_quant_engaged = quantize_text_encoders( pipe, target, @@ -1216,21 +1108,15 @@ class VideoBackend: logger = logger, ) - # ── optimisation layers, in the image backend's order: step cache FIRST - # (compile keys its fullgraph decision off an active cache: FBCache hooks - # graph-break, so compiling fullgraph before installing the cache crashes - # the first cached generation), then attention, the speed profile, and - # placement/offload last. - # A clip denoise runs minutes, so even a dense (non-GGUF) load amortises the - # one-time regional compile within a single generation: unset resolves to the - # near-lossless `default` profile for every kind. Explicit values (incl. - # "off") are honored verbatim, and `max` is never an auto choice. + # ── optimisation layers in the image backend's order: step cache FIRST (compile keys + # its fullgraph decision off an active cache; FBCache hooks graph-break), then attention, + # speed profile, placement last. A clip denoise runs minutes, so even a dense load + # amortises the compile: unset resolves to the near-lossless `default`; "off"/explicit honored. effective_speed = resolve_speed_mode( speed_mode, is_gguf = kind == "gguf", dense_default = SPEED_DEFAULT ) - # A torchao-quantised DiT must be compiled (eager dynamic quant is ~30x slower and - # would lose to the bf16 it replaced), so force at least the regional-compile - # profile when quant engaged and the effective speed was off, matching diffusion.py. + # A torchao-quantised DiT must be compiled (eager is ~30x slower), so force at least + # the regional-compile profile when quant engaged but speed was off (matches diffusion.py). if transformer_quant_engaged is not None and effective_speed == SPEED_OFF: logger.info( "video.transformer_quant: forcing speed_mode=default " @@ -1238,20 +1124,15 @@ class VideoBackend: ) effective_speed = SPEED_DEFAULT backend_flags = snapshot_backend_flags() - # Until the state commit below transfers ownership to _teardown_state, a - # failure or cancellation must restore these process-wide globals itself - # (_run_load's error handler calls _rollback_precommit_globals with this - # token). Registered BEFORE the first mutating call. + # Until the state commit transfers ownership to _teardown_state, a failure must restore + # these globals itself (via _rollback_precommit_globals). Registered BEFORE the first mutation. self._precommit_globals = (_load_token, backend_flags) - # Step cache tri-state, mirroring the image backend: unset / "auto" lets the - # step-count policy decide (engage when this model's DEFAULT schedule reaches - # FBCACHE_MIN_STEPS, re-checked against the actual step count per generation); - # explicit "off" / "fbcache" are pinned and never toggled. Run it per expert - # so both denoisers cache; the engaged mode is identical across experts. + # Step cache tri-state: unset/"auto" -> step-count policy decides (engage when the DEFAULT + # schedule reaches FBCACHE_MIN_STEPS, re-checked per generation); "off"/"fbcache" pinned. + # Run per expert so both denoisers cache. cache_request = normalize_transformer_cache(transformer_cache) cache_auto = transformer_cache is None or cache_request == TC_AUTO - # GGUF checkpoints and torchao-quantised DiTs both need the higher quantised - # threshold for the cache to still trigger over the quant noise. + # GGUF and torchao-quantised DiTs need the higher threshold to trigger over quant noise. cache_quant_active = kind == "gguf" or transformer_quant_engaged is not None default_cache_steps: Optional[int] = None if cache_auto: @@ -1263,17 +1144,14 @@ class VideoBackend: view, mode = cache_request, threshold = transformer_cache_threshold, - # A quantized transformer's block residuals are larger, so it needs the - # higher FBCache trigger threshold to cache at all. Mirror the image path - # (diffusion.py): both an engaged transformer_quant AND a GGUF checkpoint - # (quantized weights) count as quant-active here (cache_quant_active, L1172). + # A quantized transformer's residuals are larger, needing the higher FBCache + # threshold; both engaged quant and GGUF count as quant-active (cache_quant_active). quant_active = cache_quant_active, logger = logger, ) if view is pipe: cache_engaged = engaged - # The auto decision can flip at generation time, but only on a DiT that - # supports caching at all (a non-CacheMixin transformer can never engage). + # The auto decision can flip at generation time, but only on a cache-capable DiT. cache_may_toggle = cache_auto and callable( getattr(getattr(pipe, "transformer", None), "enable_cache", None) ) @@ -1295,11 +1173,9 @@ class VideoBackend: attention_engaged = None speed_optims: tuple = () for view in views: - # apply_attention_backend / apply_speed_optims both act on ``view.transformer``; - # calling them once per view sets the kernel and compiles each expert. The - # engaged values match across experts (same device/family/mode), so record the - # first pass; a dense torchao transformer on the pipeline path is not a GGUF one, - # so is_gguf keys off the load kind (gguf) AND no quant having engaged. + # Both helpers act on ``view.transformer``; call once per view to set the kernel and + # compile each expert (engaged values match, so record the first). is_gguf keys off + # kind==gguf AND no quant having engaged (a dense torchao DiT is not a GGUF one). gguf_transformer = kind == "gguf" and transformer_quant_engaged is None engaged = apply_attention_backend( view, @@ -1314,9 +1190,8 @@ class VideoBackend: is_gguf = gguf_transformer, family = fam, speed_mode = effective_speed, - # An auto cache that could still engage mid-session also drops - # fullgraph: enabling FBCache under a fullgraph-compiled DiT would - # crash the first cached generation. + # An auto cache that could still engage also drops fullgraph (FBCache under a + # fullgraph-compiled DiT crashes the first cached generation). cache_active = cache_engaged is not None or cache_may_toggle, offload_active = plan.offload_policy != "none", ) @@ -1324,23 +1199,17 @@ class VideoBackend: attention_engaged = engaged speed_optims = tuple(k for k, v in applied.items() if v) with self._generate_lock: - # A cancelled/superseded load must not place weights on the GPU the arbiter - # may already have handed to another backend; recheck right before placement - # (the commit below still does the final locked check). + # A cancelled/superseded load must not place weights on a GPU the arbiter may have + # reassigned; recheck right before placement (the commit below does the final check). if _load_token is not None and _load_token != self._load_token: del pipe clear_gpu_cache() raise RuntimeError("Video load was cancelled or superseded.") offload_policy, vae_tiling = apply_memory_plan(pipe, plan, device = device, logger = logger) - # A dual-DiT MoE pipe (Wan2.2-A14B) needs no extra per-expert offload pass here: - # apply_memory_plan's group tier (_apply_group_offload) already block-streams every - # DiT it finds on the pipe -- transformer AND transformer_2 -- and model/sequential - # offload hook every top-level module, so the second expert is covered under all tiers. - # A second _apply_group_offload on transformer_2 would re-register the group-offload - # hooks it already carries, which diffusers rejects with a duplicate-hook ValueError. + # A dual-DiT MoE needs no extra per-expert pass: apply_memory_plan already covers every + # DiT (transformer AND transformer_2) under all tiers; a second pass would duplicate-hook. if not vae_tiling: - # Decode of a whole clip is the video memory peak; tiling is near-free - # in quality and keeps the decode bounded, so it is always on. + # Whole-clip decode is the video memory peak; tiling is near-free, so always on. try: pipe.vae.enable_tiling() vae_tiling = True @@ -1413,9 +1282,7 @@ class VideoBackend: vae_tiling = vae_tiling, memory_mode = plan.requested_mode, speed_mode = effective_speed, - # Already filtered above to only the optimisations that engaged; - # apply_speed_optims returns every flag True/False and the view - # loop keeps just the True names. + # Already filtered to the engaged optimisations (True names only). speed_optims = speed_optims, backend_flags = backend_flags, attention_backend = attention_engaged, @@ -1511,9 +1378,7 @@ class VideoBackend: if self._generate_job_active: raise RuntimeError(VIDEO_GENERATION_BUSY_MSG) self._generate_job_active = True - # Register the cancel event BEFORE the worker starts so a cancel (or an - # unload) that lands in the spawn window still stops the run instead of - # returning "nothing to cancel". + # Register BEFORE the worker starts so a cancel/unload in the spawn window still stops the run. self._active_generate_cancel = cancel self._gen = { "active": True, @@ -1611,10 +1476,8 @@ class VideoBackend: with self._lock: self._generate_job_active = False if cancel_event is not None and self._active_generate_cancel is cancel_event: - # generate() clears its own registration; this covers a job whose - # worker failed before (or without) reaching generate()'s finally. - # Identity-guarded so a direct generate() that registered its own - # event in the meantime keeps its cancel handle. + # Covers a worker that failed before reaching generate()'s finally; identity-guarded + # so a direct generate() that re-registered keeps its cancel handle. self._active_generate_cancel = None if error is not None: self._gen = { @@ -1652,8 +1515,7 @@ class VideoBackend: ) -> dict[str, Any]: import torch - # begin_generate passes the event it already registered (so a cancel in the - # spawn window is honoured); a direct call makes its own. + # begin_generate passes its already-registered event; a direct call makes its own. cancel = cancel_event if cancel_event is not None else threading.Event() with self._generate_lock: with self._lock: @@ -1695,28 +1557,19 @@ class VideoBackend: "generator": generator, } if fam.guidance_via_guider: - # HunyuanVideo-1.5: __call__ has no guidance kwarg at all; the - # CFG scale is a plain attribute on the pipeline's guider - # component, set per request. Near-1 scales auto-disable CFG - # inside the guider itself (_is_cfg_enabled's is_close check). + # HunyuanVideo-1.5: __call__ has no guidance kwarg; CFG scale is a guider + # attribute set per request (near-1 scales auto-disable CFG in the guider). pipe.guider.guidance_scale = float(guidance) else: kwargs[fam.cfg_kwarg] = guidance if negative_prompt and "negative_prompt" in call_params: kwargs["negative_prompt"] = negative_prompt - # LTX-2 takes frame_rate (it shapes the audio track length); other - # pipelines fix their own rate and fps only matters at export. + # LTX-2 takes frame_rate (shapes audio length); others fix their rate, fps only at export. if "frame_rate" in call_params: kwargs["frame_rate"] = float(out_fps) - # Dual-DiT MoE (Wan2.2-A14B): the low-noise expert (transformer_2) has its - # own guidance kwarg (cfg2_kwarg = "guidance_scale_2"). Thread it only when - # the loaded family declares one AND the pipeline signature accepts it (the - # same inspect.signature gate frame_rate uses), so a single-DiT pipeline is - # never handed a kwarg its check_inputs would reject. WanPipeline raises if - # guidance_scale_2 is passed to a pipeline with boundary_ratio=None - # (pipeline_wan.py:322), so the gate is BOTH the family flag and the - # signature: TI2V-5B has no cfg2_kwarg, so it never reaches here. A None - # request lets the pipeline default it (to guidance_scale) itself. + # Dual-DiT MoE: thread the low-noise expert's guidance kwarg only when the family + # declares one AND the signature accepts it -- WanPipeline raises if guidance_scale_2 + # is passed with boundary_ratio=None (pipeline_wan.py:322); TI2V-5B never reaches here. if fam.cfg2_kwarg and fam.cfg2_kwarg in call_params and guidance_2 is not None: kwargs[fam.cfg2_kwarg] = float(guidance_2) @@ -1746,9 +1599,8 @@ class VideoBackend: return callback_kwargs def _on_scheduler_step(done: int) -> None: - # No cooperative _interrupt here: without a callback the pipeline - # never checks it, so cancellation must unwind the denoise loop - # via an exception (mapped to the cancelled sentinel below). + # No cooperative _interrupt (the pipeline never checks it), so cancellation + # must unwind the denoise loop via an exception. if cancel.is_set(): raise _VideoGenerationCancelled() _tick(done) @@ -1757,16 +1609,12 @@ class VideoBackend: kwargs["callback_on_step_end"] = _on_step progress_ctx = contextlib.nullcontext() else: - # HunyuanVideo-1.5 has no step callback; every scheduler.step - # call is exactly one denoise step, so wrap it for progress + - # cancel and restore it afterwards. + # HunyuanVideo-1.5 has no step callback; each scheduler.step is one denoise + # step, so wrap it for progress + cancel and restore afterwards. progress_ctx = _scheduler_step_progress(pipe, _on_scheduler_step) - # An AUTO cache decision is re-checked against the ACTUAL step count, - # mirroring the image backend: a many-step request gains FBCache even - # when the load's default schedule kept it off, and a few-step request - # drops it. Explicit choices never toggle. Runs per view so a dual-DiT - # MoE toggles both experts. + # Re-check an AUTO cache decision against the ACTUAL step count (a many-step + # request gains FBCache, a few-step drops it); explicit choices never toggle. Per view. if state.cache_auto: toggled = state.transformer_cache for view in _views_for(pipe, fam): @@ -1778,9 +1626,7 @@ class VideoBackend: logger = logger, ) if toggled != state.transformer_cache: - # _VideoLoadState is frozen (loads swap it as one unit); this - # tracks the pipe-level toggle that already happened so - # status() reports the true cache state. + # _VideoLoadState is frozen; record the pipe-level toggle so status() is truthful. object.__setattr__(state, "transformer_cache", toggled) entry = (state.resolved or {}).get("transformer_cache") if isinstance(entry, dict): @@ -1796,11 +1642,8 @@ class VideoBackend: with torch.inference_mode(), progress_ctx: output = pipe(**kwargs) except _VideoGenerationCancelled: - # This cancel unwinds pipe.__call__ by exception (the scheduler - # wrapper has no cooperative _interrupt), skipping the pipeline's - # end-of-call maybe_free_model_hooks(); under model/group offload - # the currently-onloaded modules would otherwise stay on the GPU - # until the next request touches them. + # Unwinding by exception skips the pipeline's end-of-call maybe_free_model_hooks(); + # under offload the onloaded modules would stay on the GPU, so free them here. free_hooks = getattr(pipe, "maybe_free_model_hooks", None) if callable(free_hooks): try: @@ -1818,9 +1661,8 @@ class VideoBackend: mp4_bytes = self._encode_mp4( video_frames, out_fps, audio_track, pipe if fam.has_audio else None ) - # A cancel that landed during the (blocking, uncancellable) export/mux must - # still discard the clip: cancel_generate() already reported success for it, - # so re-check here before it is returned and persisted to the gallery. + # A cancel during the blocking export/mux must still discard the clip; re-check + # before it is returned and persisted. if cancel.is_set(): raise RuntimeError(VIDEO_CANCELLED_MSG) duration_s = len(video_frames) / float(out_fps) if out_fps else 0.0 @@ -1877,10 +1719,8 @@ class VideoBackend: def generate_progress(self) -> dict[str, Any]: with self._lock: gen = dict(self._gen) - # generate() swaps in a bare {"active": False} on its own exit paths - # before the job worker records the terminal dict; report the job as - # still active across that gap so a poller only sees active drop - # together with a terminal phase ("completed" / "failed"). + # generate() swaps in a bare {"active": False} before the worker records the terminal + # dict; report active across that gap so a poller sees active drop only with a terminal phase. if self._generate_job_active: gen["active"] = True gen.setdefault("active", False) @@ -1903,9 +1743,8 @@ class VideoBackend: state, self._state = self._state, None if state is not None: restore_backend_flags(state.backend_flags) - # A GGUF video load may have installed the process-wide compiled GGUF - # dequantizer; restore the stock kernels so a later load that asked for - # speed_mode=off gets the bit-identical path (mirrors the image unload). + # A GGUF load may have installed the compiled GGUF dequantizer; restore the stock + # kernels so a later speed=off load gets the bit-identical path (mirrors image unload). from . import diffusion_gguf_compile diffusion_gguf_compile.uninstall_all() @@ -1919,12 +1758,8 @@ class VideoBackend: self._loading = None if self._active_generate_cancel is not None: self._active_generate_cancel.set() - # Wait for the signalled generation to actually exit before freeing the - # pipeline: the denoise loop holds its own pipe reference until the next - # step callback, so tearing down under it would report the VRAM free (and - # let the GPU arbiter start another multi-GB load) while this clip still - # occupies it. generate() holds _generate_lock for its full body, so a - # bare acquire is the exit barrier (never taken while holding _lock). + # Barrier: wait for the signalled generation to exit before freeing the pipeline, or we + # report the VRAM free (and let the arbiter start another load) while the clip still holds it. with self._generate_lock: pass self._teardown_state() diff --git a/studio/backend/core/inference/video_families.py b/studio/backend/core/inference/video_families.py index 4254210e58..9c34f42f0f 100644 --- a/studio/backend/core/inference/video_families.py +++ b/studio/backend/core/inference/video_families.py @@ -21,8 +21,7 @@ import re from dataclasses import dataclass, field from typing import Optional -# Runtime->route contract, mirroring the diffusion sentinels: the routes match -# these EXACTLY to return 409 (client-recoverable) instead of a sanitized 500. +# Runtime->route contract: routes match these EXACTLY for a 409 instead of a 500. VIDEO_NOT_LOADED_MSG = "No video model is loaded." VIDEO_CANCELLED_MSG = "Video generation was cancelled." VIDEO_GENERATION_BUSY_MSG = "A video generation is already in progress." @@ -40,26 +39,19 @@ 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 alongside frames (LTX-2): - # export must mux the audio track into the MP4 and size estimates must count - # the audio VAE + vocoder companions. + # 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 style dual-expert MoE: a second DiT (``transformer_2``) handles - # the low-noise steps, with its own guidance kwarg. None/False for single-DiT - # families. Declared now so adding the A14B family later does not churn the - # schema every module already imports. + # 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 style guidance: the pipeline __call__ takes NO guidance - # kwarg at all; CFG lives on a ``guider`` component (ClassifierFreeGuidance) - # whose ``guidance_scale`` is a plain attribute set per request. When True, - # generate() writes the scale onto ``pipe.guider`` instead of passing - # ``cfg_kwarg`` (which the pipeline would reject as an unexpected argument). + # 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 constraints. ``frame_step`` is the temporal - # compression: a valid frame count is k * frame_step + 1 (the +1 is the - # anchor frame), 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 @@ -67,38 +59,26 @@ class VideoFamily: frame_step: int = 8 # Width/height must be divisible by this (LTX-2's pipeline rejects non-/32). resolution_multiple: int = 32 - # (width, height) presets the UI offers, landscape first, including a vertical - # option. The first preset is the default. + # (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), the video analogue of the image auto-policy table. - # These are 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 family's DiT compiles cleanly with regional torch.compile - # (Wan/LTX-2 declare _repeated_blocks; set False until verified per family). + # True when the DiT compiles cleanly with regional torch.compile (declares _repeated_blocks). supports_torch_compile: bool = True - # Families whose activations overflow float16 -> the loader promotes fp16 to - # float32. Video DiTs are bf16-native, so this defaults True (fp16 is never - # the right resolution for them; bf16 or float32 only). + # Video DiTs are bf16-native, so fp16 promotes to float32; defaults True. fp16_incompatible: bool = True - # Wan's VAE decodes in float32: diffusers loads AutoencoderKLWan at torch.float32 while the - # pipe runs bf16 (WanPipeline docstring). Loading the VAE bf16 like the other components - # degrades every clip (banding / black frames), so when True the loader pins the VAE back to - # fp32 after building the pipe. The bf16_components_gb VAE term is already its fp32 size, so - # the memory plan stays consistent. + # 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. vae_force_fp32: bool = False # Curated GGUF repo for the picker (the DiT as single-file GGUF quants). gguf_repo: Optional[str] = None _FAMILIES: tuple[VideoFamily, ...] = ( - # LTX-2 (diffusers >= 0.39): a ~19B single-stream video DiT generating - # synchronized audio + video in one pass (audio VAE + vocoder + text - # connectors ride the base repo; the classes are vendored inside - # diffusers.pipelines.ltx2). The Gemma3-27B text encoder is the memory - # heavyweight: ~50 GB bf16-resident, more than the DiT itself. The diffusers - # base repo carries the dev-style config (40 steps, CFG 4); the distilled - # single-file/GGUF checkpoints run few-step (see default_video_generation_params). + # LTX-2 (diffusers >= 0.39): ~19B single-stream video DiT generating synchronized audio + + # video in one pass. The Gemma3-27B text encoder is the memory heavyweight (~50 GB bf16, + # more than the DiT). Base repo carries the dev config (40 steps, CFG 4); distilled runs few-step. VideoFamily( name = "ltx-2", pipeline_class = "LTX2Pipeline", @@ -112,70 +92,47 @@ _FAMILIES: tuple[VideoFamily, ...] = ( default_fps = 24, frame_step = 8, resolution_multiple = 32, - # The pipeline's native default is 768x512; 1216x704 is the model card's - # quality target; 704x1216 is the vertical variant. + # 768x512 native default; 1216x704 the card's quality target; 704x1216 vertical. resolution_presets = ((768, 512), (1216, 704), (704, 1216), (512, 768)), - # transformer 37.8 stored bf16; Gemma3-27B TE ~50.4; video VAE 2.4 + - # connectors 2.9 + audio VAE/vocoder 0.2 (sibling metadata, duplicates - # removed -- the repo ships the TE twice under two shard namings). + # transformer 37.8 bf16; Gemma3-27B TE ~50.4; VAE 2.4 + connectors 2.9 + audio 0.2. bf16_components_gb = (37.8, 50.4, 5.5), gguf_repo = "unsloth/LTX-2.3-GGUF", ), - # Wan2.2-TI2V-5B (diffusers >= 0.35, verified on 0.39): a ~5B single-stream - # video DiT (WanPipeline + WanTransformer3DModel + AutoencoderKLWan + a UMT5 - # text encoder). No audio, no second expert -- its model_index.json ships - # ``boundary_ratio: null`` and ``transformer_2: [null, null]``, so it is a - # plain single-DiT family (is_moe left False). The Wan VAE has a temporal - # compression of 4, so valid frame counts are 4k+1 (frame_step = 4), which - # matches the pipeline's own ``num_frames % vae_scale_factor_temporal == 1`` - # check (pipeline_wan.py:493). The pipeline defaults to 50 steps / CFG 5, but - # the 5B TI2V card ships the 720p-class few-step recipe, so the picker default - # (see _VIDEO_GENERATION_DEFAULTS) uses the pipeline's 50/5 while the UI presets - # target 720p at 24 fps (the model card's playback rate). + # Wan2.2-TI2V-5B (diffusers >= 0.35, verified on 0.39): ~5B single-stream video DiT (UMT5 + # text encoder). No audio, no second expert (boundary_ratio null, transformer_2 null), so + # single-DiT. Wan VAE temporal compression 4 -> 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", transformer_class = "WanTransformer3DModel", base_repo = "Wan-AI/Wan2.2-TI2V-5B-Diffusers", - # "wan2.2-5b" and "wan-ti2v" are the short ids the picker / GGUF filenames - # use; "wan2.2-ti2v" catches the diffusers repo stem without the "-5b". + # "wan2.2-5b"/"wan-ti2v" are the picker/GGUF short ids; "wan2.2-ti2v" catches the repo stem. aliases = ("wan2.2-5b", "wan-ti2v", "wan2.2-ti2v", "wan-ti2v-5b"), has_audio = False, default_steps = 50, default_guidance = 5.0, - # 121 frames at 24 fps is ~5s, the model card's headline clip length; on the - # 4k+1 lattice (121 = 4*30 + 1) it needs no snapping. + # 121 frames at 24 fps ~5s; on the 4k+1 lattice (121 = 4*30 + 1) it needs no snapping. default_num_frames = 121, default_fps = 24, - # Wan VAE temporal factor is 4 (autoencoder_kl_wan.py scale_factor_temporal), - # so valid counts are 4k+1, unlike LTX-2's 8k+1. + # Wan VAE temporal factor 4, so valid counts are 4k+1. frame_step = 4, - # TI2V-5B's VAE is 16x spatial (vae/config.json scale_factor_spatial=16), and the - # transformer patch is 2, so WanPipeline floors H/W to 16*2 = 32 (pipeline_wan.py:505, - # silently, with a warning). Snap to 32 so the recorded size matches the generated clip; - # a /16-but-not-/32 request (e.g. 720) would otherwise be recorded but rendered at 704. + # 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 (a /16-not-/32 request would render at 704). resolution_multiple = 32, - # 720p-class presets (all /32): 1280x704 landscape (the card's target), its vertical - # variant, and a square. The first preset is the default the loader plans memory against. + # 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 sizes. The transformer + VAE ship FP32 on disk (safetensors headers are - # F32; transformer index = 20.0 GB = 5B params x 4), so bf16-resident transformer is half - # (~10.0); the UMT5 text encoder ships bf16 (11.4). The VAE runs fp32 (vae_force_fp32), so - # its term is the fp32 size (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. - # Its model_index.json lists BOTH ``transformer`` and ``transformer_2`` as - # WanTransformer3DModel and sets ``boundary_ratio: 0.875``; the pipeline routes - # the high-noise steps (timestep >= boundary) through ``transformer`` at - # guidance_scale and the low-noise steps through ``transformer_2`` at - # guidance_scale_2 (pipeline_wan.py:584-603). ``guidance_scale_2`` exists in - # 0.39 (pipeline_wan.py:392) and is only accepted when boundary_ratio is set - # (its check_inputs raises otherwise, pipeline_wan.py:322), so cfg2_kwarg is - # threaded ONLY for this family. boundary_ratio itself lives in the pipeline - # config (loaded from model_index.json), so no per-generation plumbing is needed. + # Wan2.2-T2V-A14B (diffusers >= 0.35, verified on 0.39): the dual-expert MoE. Both + # transformer + transformer_2 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. boundary_ratio lives in the pipeline config, so no per-generation plumbing. VideoFamily( name = "wan2.2-t2v-a14b", pipeline_class = "WanPipeline", @@ -183,82 +140,57 @@ _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, - # The second expert is the same class; is_moe drives the dual-DiT optimisation - # layers (speed / attention / cache / quant apply to BOTH transformers), and - # cfg2_kwarg names the pipeline kwarg carrying 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", default_steps = 50, default_guidance = 5.0, - # 81 frames at 16 fps is ~5s (81 = 4*20 + 1), the A14B card's default clip. + # 81 frames at 16 fps ~5s (81 = 4*20 + 1), the A14B card's default clip. default_num_frames = 81, - # The A14B card runs at 16 fps (vs the 5B TI2V's 24), per its model_index / - # model card; export uses this rate. - default_fps = 16, + default_fps = 16, # A14B runs at 16 fps (vs TI2V-5B's 24) frame_step = 4, resolution_multiple = 16, - # 480p and 720p presets (landscape + vertical), the two resolutions the A14B card - # documents. 832x480 is the native 480p; 1280x720 the native 720p (true 16:9). A14B's - # VAE is 8x so resolution_multiple is 16 and 720 (= 45*16) renders exactly -- the 704 - # value belongs to TI2V-5B, whose 16x VAE floors 720 to 704 (multiple 32). + # 480p + 720p presets (landscape + vertical). A14B's VAE is 8x so multiple 16 renders + # 720 (=45*16) exactly (unlike TI2V-5B's 16x VAE, which floors 720 to 704). resolution_presets = ((1280, 720), (832, 480), (480, 832), (720, 1280)), - # bf16-RESIDENT sizes. Each expert ships FP32 on disk (safetensors headers are F32; - # transformer index = 57.15 GB = 14.3B params x 4), so bf16-resident is ~28.6 each -> - # ~57.2 for BOTH experts (the memory headline before offload), NOT the 114.3 fp32 - # on-disk sum. UMT5 text encoder ships bf16 (11.4); the VAE runs fp32 (vae_force_fp32), - # so its term is the fp32 size (0.5). + # bf16-RESIDENT. Each expert ships FP32 (index 57.15 GB = 14.3B x 4) -> ~28.6 bf16 each -> + # ~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 ship the two experts as separate files, and a - # single-file load covers only one (validate_load_request refuses it). + # No gguf_repo: community GGUFs split the experts, and a single-file load covers only one. ), - # HunyuanVideo-1.5 (diffusers >= 0.39): an 8.3B video DiT with a Qwen2.5-VL - # text encoder plus a ByT5 glyph encoder. Three quirks, all verified against - # the installed pipeline source (pipeline_hunyuan_video1_5.py): - # 1. __call__ takes NO guidance kwarg; CFG lives on the ``guider`` component - # (ClassifierFreeGuidance; the 480p t2v repo ships guidance_scale = 6.0), - # hence guidance_via_guider. - # 2. __call__ has NO callback_on_step_end; generate() falls back to the - # scheduler.step progress wrapper automatically (capability-detected). - # 3. The tencent/HunyuanVideo-1.5 repo is the ORIGINAL layout (config.json, - # no model_index.json); only the hunyuanvideo-community Diffusers repacks - # load through HunyuanVideo15Pipeline, so those are the trusted repos. - # The transformer declares _repeated_blocks and inherits CacheMixin, so the - # regional compile profile and First-Block-Cache both apply. + # 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 on the ``guider`` + # (guidance_via_guider); (2) NO callback_on_step_end (generate() uses the scheduler.step + # wrapper); (3) tencent's repo is the original layout (no model_index.json), so only the + # community Diffusers repacks load. The transformer declares _repeated_blocks + CacheMixin. VideoFamily( name = "hunyuanvideo-1.5", pipeline_class = "HunyuanVideo15Pipeline", transformer_class = "HunyuanVideo15Transformer3DModel", base_repo = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v", - # No bare "hunyuanvideo" alias: it would also claim the incompatible 1.0 - # repos (HunyuanVideoPipeline), which this family cannot load. + # No bare "hunyuanvideo" alias: it would also claim the incompatible 1.0 repos. aliases = ("hunyuanvideo-1-5", "hunyuanvideo1.5", "hunyuanvideo1-5", "hv15"), has_audio = False, guidance_via_guider = True, default_steps = 50, default_guidance = 6.0, - # 121 frames at 24 fps is ~5s, the pipeline's own num_frames default. + # 121 frames at 24 fps ~5s, the pipeline's own default. default_num_frames = 121, default_fps = 24, - # The HV15 VAE compresses 16x spatial / 4x temporal (vae config: - # spatial_compression_ratio 16, temporal_compression_ratio 4) with a - # patch-1 transformer, so sizes snap to /16 and frames to 4k+1. + # HV15 VAE compresses 16x spatial / 4x temporal, patch-1, so sizes snap /16, frames 4k+1. frame_step = 4, resolution_multiple = 16, - # 480p-class presets (the base repo is the 480p t2v variant): landscape, - # vertical, square. + # 480p-class presets (the base is the 480p variant): landscape, vertical, square. resolution_presets = ((832, 480), (480, 832), (624, 624)), - # Disk shards are fp32 for the DiT (32.0 GB -> 16.6 bf16-resident) and the - # VAE (4.7 -> 2.4); the Qwen2.5-VL TE is stored bf16 (14.0) plus ByT5 0.8. + # DiT fp32 on disk (32.0 -> 16.6 bf16); VAE (4.7 -> 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, pipeline quirks, guider config - # (guidance 6.0) and shard footprint as the 480p entry above; only the - # trained resolution class differs. Kept as its OWN family so a 720p load - # defaults to 720p-class sizes instead of silently rendering at 832x480. - # The repo-id alias is the full path segment, so it out-lengths (and thus - # 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", @@ -336,21 +268,14 @@ 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 against -# the picked id (then the base repo), most specific first: the distilled LTX-2.3 -# checkpoints run few-step with CFG off, while the dev-config base repo wants the -# full 40-step CFG schedule. Mirrors default_generation_params on the image side. +# 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), - # Wan2.2 pipelines default to 50 steps at CFG 5.0 (WanPipeline.__call__: - # num_inference_steps = 50, guidance_scale = 5.0, verified in diffusers 0.39). - # Both TI2V-5B and A14B share these; the substring "wan" catches the picked id - # and the base repo. A future distilled Wan GGUF is caught by the "distilled" - # row above (listed first), exactly as the LTX-2.3 distilled checkpoints are. + # Wan2.2 pipelines default to 50 steps / CFG 5.0; both TI2V-5B and A14B share these. ("wan", 50, 5.0), - # HunyuanVideo-1.5 runs the pipeline's 50 steps with the guider's shipped - # CFG 6.0 (guider_config.json in the community Diffusers repacks). + # HunyuanVideo-1.5: 50 steps with the guider's shipped CFG 6.0. ("hunyuanvideo", 50, 6.0), ) @@ -366,11 +291,8 @@ 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, not a raw substring: reject a - # preceding ASCII letter so an opaque path/repo like "user/swan-video" - # or "taiwan-clips" does not false-match "wan" and silently apply Wan's - # 50-step/CFG-5 schedule to a non-Wan model. Trailing chars stay free so - # "wan2.2-ti2v", "ltxv-2.3" and "...-distilled-..." still match. + # 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. if re.search(r"(? None: diffusion_active = get_diffusion_training_service().is_active() except Exception: # noqa: BLE001 diffusion_active = False - # An SDXL LoRA trainer runs in its own subprocess on the same GPU, so a video - # load must be refused while one is active too -- otherwise the resident pipeline - # competes with the trainer for VRAM. Symmetric with the image-load interlock. + # An SDXL LoRA trainer runs in its own subprocess on the same GPU, so refuse a video + # load while one is active too (VRAM competition). Symmetric with the image-load interlock. if not llm_active and not diffusion_active: return raise HTTPException( @@ -85,9 +84,7 @@ async def load_video_model( backend = get_video_backend() try: - # Validate cheaply BEFORE touching the GPU: an unloadable pick (bad family, - # missing local checkpoint, a non-trusted non-GGUF repo) must not evict a - # working chat model and then 400. + # Validate cheaply BEFORE touching the GPU so an unloadable pick can't evict chat then 400. await asyncio.to_thread( backend.validate_load_request, request.model_path, @@ -98,15 +95,10 @@ async def load_video_model( transformer_quant = request.transformer_quant, text_encoder_quant = request.text_encoder_quant, ) - # Refuse while training is running: a multi-GB video pipeline would compete - # with the training subprocess for VRAM. Mirrors the image-load guard. + # Refuse while training is running (VRAM competition). Mirrors the image-load guard. _guard_video_load_against_training() - # Take the GPU from the chat backend only when this load will actually use it, - # which is exactly the resolved device being non-CPU. A CPU-only load never - # touches GPU memory, so keying off the device (not the load) avoids wrongly - # evicting a resident chat model. Release any stale VIDEO ownership on a CPU - # load -- release() is owner-guarded, so it is a no-op when video never owned - # the GPU. + # Take the GPU from chat only for a non-CPU load; a CPU load never touches GPU memory, + # so key off the device. Release stale VIDEO ownership on a CPU load (owner-guarded no-op). device = await asyncio.to_thread(lambda: resolve_diffusion_device_target().device) if device != "cpu": await asyncio.to_thread(acquire_for, VIDEO) @@ -173,9 +165,8 @@ async def generate_video( # Bad client input -- a 400 with the reason, not a generic 500. raise HTTPException(status_code = 400, detail = str(exc)) except RuntimeError as exc: - # Only "no model loaded" / "already generating" are client-state (409). - # Match the sentinels exactly, not as a substring, so an unrelated failure - # can't misroute to 409 and leak its message. + # Only the not-loaded / busy sentinels are client-state (409); match exactly so an + # unrelated failure can't misroute and leak its message. msg = str(exc) if msg in (VIDEO_NOT_LOADED_MSG, VIDEO_GENERATION_BUSY_MSG): raise HTTPException(status_code = 409, detail = msg) @@ -212,11 +203,8 @@ async def unload_video_model(current_subject: str = Depends(get_current_subject) backend = get_video_backend() status_dict = await asyncio.to_thread(backend.unload) # Drop VIDEO ownership only if nothing is resident AND no new load is in flight: a concurrent - # /video/load that re-acquired VIDEO while this (slow) unload ran must keep ownership, or a - # later chat/image load would see no owner, skip eviction, and OOM against the newly resident - # (or still in-flight) video pipeline. release() is owner-guarded and identity-less, so an - # unconditional release here would clear the newer load's claim. Mirrors the images-route - # guard (inference.py), plus the in-flight check the committed-loaded state cannot cover. + # /video/load that re-acquired VIDEO must keep ownership (release() is owner-guarded but + # identity-less, so an unconditional release would clear the newer claim). Mirrors the images route. if not backend.loading_repo_ids() and not backend.status()["loaded"]: release(VIDEO) return VideoStatusResponse(**status_dict) @@ -235,10 +223,8 @@ async def list_gallery_videos( # Fetch one extra to learn whether more remain, without a second scan. records = await asyncio.to_thread(video_gallery.list_videos, limit + 1, offset) has_more = len(records) > limit - # Build the response per record and drop any that fail schema validation: a - # sidecar with all required keys but a wrong value type (a hand-dropped or - # older-schema file) passes the presence-only read but would raise inside - # GalleryVideo(**r). Skipping it keeps one bad file from 500-ing the listing. + # Build per record, dropping any that fail schema validation, so one bad sidecar + # (wrong value type) doesn't 500 the whole listing. videos = [] for r in records[:limit]: try: @@ -259,9 +245,8 @@ async def get_gallery_video_file( raise HTTPException(status_code = 404, detail = "Video not found.") from fastapi.responses import FileResponse - # FileResponse streams from disk (no whole-clip buffering per request) and - # serves HTTP range requests so a direct URL can seek without a full fetch. - # Immutable content (id is unique per video), so let the browser cache it. + # FileResponse streams from disk and serves range requests (seek without a full fetch). + # Immutable per id, so let the browser cache it. return FileResponse( path, media_type = "video/mp4",