diff --git a/scripts/video_quality.py b/scripts/video_quality.py index 515700c370..ebdc731ae9 100644 --- a/scripts/video_quality.py +++ b/scripts/video_quality.py @@ -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) diff --git a/studio/backend/core/inference/diffusion_ideogram4.py b/studio/backend/core/inference/diffusion_ideogram4.py index 999a38e37c..448240eecf 100644 --- a/studio/backend/core/inference/diffusion_ideogram4.py +++ b/studio/backend/core/inference/diffusion_ideogram4.py @@ -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 diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index 576124d7c9..300ee82a8d 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -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. diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index c949491b38..2589f74ff3 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -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: diff --git a/studio/backend/core/inference/video_ltx2.py b/studio/backend/core/inference/video_ltx2.py index 93cc88506c..c549d84b0c 100644 --- a/studio/backend/core/inference/video_ltx2.py +++ b/studio/backend/core/inference/video_ltx2.py @@ -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 diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index ba4adb630d..2c07dcfe4d 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -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" diff --git a/studio/backend/routes/video.py b/studio/backend/routes/video.py index f6b6fc4b4c..6ad09af9de 100644 --- a/studio/backend/routes/video.py +++ b/studio/backend/routes/video.py @@ -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} diff --git a/studio/backend/tests/test_diffusion_more_families.py b/studio/backend/tests/test_diffusion_more_families.py index 738c8b90a5..2e21e3023f 100644 --- a/studio/backend/tests/test_diffusion_more_families.py +++ b/studio/backend/tests/test_diffusion_more_families.py @@ -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") diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 1862e9bce7..2c7bbf83f2 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -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") diff --git a/studio/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index 4ef7d3000f..9abe77c529 100644 --- a/studio/backend/tests/test_video_routes.py +++ b/studio/backend/tests/test_video_routes.py @@ -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"]