diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index dcdbb863bb..79d19c9306 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -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() diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index 19322be2b4..09a0a1ae51 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -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, ) diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index edb0a41bbb..2c3a027e54 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -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) diff --git a/studio/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index 5928d6523b..d1fd6b3ba4 100644 --- a/studio/backend/tests/test_video_routes.py +++ b/studio/backend/tests/test_video_routes.py @@ -68,6 +68,7 @@ class _FakeBackend: model_path, *, gguf_filename = None, + base_repo = None, family_override = None, model_kind = None, ):