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

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2026-07-04 14:47:22 +00:00
commit 3ebec8e037
10 changed files with 230 additions and 130 deletions
Download patch file Download diff file Expand all files Collapse all files

77
scripts/video_quality.py
View file

@ -79,7 +79,6 @@ _PERFECT_MATCH_PSNR = 100.0
def _gray(frame: Any) -> Any:
import numpy as np
f = np.asarray(frame, dtype = np.float64)
return f @ np.array([0.299, 0.587, 0.114])
@ -109,7 +108,11 @@ def _box_mean(x: Any, w: int) -> Any:
return total / float(w * w)
def frame_ssim(a: Any, b: Any, window: int = 7) -> float:
def frame_ssim(
a: Any,
b: Any,
window: int = 7,
) -> float:
"""Pure numpy box-window SSIM on luminance (Wang et al. constants); identical
math to scripts/diffusion_quality.py so image and video budgets compare."""
ga, gb = _gray(a), _gray(b)
@ -148,7 +151,11 @@ def temporal_deviation(ref_frames: Any, cand_frames: Any) -> float:
return sum(abs(r - c) for r, c in zip(ref_series, cand_series)) / denom
def clip_metrics(ref_frames: Any, cand_frames: Any, sample_count: int = 5) -> dict[str, Any]:
def clip_metrics(
ref_frames: Any,
cand_frames: Any,
sample_count: int = 5,
) -> dict[str, Any]:
"""All frame metrics for one candidate clip vs the reference clip."""
import numpy as np
@ -179,7 +186,7 @@ def audio_metrics(ref_audio: Optional[Any], cand_audio: Optional[Any]) -> dict[s
if a is None:
return None
arr = np.asarray(a, dtype = np.float64)
return float(np.sqrt((arr ** 2).mean())) if arr.size else 0.0
return float(np.sqrt((arr**2).mean())) if arr.size else 0.0
ref_rms, cand_rms = _rms(ref_audio), _rms(cand_audio)
silent_collapse = (
@ -243,10 +250,14 @@ def parse_spec(spec: str) -> dict[str, str]:
def spec_label(spec: dict[str, str]) -> str:
if not spec:
return "base"
return ",".join(f"{k}={Path(v).name if k == 'gguf_filename' else v}" for k, v in sorted(spec.items()))
return ",".join(
f"{k}={Path(v).name if k == 'gguf_filename' else v}" for k, v in sorted(spec.items())
)
def run_config(backend: Any, args: Any, spec: dict[str, str], workdir: Path, name: str) -> dict[str, Any]:
def run_config(
backend: Any, args: Any, spec: dict[str, str], workdir: Path, name: str
) -> dict[str, Any]:
"""Load per spec, generate the fixed clip, unload. Returns frames/audio/cost."""
import torch
@ -276,9 +287,7 @@ def run_config(backend: Any, args: Any, spec: dict[str, str], workdir: Path, nam
seed = args.seed,
)
generate_s = time.monotonic() - t0
peak_gib = (
torch.cuda.max_memory_allocated() / 2**30 if torch.cuda.is_available() else 0.0
)
peak_gib = torch.cuda.max_memory_allocated() / 2**30 if torch.cuda.is_available() else 0.0
backend.unload()
frames, audio = decode_mp4(result["mp4_bytes"], workdir, name)
return {
@ -287,10 +296,19 @@ def run_config(backend: Any, args: Any, spec: dict[str, str], workdir: Path, nam
"load_s": round(load_s, 1),
"generate_s": round(generate_s, 1),
"peak_vram_gib": round(peak_gib, 2),
"resolved": {k: v for k, v in status.items() if k in (
"speed_mode", "attention_backend", "transformer_cache", "transformer_quant",
"offload_policy", "model_kind",
)},
"resolved": {
k: v
for k, v in status.items()
if k
in (
"speed_mode",
"attention_backend",
"transformer_cache",
"transformer_quant",
"offload_policy",
"model_kind",
)
},
}
@ -319,8 +337,10 @@ def run_gate(args: Any) -> int:
audio = audio_metrics(ref["audio"], cand["audio"])
row = {
"candidate": label,
**{k: (round(v, 4) if isinstance(v, float) and math.isfinite(v) else v)
for k, v in metrics.items()},
**{
k: (round(v, 4) if isinstance(v, float) and math.isfinite(v) else v)
for k, v in metrics.items()
},
**{f"audio_{k}": v for k, v in audio.items()},
"load_s": cand["load_s"],
"generate_s": cand["generate_s"],
@ -363,7 +383,11 @@ def selftest() -> int:
rng = np.random.default_rng(0)
h, w, n = 64, 96, 12
def make_clip(offset = 0.0, noise = 0.0, black = False):
def make_clip(
offset = 0.0,
noise = 0.0,
black = False,
):
frames = []
for t in range(n):
x = np.linspace(0, 1, w)[None, :] + t * 0.05 + offset
@ -385,24 +409,29 @@ def selftest() -> int:
ok = ok and cond
same = clip_metrics(ref, make_clip())
check(same["ssim_mean"] > 0.99 and same["temporal_deviation"] < 0.01,
f"identical clip scores ~1 (ssim {same['ssim_mean']:.3f})")
check(
same["ssim_mean"] > 0.99 and same["temporal_deviation"] < 0.01,
f"identical clip scores ~1 (ssim {same['ssim_mean']:.3f})",
)
check(verdict(same, {"silent_collapse": False}) == "PASS", "identical clip verdict PASS")
noisy = clip_metrics(ref, make_clip(noise = 12.0))
check(0.3 < noisy["ssim_mean"] < 0.99, f"noisy clip degrades ssim ({noisy['ssim_mean']:.3f})")
black = clip_metrics(ref, make_clip(black = True))
check(verdict(black, {"silent_collapse": False}) == "FAIL",
f"black clip verdict FAIL (min_luma {black['min_luma']:.3f})")
check(
verdict(black, {"silent_collapse": False}) == "FAIL",
f"black clip verdict FAIL (min_luma {black['min_luma']:.3f})",
)
shifted = clip_metrics(ref, make_clip(offset = 0.5))
check(shifted["ssim_mean"] < same["ssim_mean"], "content shift lowers ssim")
audio = audio_metrics(np.sin(np.linspace(0, 100, 16000)), np.zeros(16000))
check(audio["silent_collapse"] is True, "silent audio collapse detected")
audio_ok = audio_metrics(np.sin(np.linspace(0, 100, 16000)),
np.sin(np.linspace(0, 100, 16000)) * 0.8)
audio_ok = audio_metrics(
np.sin(np.linspace(0, 100, 16000)), np.sin(np.linspace(0, 100, 16000)) * 0.8
)
check(audio_ok["silent_collapse"] is False, "attenuated audio is not a collapse")
print("VIDEO-QUALITY-SELFTEST", "PASS" if ok else "FAIL")
@ -414,7 +443,9 @@ def main() -> int:
parser.add_argument("--selftest", action = "store_true", help = "CPU metric sanity check")
parser.add_argument("--model", help = "Repo id handed to the video backend")
parser.add_argument("--model-kind", default = None, help = "pipeline | gguf | single_file")
parser.add_argument("--reference", default = "", help = "Reference spec 'k=v;k=v' ('' = plain base load)")
parser.add_argument(
"--reference", default = "", help = "Reference spec 'k=v;k=v' ('' = plain base load)"
)
parser.add_argument("--candidates", nargs = "+", default = [], help = "Candidate specs 'k=v;k=v'")
parser.add_argument("--prompt", default = DEFAULT_PROMPT)
parser.add_argument("--width", type = int, default = 768)

24
studio/backend/core/inference/diffusion_ideogram4.py
View file

@ -104,6 +104,7 @@ def _patch_create_causal_mask() -> None:
pipe_mod.create_causal_mask = create_causal_mask_compat
_CAUSAL_MASK_PATCHED = True
# The fp8 attention is stored as a single fused ``qkv`` matrix with the Q, K and V
# rows stacked in that order; each block is ``hidden_size`` rows tall. hidden_size =
# attention_head_dim * num_attention_heads, read from the transformer config so a
@ -207,7 +208,9 @@ def _text_encoder_shard_paths(repo_id: str, token: Optional[str]) -> list[str]:
raise FileNotFoundError(f"no text_encoder safetensors under {sub}")
try:
index_path = hf_hub_download(repo_id, "text_encoder/model.safetensors.index.json", token = token)
index_path = hf_hub_download(
repo_id, "text_encoder/model.safetensors.index.json", token = token
)
weight_map = json.loads(Path(index_path).read_text())["weight_map"]
shards = sorted(set(weight_map.values()))
except Exception: # noqa: BLE001 -- single-file text encoder has no index
@ -243,7 +246,11 @@ def _text_encoder_is_fp8(repo_id: str, token: Optional[str]) -> bool:
return False
def load_ideogram4_text_encoder(repo_id: str, dtype, hf_token: Optional[str] = None):
def load_ideogram4_text_encoder(
repo_id: str,
dtype,
hf_token: Optional[str] = None,
):
"""The Qwen3-VL text encoder for ``repo_id``.
The ``-fp8`` repo stores this encoder in the SAME float8-plus-per-channel-scale
@ -302,7 +309,12 @@ def load_ideogram4_text_encoder(repo_id: str, dtype, hf_token: Optional[str] = N
return model
def load_ideogram4_transformer(repo_id: str, subfolder: str, dtype, hf_token: Optional[str] = None):
def load_ideogram4_transformer(
repo_id: str,
subfolder: str,
dtype,
hf_token: Optional[str] = None,
):
"""An ``Ideogram4Transformer2DModel`` for ``repo_id/subfolder`` (still on CPU).
Reads the transformer config, and if the shards carry the vendor fp8 layout
@ -356,7 +368,11 @@ def load_ideogram4_transformer(repo_id: str, subfolder: str, dtype, hf_token: Op
return model
def load_ideogram4_pipeline(repo_id: str, dtype, hf_token: Optional[str] = None):
def load_ideogram4_pipeline(
repo_id: str,
dtype,
hf_token: Optional[str] = None,
):
"""Assemble Ideogram4Pipeline from ``repo_id`` per-component (see module doc)."""
import diffusers

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

@ -270,10 +270,7 @@ def estimate_image_runtime_mib(
def estimate_video_runtime_mib(
*,
width: Optional[int],
height: Optional[int],
num_frames: Optional[int],
*, width: Optional[int], height: Optional[int], num_frames: Optional[int]
) -> int:
"""Per-call activation / latent / decode headroom for a video generation.

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

@ -408,7 +408,6 @@ class VideoBackend:
# 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"],
@ -472,7 +471,6 @@ class VideoBackend:
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:
@ -556,9 +554,7 @@ class VideoBackend:
components = fam.bf16_components_gb
mib_per_gb = 1000.0**3 / (1024.0 * 1024.0)
if kind == "pipeline":
model_dense_mib = (
int(sum(components) * mib_per_gb) if components is not None else None
)
model_dense_mib = int(sum(components) * mib_per_gb) if components is not None else None
companion_mib = None
else:
checkpoint_path = self._resolve_checkpoint_path(repo_id, gguf_filename, hf_token)
@ -622,12 +618,8 @@ class VideoBackend:
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
)
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
@ -736,9 +728,7 @@ class VideoBackend:
if view is pipe:
attention_engaged = engaged
speed_optims = tuple(k for k, v in applied.items() if v)
offload_policy, vae_tiling = apply_memory_plan(
pipe, plan, device = device, logger = logger
)
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.
@ -826,7 +816,11 @@ class VideoBackend:
)
logger.info(
"video.loaded: %s (%s, %s, offload=%s, speed=%s, quant=%s)",
repo_id, fam.name, kind, offload_policy, effective_speed,
repo_id,
fam.name,
kind,
offload_policy,
effective_speed,
transformer_quant_engaged or "off",
)
return self.status()
@ -845,9 +839,7 @@ class VideoBackend:
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)
)
return Path(hf_hub_download_with_xet_fallback(repo_id, gguf_filename or "", hf_token))
# ── generation ───────────────────────────────────────────────────────────
@ -866,7 +858,6 @@ class VideoBackend:
seed: Optional[int] = None,
) -> dict[str, Any]:
import torch
cancel = threading.Event()
with self._generate_lock:
with self._lock:
@ -877,7 +868,8 @@ class VideoBackend:
try:
fam = state.family
width, height = snap_video_size(
fam, width or fam.resolution_presets[0][0],
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)
@ -926,17 +918,18 @@ class VideoBackend:
# (pipeline_wan.py:322), so the gate is BOTH the family flag and the
# signature: TI2V-5B has no cfg2_kwarg, so it never reaches here. A None
# request lets the pipeline default it (to guidance_scale) itself.
if (
fam.cfg2_kwarg
and fam.cfg2_kwarg in call_params
and guidance_2 is not None
):
if fam.cfg2_kwarg and fam.cfg2_kwarg in call_params and guidance_2 is not None:
kwargs[fam.cfg2_kwarg] = float(guidance_2)
started = time.monotonic()
self._gen = {
"active": True, "phase": "denoise", "step": 0, "total": steps,
"started": started, "eta_seconds": None, "error": None,
"active": True,
"phase": "denoise",
"step": 0,
"total": steps,
"started": started,
"eta_seconds": None,
"error": None,
}
def _tick(done: int) -> None:

128
studio/backend/core/inference/video_ltx2.py
View file

@ -270,12 +270,10 @@ def read_checkpoint_header(checkpoint_path: Path | str) -> dict[str, tuple[int,
path = str(checkpoint_path)
if path.lower().endswith(".gguf"):
from gguf import GGUFReader
for tensor in GGUFReader(path).tensors:
names_shapes[str(tensor.name)] = tuple(int(x) for x in tensor.shape)
else:
from safetensors import safe_open
with safe_open(path, framework = "pt") as handle:
for name in handle.keys():
names_shapes[name] = tuple(handle.get_slice(name).get_shape())
@ -344,16 +342,20 @@ def _split_checkpoint(state: dict[str, Any]) -> dict[str, dict[str, Any]]:
connector keys, nothing else).
"""
groups: dict[str, dict[str, Any]] = {
"dit": {}, "connectors": {}, "vae": {}, "audio_vae": {}, "vocoder": {}
"dit": {},
"connectors": {},
"vae": {},
"audio_vae": {},
"vocoder": {},
}
for key, value in state.items():
bare = key[len(_DIT_PREFIX):] if key.startswith(_DIT_PREFIX) else key
bare = key[len(_DIT_PREFIX) :] if key.startswith(_DIT_PREFIX) else key
if bare.startswith("vae."):
groups["vae"][bare[len("vae."):]] = value
groups["vae"][bare[len("vae.") :]] = value
elif bare.startswith("audio_vae."):
groups["audio_vae"][bare[len("audio_vae."):]] = value
groups["audio_vae"][bare[len("audio_vae.") :]] = value
elif bare.startswith("vocoder."):
groups["vocoder"][bare[len("vocoder."):]] = value
groups["vocoder"][bare[len("vocoder.") :]] = value
elif bare.startswith(_CONNECTOR_KEY_PREFIXES):
groups["connectors"][bare] = value
else:
@ -382,9 +384,14 @@ def checkpoint_variant(checkpoint_path: Path | str) -> str:
# ── component builders ───────────────────────────────────────────────────────
def _build_from_config(model_cls: Any, config: dict[str, Any], state: dict[str, Any],
rename: dict[str, str], torch_dtype: Any,
remove_suffixes: tuple[str, ...] = ()) -> Any:
def _build_from_config(
model_cls: Any,
config: dict[str, Any],
state: dict[str, Any],
rename: dict[str, str],
torch_dtype: Any,
remove_suffixes: tuple[str, ...] = (),
) -> Any:
from accelerate import init_empty_weights
state = _apply_rename(_to_plain_dtype(state, torch_dtype), rename)
@ -396,8 +403,14 @@ def _build_from_config(model_cls: Any, config: dict[str, Any], state: dict[str,
return model.to(torch_dtype)
def load_ltx23_transformer(dit_state: dict[str, Any], *, base_repo: str, torch_dtype: Any,
is_gguf: bool, hf_token: Optional[str]) -> Any:
def load_ltx23_transformer(
dit_state: dict[str, Any],
*,
base_repo: str,
torch_dtype: Any,
is_gguf: bool,
hf_token: Optional[str],
) -> Any:
import diffusers
from diffusers import LTX2VideoTransformer3DModel
@ -406,7 +419,7 @@ def load_ltx23_transformer(dit_state: dict[str, Any], *, base_repo: str, torch_d
# base repo's 2.0 transformer config and runs the stock 2.0 key conversion.
for old, new in _TRANSFORMER_PRERENAME:
for key in [k for k in dit_state if k.startswith(old)]:
dit_state[new + key[len(old):]] = dit_state.pop(key)
dit_state[new + key[len(old) :]] = dit_state.pop(key)
kwargs: dict[str, Any] = {
"config": base_repo,
"subfolder": "transformer",
@ -415,44 +428,54 @@ def load_ltx23_transformer(dit_state: dict[str, Any], *, base_repo: str, torch_d
**LTX_2_3_TRANSFORMER_CONFIG_OVERRIDES,
}
if is_gguf:
kwargs["quantization_config"] = diffusers.GGUFQuantizationConfig(
compute_dtype = torch_dtype
)
kwargs["quantization_config"] = diffusers.GGUFQuantizationConfig(compute_dtype = torch_dtype)
return LTX2VideoTransformer3DModel.from_single_file(dit_state, **kwargs)
def load_ltx23_connectors(connector_state: dict[str, Any], *, variant: str,
torch_dtype: Any, hf_token: Optional[str]) -> Any:
def load_ltx23_connectors(
connector_state: dict[str, Any], *, variant: str, torch_dtype: Any, hf_token: Optional[str]
) -> Any:
from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors
# Transformer-only checkpoints carry the connector stacks but not the huge
# per-modality text projections; fetch those from the companion file.
if not any(k.startswith("text_embedding_projection") for k in connector_state):
connector_state = dict(connector_state)
connector_state.update(_load_extras_file(
_EXTRAS_TEXT_PROJ.format(variant = variant), hf_token
))
connector_state.update(
_load_extras_file(_EXTRAS_TEXT_PROJ.format(variant = variant), hf_token)
)
return _build_from_config(
LTX2TextConnectors, _CONNECTORS_CONFIG, connector_state, _CONNECTORS_RENAME,
LTX2TextConnectors,
_CONNECTORS_CONFIG,
connector_state,
_CONNECTORS_RENAME,
torch_dtype,
)
def load_ltx23_vae(vae_state: dict[str, Any], *, variant: str, torch_dtype: Any,
hf_token: Optional[str]) -> Any:
def load_ltx23_vae(
vae_state: dict[str, Any], *, variant: str, torch_dtype: Any, hf_token: Optional[str]
) -> Any:
from diffusers import AutoencoderKLLTX2Video
if not vae_state:
vae_state = _load_extras_file(_EXTRAS_VIDEO_VAE.format(variant = variant), hf_token)
return _build_from_config(
AutoencoderKLLTX2Video, _VIDEO_VAE_CONFIG, vae_state, _VIDEO_VAE_RENAME,
torch_dtype, remove_suffixes = _VIDEO_VAE_REMOVE_SUFFIXES,
AutoencoderKLLTX2Video,
_VIDEO_VAE_CONFIG,
vae_state,
_VIDEO_VAE_RENAME,
torch_dtype,
remove_suffixes = _VIDEO_VAE_REMOVE_SUFFIXES,
)
def load_ltx23_audio_vae_and_vocoder(
audio_vae_state: dict[str, Any], vocoder_state: dict[str, Any], *, variant: str,
torch_dtype: Any, hf_token: Optional[str],
audio_vae_state: dict[str, Any],
vocoder_state: dict[str, Any],
*,
variant: str,
torch_dtype: Any,
hf_token: Optional[str],
) -> tuple[Any, Any]:
from diffusers import AutoencoderKLLTX2Audio
from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE
@ -460,13 +483,16 @@ def load_ltx23_audio_vae_and_vocoder(
if not audio_vae_state or not vocoder_state:
combined = _load_extras_file(_EXTRAS_AUDIO_VAE.format(variant = variant), hf_token)
audio_vae_state = {
k[len("audio_vae."):]: v for k, v in combined.items() if k.startswith("audio_vae.")
k[len("audio_vae.") :]: v for k, v in combined.items() if k.startswith("audio_vae.")
}
vocoder_state = {
k[len("vocoder."):]: v for k, v in combined.items() if k.startswith("vocoder.")
k[len("vocoder.") :]: v for k, v in combined.items() if k.startswith("vocoder.")
}
audio_vae = _build_from_config(
AutoencoderKLLTX2Audio, _AUDIO_VAE_CONFIG, audio_vae_state, _AUDIO_VAE_RENAME,
AutoencoderKLLTX2Audio,
_AUDIO_VAE_CONFIG,
audio_vae_state,
_AUDIO_VAE_RENAME,
torch_dtype,
)
# The 2.3 vocoder is a composite (base vocoder + bandwidth-extension stack +
@ -485,8 +511,14 @@ def load_ltx23_audio_vae_and_vocoder(
# ── pipeline assembly ────────────────────────────────────────────────────────
def load_ltx23_pipeline(checkpoint_path: Path | str, *, base_repo: str, torch_dtype: Any,
is_gguf: bool, hf_token: Optional[str] = None) -> Any:
def load_ltx23_pipeline(
checkpoint_path: Path | str,
*,
base_repo: str,
torch_dtype: Any,
is_gguf: bool,
hf_token: Optional[str] = None,
) -> Any:
"""Full LTX-2.3 pipeline from a single-file/GGUF checkpoint.
Assembled per-component (constructor, not from_pretrained) because the base
@ -500,7 +532,9 @@ def load_ltx23_pipeline(checkpoint_path: Path | str, *, base_repo: str, torch_dt
variant = checkpoint_variant(checkpoint_path)
logger.info(
"video.ltx23_assembly: variant=%s gguf=%s extras=%s",
variant, is_gguf, LTX23_EXTRAS_REPO,
variant,
is_gguf,
LTX23_EXTRAS_REPO,
)
state = load_single_file_checkpoint(str(checkpoint_path))
groups = _split_checkpoint(state)
@ -518,19 +552,25 @@ def load_ltx23_pipeline(checkpoint_path: Path | str, *, base_repo: str, torch_dt
)
transformer = load_ltx23_transformer(
groups["dit"], base_repo = base_repo, torch_dtype = torch_dtype,
is_gguf = is_gguf, hf_token = hf_token,
)
connectors = load_ltx23_connectors(
groups["connectors"], variant = variant, torch_dtype = torch_dtype,
groups["dit"],
base_repo = base_repo,
torch_dtype = torch_dtype,
is_gguf = is_gguf,
hf_token = hf_token,
)
vae = load_ltx23_vae(
groups["vae"], variant = variant, torch_dtype = torch_dtype, hf_token = hf_token
connectors = load_ltx23_connectors(
groups["connectors"],
variant = variant,
torch_dtype = torch_dtype,
hf_token = hf_token,
)
vae = load_ltx23_vae(groups["vae"], variant = variant, torch_dtype = torch_dtype, hf_token = hf_token)
audio_vae, vocoder = load_ltx23_audio_vae_and_vocoder(
groups["audio_vae"], groups["vocoder"], variant = variant,
torch_dtype = torch_dtype, hf_token = hf_token,
groups["audio_vae"],
groups["vocoder"],
variant = variant,
torch_dtype = torch_dtype,
hf_token = hf_token,
)
# Shared 2.0/2.3 components from the base repo, resolved through model_index

4
studio/backend/models/inference.py
View file

@ -2461,7 +2461,9 @@ class VideoGenerationDefaults(BaseModel):
guidance: float = Field(..., description = "Default guidance scale")
num_frames: int = Field(..., description = "Default frame count")
fps: int = Field(..., description = "Default playback frame rate")
frame_step: int = Field(..., description = "Temporal lattice: valid counts are k * frame_step + 1")
frame_step: int = Field(
..., description = "Temporal lattice: valid counts are k * frame_step + 1"
)
resolution_multiple: int = Field(..., description = "Width/height must be divisible by this")
resolution_presets: list[list[int]] = Field(
default_factory = list, description = "(width, height) presets the UI offers, default first"

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

@ -55,7 +55,6 @@ def _guard_video_load_against_training() -> None:
diffusion_active = False
try:
from core.training.diffusion_training_service import get_diffusion_training_service
diffusion_active = get_diffusion_training_service().is_active()
except Exception: # noqa: BLE001
diffusion_active = False
@ -136,7 +135,6 @@ async def load_video_model(
@router.get("/video/load-progress", response_model = VideoLoadProgressResponse)
async def video_load_progress(current_subject: str = Depends(get_current_subject)):
from core.inference.video import get_video_backend
return VideoLoadProgressResponse(**get_video_backend().load_progress())
@ -215,14 +213,12 @@ async def generate_video(
@router.get("/video/generate-progress", response_model = VideoGenerateProgressResponse)
async def video_generate_progress(current_subject: str = Depends(get_current_subject)):
from core.inference.video import get_video_backend
return VideoGenerateProgressResponse(**get_video_backend().generate_progress())
@router.post("/video/generate/cancel")
async def cancel_video_generation(current_subject: str = Depends(get_current_subject)):
from core.inference.video import get_video_backend
cancelled = await asyncio.to_thread(get_video_backend().cancel_generate)
return {"cancelled": cancelled}
@ -230,7 +226,6 @@ async def cancel_video_generation(current_subject: str = Depends(get_current_sub
@router.get("/video/status", response_model = VideoStatusResponse)
async def video_status(current_subject: str = Depends(get_current_subject)):
from core.inference.video import get_video_backend
return VideoStatusResponse(**get_video_backend().status())
@ -301,6 +296,5 @@ async def delete_gallery_video(video_id: str, current_subject: str = Depends(get
@router.delete("/video/gallery")
async def clear_gallery_videos(current_subject: str = Depends(get_current_subject)):
from core.inference import video_gallery
removed = await asyncio.to_thread(video_gallery.clear)
return {"removed": removed}

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

@ -85,7 +85,6 @@ def test_excluded_model_reason_none_for_supported_and_unknown():
def test_validate_load_request_surfaces_exclusion_reason():
from core.inference.diffusion import DiffusionBackend
backend = DiffusionBackend()
with pytest.raises(ValueError, match = "trust_remote_code"):
backend.validate_load_request("tencent/HunyuanImage-3.0")

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

@ -479,17 +479,21 @@ def test_is_ltx23_checkpoint_gguf(monkeypatch, tmp_path):
gguf = types.ModuleType("gguf")
# GGUF headers store dims in GGML (reversed) order.
gguf.GGUFReader = _reader_for({
"model.diffusion_model.transformer_blocks.0.scale_shift_table": (4096, 9),
})
gguf.GGUFReader = _reader_for(
{
"model.diffusion_model.transformer_blocks.0.scale_shift_table": (4096, 9),
}
)
monkeypatch.setitem(sys.modules, "gguf", gguf)
path = tmp_path / "ltx23.gguf"
path.write_bytes(b"x")
assert is_ltx23_checkpoint(path) is True
gguf.GGUFReader = _reader_for({
"model.diffusion_model.transformer_blocks.0.scale_shift_table": (4096, 6),
})
gguf.GGUFReader = _reader_for(
{
"model.diffusion_model.transformer_blocks.0.scale_shift_table": (4096, 6),
}
)
assert is_ltx23_checkpoint(path) is False
def _boom(path):
@ -643,9 +647,7 @@ def test_hv15_cancel_unwinds_scheduler_loop(fake_runtime):
pipe = _FakeHV15Pipeline.instance
# Cancel lands during the FIRST real step; the next wrapped call must raise out
# of the denoise loop and generate() must surface the cancelled sentinel.
pipe.scheduler.on_step = (
lambda n: backend.cancel_generate() if n == 1 else None
)
pipe.scheduler.on_step = lambda n: backend.cancel_generate() if n == 1 else None
with pytest.raises(RuntimeError, match = VIDEO_CANCELLED_MSG):
backend.generate(prompt = "a fox", steps = 4)
assert pipe.scheduler.calls == 1
@ -785,7 +787,14 @@ def test_wan_a14b_dense_quant_applies_to_both_dits(fake_runtime, monkeypatch):
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
quantised = []
def _fake_quant(view, target, *, mode, family, logger = None):
def _fake_quant(
view,
target,
*,
mode,
family,
logger = None,
):
# The helper reads view.transformer; record the object it would quantise so the
# test proves the second expert was reached through the proxy.
quantised.append(view.transformer)
@ -815,7 +824,14 @@ def test_dense_quant_skipped_under_offload(fake_runtime, monkeypatch):
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
quantised = []
def _fake_quant(view, target, *, mode, family, logger = None):
def _fake_quant(
view,
target,
*,
mode,
family,
logger = None,
):
quantised.append(view.transformer)
return "int8"
@ -856,7 +872,14 @@ def test_wan_ti2v_dense_quant_applies_to_single_dit(fake_runtime, monkeypatch):
monkeypatch.setattr(video_mod, "dense_transformer_supported", lambda target: True)
quantised = []
def _fake_quant(view, target, *, mode, family, logger = None):
def _fake_quant(
view,
target,
*,
mode,
family,
logger = None,
):
quantised.append(view.transformer)
return "fp8"
@ -876,13 +899,9 @@ def test_wan_validate_trusted_repos(fake_runtime):
# The two Wan base repos are trusted for non-GGUF (pipeline) loads; an unrelated
# repo carrying the family name is not.
backend = VideoBackend()
fam = backend.validate_load_request(
"Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline"
)
fam = backend.validate_load_request("Wan-AI/Wan2.2-TI2V-5B-Diffusers", model_kind = "pipeline")
assert fam.name == "wan2.2-ti2v-5b"
fam2 = backend.validate_load_request(
"Wan-AI/Wan2.2-T2V-A14B-Diffusers", model_kind = "pipeline"
)
fam2 = backend.validate_load_request("Wan-AI/Wan2.2-T2V-A14B-Diffusers", model_kind = "pipeline")
assert fam2.name == "wan2.2-t2v-a14b"
with pytest.raises(ValueError, match = "limited to"):
backend.validate_load_request("evil/wan2.2-ti2v-5b-repack", model_kind = "pipeline")

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

@ -101,7 +101,8 @@ class _FakeBackend:
"base_repo": kwargs.get("base_repo") or "Lightricks/LTX-2",
"device": "cpu",
"dtype": "float32",
"model_kind": kwargs.get("model_kind") or ("gguf" if kwargs.get("gguf_filename") else "pipeline"),
"model_kind": kwargs.get("model_kind")
or ("gguf" if kwargs.get("gguf_filename") else "pipeline"),
"memory_mode": kwargs.get("memory_mode") or "auto",
"has_audio": True,
"defaults": _defaults(),
@ -115,7 +116,13 @@ class _FakeBackend:
"error": None,
}
def generate(self, *, prompt, seed = None, **kwargs):
def generate(
self,
*,
prompt,
seed = None,
**kwargs,
):
if not self.loaded:
raise RuntimeError(VIDEO_NOT_LOADED_MSG)
return {
@ -297,7 +304,9 @@ def test_generate_happy_path_persists_and_returns_record(client):
"/api/inference/video/load",
json = {"model_path": "unsloth/LTX-2.3-GGUF", "gguf_filename": "q.gguf"},
)
gen = client.post("/api/inference/video/generate", json = {"prompt": "a sloth surfing", "seed": 7})
gen = client.post(
"/api/inference/video/generate", json = {"prompt": "a sloth surfing", "seed": 7}
)
assert gen.status_code == 200
video = gen.json()["video"]
assert video["seed"] == 7 and video["prompt"] == "a sloth surfing" and video["id"]