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

Harden video load path: early GGUF-repo rejection, family fallback parity, rollback and teardown fixes

Browse source 
Review follow-ups on the video inference backend:

- validate_load_request now rejects a -GGUF repo picked as a diffusers
  pipeline (no gguf_filename) up front, instead of failing minutes later
  in from_pretrained after the GPU owner was already evicted.
- New _detect_load_family helper shared by validate_load_request and
  _run_load: when the repo id alone does not carry the family, fall back
  to detecting it from the picked GGUF filename, so both paths agree.
- routes/video.py now threads base_repo into validate_load_request so an
  untrusted companion repo is refused before the arbiter handoff.
- unload() now drains _generate_lock before _teardown_state so a
  cancelled clip actually exits the denoise loop before the VRAM is
  reported free.
- load_pipeline re-checks the load token after the generate-lock barrier
  and raises if the load was superseded while waiting.
- Pre-commit global mutations (backend flags, gguf compile installs) are
  registered per load token and rolled back in _run_load's error path
  via _rollback_precommit_globals, so a failed load no longer leaks
  process-wide state.
- fp32 memory estimates now apply a 2x dtype scale on non-CPU devices
  for pipeline, single-file and companion sizes (bf16 tables assume
  2 bytes/param); GGUF quant estimates stay unscaled.

