870 lines
35 KiB
Python
870 lines
35 KiB
Python
# 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
|
|
)
|
|
from .video_ltx2 import is_ltx23_checkpoint, load_ltx23_pipeline
|
|
|
|
if fam.name == "ltx-2" and is_ltx23_checkpoint(checkpoint_path):
|
|
# 2.3 checkpoints need the full assembly: new transformer config
|
|
# flags, key renames the stock converter lacks, and the 2.3
|
|
# connectors/VAEs/vocoder the 2.0 base repo does not carry.
|
|
pipe = load_ltx23_pipeline(
|
|
checkpoint_path,
|
|
base_repo = base,
|
|
torch_dtype = dtype,
|
|
is_gguf = kind == "gguf",
|
|
hf_token = hf_token,
|
|
)
|
|
else:
|
|
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
|