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

Tighten comments across the remaining image stack files

Browse source
This commit is contained in:
Daniel Han 2026-07-12 11:40:05 +00:00
commit 9fde4b9991
6 changed files with 722 additions and 1500 deletions
Download patch file Download diff file Expand all files Collapse all files

1145
studio/backend/core/inference/diffusion.py
View file

File diff suppressed because it is too large Load diff

227
studio/backend/core/inference/sd_cpp_backend.py
View file

@ -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/<user>/...), 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: <lora:ALIAS:w> 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 <lora:ALIAS:w> tags into the prompt (deduped against any the user
# typed). Empty -> prompt/dir unchanged.
# Materialize LoRAs into a scan dir and inject <lora:ALIAS:w> 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"],
}

94
studio/backend/core/inference/sd_cpp_server.py
View file

@ -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 []

503
studio/backend/core/inference/video.py
View file

@ -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()

210
studio/backend/core/inference/video_families.py
View file

@ -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"(?<![a-z])" + re.escape(key), needle):
return steps, guidance
return fallback

43
studio/backend/routes/video.py
View file

@ -59,9 +59,8 @@ def _guard_video_load_against_training() -> 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",