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

Video inference engine: LTX-2 family registry, VideoBackend, MP4 gallery

Browse source 
Text-to-video lands as a SIBLING of the image diffusion backend, not a mode of
it: video pipelines take frame/fps arguments, return frame stacks plus, for
LTX-2, synchronized audio, and persist MP4s -- none of the image module's
img2img/inpaint/ControlNet/LoRA surface applies. The image backend's hardware
and optimisation layers are imported unchanged (device/dtype resolution, memory
planning + offload tiers, attention backends, speed profiles, FBCache), and the
load-token/cancel-event concurrency skeleton is copied verbatim so lifecycle
behaviour cannot diverge.

core/inference/video_families.py: a pure VideoFamily registry (no torch) with
the ltx-2 entry -- LTX2Pipeline + LTX2VideoTransformer3DModel, base
Lightricks/LTX-2, unsloth/LTX-2.3-GGUF as the curated GGUF source, audio on,
frame lattice k*8+1, /32 resolutions with a vertical preset, and measured bf16
component sizes (the Gemma3-27B text encoder outweighs the 19B DiT itself).
MoE fields (transformer_2, guidance_scale_2) are declared now so the Wan2.2
A14B family lands later without churning the schema.

core/inference/video.py: VideoBackend with async begin_load + cache-scan
download progress, GGUF / single-file / full-pipeline loads (the GGUF DiT
assembles onto the base repo exactly like the image path), generation with
frame/size snapping BEFORE latents allocate, per-step progress + ETA and
cooperative cancel via the standard diffusers callback, and MP4 (H.264) export
through diffusers' PyAV encoder with the audio track muxed when the family
produces one. VAE tiling is always on: decoding a 100+ frame clip is the
memory peak, and the frames-aware estimate_video_runtime_mib (new, in
diffusion_memory) feeds the planner where the pixel-area image estimate would
badly undershoot. Loads are gated to unsloth/*, the official Lightricks base
repos, or local paths; PyAV availability is checked at load time so a missing
encoder cannot fail a clip after a multi-minute denoise.

core/inference/video_gallery.py: {id}.mp4 + {id}.json recipe sidecar pairs
under studio_root()/videos (an MP4 has no PNG text chunk to embed the recipe
in), with the image gallery's id/containment guards, newest-first listing that
skips orphans, delete/clear.

gpu_arbiter gains the VIDEO owner: ownership is exclusive, so the existing
evict-the-current-owner already generalises to chat/image/video all evicting
each other. The av (PyAV) dependency joins requirements/studio.txt.

Tests: video family detection/snapping/defaults, backend lifecycle on a faked
torch/diffusers runtime (GGUF assembly, shape snapping, distilled defaults,
cancel/progress, sentinel), gallery roundtrip/containment/orphans. 52 new
tests green plus the arbiter suite.
This commit is contained in:
Daniel Han 2026-07-04 13:08:43 +00:00
commit 1b4a66dcca
9 changed files with 1757 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

23
studio/backend/core/inference/diffusion_memory.py
View file

@ -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."""

13
studio/backend/core/inference/gpu_arbiter.py
View file

@ -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:

855
studio/backend/core/inference/video.py Normal file
View file

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

189
studio/backend/core/inference/video_families.py Normal file
View file

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

162
studio/backend/core/inference/video_gallery.py Normal file
View file

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

5
studio/backend/requirements/studio.txt
View file

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

243
studio/backend/tests/test_video_backend.py Normal file
View file

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

102
studio/backend/tests/test_video_families.py Normal file
View file

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

167
studio/backend/tests/test_video_gallery.py Normal file
View file

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