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 73f62515b9..10a10a57e2 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -290,7 +290,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"], @@ -354,7 +353,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: @@ -436,9 +434,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) @@ -502,12 +498,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 @@ -534,9 +526,7 @@ class VideoBackend: speed_mode = effective_speed, offload_active = plan.offload_policy != "none", ) - 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. @@ -603,7 +593,11 @@ class VideoBackend: ) logger.info( "video.loaded: %s (%s, %s, offload=%s, speed=%s)", - repo_id, fam.name, kind, offload_policy, effective_speed, + repo_id, + fam.name, + kind, + offload_policy, + effective_speed, ) return self.status() @@ -621,9 +615,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 ─────────────────────────────────────────────────────────── @@ -641,7 +633,6 @@ class VideoBackend: seed: Optional[int] = None, ) -> dict[str, Any]: import torch - cancel = threading.Event() with self._generate_lock: with self._lock: @@ -652,7 +643,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) @@ -688,8 +680,13 @@ class VideoBackend: 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 _on_step(p, step_index, timestep, callback_kwargs): 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 4d37815cb1..99d748da91 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -2439,7 +2439,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 7ff295edad..080831103f 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 @@ -134,7 +133,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()) @@ -212,14 +210,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} @@ -227,7 +223,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()) @@ -298,6 +293,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 9e93ed39ba..f44ee0151a 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -240,17 +240,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): diff --git a/studio/backend/tests/test_video_routes.py b/studio/backend/tests/test_video_routes.py index 355130646f..5928d6523b 100644 --- a/studio/backend/tests/test_video_routes.py +++ b/studio/backend/tests/test_video_routes.py @@ -99,7 +99,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(), @@ -113,7 +114,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 { @@ -260,7 +267,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"]