diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index c56d1592ce..576124d7c9 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -269,6 +269,29 @@ def estimate_image_runtime_mib( return max(1024, int(8192 * max(0.25, pixel_scale) * multiplier)) +def estimate_video_runtime_mib( + *, + width: Optional[int], + height: Optional[int], + num_frames: Optional[int], +) -> int: + """Per-call activation / latent / decode headroom for a video generation. + + The image estimator is pixel-area only and badly undershoots video: latents + carry a frames dimension and the VAE DECODE is the peak -- the decoded clip + materialises as num_frames full-resolution fp32 frames plus the decoder's + intermediates, typically dwarfing the denoise activations. Scale by the + decoded-clip footprint (frames x H x W x 3 x 4 bytes) with a 3x factor for + decoder intermediates + the PIL/tensor copy held during export, on top of a + fixed denoise-side base. + """ + w = max(64, int(width or 768)) + h = max(64, int(height or 512)) + frames = max(1, int(num_frames or 121)) + decoded_mib = (frames * w * h * 3 * 4) / float(1024 * 1024) + return max(3072, int(4096 + 3.0 * decoded_mib)) + + def _safe_device_budget_mib(memory: DeviceMemory) -> Optional[int]: """Free memory minus a headroom reserve, so a plan that "fits" leaves room for fragmentation and other tenants. None when free memory is unknown.""" diff --git a/studio/backend/core/inference/gpu_arbiter.py b/studio/backend/core/inference/gpu_arbiter.py index a6c79bf1e8..f4d93203af 100644 --- a/studio/backend/core/inference/gpu_arbiter.py +++ b/studio/backend/core/inference/gpu_arbiter.py @@ -24,6 +24,7 @@ logger = get_logger(__name__) CHAT = "chat" DIFFUSION = "diffusion" +VIDEO = "video" _lock = threading.Lock() _owner: Optional[str] = None @@ -61,8 +62,16 @@ def _evict_diffusion() -> None: get_active_diffusion_engine().unload() -# Patchable in tests via monkeypatch.setitem. -_EVICTORS = {CHAT: _evict_chat, DIFFUSION: _evict_diffusion} +def _evict_video() -> None: + from core.inference.video import get_video_backend + get_video_backend().unload() + + +# Patchable in tests via monkeypatch.setitem. Ownership is exclusive -- only one +# owner holds the GPU at a time -- so acquire_for's evict-the-current-owner +# already generalises to any number of registered owners (chat / image / video +# all evict whichever of the others currently holds the GPU). +_EVICTORS = {CHAT: _evict_chat, DIFFUSION: _evict_diffusion, VIDEO: _evict_video} def acquire_for(owner: str) -> None: diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py new file mode 100644 index 0000000000..88908f0d9d --- /dev/null +++ b/studio/backend/core/inference/video.py @@ -0,0 +1,855 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Local text-to-video inference backend (diffusers). + +A deliberate sibling of ``DiffusionBackend`` rather than a mode of it: video +pipelines take frame/fps arguments, return frame stacks (plus synchronized audio +for LTX-2) instead of PIL images, and persist MP4s -- none of the image module's +img2img/inpaint/ControlNet/LoRA surface applies. The concurrency skeleton +(load token, per-generation cancel event, split status/generate locks) is copied +from the image backend so the two cannot diverge in lifecycle behaviour, and the +hardware/optimisation layers are IMPORTED from the image stack unchanged: +device/dtype resolution, memory planning + offload tiers, attention backends, +speed profiles (regional torch.compile), and FBCache step caching all operate on +``pipe.transformer`` generically. + +Video-specific behaviour lives here: +- the runtime headroom estimate is frames-aware (``estimate_video_runtime_mib``): + the VAE decode of a whole clip is the memory peak, not the denoise; +- VAE tiling is always enabled (decode of 100+ frames at 720p-class resolutions + spikes far beyond the image case, and tiling's quality cost is negligible); +- generation snaps num_frames to the family's temporal lattice (k * step + 1) + and width/height to its required multiple BEFORE latents are allocated; +- the result is encoded to MP4 (H.264) via diffusers' PyAV-backed exporter, + muxing the audio track for families that produce one. + +Loads are gated to trusted repos exactly like the image backend: unsloth/*, the +family's official base repos, or a local path the user explicitly picked. +""" + +from __future__ import annotations + +import inspect +import os +import tempfile +import threading +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional + +from loggers import get_logger + +from .diffusion_attention import apply_attention_backend, select_attention_backend +from .diffusion_cache import apply_step_cache, normalize_transformer_cache +from .diffusion_device import resolve_diffusion_device_target +from .diffusion_memory import ( + apply_memory_plan, + estimate_gguf_resident_mib, + estimate_safetensors_dense_mib, + estimate_video_runtime_mib, + file_size_mib, + normalize_memory_mode, + plan_diffusion_memory, + snapshot_device_memory, +) +from .diffusion_speed import ( + SPEED_OFF, + apply_speed_optims, + resolve_speed_mode, + restore_backend_flags, + snapshot_backend_flags, +) +from .diffusion_auto_policy import build_resolved_record +from .video_families import ( + VIDEO_CANCELLED_MSG, + VIDEO_NOT_LOADED_MSG, + VideoFamily, + default_video_generation_params, + detect_video_family, + resolve_video_base_repo, + snap_num_frames, + snap_video_size, + supported_video_family_names, +) +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). +_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. +_TRUSTED_NON_GGUF_VIDEO_REPOS = frozenset( + { + "lightricks/ltx-2", + "lightricks/ltx-2.3", + "lightricks/ltx-2.3-fp8", + } +) + + +def resolve_video_model_kind(gguf_filename: Optional[str], model_kind: Optional[str]) -> str: + """Classify a load request; explicit model_kind wins, else the filename decides.""" + if model_kind: + kind = model_kind.strip().lower() + if kind not in _MODEL_KINDS: + raise ValueError( + f"Unknown model_kind '{model_kind}'. Expected one of {sorted(_MODEL_KINDS)}." + ) + return kind + if not gguf_filename: + return "pipeline" + return "gguf" if gguf_filename.strip().lower().endswith(".gguf") else "single_file" + + +def _is_trusted_video_repo(repo_id: str) -> bool: + """Whether a NON-GGUF load may deserialise this repo (see the image twin).""" + try: + if Path(repo_id).expanduser().exists(): + return True + except OSError: + pass + rid = repo_id.strip().lower() + return rid.startswith("unsloth/") or rid in _TRUSTED_NON_GGUF_VIDEO_REPOS + + +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.""" + try: + import av # noqa: F401 + except Exception as exc: # noqa: BLE001 -- any import failure means no encoder + raise ValueError( + "Video generation needs the 'av' package (PyAV) to encode MP4s. " + "Install it with: pip install av" + ) from exc + + +@dataclass(frozen = True) +class _VideoLoadState: + """Everything about the currently-loaded video pipeline, swapped as one unit.""" + + pipe: Any + family: VideoFamily + repo_id: str + base_repo: str + device: str + dtype: str + kind: str + gguf_filename: Optional[str] = None + offload_policy: str = "none" + vae_tiling: bool = True + memory_mode: str = "auto" + speed_mode: str = SPEED_OFF + speed_optims: tuple = () + backend_flags: Optional[dict] = None + attention_backend: Optional[str] = None + transformer_cache: Optional[str] = None + resolved: Optional[dict] = None + + +@dataclass +class _VideoLoadingState: + repo_id: str + base_repo: str + expected_bytes: Optional[int] = None + error: Optional[str] = None + + +def _progress(phase: Optional[str], **extra: Any) -> dict[str, Any]: + return {"phase": phase, **extra} + + +class VideoBackend: + """One loaded video pipeline; loads swap it atomically (same model as images).""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._generate_lock = threading.Lock() + self._state: Optional[_VideoLoadState] = None + self._loading: Optional[_VideoLoadingState] = None + self._load_token = 0 + self._cancel_event = threading.Event() + self._active_generate_cancel: Optional[threading.Event] = None + # Generation progress, written by the step callback / phase transitions. + self._gen: dict[str, Any] = {"active": False} + + # ── validation ─────────────────────────────────────────────────────────── + + def validate_load_request( + self, + repo_id: str, + *, + gguf_filename: Optional[str] = None, + family_override: Optional[str] = None, + model_kind: Optional[str] = None, + ) -> 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 + ) + if fam is None: + raise ValueError( + f"'{repo_id}' is not a supported text-to-video model. Supported families: " + f"{', '.join(supported_video_family_names())}. If this is a variant of one " + f"of them, pass family_override with that family name." + ) + if kind != "gguf" and not _is_trusted_video_repo(repo_id): + raise ValueError( + f"Non-GGUF video loads are limited to unsloth/* repos, the official " + f"family base repos, and local paths; '{repo_id}' is neither." + ) + if kind in ("gguf", "single_file") and not gguf_filename: + raise ValueError("A gguf/single_file load needs the checkpoint filename.") + _ensure_mp4_encoder_available() + return fam + + # ── background load + progress ─────────────────────────────────────────── + + def begin_load( + self, + repo_id: str, + *, + gguf_filename: Optional[str] = None, + base_repo: Optional[str] = None, + family_override: Optional[str] = None, + hf_token: Optional[str] = None, + memory_mode: Optional[str] = None, + speed_mode: Optional[str] = None, + attention_backend: Optional[str] = None, + transformer_cache: Optional[str] = None, + transformer_cache_threshold: Optional[float] = None, + model_kind: Optional[str] = None, + ) -> dict[str, Any]: + """Validate, then run the (slow) load on a daemon thread. Returns at once.""" + hf_token = (hf_token.strip() if isinstance(hf_token, str) else hf_token) or None + fam = self.validate_load_request( + repo_id, + gguf_filename = gguf_filename, + family_override = family_override, + model_kind = model_kind, + ) + with self._lock: + if self._loading is not None and self._loading.error is None: + raise RuntimeError("A video load is already in progress.") + self._load_token += 1 + token = self._load_token + self._cancel_event.clear() + self._loading = _VideoLoadingState(repo_id = repo_id, base_repo = fam.base_repo) + + threading.Thread( + target = self._run_load, + kwargs = dict( + repo_id = repo_id, + gguf_filename = gguf_filename, + base_repo = base_repo, + family_override = family_override, + hf_token = hf_token, + memory_mode = memory_mode, + speed_mode = speed_mode, + attention_backend = attention_backend, + transformer_cache = transformer_cache, + transformer_cache_threshold = transformer_cache_threshold, + model_kind = model_kind, + _load_token = token, + ), + daemon = True, + ).start() + return self.status() + + def _run_load(self, **kwargs: Any) -> None: + token = kwargs.get("_load_token") + try: + fam = detect_video_family(kwargs["repo_id"], kwargs.get("family_override")) + kind = resolve_video_model_kind(kwargs.get("gguf_filename"), kwargs.get("model_kind")) + base = ( + kwargs["repo_id"] + if kind == "pipeline" + else resolve_video_base_repo(fam, kwargs.get("base_repo")) + ) + kwargs["base_repo"] = base + expected = self._estimate_download_bytes( + kwargs["repo_id"], kwargs.get("gguf_filename"), base, kwargs.get("hf_token"), kind + ) + with self._lock: + 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 download inside from_pretrained (which resumes from the + # cache, so a cancelled pull costs nothing). + 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 + + hf_hub_download_with_xet_fallback( + kwargs["repo_id"], + kwargs["gguf_filename"], + kwargs.get("hf_token"), + cancel_event = self._cancel_event, + ) + 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 + if self._load_token != token: + return + logger.error("video.load_failed: %s", exc) + from utils.native_path_leases import redact_native_paths + + with self._lock: + if self._load_token == token and self._loading is not None: + self._loading.error = redact_native_paths(str(exc)) + + def _estimate_download_bytes( + self, + repo_id: str, + gguf_filename: Optional[str], + base: str, + hf_token: Optional[str], + kind: str, + ) -> Optional[int]: + """Total bytes this load will pull (checkpoint + companions), or None.""" + try: + from huggingface_hub import HfApi + + total = 0 + api = HfApi(token = hf_token or None) + if gguf_filename and not Path(repo_id).expanduser().exists(): + info = api.model_info(repo_id, files_metadata = True) + for sibling in info.siblings or []: + if sibling.rfilename == gguf_filename and sibling.size: + total += int(sibling.size) + if base and not Path(base).expanduser().exists(): + info = api.model_info(base, files_metadata = True) + for sibling in info.siblings or []: + name, size = sibling.rfilename, sibling.size or 0 + if not name.endswith((".safetensors", ".json", ".model", ".txt")): + continue + # A GGUF/single-file load replaces the base repo's DiT, and the + # LTX-2 base repo ships its text encoder TWICE (two shard + # namings); count only what from_pretrained will pull. + if kind != "pipeline" and name.startswith("transformer/"): + continue + if name.startswith("text_encoder/diffusion_pytorch_model"): + continue + total += int(size) + return total or None + except Exception: # noqa: BLE001 -- progress totals are best-effort only + return None + + def _cache_bytes(self, repo_id: Optional[str]) -> int: + """Bytes of ``repo_id`` currently in the HF blob cache (progress polling).""" + if not repo_id: + return 0 + try: + from huggingface_hub import scan_cache_dir + + rid = repo_id.strip() + for repo in scan_cache_dir().repos: + if repo.repo_id == rid: + return int(repo.size_on_disk) + except Exception: # noqa: BLE001 -- cache scan is best-effort + return 0 + return 0 + + def load_progress(self) -> dict[str, Any]: + """Phase + downloaded/total bytes for the in-flight load (cache-scan based).""" + loading = self._loading + if loading is not None and loading.error: + return _progress("error", error = loading.error) + if loading is None: + return _progress("ready" if self._state is not None else None) + downloaded = self._cache_bytes(loading.repo_id) + if loading.base_repo and loading.base_repo != loading.repo_id: + downloaded += self._cache_bytes(loading.base_repo) + expected = loading.expected_bytes + phase = "downloading" + if expected and downloaded >= expected: + phase = "finalizing" + return _progress( + phase, + downloaded_bytes = int(downloaded), + expected_bytes = int(expected) if expected else None, + ) + + # ── the load itself ────────────────────────────────────────────────────── + + def load_pipeline( + self, + repo_id: str, + *, + gguf_filename: Optional[str] = None, + base_repo: Optional[str] = None, + family_override: Optional[str] = None, + hf_token: Optional[str] = None, + memory_mode: Optional[str] = None, + speed_mode: Optional[str] = None, + attention_backend: Optional[str] = None, + transformer_cache: Optional[str] = None, + transformer_cache_threshold: Optional[float] = None, + model_kind: Optional[str] = None, + _load_token: Optional[int] = None, + ) -> dict[str, Any]: + import diffusers + import torch + + fam = self.validate_load_request( + repo_id, + gguf_filename = gguf_filename, + family_override = family_override, + model_kind = model_kind, + ) + kind = resolve_video_model_kind(gguf_filename, model_kind) + base = repo_id if kind == "pipeline" else resolve_video_base_repo(fam, base_repo) + + 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. + if self._active_generate_cancel is not None: + self._active_generate_cancel.set() + 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. + dtype = target.dtype + if fam.fp16_incompatible and dtype is torch.float16: + dtype = torch.float32 + + # ── 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 + ) + companion_mib = None + else: + checkpoint_path = self._resolve_checkpoint_path(repo_id, gguf_filename, hf_token) + size_mib = file_size_mib(str(checkpoint_path)) + model_dense_mib = None + if kind == "gguf": + transformer_mib = estimate_gguf_resident_mib(size_mib) + else: + transformer_mib = estimate_safetensors_dense_mib(size_mib) + companion_mib = ( + int((components[1] + components[2]) * mib_per_gb) + if components is not None + else None + ) + model_dense_mib = transformer_mib + runtime_mib = estimate_video_runtime_mib( + width = fam.resolution_presets[0][0], + height = fam.resolution_presets[0][1], + num_frames = fam.default_num_frames, + ) + plan = plan_diffusion_memory( + target = target, + device_memory = device_memory, + model_dense_mib = model_dense_mib, + runtime_headroom_mib = runtime_mib, + companion_dense_mib = companion_mib, + requested_mode = normalize_memory_mode(memory_mode), + ) + + # ── build the pipeline. + pipeline_cls = getattr(diffusers, fam.pipeline_class) + pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype} + if hf_token: + pipe_kwargs["token"] = hf_token + if kind == "pipeline": + pipe = pipeline_cls.from_pretrained(repo_id, **pipe_kwargs) + else: + transformer_cls = getattr(diffusers, fam.transformer_class) + checkpoint_path = self._resolve_checkpoint_path(repo_id, gguf_filename, hf_token) + sf_kwargs: dict[str, Any] = { + "torch_dtype": dtype, + "config": base, + "subfolder": "transformer", + "token": hf_token, + } + if kind == "gguf": + sf_kwargs["quantization_config"] = diffusers.GGUFQuantizationConfig( + compute_dtype = dtype + ) + transformer = transformer_cls.from_single_file(str(checkpoint_path), **sf_kwargs) + pipe = pipeline_cls.from_pretrained(base, transformer = transformer, **pipe_kwargs) + + 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.") + + # ── optimisation layers, in the image backend's order: speed profile + # (compile must precede placement), attention (compile traces it), + # placement/offload, then the step cache. + effective_speed = resolve_speed_mode(speed_mode, is_gguf = kind == "gguf") + backend_flags = snapshot_backend_flags() + attention_engaged = apply_attention_backend( + pipe, + select_attention_backend( + target, attention_backend, speed_active = effective_speed != SPEED_OFF + ), + logger = logger, + ) + speed_optims = apply_speed_optims( + pipe, + target, + is_gguf = kind == "gguf", + family = fam, + speed_mode = effective_speed, + offload_active = plan.offload_policy != "none", + ) + offload_policy, vae_tiling = apply_memory_plan( + pipe, plan, device = device, logger = logger + ) + 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. + try: + pipe.vae.enable_tiling() + vae_tiling = True + except Exception as exc: # noqa: BLE001 -- tiling is an optimisation only + logger.warning("video.vae_tiling_failed: %s", exc) + cache_engaged = apply_step_cache( + pipe, + mode = normalize_transformer_cache(transformer_cache), + threshold = transformer_cache_threshold, + logger = logger, + ) + + resolved = build_resolved_record( + { + "memory_mode": ( + memory_mode, + plan.requested_mode, + f"planned '{plan.offload_policy}' offload from the family size table", + ), + "speed_mode": ( + speed_mode, + effective_speed, + "GGUF video loads default to the near-lossless compile profile", + ), + "attention_backend": ( + attention_backend, + attention_engaged or "native", + "cuDNN fused attention on NVIDIA when a speed profile is active", + ), + "transformer_cache": ( + transformer_cache, + cache_engaged or "off", + "step cache engages on many-step schedules only", + ), + } + ) + + with self._lock: + 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.") + self._state = _VideoLoadState( + pipe = pipe, + family = fam, + repo_id = repo_id, + base_repo = base, + device = device, + dtype = str(dtype).replace("torch.", ""), + kind = kind, + gguf_filename = gguf_filename, + offload_policy = offload_policy, + vae_tiling = vae_tiling, + memory_mode = plan.requested_mode, + speed_mode = effective_speed, + speed_optims = tuple(speed_optims or ()), + backend_flags = backend_flags, + attention_backend = attention_engaged, + transformer_cache = cache_engaged, + resolved = resolved, + ) + logger.info( + "video.loaded: %s (%s, %s, offload=%s, speed=%s)", + repo_id, fam.name, kind, offload_policy, effective_speed, + ) + return self.status() + + @staticmethod + def _resolve_checkpoint_path( + repo_id: str, gguf_filename: Optional[str], hf_token: Optional[str] + ) -> Path: + """The local checkpoint file for a gguf/single_file load (downloads if hub).""" + from .diffusion_families import resolve_local_gguf_child + + root = Path(repo_id).expanduser() + if root.is_dir(): + return resolve_local_gguf_child(root, gguf_filename or "") + if root.is_file(): + return root + from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback + + return Path( + hf_hub_download_with_xet_fallback(repo_id, gguf_filename or "", hf_token) + ) + + # ── generation ─────────────────────────────────────────────────────────── + + def generate( + self, + *, + prompt: str, + negative_prompt: Optional[str] = None, + width: Optional[int] = None, + height: Optional[int] = None, + num_frames: Optional[int] = None, + fps: Optional[int] = None, + steps: Optional[int] = None, + guidance: Optional[float] = None, + seed: Optional[int] = None, + ) -> dict[str, Any]: + import torch + + cancel = threading.Event() + with self._generate_lock: + with self._lock: + state = self._state + if state is None: + raise RuntimeError(VIDEO_NOT_LOADED_MSG) + self._active_generate_cancel = cancel + try: + fam = state.family + width, height = snap_video_size( + fam, width or fam.resolution_presets[0][0], + height or fam.resolution_presets[0][1], + ) + frames = snap_num_frames(fam, num_frames or fam.default_num_frames) + out_fps = int(fps or fam.default_fps) + default_steps, default_guidance = default_video_generation_params( + state.gguf_filename, state.repo_id, state.base_repo + ) + steps = int(steps or default_steps) + guidance = float(default_guidance if guidance is None else guidance) + + generator = torch.Generator(device = state.device) + if seed is None: + seed = int(generator.seed()) % (2**53) + generator = generator.manual_seed(int(seed)) + + pipe = state.pipe + call_params = inspect.signature(pipe.__call__).parameters + kwargs: dict[str, Any] = { + "prompt": prompt, + "num_inference_steps": steps, + fam.cfg_kwarg: guidance, + "width": width, + "height": height, + "num_frames": frames, + "generator": generator, + } + 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. + if "frame_rate" in call_params: + kwargs["frame_rate"] = float(out_fps) + + started = time.monotonic() + self._gen = { + "active": True, "phase": "denoise", "step": 0, "total": steps, + "started": started, "eta_seconds": None, "error": None, + } + + def _on_step(p, step_index, timestep, callback_kwargs): + if cancel.is_set(): + p._interrupt = True + return callback_kwargs + done = step_index + 1 + elapsed = time.monotonic() - started + self._gen.update( + step = done, + eta_seconds = (elapsed / max(1, done)) * max(0, steps - done), + ) + return callback_kwargs + + if "callback_on_step_end" in call_params: + kwargs["callback_on_step_end"] = _on_step + + with torch.inference_mode(): + output = pipe(**kwargs) + if cancel.is_set(): + raise RuntimeError(VIDEO_CANCELLED_MSG) + + self._gen.update(phase = "export", eta_seconds = None) + video_frames = output.frames[0] + audio = getattr(output, "audio", None) + audio_track = audio[0] if fam.has_audio and audio is not None else None + mp4_bytes = self._encode_mp4( + video_frames, out_fps, audio_track, pipe if fam.has_audio else None + ) + duration_s = len(video_frames) / float(out_fps) if out_fps else 0.0 + self._gen = {"active": False} + return { + "mp4_bytes": mp4_bytes, + "seed": int(seed), + "repo_id": state.repo_id, + "width": width, + "height": height, + "num_frames": len(video_frames), + "fps": out_fps, + "duration_s": duration_s, + "has_audio": bool(audio_track is not None), + "steps": steps, + "guidance": guidance, + } + except Exception: + self._gen = {"active": False} + raise + finally: + with self._lock: + if self._active_generate_cancel is cancel: + self._active_generate_cancel = None + + @staticmethod + def _encode_mp4(video_frames, fps: int, audio, pipe) -> bytes: + """Encode frames (+ optional audio) to H.264 MP4 bytes via diffusers' PyAV + exporter. A temp file bridges the exporter's path-based API; the bytes are + what the gallery persists.""" + from diffusers.utils.export_utils import encode_video + + tmp = tempfile.NamedTemporaryFile(suffix = ".mp4", delete = False) + tmp.close() + try: + encode_kwargs: dict[str, Any] = {} + if audio is not None and pipe is not None: + encode_kwargs["audio"] = audio + sample_rate = getattr( + getattr(getattr(pipe, "vocoder", None), "config", None), + "output_sampling_rate", + None, + ) + if sample_rate: + encode_kwargs["audio_sample_rate"] = int(sample_rate) + encode_video(video_frames, fps, tmp.name, **encode_kwargs) + return Path(tmp.name).read_bytes() + finally: + try: + os.unlink(tmp.name) + except OSError: + pass + + def generate_progress(self) -> dict[str, Any]: + gen = dict(self._gen) + gen.setdefault("active", False) + return gen + + def cancel_generate(self) -> bool: + """Signal the in-flight generation to stop at its next step callback.""" + with self._lock: + cancel = self._active_generate_cancel + if cancel is None: + return False + cancel.set() + return True + + # ── teardown + status ──────────────────────────────────────────────────── + + def _teardown_state(self) -> None: + state = None + with self._lock: + state, self._state = self._state, None + if state is not None: + restore_backend_flags(state.backend_flags) + del state + clear_gpu_cache() + + def unload(self) -> dict[str, Any]: + with self._lock: + self._load_token += 1 + self._cancel_event.set() + self._loading = None + if self._active_generate_cancel is not None: + self._active_generate_cancel.set() + self._teardown_state() + logger.info("video.unloaded") + return self.status() + + def status(self) -> dict[str, Any]: + state = self._state + if state is None: + return { + "loaded": False, + "repo_id": None, + "family": None, + "base_repo": None, + "device": None, + "dtype": None, + "model_kind": None, + "offload_policy": None, + "vae_tiling": False, + "memory_mode": None, + "speed_mode": None, + "speed_optims": [], + "attention_backend": None, + "transformer_cache": None, + "has_audio": False, + "defaults": None, + "resolved": None, + } + fam = state.family + default_steps, default_guidance = default_video_generation_params( + state.gguf_filename, state.repo_id, state.base_repo + ) + return { + "loaded": True, + "repo_id": state.repo_id, + "family": fam.name, + "base_repo": state.base_repo, + "device": state.device, + "dtype": state.dtype, + "model_kind": state.kind, + "offload_policy": state.offload_policy, + "vae_tiling": state.vae_tiling, + "memory_mode": state.memory_mode, + "speed_mode": state.speed_mode, + "speed_optims": list(state.speed_optims), + "attention_backend": state.attention_backend, + "transformer_cache": state.transformer_cache, + "has_audio": fam.has_audio, + "defaults": { + "steps": default_steps, + "guidance": default_guidance, + "num_frames": fam.default_num_frames, + "fps": fam.default_fps, + "frame_step": fam.frame_step, + "resolution_multiple": fam.resolution_multiple, + "resolution_presets": [list(p) for p in fam.resolution_presets], + }, + "resolved": state.resolved, + } + + +_backend: Optional[VideoBackend] = None +_backend_lock = threading.Lock() + + +def get_video_backend() -> VideoBackend: + global _backend + with _backend_lock: + if _backend is None: + _backend = VideoBackend() + return _backend diff --git a/studio/backend/core/inference/video_families.py b/studio/backend/core/inference/video_families.py new file mode 100644 index 0000000000..c941b6ce29 --- /dev/null +++ b/studio/backend/core/inference/video_families.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Pure helpers for text-to-video model identification. + +The video registry mirrors ``diffusion_families`` (no torch/diffusers imports, so +everything unit-tests without the heavy runtime) but is a SEPARATE registry with a +separate backend: video pipelines take frame/fps arguments, return frame stacks +(and, for LTX-2, synchronized audio) instead of PIL images, and their artifacts are +MP4s. Keeping the registries apart means neither picker can mis-route a checkpoint +to the wrong engine. + +A video checkpoint published as a single-file GGUF only carries the DiT weights; +the VAE / text encoder / connectors / vocoder come from the companion diffusers +base repo, exactly like the image GGUF path. +""" + +from __future__ import annotations + +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. +VIDEO_NOT_LOADED_MSG = "No video model is loaded." +VIDEO_CANCELLED_MSG = "Video generation was cancelled." + + +@dataclass(frozen = True) +class VideoFamily: + name: str + pipeline_class: str + transformer_class: str + base_repo: str + # Pipeline kwarg carrying the guidance value. + cfg_kwarg: str = "guidance_scale" + # The pipe attribute holding the denoiser (all current video families are DiTs). + 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. + 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. + transformer2_class: Optional[str] = None + is_moe: bool = False + cfg2_kwarg: Optional[str] = None + # 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. + default_steps: int = 40 + default_guidance: float = 4.0 + default_num_frames: int = 121 + default_fps: int = 24 + 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. + 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. + 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). + 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). + fp16_incompatible: bool = True + # 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). + VideoFamily( + name = "ltx-2", + pipeline_class = "LTX2Pipeline", + transformer_class = "LTX2VideoTransformer3DModel", + base_repo = "Lightricks/LTX-2", + aliases = ("ltx-2.3", "ltx2", "ltx-video", "ltxv", "ltx"), + has_audio = True, + default_steps = 40, + default_guidance = 4.0, + default_num_frames = 121, + 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. + 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). + bf16_components_gb = (37.8, 50.4, 5.5), + gguf_repo = "unsloth/LTX-2.3-GGUF", + ), +) + + +def _token_in_needle(token: str, needle: str) -> bool: + """Whole path/name segment match, as in diffusion_families (a short alias like + 'ltx' must not match inside an unrelated word).""" + return re.search(r"(?:^|[-_./\\])" + re.escape(token) + r"(?:$|[-_./\\])", needle) is not None + + +def detect_video_family(repo_id: str, override: Optional[str] = None) -> Optional[VideoFamily]: + """Resolve a ``VideoFamily`` from a repo id, or an explicit override. + + Same contract as ``diffusion_families.detect_family``: an override matches a + name/alias exactly; otherwise the longest name/alias appearing as a whole + segment of the repo id wins. + """ + if override: + key = override.strip().lower() + for fam in _FAMILIES: + if key == fam.name or key in fam.aliases: + return fam + return None + needle = repo_id.lower() + best: Optional[tuple[VideoFamily, int]] = None + for fam in _FAMILIES: + for token in (fam.name, *fam.aliases): + if _token_in_needle(token, needle) and (best is None or len(token) > best[1]): + best = (fam, len(token)) + return best[0] if best else None + + +def supported_video_family_names() -> tuple[str, ...]: + return tuple(fam.name for fam in _FAMILIES) + + +def resolve_video_base_repo(fam: VideoFamily, base_repo: Optional[str]) -> str: + """The companion diffusers repo: caller-supplied if given, else the family fallback.""" + base = (base_repo or "").strip() + return base or fam.base_repo + + +def snap_num_frames(fam: VideoFamily, num_frames: int) -> int: + """The nearest valid frame count at or below the request (k * frame_step + 1). + + Video latents are allocated as (num_frames - 1) / temporal_compression + 1, so + an off-lattice count wastes a partial latent frame at best and trips shape + checks at worst; snapping mirrors the image path's silent /16 size snap. + """ + step = max(1, fam.frame_step) + return max(1, ((max(1, num_frames) - 1) // step) * step + 1) + + +def snap_video_size(fam: VideoFamily, width: int, height: int) -> tuple[int, int]: + """Width/height floored to the family's required multiple (minimum one unit).""" + multiple = max(1, fam.resolution_multiple) + snap = lambda v: max(multiple, (max(1, v) // multiple) * multiple) # noqa: E731 + 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. +_VIDEO_GENERATION_DEFAULTS: tuple[tuple[str, int, float], ...] = ( + ("distilled", 8, 1.0), + ("ltx", 40, 4.0), +) + + +def default_video_generation_params(*identifiers: Optional[str]) -> tuple[int, float]: + """Default ``(steps, guidance)`` for a loaded video model; the first identifier + naming a known variant wins, so a GGUF filename ('...distilled...Q4_K_M.gguf') + beats the family base repo.""" + for identifier in identifiers: + needle = (identifier or "").lower() + for key, steps, guidance in _VIDEO_GENERATION_DEFAULTS: + if key in needle: + return steps, guidance + return 40, 4.0 diff --git a/studio/backend/core/inference/video_gallery.py b/studio/backend/core/inference/video_gallery.py new file mode 100644 index 0000000000..06753c5384 --- /dev/null +++ b/studio/backend/core/inference/video_gallery.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Disk-backed persistence for generated videos. + +Each video is an MP4 under ``studio_root()/videos``. Unlike a PNG, an MP4 has no +portable text-chunk we can embed the generation recipe into, so every video is +stored as a pair: ``{id}.mp4`` holds the encoded bytes and ``{id}.json`` holds +the full recipe as a UTF-8 JSON sidecar (the source of truth the gallery reads +back). The pair travels together on disk; a lone file is not a valid record. + +The gallery is intentionally dumb storage: the route owns the metadata schema +and passes a plain dict; this module only writes/reads/sorts files. +""" + +from __future__ import annotations + +import json +import os +import re +import uuid +from pathlib import Path +from typing import Any, Optional + +from loggers import get_logger +from utils.paths import ensure_dir, studio_root + +logger = get_logger(__name__) + +# Video ids are file stems; restrict to filename-safe chars so a crafted id +# can't escape the gallery directory. +_ID_RE = re.compile(r"^[A-Za-z0-9_-]{1,128}$") + + +def gallery_dir() -> Path: + return ensure_dir(studio_root() / "videos") + + +def save(mp4_bytes: bytes, meta: dict[str, Any]) -> dict[str, Any]: + """Persist encoded MP4 bytes plus their recipe sidecar; return the record.""" + video_id = uuid.uuid4().hex + directory = gallery_dir() + (directory / f"{video_id}.mp4").write_bytes(mp4_bytes) + # Write the sidecar via a tmp file + os.replace so a reader never sees a + # half-written recipe (an incomplete json would make the pair unlistable). + sidecar = directory / f"{video_id}.json" + tmp = directory / f"{video_id}.json.tmp" + tmp.write_text(json.dumps(meta), encoding = "utf-8") + os.replace(tmp, sidecar) + return _record(video_id, meta) + + +def _record(video_id: str, meta: dict[str, Any]) -> dict[str, Any]: + return { + **meta, + "id": video_id, + "url": f"/api/inference/video/gallery/{video_id}/file", + } + + +def video_path(video_id: str) -> Optional[Path]: + """Resolve an id to its on-disk MP4, or None if missing / unsafe.""" + if not _ID_RE.match(video_id): + return None + path = gallery_dir() / f"{video_id}.mp4" + # Defence in depth: confirm the resolved path is still inside the gallery. + try: + path.resolve().relative_to(gallery_dir().resolve()) + except ValueError: + return None + return path if path.is_file() else None + + +def _sidecar_path(video_id: str) -> Path: + return gallery_dir() / f"{video_id}.json" + + +def _read_meta(sidecar: Path) -> Optional[dict[str, Any]]: + try: + raw = sidecar.read_text(encoding = "utf-8") + except OSError: + return None + try: + meta = json.loads(raw) + except (ValueError, TypeError): + return None + return meta if isinstance(meta, dict) else None + + +def _mtime(path: Path) -> float: + try: + return path.stat().st_mtime + except OSError: + return 0.0 + + +def list_videos(limit: Optional[int] = None, offset: int = 0) -> list[dict[str, Any]]: + """A newest-first window of videos for infinite scroll. + + Ordered by MP4 mtime (a cheap stat, ~= generation order) so a months-old + gallery isn't opened in full just to sort it; only the window's sidecars are + read. limit=None returns everything from ``offset`` on. A file without its + pair (an MP4 with no readable json sidecar, or a sidecar with no MP4) is not + a valid record and is skipped.""" + try: + paths = list(gallery_dir().glob("*.mp4")) + except OSError: + return [] + paths.sort(key = _mtime, reverse = True) + # Page over READABLE records, not raw files: filtering an orphan MP4 out of an + # already-sliced window would drop valid videos that sort after it and make the + # route's has_more wrong. Read only as far as needed to fill the requested window. + want = None if limit is None else offset + limit + records = [] + for path in paths: + meta = _read_meta(_sidecar_path(path.stem)) + if meta is None: # no readable sidecar (orphan mp4) — skip + continue + records.append(_record(path.stem, meta)) + if want is not None and len(records) >= want: + break + return records[offset:] if limit is None else records[offset : offset + limit] + + +def delete(video_id: str) -> bool: + """Remove both files of a pair; True if the MP4 existed.""" + path = video_path(video_id) + if path is None: + return False + # Best-effort on the sidecar: a leftover json without its mp4 is skipped by + # list_videos anyway, so a failed sidecar unlink must not fail the delete. + try: + _sidecar_path(video_id).unlink() + except OSError: + pass + try: + path.unlink() + return True + except OSError as exc: + logger.warning("video_gallery.delete_failed: %s", exc) + return False + + +def clear() -> int: + """Delete every gallery pair; return how many videos were removed.""" + removed = 0 + try: + paths = list(gallery_dir().glob("*.mp4")) + except OSError: + return 0 + for path in paths: + # Drop the sidecar first; an orphaned json is harmless (skipped by list). + try: + _sidecar_path(path.stem).unlink() + except OSError: + pass + try: + path.unlink() + removed += 1 + except OSError: + continue + return removed diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 2a7bc9f47a..4c28554562 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -22,6 +22,11 @@ fastmcp>=3.0.2 # the gguf reader. diffusers itself is torch-bound and comes via the ML stack; # gguf is torch-free so it lives here, where it's installed in every mode. gguf +# Local video generation (Video page): MP4 encode (H.264 + AAC mux for LTX-2's +# synchronized audio) runs through PyAV, which bundles its own ffmpeg libs. The +# video load path fails fast with a clear error when this is missing, so keep it +# in the base studio set rather than an extra. +av>=12.0.0 # RAG (knowledge bases, hybrid retrieval). sentence-transformers lives in # extras-no-deps.txt; these add the lexical+dense store and document parsing. sqlite-vec==0.1.9 diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py new file mode 100644 index 0000000000..1819ec80f4 --- /dev/null +++ b/studio/backend/tests/test_video_backend.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""VideoBackend lifecycle on a faked torch/diffusers runtime (CPU-only, offline). +Mirrors test_diffusion_backend's fake_runtime pattern: explicit fake signatures so +the signature-gated kwargs actually exercise, sys.modules stubs so no real ML +stack loads.""" + +import contextlib +import sys +import types + +import pytest + +from core.inference.video import VideoBackend, get_video_backend, resolve_video_model_kind +from core.inference.video_families import VIDEO_NOT_LOADED_MSG + + +class _FakeDtype: + def __init__(self, name: str) -> None: + self._name = name + + def __repr__(self) -> str: + return f"torch.{self._name}" + + __str__ = __repr__ + + +class _FakeGenerator: + def __init__(self, device = None) -> None: + self.device = device + self.manual = None + + def seed(self) -> int: + return 4242 + + def manual_seed(self, value: int): + self.manual = value + return self + + +class _FakeVae: + def __init__(self) -> None: + self.tiled = False + + def enable_tiling(self) -> None: + self.tiled = True + + +class _FakePipe: + def __init__(self) -> None: + self.moved_to = None + self.vae = _FakeVae() + self.last_kwargs = None + self._interrupt = False + + def to(self, device): + self.moved_to = device + return self + + def enable_vae_tiling(self) -> None: + self.vae.tiled = True + + # Explicit signature so generate()'s signature-gated kwargs (negative_prompt, + # frame_rate, callback) actually engage; **kwargs would defeat the gates. + def __call__( + self, + *, + prompt = None, + negative_prompt = None, + num_inference_steps = None, + guidance_scale = None, + width = None, + height = None, + num_frames = None, + frame_rate = None, + generator = None, + callback_on_step_end = None, + **kwargs, + ): + self.last_kwargs = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + "width": width, + "height": height, + "num_frames": num_frames, + "frame_rate": frame_rate, + **kwargs, + } + if callback_on_step_end is not None: + for step in range(int(num_inference_steps or 1)): + callback_on_step_end(self, step, 0, {}) + if self._interrupt: + break + frames = [[object() for _ in range(int(num_frames or 1))]] + return types.SimpleNamespace(frames = frames, audio = None) + + +class _FakePipeline: + last: dict = {} + + @classmethod + def from_pretrained(cls, base, **kwargs): + _FakePipeline.last = {"base": base, **kwargs} + return _FakePipe() + + +class _FakeTransformer: + last: dict = {} + + @classmethod + def from_single_file(cls, path, **kwargs): + _FakeTransformer.last = {"path": path, **kwargs} + return object() + + +@pytest.fixture +def fake_runtime(monkeypatch): + torch = types.ModuleType("torch") + torch.bfloat16 = _FakeDtype("bfloat16") + torch.float16 = _FakeDtype("float16") + torch.float32 = _FakeDtype("float32") + torch.Generator = _FakeGenerator + torch.cuda = types.SimpleNamespace(is_available = lambda: False) + torch.backends = types.SimpleNamespace(mps = None) + torch.inference_mode = lambda: contextlib.nullcontext() + + diffusers = types.ModuleType("diffusers") + diffusers.GGUFQuantizationConfig = lambda compute_dtype = None: ("quant", compute_dtype) + diffusers.LTX2Pipeline = _FakePipeline + diffusers.LTX2VideoTransformer3DModel = _FakeTransformer + + monkeypatch.setitem(sys.modules, "torch", torch) + monkeypatch.setitem(sys.modules, "diffusers", diffusers) + monkeypatch.setattr("core.inference.video.clear_gpu_cache", lambda: None) + # MP4 encode needs real frames + PyAV; the backend contract under test is the + # byte handoff, so stub the encoder. + monkeypatch.setattr( + VideoBackend, "_encode_mp4", staticmethod(lambda frames, fps, audio, pipe: b"MP4") + ) + _FakePipeline.last = {} + _FakeTransformer.last = {} + yield + + +def _load_gguf(backend, tmp_path): + (tmp_path / "model.gguf").write_bytes(b"weights") + return backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "Lightricks/LTX-2", + family_override = "ltx-2", + ) + + +def test_resolve_kind(): + assert resolve_video_model_kind("x.gguf", None) == "gguf" + assert resolve_video_model_kind("x.safetensors", None) == "single_file" + assert resolve_video_model_kind(None, None) == "pipeline" + with pytest.raises(ValueError): + resolve_video_model_kind(None, "bogus") + + +def test_validate_rejects_unknown_and_untrusted(): + backend = VideoBackend() + with pytest.raises(ValueError, match = "not a supported"): + backend.validate_load_request("someorg/some-image-model") + # A known family but an untrusted repo id must not open from_pretrained. + with pytest.raises(ValueError, match = "limited to"): + backend.validate_load_request("evil/ltx-2-repack") + # GGUF loads stay open to any repo (single-file read, no pickle). + fam = backend.validate_load_request( + "anyorg/ltx-2-GGUF", gguf_filename = "x.gguf", model_kind = "gguf" + ) + assert fam.name == "ltx-2" + with pytest.raises(ValueError, match = "filename"): + backend.validate_load_request("unsloth/LTX-2.3-GGUF", model_kind = "gguf") + + +def test_load_generate_unload_gguf(fake_runtime, tmp_path): + backend = VideoBackend() + status = _load_gguf(backend, tmp_path) + assert status["loaded"] is True and status["family"] == "ltx-2" + assert status["model_kind"] == "gguf" + assert status["has_audio"] is True + # The GGUF transformer is dequant-configured and assembled onto the base repo. + assert _FakeTransformer.last["path"].endswith("model.gguf") + assert _FakeTransformer.last["quantization_config"][0] == "quant" + assert _FakePipeline.last["base"] == "Lightricks/LTX-2" + assert "transformer" in _FakePipeline.last + # Video decode is the memory peak: tiling is always on. + assert status["vae_tiling"] is True + assert status["defaults"]["frame_step"] == 8 + + result = backend.generate( + prompt = "a sloth surfing", width = 1000, height = 700, num_frames = 120, fps = 24 + ) + call = backend._state.pipe.last_kwargs + # Shape snapping happened BEFORE the pipe call: /32 sizes, 8k+1 frames. + assert (call["width"], call["height"]) == (992, 672) + assert call["num_frames"] == 113 + assert call["frame_rate"] == 24.0 + assert result["mp4_bytes"] == b"MP4" + assert result["num_frames"] == 113 and result["fps"] == 24 + assert result["has_audio"] is False # fake pipe returned no audio track + assert 0 <= result["seed"] < 2**53 + + status = backend.unload() + assert status["loaded"] is False + + +def test_generate_defaults_from_variant(fake_runtime, tmp_path): + # A distilled GGUF pick defaults to the few-step no-CFG schedule. + (tmp_path / "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf").write_bytes(b"w") + backend = VideoBackend() + backend.load_pipeline( + str(tmp_path), + gguf_filename = "ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf", + base_repo = "Lightricks/LTX-2", + family_override = "ltx-2", + ) + backend.generate(prompt = "a sloth") + call = backend._state.pipe.last_kwargs + assert call["num_inference_steps"] == 8 + assert call["guidance_scale"] == 1.0 + + +def test_generate_without_load_raises(fake_runtime): + backend = VideoBackend() + with pytest.raises(RuntimeError, match = VIDEO_NOT_LOADED_MSG): + backend.generate(prompt = "x") + + +def test_generate_progress_and_cancel_idle(fake_runtime): + backend = VideoBackend() + assert backend.generate_progress() == {"active": False} + assert backend.cancel_generate() is False + + +def test_singleton(): + assert get_video_backend() is get_video_backend() diff --git a/studio/backend/tests/test_video_families.py b/studio/backend/tests/test_video_families.py new file mode 100644 index 0000000000..852a9f5bb3 --- /dev/null +++ b/studio/backend/tests/test_video_families.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Video family registry: detection, shape snapping, generation defaults. +Pure-module tests: no torch, no network.""" + +import pytest + +from core.inference.video_families import ( + VIDEO_CANCELLED_MSG, + VIDEO_NOT_LOADED_MSG, + default_video_generation_params, + detect_video_family, + resolve_video_base_repo, + snap_num_frames, + snap_video_size, + supported_video_family_names, +) + + +@pytest.mark.parametrize( + "repo_id", + [ + "unsloth/LTX-2.3-GGUF", + "Lightricks/LTX-2", + "Lightricks/LTX-2.3-fp8", + "lightricks/ltx-2.3", + "some/dir/ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf", + ], +) +def test_detect_ltx2(repo_id): + fam = detect_video_family(repo_id) + assert fam is not None and fam.name == "ltx-2" + assert fam.pipeline_class == "LTX2Pipeline" + assert fam.has_audio is True + + +def test_detect_override_and_unknown(): + assert detect_video_family("x", override = "ltx-2").name == "ltx-2" + assert detect_video_family("x", override = "ltx2").name == "ltx-2" + assert detect_video_family("x", override = "nope") is None + assert detect_video_family("Wan-AI/Wan2.2-TI2V-5B-Diffusers") is None # V3 family + # A short alias must not match inside an unrelated word. + assert detect_video_family("someorg/deluxtreme-model") is None + + +def test_sentinels_are_video_specific(): + # The routes match these EXACTLY for 409s; they must not collide with the + # image sentinels or a video 409 would be mis-attributed. + assert "video" in VIDEO_NOT_LOADED_MSG.lower() + assert "video" in VIDEO_CANCELLED_MSG.lower() + + +def test_resolve_base_repo(): + fam = detect_video_family("unsloth/LTX-2.3-GGUF") + assert resolve_video_base_repo(fam, None) == "Lightricks/LTX-2" + assert resolve_video_base_repo(fam, " ") == "Lightricks/LTX-2" + assert resolve_video_base_repo(fam, "other/base") == "other/base" + + +def test_snap_num_frames_lattice(): + fam = detect_video_family("unsloth/LTX-2.3-GGUF") + # Valid counts are k * 8 + 1: on-lattice values pass through, everything + # else floors to the previous lattice point, never below 1. + assert snap_num_frames(fam, 121) == 121 + assert snap_num_frames(fam, 120) == 113 + assert snap_num_frames(fam, 122) == 121 + assert snap_num_frames(fam, 1) == 1 + assert snap_num_frames(fam, 0) == 1 + assert snap_num_frames(fam, 9) == 9 + + +def test_snap_video_size_multiple(): + fam = detect_video_family("unsloth/LTX-2.3-GGUF") + assert snap_video_size(fam, 768, 512) == (768, 512) + assert snap_video_size(fam, 1000, 700) == (992, 672) + # Never snaps to zero: the floor is one multiple. + assert snap_video_size(fam, 1, 1) == (32, 32) + + +def test_generation_defaults_distilled_vs_dev(): + # The distilled checkpoints run few-step with CFG off; the dev-config base + # repo wants the full schedule. The picked filename wins over the base repo. + assert default_video_generation_params( + "distilled-1.1/ltx-2.3-22b-distilled-1.1-Q4_K_M.gguf", "Lightricks/LTX-2" + ) == (8, 1.0) + assert default_video_generation_params(None, "Lightricks/LTX-2") == (40, 4.0) + assert default_video_generation_params("unknown/thing") == (40, 4.0) + + +def test_supported_names(): + assert supported_video_family_names() == ("ltx-2",) + + +def test_family_size_table_present(): + fam = detect_video_family("unsloth/LTX-2.3-GGUF") + assert fam.bf16_components_gb is not None + transformer_gb, text_encoder_gb, companions_gb = fam.bf16_components_gb + # The Gemma3-27B text encoder outweighs the DiT itself; a table that lost + # that would let auto planning under-reserve by ~50 GB. + assert text_encoder_gb > transformer_gb > 20.0 + assert companions_gb > 0.0 diff --git a/studio/backend/tests/test_video_gallery.py b/studio/backend/tests/test_video_gallery.py new file mode 100644 index 0000000000..7248ec2c63 --- /dev/null +++ b/studio/backend/tests/test_video_gallery.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the disk-backed video gallery: MP4 + JSON-sidecar round-trips, +listing order, safe id handling, orphan-pair skipping, and delete/clear.""" + +from __future__ import annotations + +import json +import os + +import core.inference.video_gallery as gallery + + +import pytest + + +@pytest.fixture(autouse = True) +def _tmp_gallery(monkeypatch, tmp_path): + # Point the gallery at a throwaway root instead of ~/.unsloth/studio. + monkeypatch.setattr(gallery, "studio_root", lambda: tmp_path) + + +def _mp4(tag = b"\x00\x00\x00\x18ftypmp42"): + # Not a real container; the gallery treats the bytes as opaque payload. + return tag + + +def _meta(**over): + base = { + "prompt": "a sloth surfing", + "negative_prompt": None, + "width": 1024, + "height": 576, + "num_frames": 49, + "steps": 30, + "guidance": 6.0, + "seed": 7, + "model": "unsloth/some-video-model", + "created_at": 100.0, + } + base.update(over) + return base + + +def test_save_writes_pair_and_round_trips(): + record = gallery.save(_mp4(), _meta()) + assert record["id"] and record["url"].endswith(f"{record['id']}/file") + + # Both files of the pair exist: the mp4 payload and the json recipe sidecar. + directory = gallery.gallery_dir() + assert (directory / f"{record['id']}.mp4").is_file() + sidecar = directory / f"{record['id']}.json" + assert json.loads(sidecar.read_text(encoding = "utf-8"))["prompt"] == "a sloth surfing" + + listed = gallery.list_videos() + assert len(listed) == 1 + assert listed[0]["prompt"] == "a sloth surfing" and listed[0]["seed"] == 7 + # Meta fields survive the sidecar round-trip untouched. + assert listed[0]["num_frames"] == 49 and listed[0]["model"] == "unsloth/some-video-model" + + +def test_url_shape(): + record = gallery.save(_mp4(), _meta()) + assert record["url"] == f"/api/inference/video/gallery/{record['id']}/file" + + +def _save_with_mtime(prompt: str, t: float) -> dict: + record = gallery.save(_mp4(), _meta(prompt = prompt, created_at = t)) + # Listing orders by mp4 mtime; set it explicitly so a tight test loop can't tie it. + os.utime(gallery.gallery_dir() / f"{record['id']}.mp4", (t, t)) + return record + + +def test_list_is_newest_first(): + old = _save_with_mtime("old", 100.0) + new = _save_with_mtime("new", 200.0) + assert [r["id"] for r in gallery.list_videos()] == [new["id"], old["id"]] + + +def test_list_paginates_with_limit_offset(): + # 5 videos, newest (t=4) first. + for i in range(5): + _save_with_mtime(f"p{i}", float(i)) + page1 = gallery.list_videos(limit = 2, offset = 0) + page2 = gallery.list_videos(limit = 2, offset = 2) + assert [r["prompt"] for r in page1] == ["p4", "p3"] + assert [r["prompt"] for r in page2] == ["p2", "p1"] + # limit=None still returns everything from the offset. + assert len(gallery.list_videos()) == 5 + assert len(gallery.list_videos(offset = 4)) == 1 + + +def test_video_path_rejects_unsafe_ids(): + # Traversal / bad chars / absolute paths never resolve to a path. + assert gallery.video_path("../../etc/passwd") is None + assert gallery.video_path("/etc/passwd") is None + assert gallery.video_path("a/b") is None + assert gallery.video_path("missing") is None + + +def test_video_path_returns_mp4_for_saved_id(): + record = gallery.save(_mp4(), _meta()) + path = gallery.video_path(record["id"]) + assert path is not None and path.name == f"{record['id']}.mp4" + + +def test_delete_removes_both_files(): + record = gallery.save(_mp4(), _meta(prompt = "a")) + gallery.save(_mp4(), _meta(prompt = "b")) + directory = gallery.gallery_dir() + assert gallery.delete(record["id"]) is True + # Both halves of the pair are gone. + assert not (directory / f"{record['id']}.mp4").exists() + assert not (directory / f"{record['id']}.json").exists() + assert gallery.delete(record["id"]) is False # already gone + assert len(gallery.list_videos()) == 1 + + +def test_clear_returns_count(): + gallery.save(_mp4(), _meta(prompt = "a")) + gallery.save(_mp4(), _meta(prompt = "b")) + assert gallery.clear() == 2 + assert gallery.list_videos() == [] + # No stray sidecars left behind after a clear. + assert list(gallery.gallery_dir().glob("*.json")) == [] + + +def test_list_skips_orphan_mp4_without_sidecar(): + # An MP4 with no readable json sidecar (a hand-dropped file) is not a record. + orphan = gallery.gallery_dir() / "orphan.mp4" + orphan.write_bytes(_mp4()) + gallery.save(_mp4(), _meta(prompt = "ours")) + listed = gallery.list_videos() + assert [r["prompt"] for r in listed] == ["ours"] + + +def test_list_skips_orphan_sidecar_without_mp4(): + # A json sidecar with no MP4 alongside it is never surfaced (listing globs mp4s). + orphan = gallery.gallery_dir() / "lonely.json" + orphan.write_text(json.dumps(_meta(prompt = "no video")), encoding = "utf-8") + gallery.save(_mp4(), _meta(prompt = "ours")) + listed = gallery.list_videos() + assert [r["prompt"] for r in listed] == ["ours"] + + +def test_orphan_mp4_in_window_does_not_drop_valid_videos(): + # An orphan MP4 sorting INTO the requested page must not consume a window slot + # and drop a valid video that sorts after it: paging is over readable records. + _save_with_mtime("p2", 100.0) + orphan = gallery.gallery_dir() / "zzz_orphan.mp4" + orphan.write_bytes(_mp4()) # newest by mtime (set below), sorts first + os.utime(orphan, (300.0, 300.0)) + _save_with_mtime("p1", 200.0) + # First page of 2 must still return both real videos, not [p1] (orphan eating a slot). + page1 = gallery.list_videos(limit = 2, offset = 0) + assert [r["prompt"] for r in page1] == ["p1", "p2"] + + +def test_list_skips_corrupt_sidecar(): + # A sidecar that is not valid JSON is treated as a foreign/orphan mp4 and skipped. + directory = gallery.gallery_dir() + (directory / "broken.mp4").write_bytes(_mp4()) + (directory / "broken.json").write_text("{not json", encoding = "utf-8") + gallery.save(_mp4(), _meta(prompt = "ours")) + listed = gallery.list_videos() + assert [r["prompt"] for r in listed] == ["ours"]