Tests: GGUF-repo-as-pipeline rejection, _detect_load_family fallback and
override semantics; fake route backend accepts base_repo. 66 passed
across test_video_backend, test_video_routes, test_video_families,
test_video_gallery.
This commit is contained in:
Daniel Han 2026-07-05 07:55:26 +00:00
commit 6485e68147
4 changed files with 123 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -119,6 +119,22 @@ def _is_trusted_video_repo(repo_id: str) -> bool:
return rid.startswith("unsloth/") or rid in _TRUSTED_NON_GGUF_VIDEO_REPOS
def _detect_load_family(
repo_id: str,
gguf_filename: Optional[str],
family_override: Optional[str],
) -> Optional[VideoFamily]:
"""Family detection shared by validate_load_request and the load worker: the
repo id first, then the picked filename -- a local directory or generically
named repo often carries the family token only in the checkpoint filename,
and the worker must resolve the same family the validator accepted."""
return detect_video_family(repo_id, family_override) or (
detect_video_family(f"{repo_id}/{gguf_filename}")
if gguf_filename and not family_override
else None
)
def _ensure_mp4_encoder_available() -> None:
"""Fail a load fast when PyAV is missing: the export otherwise dies AFTER a
multi-minute denoise, which is the worst possible time to learn about it."""
@ -193,11 +209,16 @@ class VideoBackend:
) -> VideoFamily:
"""Cheap, network-free validation shared by the route and the load path."""
kind = resolve_video_model_kind(gguf_filename, model_kind)
fam = detect_video_family(repo_id, family_override) or (
detect_video_family(f"{repo_id}/{gguf_filename}")
if gguf_filename and not family_override
else None
)
# 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.
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 "
"(gguf_filename) instead of loading it as a diffusers pipeline."
)
fam = _detect_load_family(repo_id, gguf_filename, family_override)
if fam is None:
raise ValueError(
f"'{repo_id}' is not a supported text-to-video model. Supported families: "
@ -291,7 +312,9 @@ class VideoBackend:
def _run_load(self, **kwargs: Any) -> None:
token = kwargs.get("_load_token")
try:
fam = detect_video_family(kwargs["repo_id"], kwargs.get("family_override"))
fam = _detect_load_family(
kwargs["repo_id"], kwargs.get("gguf_filename"), kwargs.get("family_override")
)
kind = resolve_video_model_kind(kwargs.get("gguf_filename"), kwargs.get("model_kind"))
base = (
kwargs["repo_id"]
@ -369,6 +392,11 @@ class VideoBackend:
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).
self._rollback_precommit_globals(token)
if self._load_token != token:
return
logger.error("video.load_failed: %s", exc)
@ -378,6 +406,25 @@ class VideoBackend:
if self._load_token == token and self._loading is not None:
self._loading.error = redact_native_paths(str(exc))
def _rollback_precommit_globals(self, token: Optional[int]) -> None:
"""Restore process-wide speed globals (cudnn.benchmark / TF32 / the compiled
GGUF dequantizer) for a load that died BEFORE committing _VideoLoadState.
_teardown_state only restores from the committed state's snapshot, so an
uncommitted load would otherwise leak its profile into the next speed=off
load. Token-scoped: when a newer load has already taken the snapshot slot,
the stale worker must leave the globals alone."""
stored = getattr(self, "_precommit_globals", None)
if stored is None:
return
stored_token, flags = stored
if token is not None and stored_token is not None and stored_token != token:
return
self._precommit_globals = None
restore_backend_flags(flags)
from . import diffusion_gguf_compile
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.
@ -584,6 +631,12 @@ class VideoBackend:
# body, so a bare acquire is the exit barrier (never while holding _lock).
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.
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()
@ -594,13 +647,23 @@ class VideoBackend:
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.
dtype_scale = 2.0 if device != "cpu" and dtype is torch.float32 else 1.0
# ── memory plan: family-table resident estimate + frames-aware headroom.
device_memory = snapshot_device_memory(target)
components = fam.bf16_components_gb
mib_per_gb = 1000.0**3 / (1024.0 * 1024.0)
if kind == "pipeline":
model_dense_mib = int(sum(components) * mib_per_gb) if components is not None else None
model_dense_mib = (
int(sum(components) * mib_per_gb * dtype_scale)
if components is not None
else None
)
companion_mib = None
else:
checkpoint_path = self._resolve_checkpoint_path(repo_id, gguf_filename, hf_token)
@ -610,8 +673,10 @@ class VideoBackend:
transformer_mib = estimate_gguf_resident_mib(size_mib)
else:
transformer_mib = estimate_safetensors_dense_mib(size_mib)
if transformer_mib is not None:
transformer_mib = int(transformer_mib * dtype_scale)
companion_mib = (
int((components[1] + components[2]) * mib_per_gb)
int((components[1] + components[2]) * mib_per_gb * dtype_scale)
if components is not None
else None
)
@ -691,6 +756,11 @@ class VideoBackend:
# placement/offload last.
effective_speed = resolve_speed_mode(speed_mode, is_gguf = kind == "gguf")
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.
self._precommit_globals = (_load_token, backend_flags)
cache_engaged = apply_step_cache(
pipe,
mode = normalize_transformer_cache(transformer_cache),
@ -782,6 +852,8 @@ class VideoBackend:
transformer_cache = cache_engaged,
resolved = resolved,
)
# Ownership of the globals transferred to _state / _teardown_state.
self._precommit_globals = None
logger.info(
"video.loaded: %s (%s, %s, offload=%s, speed=%s)",
repo_id,
@ -1019,6 +1091,14 @@ 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).
with self._generate_lock:
pass
self._teardown_state()
logger.info("video.unloaded")
return self.status()

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

@ -91,6 +91,7 @@ async def load_video_model(
backend.validate_load_request,
request.model_path,
gguf_filename = request.gguf_filename,
base_repo = request.base_repo,
family_override = request.family_override,
model_kind = request.model_kind,
)

34
studio/backend/tests/test_video_backend.py
View file

@ -12,7 +12,12 @@ import types
import pytest
from core.inference.video import VideoBackend, get_video_backend, resolve_video_model_kind
from core.inference.video import (
VideoBackend,
_detect_load_family,
get_video_backend,
resolve_video_model_kind,
)
from core.inference.video_families import VIDEO_NOT_LOADED_MSG
@ -211,6 +216,33 @@ def test_validate_gates_base_repo_and_local_paths(tmp_path):
)
def test_validate_rejects_gguf_repo_as_pipeline():
backend = VideoBackend()
# A -GGUF repo with no quant filename resolves to the pipeline kind and would
# only fail minutes later in from_pretrained, AFTER evicting the GPU owner.
with pytest.raises(ValueError, match = "pick one of its .gguf files"):
backend.validate_load_request("unsloth/LTX-2.3-GGUF")
with pytest.raises(ValueError, match = "pick one of its .gguf files"):
backend.validate_load_request("unsloth/Wan2.2-TI2V-5B-GGUF/")
def test_detect_load_family_filename_fallback():
# Repo id alone carries the family.
fam = _detect_load_family("Lightricks/LTX-2", None, None)
assert fam is not None and fam.name == "ltx-2"
# Repo id is opaque but the picked filename carries it: fall back to the
# combined path so validate and _run_load agree on the family.
fam = _detect_load_family("someorg/quants", "ltx-2-19b-Q4_K_M.gguf", None)
assert fam is not None and fam.name == "ltx-2"
# No filename and no recognisable repo id: no family.
assert _detect_load_family("someorg/quants", None, None) is None
# An explicit override resolves by name/alias and skips the filename fallback:
# a bogus override stays None even when the filename would have matched.
fam = _detect_load_family("someorg/quants", "ltx-2-19b-Q4_K_M.gguf", "ltxv")
assert fam is not None and fam.name == "ltx-2"
assert _detect_load_family("someorg/quants", "ltx-2-19b-Q4_K_M.gguf", "bogus") is None
def test_load_generate_unload_gguf(fake_runtime, tmp_path):
backend = VideoBackend()
status = _load_gguf(backend, tmp_path)

1
studio/backend/tests/test_video_routes.py
View file

@ -68,6 +68,7 @@ class _FakeBackend:
model_path,
*,
gguf_filename = None,
base_repo = None,
family_override = None,
model_kind = None,
):