From a2743083309156c5e729c2e55aa7b2bee529e3f4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 05:25:03 +0000 Subject: [PATCH 1/6] Check the cancel event between predownload files and probe local dirs for LTX-2.3 A warm-cache predownload sweep never consults the cancel event (each cached file returns instantly), so an unload during it was ignored until a cold file hit the network. The 2.3 detection also only probed bare-file local repos; resolve directory repos through the same child resolver the loader uses so their base pull is scoped too. --- studio/backend/core/inference/video.py | 20 +++++++++++++- studio/backend/tests/test_video_backend.py | 31 ++++++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 4bd7751f7c..3b64e37dca 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -331,8 +331,21 @@ class VideoBackend: probe = checkpoint_local if probe is None: + # Local repos: a bare file, or a directory whose child the same + # resolver load_pipeline uses picks out. Unresolvable here means + # load_pipeline will surface the real error; keep the wide pull. root = Path(kwargs["repo_id"]).expanduser() - probe = root if root.is_file() else None + if root.is_file(): + probe = root + elif root.is_dir(): + try: + probe = self._resolve_checkpoint_path( + kwargs["repo_id"], + kwargs.get("gguf_filename"), + kwargs.get("hf_token"), + ) + except Exception: # noqa: BLE001 -- surfaced by load_pipeline + probe = None ltx23 = probe is not None and is_ltx23_checkpoint(probe) if ltx23: expected = self._estimate_download_bytes( @@ -463,6 +476,11 @@ class VideoBackend: snapshot_root: Optional[Path] = None for name, _ in files: + # Explicit per-file check: a fully-cached file returns without ever + # consulting the event, so a warm-cache sweep would otherwise run to + # completion after an unload already cancelled this load. + if self._cancel_event.is_set(): + raise RuntimeError(VIDEO_CANCELLED_MSG) local = Path( hf_hub_download_with_xet_fallback( base, name, hf_token, cancel_event = self._cancel_event diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index c34cee5bf9..209983e4f6 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -528,3 +528,34 @@ def test_base_download_files_ltx23_keeps_only_shared_components(): assert not any( n.startswith(("vae/", "connectors/", "latent_upsampler/", "transformer/")) for n in names ) + + +def test_predownload_base_honors_cancel_between_files(monkeypatch): + # A warm-cache sweep returns each file instantly without consulting the event, + # so the loop must check it explicitly or an unload mid-predownload is ignored. + backend = VideoBackend() + backend._cancel_event.set() + calls: list = [] + monkeypatch.setattr( + "utils.hf_xet_fallback.hf_hub_download_with_xet_fallback", + lambda repo, fn, tok, **kw: (calls.append(fn), f"/cache/{fn}")[1], + ) + + class _Api: + def __init__(self, token = None): + pass + + def model_info(self, repo, files_metadata = True): + return types.SimpleNamespace( + siblings = [ + _sibling("model_index.json", 1), + _sibling("vae/diffusion_pytorch_model.safetensors", 2), + ] + ) + + import huggingface_hub + + monkeypatch.setattr(huggingface_hub, "HfApi", _Api) + with pytest.raises(RuntimeError, match = "cancelled"): + backend._predownload_base("base/repo", None, "pipeline") + assert calls == [] From de2cddc60600ce68147fa0d8fb8d04d34e6c7e9e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 5 Jul 2026 05:25:36 +0000 Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_video_backend.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 209983e4f6..30cbacdee0 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -545,7 +545,11 @@ def test_predownload_base_honors_cancel_between_files(monkeypatch): def __init__(self, token = None): pass - def model_info(self, repo, files_metadata = True): + def model_info( + self, + repo, + files_metadata = True, + ): return types.SimpleNamespace( siblings = [ _sibling("model_index.json", 1), From d6795ed077e42c2984d8e75e7c93cc682b26aa81 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 05:27:58 +0000 Subject: [PATCH 3/6] Restore pre-Ampere bf16 fail-fast in the DiT trainer The perf rewrite dropped the bf16 capability guard, so a pre-Ampere CUDA device (T4/V100/RTX 20xx) would die deep in model load with an opaque dtype error instead of a clear message. Restores parity with the SDXL trainer. --- studio/backend/core/training/diffusion_dit_trainer.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/studio/backend/core/training/diffusion_dit_trainer.py b/studio/backend/core/training/diffusion_dit_trainer.py index af443059aa..c27f1f6fef 100644 --- a/studio/backend/core/training/diffusion_dit_trainer.py +++ b/studio/backend/core/training/diffusion_dit_trainer.py @@ -812,6 +812,13 @@ def run_dit_lora_training( device = "cuda" if torch.cuda.is_available() else "cpu" # The flow-matching + 4-bit path is bf16 throughout (fp32 on a CPU-only box, which is # unsupported for real runs but keeps import/unit tests architecture-agnostic). + # Fail fast on pre-Ampere CUDA (T4/V100/RTX 20xx): bf16 compute is required and the run + # would otherwise die deep in model load with an opaque dtype error. + if device == "cuda" and not torch.cuda.is_bf16_supported(): + raise ValueError( + "This trainer requires a bfloat16-capable GPU (Ampere or newer); " + "this CUDA device does not support bf16." + ) weight_dtype = torch.bfloat16 if device == "cuda" else torch.float32 _assert_trusted_base_model(cfg.base_model) From 7a05d8655b8b2e29676b7e7fdbbb62802f8a36f5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 05:30:58 +0000 Subject: [PATCH 4/6] Stream every DiT through group offload, not just the primary transformer A dual-DiT pipeline (Ideogram 4's unconditional tower) placed its second denoiser resident under the group tier, which defeats the tier since the pair rarely fits where one alone did not. Stream transformer_2 and unconditional_transformer alongside the transformer and keep only the smaller companions resident. --- studio/backend/core/inference/diffusion_memory.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index c56d1592ce..4c4a3cca2c 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -502,6 +502,16 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool: import torch from diffusers.hooks import apply_group_offloading + # A dual-DiT pipeline (e.g. Ideogram 4's unconditional tower) carries a second + # denoiser as large as the first; leaving it resident would defeat this tier + # (the pair rarely fits where one alone did not). Stream every DiT and keep + # only the genuinely smaller companions resident. + streamed: dict[str, Any] = {"transformer": transformer} + for extra in ("transformer_2", "unconditional_transformer"): + module = getattr(pipe, extra, None) + if isinstance(module, torch.nn.Module): + streamed[extra] = module + onload = torch.device(device) use_stream = onload.type == "cuda" # overlap H2D copies with compute on CUDA gkwargs: dict[str, Any] = { @@ -531,11 +541,12 @@ def _apply_group_offload(pipe: Any, device: str, logger: Any) -> bool: # load-time crash. The streamed transformer manages its own placement via the # offloading hooks applied next. for name, comp in getattr(pipe, "components", {}).items(): - if name == "transformer": + if name in streamed: continue if isinstance(comp, torch.nn.Module): comp.to(onload) - apply_group_offloading(transformer, **gkwargs) + for module in streamed.values(): + apply_group_offloading(module, **gkwargs) return True except Exception as exc: # noqa: BLE001 — fall back to whole-module offload if logger is not None: From 14f4e7ef7ccfec44121dac7a19afda9a3cd3be99 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 05:32:24 +0000 Subject: [PATCH 5/6] Keep standalone chat templates in the scoped video download tokenizer/chat_template.jinja ships as its own file in the LTX-2 and HunyuanVideo-1.5 repos and apply_chat_template reads it at generation time, so a scoped snapshot without it loads fine and then crashes the first generation. --- studio/backend/core/inference/video.py | 6 +++++- studio/backend/tests/test_video_backend.py | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 3b64e37dca..dcdbb863bb 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -406,7 +406,11 @@ class VideoBackend: files: list[tuple[str, int]] = [] for sibling in info.siblings or []: name, size = sibling.rfilename, sibling.size or 0 - if not name.endswith((".safetensors", ".json", ".model", ".txt")): + # .jinja: tokenizer/chat_template.jinja ships as a standalone file in the + # LTX-2 and HunyuanVideo-1.5 repos (not embedded in tokenizer_config.json) + # and apply_chat_template needs it at generation time, so a snapshot + # without it loads fine and then crashes the first generation. + if not name.endswith((".safetensors", ".json", ".model", ".txt", ".jinja")): continue if "/" not in name and name.endswith(".safetensors"): continue diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index 30cbacdee0..edb0a41bbb 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -459,6 +459,7 @@ _LTX2_SIBLINGS = [ _sibling("text_encoder/diffusion_pytorch_model-00002-of-00002.safetensors", 25), _sibling("vae/diffusion_pytorch_model.safetensors", 3), _sibling("tokenizer/tokenizer.model", 1), + _sibling("tokenizer/chat_template.jinja", 1), _sibling("assets/example.mp4", 500), ] @@ -473,7 +474,10 @@ def test_base_download_files_scopes_pipeline_pull(): assert "assets/example.mp4" not in files assert files["text_encoder/model-00001-of-00002.safetensors"] == 25 assert files["transformer/diffusion_pytorch_model-00001-of-00002.safetensors"] == 20 - assert sum(files.values()) == 10 + 1 + 20 + 18 + 25 + 25 + 3 + 1 + # The standalone chat template must survive the whitelist: apply_chat_template + # reads it at generation time and it is not embedded in tokenizer_config.json. + assert "tokenizer/chat_template.jinja" in files + assert sum(files.values()) == 10 + 1 + 20 + 18 + 25 + 25 + 3 + 1 + 1 def test_base_download_files_gguf_drops_transformer(): From fc35c40c18c5f58f4b7acb4c354eaca6ea9e9172 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 5 Jul 2026 05:34:18 +0000 Subject: [PATCH 6/6] Return 404 for malformed diffusion run records in the detail route The run detail route built DiffusionTrainingRunDetail(**rec) unguarded, so a valid-JSON-but-wrong-shape record (hand-edited or an older schema) would 500 instead of reading as absent. Catch ValidationError and 404, matching how the list route skips malformed records. --- studio/backend/routes/training.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 74088ab886..35d6921605 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1339,7 +1339,12 @@ async def get_diffusion_training_run( rec = get_diffusion_run(job_id) if rec is None: raise HTTPException(status_code = 404, detail = "No such training run.") - return DiffusionTrainingRunDetail(**rec) + try: + return DiffusionTrainingRunDetail(**rec) + except ValidationError: + # A malformed on-disk record (hand-edited / older shape) should read as absent + # rather than 500 the endpoint, mirroring how the list route skips bad records. + raise HTTPException(status_code = 404, detail = "No such training run.") # Extensions accepted into an image-training dataset folder: images the trainer reads,