From 62cae5fe71cacff56dd2ec72b881a6697b028cce Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 04:41:43 +0000
Subject: [PATCH 1/8] Load image pipelines from the prefetched snapshot instead
 of re-sweeping the hub

The prefetch already scopes the file list (no packaged root singles, no
dtype-variant twins, no ONNX/Flax exports), but from_pretrained was then
called with the hub id, and its own snapshot sweep re-downloaded the
skipped files anyway: 24 GB per FLUX.1 repo and 65 GB on FLUX.2-dev, as
found in the blob cache. Return the snapshot dir from the prefetch (keyed
on the pipeline manifest) and hand it to every pipeline-assembly
from_pretrained site; any prefetch failure keeps the hub id and the old
behavior.
---
 studio/backend/core/inference/diffusion.py    | 41 +++++++++++++++----
 .../backend/tests/test_diffusion_backend.py   | 32 +++++++++++++++
 2 files changed, 64 insertions(+), 9 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 09b2fa1b42..2d69c9bb56 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -439,11 +439,17 @@ class DiffusionBackend:
         base: str,
         base_files: list[str],
         hf_token: Optional[str],
-    ) -> None:
+    ) -> Optional[str]:
         """Pre-download the GGUF + the given ``base_files`` into the HF cache,
         WITHOUT the lock and honoring ``_cancel_event``, so load_pipeline's
         from_single_file / from_pretrained hit the cache and the heavy download can
-        be preempted by an unload/eviction. Raises ``RuntimeError("Cancelled")``."""
+        be preempted by an unload/eviction. Raises ``RuntimeError("Cancelled")``.
+
+        Returns the base repo's local snapshot dir when the prefetched set includes
+        the pipeline manifest, so from_pretrained can load from disk instead of
+        re-sweeping the hub (its own sweep also pulls files the scoped list skips,
+        e.g. the 24 GB packaged root singles in each FLUX.1 repo); None otherwise
+        (estimate failure, config-only base, local repo) -> hub id as before."""
         from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback
 
         # GGUF transformer (hub repos only; a local path is already on disk).
@@ -452,12 +458,16 @@ class DiffusionBackend:
                 repo_id, gguf_filename, hf_token, cancel_event = self._cancel_event
             )
         # Base repo (VAE / text-encoder / scheduler); list comes from the estimate.
+        snapshot_root: Optional[str] = None
         for rfilename in base_files:
             if self._cancel_event.is_set():
                 raise RuntimeError("Cancelled")
-            hf_hub_download_with_xet_fallback(
+            local = hf_hub_download_with_xet_fallback(
                 base, rfilename, hf_token, cancel_event = self._cancel_event
             )
+            if rfilename == "model_index.json":
+                snapshot_root = str(Path(local).parent)
+        return snapshot_root
 
     def validate_load_request(
         self,
@@ -674,7 +684,7 @@ class DiffusionBackend:
                     self._loading.expected_bytes = expected
             # Download outside the lock so unload()/an eviction can preempt the
             # multi-GB pull; load_pipeline below then assembles from the cache.
-            self._prefetch_files(
+            kwargs["_base_local_dir"] = self._prefetch_files(
                 kwargs["repo_id"],
                 kwargs.get("gguf_filename"),
                 base,
@@ -887,6 +897,7 @@ class DiffusionBackend:
         transformer_cache_threshold: Optional[float] = None,
         model_kind: Optional[str] = None,
         _load_token: Optional[int] = None,
+        _base_local_dir: Optional[str] = None,
     ) -> dict[str, Any]:
         # A blank / whitespace-only token must degrade to anonymous access, not be passed
         # as an explicit credential (from_single_file / from_pretrained / the Hub client
@@ -1003,6 +1014,7 @@ class DiffusionBackend:
                             transformer_quant,
                             transformer_quant_fast_accum,
                             fam = fam,
+                            base_local_dir = _base_local_dir,
                             prequant_path = transformer_prequant_path,
                         )
                     except Exception as exc:  # noqa: BLE001 — fall back to the GGUF build
@@ -1033,7 +1045,12 @@ class DiffusionBackend:
                         pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype}
                         if hf_token:
                             pipe_kwargs["token"] = hf_token
-                        pipe = pipeline_cls.from_pretrained(repo_id, **pipe_kwargs)
+                        # The prefetched snapshot dir keeps from_pretrained off the hub:
+                        # its own snapshot sweep re-downloads files the scoped prefetch
+                        # skipped (root packaged singles, e.g. 24 GB per FLUX.1 repo).
+                        pipe = pipeline_cls.from_pretrained(
+                            _base_local_dir or repo_id, **pipe_kwargs
+                        )
                     elif kind == "single_file" and fam.single_file_is_pipeline:
                         # A single-file SDXL-style checkpoint is the WHOLE pipeline
                         # (U-Net + VAE + both text encoders), not a transformer-only file,
@@ -1069,7 +1086,9 @@ class DiffusionBackend:
                         pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer}
                         if hf_token:
                             pipe_kwargs["token"] = hf_token
-                        pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
+                        pipe = pipeline_cls.from_pretrained(
+                            _base_local_dir or base, **pipe_kwargs
+                        )
 
                 # Resolve the effective speed mode: GGUF models default to the
                 # near-lossless `default` profile (compile is ~2.2x and sits below
@@ -1301,6 +1320,7 @@ class DiffusionBackend:
         *,
         fam: Optional[DiffusionFamily] = None,
         prequant_path: Optional[str] = None,
+        base_local_dir: Optional[str] = None,
     ) -> tuple[Any, str]:
         """Build the opt-in fast pipeline and return ``(pipe, engaged_scheme)``.
 
@@ -1348,7 +1368,7 @@ class DiffusionBackend:
                 )
                 if transformer is not None:
                     pipe = self._assemble_pipe(
-                        pipeline_cls, base, transformer, dtype, hf_token, device
+                        pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir
                     )
                     return pipe, scheme
 
@@ -1356,7 +1376,9 @@ class DiffusionBackend:
         transformer = transformer_cls.from_pretrained(
             base, subfolder = "transformer", torch_dtype = dtype, token = hf_token
         )
-        pipe = self._assemble_pipe(pipeline_cls, base, transformer, dtype, hf_token, device)
+        pipe = self._assemble_pipe(
+            pipeline_cls, base, transformer, dtype, hf_token, device, base_local_dir
+        )
         scheme = quantize_transformer(
             pipe,
             target,
@@ -1377,13 +1399,14 @@ class DiffusionBackend:
         dtype: Any,
         hf_token: Optional[str],
         device: str,
+        base_local_dir: Optional[str] = None,
     ) -> Any:
         """Assemble the diffusers pipeline around ``transformer`` and place it on ``device``
         (a no-op for an already-placed pre-quantized transformer; it moves the companions)."""
         pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer}
         if hf_token:
             pipe_kwargs["token"] = hf_token
-        pipe = pipeline_cls.from_pretrained(base, **pipe_kwargs)
+        pipe = pipeline_cls.from_pretrained(base_local_dir or base, **pipe_kwargs)
         pipe.to(device)
         return pipe
 
diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py
index 6baa0a8a0a..1dccc2d0ac 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -2121,3 +2121,35 @@ def test_generate_resets_step_cache_only_when_engaged(fake_runtime, tmp_path):
     backend.generate(prompt = "a sloth")
     backend.generate(prompt = "another sloth")
     assert resets == [True, True]
+
+
+def test_prefetch_returns_snapshot_dir_for_manifest(monkeypatch):
+    # The prefetched pipeline manifest's directory is the local snapshot root; a
+    # config-only base list (no manifest) returns None so the hub id stays in use.
+    backend = DiffusionBackend()
+    monkeypatch.setattr(
+        "utils.hf_xet_fallback.hf_hub_download_with_xet_fallback",
+        lambda repo, fn, tok, **k: f"/cache/snap/{fn}",
+    )
+    root = backend._prefetch_files(
+        "base/repo", None, "base/repo", ["model_index.json", "vae/x.safetensors"], None
+    )
+    assert root == "/cache/snap"
+    assert (
+        backend._prefetch_files("base/repo", None, "base/repo", ["vae/x.safetensors"], None)
+        is None
+    )
+
+
+def test_pipeline_load_uses_predownloaded_dir(fake_runtime, tmp_path):
+    # With a prefetched snapshot, from_pretrained must receive the local dir --
+    # its own hub sweep would re-download the root packaged singles the scoped
+    # prefetch skips (24 GB per FLUX.1 repo).
+    backend = DiffusionBackend()
+    backend.load_pipeline(
+        "unsloth/Qwen-Image-2512-bnb-4bit",
+        model_kind = "pipeline",
+        _base_local_dir = str(tmp_path),
+    )
+    assert _FakePipeline.last["base"] == str(tmp_path)
+    backend.unload()

From 919661ddf084f63b4df31ae0f598f86f9f24b27b 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 04:42:14 +0000
Subject: [PATCH 2/8] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/inference/diffusion.py     | 4 +---
 studio/backend/tests/test_diffusion_backend.py | 3 +--
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 2d69c9bb56..3b47287722 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -1086,9 +1086,7 @@ class DiffusionBackend:
                         pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer}
                         if hf_token:
                             pipe_kwargs["token"] = hf_token
-                        pipe = pipeline_cls.from_pretrained(
-                            _base_local_dir or base, **pipe_kwargs
-                        )
+                        pipe = pipeline_cls.from_pretrained(_base_local_dir or base, **pipe_kwargs)
 
                 # Resolve the effective speed mode: GGUF models default to the
                 # near-lossless `default` profile (compile is ~2.2x and sits below
diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py
index 1dccc2d0ac..65516fb5e9 100644
--- a/studio/backend/tests/test_diffusion_backend.py
+++ b/studio/backend/tests/test_diffusion_backend.py
@@ -2136,8 +2136,7 @@ def test_prefetch_returns_snapshot_dir_for_manifest(monkeypatch):
     )
     assert root == "/cache/snap"
     assert (
-        backend._prefetch_files("base/repo", None, "base/repo", ["vae/x.safetensors"], None)
-        is None
+        backend._prefetch_files("base/repo", None, "base/repo", ["vae/x.safetensors"], None) is None
     )
 
 

From c53a6cb65e05e231a8bb007d7f5a6463dd6d19e4 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 04:53:37 +0000
Subject: [PATCH 3/8] Address review findings: dataset preflight, sd.cpp unload
 barrier, caption tombstone, ControlNet cache race

- Run the trainer's caption discovery in the start route BEFORE freeing GPU
  residents, so a missing or uncaptionable dataset 400s without evicting the
  loaded chat/Images model.
- sd.cpp unload now waits out a cancelled one-shot generation on the generate
  lock before reporting the device free, matching the diffusers backend.
- Clearing a caption that came from metadata.jsonl writes an empty sidecar
  tombstone instead of unlinking (both readers treat an existing sidecar as
  authoritative), so the cleared label cannot resurface.
- The ControlNet wrapper pipe is only cached while its load is still current,
  closing the unload race the model cache already handled.
---
 studio/backend/core/inference/diffusion.py    |  9 ++++-
 .../backend/core/inference/sd_cpp_backend.py  |  7 ++++
 studio/backend/routes/training.py             | 37 ++++++++++++++++---
 .../tests/test_diffusion_controlnet.py        | 18 +++++++++
 .../backend/tests/test_diffusion_training.py  |  7 ++++
 5 files changed, 72 insertions(+), 6 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 2d69c9bb56..ddf31db97a 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -1582,7 +1582,14 @@ class DiffusionBackend:
             pipe = getattr(diffusers, pipe_cls_name).from_pipe(
                 state.pipe, controlnet = cn_model, torch_dtype = None
             )
-            self._cn_pipes[key] = pipe
+            with self._lock:
+                # Same race as the model cache above: an unload/superseding load may
+                # have cleared _cn_pipes while from_pipe ran; caching now would pin a
+                # pipeline built around the UNLOADED base and hand it to the next load.
+                if cancel.is_set() or self._state is not state:
+                    del pipe
+                    raise RuntimeError(DIFFUSION_CANCELLED_MSG)
+                self._cn_pipes[key] = pipe
         return pipe
 
     @staticmethod
diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py
index 75572b500a..baee6605e0 100644
--- a/studio/backend/core/inference/sd_cpp_backend.py
+++ b/studio/backend/core/inference/sd_cpp_backend.py
@@ -1096,6 +1096,13 @@ class SdCppDiffusionBackend:
             state.server.stop()
         if pending is not None and pending is not (state.server if state else None):
             pending.stop()
+        # Wait for a signalled one-shot generation to actually exit before reporting
+        # unloaded: callers (the GPU arbiter, training cleanup) treat this return as
+        # "the device is free", but a one-shot sd-cli child killed by the cancel above
+        # unwinds under _generate_lock. A bare acquire is the exit barrier (never taken
+        # while holding _lock; same pattern as DiffusionBackend.unload).
+        with self._generate_lock:
+            pass
         return self.status()
 
     def status(self) -> dict[str, Any]:
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index a89a8d01c2..56514d2272 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -1234,6 +1234,21 @@ async def start_diffusion_training(
     # user's loaded chat/Images model, and never surfaces as a confusing mid-load 401.
     _preflight_gated_base(config.get("base_model", ""), config.get("hf_token"))
 
+    # Preflight the dataset too: a missing/empty/uncaptionable data_dir otherwise
+    # fails inside the spawned trainer AFTER the user's chat/Images model was
+    # evicted. Same discovery the trainer runs, so the two cannot disagree.
+    from core.training import diffusion_train_common as _dtc
+
+    try:
+        await asyncio.to_thread(
+            _dtc.discover_image_caption_pairs,
+            config["data_dir"],
+            instance_prompt = config.get("instance_prompt") or None,
+            caption_column = config.get("caption_column") or "text",
+        )
+    except (FileNotFoundError, ValueError) as e:
+        raise HTTPException(status_code = 400, detail = str(e))
+
     # Free resident GPU workloads (export / Images pipeline / chat) before the trainer
     # loads its own pipeline.
     _free_gpu_for_diffusion_training()
@@ -1636,12 +1651,24 @@ async def set_diffusion_dataset_caption(
         sidecar = image_path.with_suffix(".txt")
         if caption:
             sidecar.write_text(caption, encoding = "utf-8")
-        else:
-            # Blank clears the sidecar; also drop a stale .caption so the image reads as
-            # uncaptioned afterwards.
-            sidecar.unlink(missing_ok = True)
             image_path.with_suffix(".caption").unlink(missing_ok = True)
-        return _image_record(folder, image_path, _load_metadata_captions(folder))
+            return _image_record(folder, image_path, _load_metadata_captions(folder))
+        # Blank must actually clear. Unlinking alone would resurface this image's
+        # metadata.jsonl / captions.jsonl caption (the fallback source), so when one
+        # exists write an EMPTY sidecar instead: both the record reader and the
+        # trainer's discovery treat an existing sidecar as authoritative even when
+        # empty, which makes it a tombstone. No metadata caption -> plain cleanup.
+        meta = _load_metadata_captions(folder)
+        try:
+            rel = image_path.relative_to(folder).as_posix()
+        except ValueError:
+            rel = image_path.name
+        if image_path.name in meta or rel in meta:
+            sidecar.write_text("", encoding = "utf-8")
+        else:
+            sidecar.unlink(missing_ok = True)
+        image_path.with_suffix(".caption").unlink(missing_ok = True)
+        return _image_record(folder, image_path, meta)
 
     return await asyncio.to_thread(write)
 
diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py
index 5708c3a3ae..3d4f7c98cf 100644
--- a/studio/backend/tests/test_diffusion_controlnet.py
+++ b/studio/backend/tests/test_diffusion_controlnet.py
@@ -239,6 +239,9 @@ def test_controlnet_pipe_loads_once_and_caches(monkeypatch):
     monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
     b = DiffusionBackend()
     st = _state()
+    # The pipe cache only commits while ``st`` is the CURRENT load (an unload racing
+    # from_pipe must not repopulate the cache), so mirror the loaded invariant.
+    b._state = st
     resolved = dc.ResolvedControlNet("flux-union-pro", "repo/id", is_local = False)
     p1 = b._controlnet_pipe(st, resolved, threading.Event())
     assert isinstance(p1, _FakeCNPipe) and isinstance(p1.controlnet, _FakeCNModel)
@@ -259,3 +262,18 @@ def test_controlnet_pipe_rejects_family_without_classes():
     st.family.controlnet_pipeline_class = None
     with pytest.raises(ValueError, match = "not supported"):
         b._controlnet_pipe(st, dc.ResolvedControlNet("x", "y", False), threading.Event())
+
+def test_controlnet_pipe_not_cached_after_unload_race(monkeypatch):
+    # An unload that lands while from_pipe is assembling must not let the wrapper
+    # repopulate the cache around the torn-down base pipe.
+    import threading
+
+    from core.inference.diffusion import DiffusionBackend
+
+    monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
+    b = DiffusionBackend()
+    st = _state()  # never committed to b._state: the load is already gone
+    resolved = dc.ResolvedControlNet("flux-union-pro", "repo/id", is_local = False)
+    with pytest.raises(RuntimeError, match = "cancelled"):
+        b._controlnet_pipe(st, resolved, threading.Event())
+    assert b._cn_pipes == {}
diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py
index 1f85b08ccb..b5acaa7bb9 100644
--- a/studio/backend/tests/test_diffusion_training.py
+++ b/studio/backend/tests/test_diffusion_training.py
@@ -265,6 +265,13 @@ def client(monkeypatch):
 
     monkeypatch.setattr(tr, "get_training_backend", lambda: _FakeLLMBackend(active = False))
     monkeypatch.setattr(tr, "_free_gpu_for_diffusion_training", lambda: None)
+    # The dataset preflight runs the trainer's discovery against _BODY's fake
+    # data_dir; stub it here so wiring tests pass, and let the dedicated preflight
+    # tests below re-point it at a real tmp dataset.
+    monkeypatch.setattr(
+        "core.training.diffusion_train_common.discover_image_caption_pairs",
+        lambda data_dir, **kw: [("img.png", "caption")],
+    )
     app = FastAPI()
     app.include_router(training_router, prefix = "/api/train")
     app.dependency_overrides[get_current_subject] = lambda: "test-user"

From bcf25ca569d4a98b28d629915ba6097216a8a049 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 04:56:25 +0000
Subject: [PATCH 4/8] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/tests/test_diffusion_controlnet.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py
index 3d4f7c98cf..50bf7159b0 100644
--- a/studio/backend/tests/test_diffusion_controlnet.py
+++ b/studio/backend/tests/test_diffusion_controlnet.py
@@ -263,6 +263,7 @@ def test_controlnet_pipe_rejects_family_without_classes():
     with pytest.raises(ValueError, match = "not supported"):
         b._controlnet_pipe(st, dc.ResolvedControlNet("x", "y", False), threading.Event())
 
+
 def test_controlnet_pipe_not_cached_after_unload_race(monkeypatch):
     # An unload that lands while from_pipe is assembling must not let the wrapper
     # repopulate the cache around the torn-down base pipe.

From 7a05d8655b8b2e29676b7e7fdbbb62802f8a36f5 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 05:30:58 +0000
Subject: [PATCH 5/8] 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 f313cfd7e50121ef08104a227f21520c42e8f393 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 07:40:45 +0000
Subject: [PATCH 6/8] Security audit: baseline the new huggingface-hub
 Sandboxes findings

The latest huggingface-hub release added the Sandboxes feature. Its
bootstrap (_sandbox.py) fetches the static sbx-server binary into /tmp with
an Authorization header and marks it executable, which is exactly the
staged-dropper pattern the scanner hunts, and three while True polling loops
in _sandbox.py / hf_api.py / utils/_http.py match the beaconing heuristic.
All four verified against the official huggingface/huggingface_hub
repository: the snippet is the documented sandbox server injection and the
loops are deadline-style job and sandbox polling. Entries generated with
--write-baseline and reviewed line by line; scan_packages.py huggingface-hub
now exits 0 with the four findings suppressed.
---
 scripts/scan_packages_baseline.json | 32 +++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json
index 8e7cd4ea5a..3b33a7ce84 100644
--- a/scripts/scan_packages_baseline.json
+++ b/scripts/scan_packages_baseline.json
@@ -226,6 +226,22 @@
    "evidence": "FS: L528:         history: list[Response] | None = None, sha256:f56272dccd651b2644aa41ef6e688e211462427aad07fef5150240ec7347446e\nNetwork: L9: import urllib.request | L1243:     class _CookieCompatRequest(urllib.request.Request):",
    "evidence_hash": "b32f79e58c938680d89efa74113eeba76c9fc5aedf5de18086f93bef274c4bda"
   },
+  {
+   "package": "huggingface-hub",
+   "file": "huggingface_hub/_sandbox.py",
+   "check": "C2 polling/beaconing loop detected",
+   "severity": "CRITICAL",
+   "evidence": "L1179:             while True: sha256:33ceddf9e42aae207e891e97808c518e92a0b27ab60e4326256717bfb25a3a38",
+   "evidence_hash": "802fd41d8bb17bf425e99d128c0351c820103a5efb74690a4086e542a71437b8"
+  },
+  {
+   "package": "huggingface-hub",
+   "file": "huggingface_hub/_sandbox.py",
+   "check": "Writes to /tmp and executes (staged dropper)",
+   "severity": "CRITICAL",
+   "evidence": "L83: d=/tmp/.sbx-server\nL84: if command -v wget >/dev/null 2>&1; then wget -q --header \"Authorization: Bearer $SBX_DL_TOKEN\" -O \"$d\" \"$SBX_SERVER_URL\"\nL85: elif command -v curl >/dev/null 2>&1; then curl -fsSL -H \"Authorization: Bearer $SBX_DL_TOKEN\" -o \"$d\" \"$SBX_SERVER_URL\"\nL86: else cp \"$SBX_SERVER_MOUNT/sbx-server\" \"$d\"; fi\nL87: chmod +x \"$d\"",
+   "evidence_hash": "6908a3fe328fa94ee22a119998d6ad07cfa1ba4efa2628acf240f4204fd76e22"
+  },
   {
    "package": "huggingface-hub",
    "file": "huggingface_hub/hf_api.py",
@@ -242,6 +258,14 @@
    "evidence": "L4600:         while True: sha256:f4a851312a1832efe1b435aa1275a82184e19cc3f47e2cd244373d56c11de272",
    "evidence_hash": "dc8fcf44788e32f42d1cc2eb0e2deb55eb2dbf2c3a55909a7d503e450f45e602"
   },
+  {
+   "package": "huggingface-hub",
+   "file": "huggingface_hub/hf_api.py",
+   "check": "C2 polling/beaconing loop detected",
+   "severity": "CRITICAL",
+   "evidence": "L4613:         while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6",
+   "evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4"
+  },
   {
    "package": "huggingface-hub",
    "file": "huggingface_hub/hf_api.py",
@@ -266,6 +290,14 @@
    "evidence": "L298:     while True: sha256:6b8e5e569594caf7c4eca6137646dae471a7c3aae7294096cf876f30b5f90306",
    "evidence_hash": "c066cc27bce31ee7b6ce07411ee7a7d9ecfbf3aafc8848f6641fabfe522a7703"
   },
+  {
+   "package": "huggingface-hub",
+   "file": "huggingface_hub/utils/_http.py",
+   "check": "C2 polling/beaconing loop detected",
+   "severity": "CRITICAL",
+   "evidence": "L462:     while True: sha256:c75d1ee228cf7703a8c28551d649395a1f89f69a3aba69413f5bbcbd10c31958",
+   "evidence_hash": "d4d5f83fed39b87898cf776d5dad0bf1a6388a932f5fb7997d1070b50e46213e"
+  },
   {
    "package": "ipython",
    "file": "IPython/core/debugger.py",

From fc1e0991244f07045ac06152a5a82f4cc79685d3 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 5 Jul 2026 11:36:58 +0000
Subject: [PATCH 7/8] Harden ControlNet loads, thumbnail cache keys, API
 training guard, and picker roving keys

Review follow-ups on the image-generation PR:

- ControlNet: resolve_controlnet accepts a bare owner/name repo without the
  non-GGUF base trust gate, and _controlnet_pipe hands it straight to
  from_pretrained. A malicious pickle .bin would deserialize on load, so run
  the same Hugging Face malware preflight (evaluate_file_security) the chat and
  export loaders use before any remote ControlNet load; local dirs are exempt.
- Dataset thumbnails: key the cache on the full filename instead of the stem so
  sample.png and sample.jpg no longer collide on one .thumbs file (which could
  serve or delete the wrong image); the delete cleanup globs the same key.
- Diffusion training start: mirror start_training's API-key guard so an API
  client cannot start training (which frees VRAM by unloading chat) while an
  inference request is streaming; it now returns 409 before any GPU is freed.
- Model picker: include the curated safetensors row keys in the recommended
  roving key list so arrow-key navigation reaches those rows instead of hitting
  the duplicate option-missing id.

Tests: ControlNet malware gate (remote blocked before from_pretrained, local
skipped), thumbnail same-stem cache separation, API-key diffusion-start 409
before GPU free. Full diffusion suites green.
---
 studio/backend/core/inference/diffusion.py    | 11 +++
 studio/backend/routes/training.py             | 31 +++++++-
 .../tests/test_diffusion_controlnet.py        | 70 +++++++++++++++++++
 .../tests/test_diffusion_dataset_api.py       | 20 +++++-
 .../backend/tests/test_diffusion_training.py  | 37 +++++++++-
 .../assistant-ui/model-selector/pickers.tsx   |  9 +++
 6 files changed, 171 insertions(+), 7 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 007d6d349d..603e4ca637 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -1537,6 +1537,17 @@ class DiffusionBackend:
         if cn_model is None:
             if cancel.is_set():
                 raise RuntimeError(DIFFUSION_CANCELLED_MSG)
+            # resolve_controlnet accepts a bare owner/name repo without the non-GGUF base
+            # trust gate, and from_pretrained below downloads and deserializes it. A
+            # malicious pickle .bin would execute on load, so run the same Hub malware
+            # preflight the chat/export loaders use before any remote ControlNet load. A
+            # local dir the user picked has no Hub scan and is exempt (fail-open there).
+            if not getattr(resolved_cn, "is_local", False):
+                from utils.security import evaluate_file_security
+
+                _cn_fs = evaluate_file_security(resolved_cn.path, hf_token = state.hf_token or None)
+                if _cn_fs.blocked:
+                    raise ValueError(_cn_fs.reason)
             import torch
 
             # state.dtype is the display string saved at load ("bfloat16"), NOT a
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index 56514d2272..07dbd5c812 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -1177,11 +1177,29 @@ def _preflight_gated_base(base_model: str, hf_token: Optional[str]) -> None:
 
 @router.post("/diffusion/start", response_model = DiffusionTrainingStartResponse)
 async def start_diffusion_training(
-    body: DiffusionTrainingStartRequest, current_subject: str = Depends(get_current_subject)
+    body: DiffusionTrainingStartRequest,
+    current_subject: str = Depends(get_current_subject),
+    via_api_key: bool = Depends(authenticated_via_api_key),
 ):
     """Start an SDXL LoRA training job from an image + caption dataset."""
     from core.training.diffusion_training_service import get_diffusion_training_service
 
+    # When Studio is driven as an inference API (API-key auth), refuse to start training
+    # while a request is in flight: _free_gpu_for_diffusion_training() below unloads the
+    # chat backends to reclaim VRAM, which would kill the stream. Mirrors start_training so
+    # a diffusion start cannot silently drop an active API inference request.
+    if via_api_key is True:
+        from core.inference.llama_keepwarm import other_inference_request_count
+        if other_inference_request_count(current_request_counted = False) > 0:
+            raise HTTPException(
+                status_code = 409,
+                detail = (
+                    "Cannot start diffusion (Images) training over the API while an inference "
+                    "request is in progress. Wait for it to finish, or start training from the "
+                    "Studio UI."
+                ),
+            )
+
     # Interlock: refuse while an LLM training run holds the GPU (symmetric with the
     # diffusion check in start_training), so the two trainers never contend for VRAM.
     try:
@@ -1606,7 +1624,11 @@ async def get_diffusion_dataset_image(
 
         thumbs_dir = folder / _THUMBS_DIRNAME
         thumbs_dir.mkdir(exist_ok = True)
-        thumb_path = thumbs_dir / f"{image_path.stem}_{size}.jpg"
+        # Key on the full filename (stem + extension), not the stem: two images that
+        # share a stem but differ by extension (sample.png / sample.jpg) would otherwise
+        # collide on one cache file, and an mtime-newer cache built for the first would
+        # be served for the second, showing the wrong image in the labeling grid.
+        thumb_path = thumbs_dir / f"{image_path.name}_{size}.jpg"
         src_mtime = image_path.stat().st_mtime
         if thumb_path.is_file() and thumb_path.stat().st_mtime >= src_mtime:
             return thumb_path
@@ -1691,7 +1713,10 @@ async def delete_diffusion_dataset_image(
             image_path.with_suffix(ext).unlink(missing_ok = True)
         thumbs_dir = folder / _THUMBS_DIRNAME
         if thumbs_dir.is_dir():
-            for t in thumbs_dir.glob(f"{image_path.stem}_*.jpg"):
+            # Thumbs are keyed on the full filename (stem + extension), so match that
+            # here too; a stem-only glob would leave this image's thumbs behind and
+            # could delete a same-stem sibling's (sample.png vs sample.jpg).
+            for t in thumbs_dir.glob(f"{image_path.name}_*.jpg"):
                 t.unlink(missing_ok = True)
         return {"deleted": image_path.name}
 
diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py
index 50bf7159b0..b7ae037197 100644
--- a/studio/backend/tests/test_diffusion_controlnet.py
+++ b/studio/backend/tests/test_diffusion_controlnet.py
@@ -231,12 +231,24 @@ def _state():
     )
 
 
+def _allow_cn_security(monkeypatch):
+    """Stub the Hub malware preflight to allow the load (hermetic, no network)."""
+    import utils.security
+
+    monkeypatch.setattr(
+        utils.security,
+        "evaluate_file_security",
+        lambda name, hf_token = None, **kw: types.SimpleNamespace(blocked = False, reason = ""),
+    )
+
+
 def test_controlnet_pipe_loads_once_and_caches(monkeypatch):
     import threading
 
     from core.inference.diffusion import DiffusionBackend
 
     monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
+    _allow_cn_security(monkeypatch)
     b = DiffusionBackend()
     st = _state()
     # The pipe cache only commits while ``st`` is the CURRENT load (an unload racing
@@ -252,6 +264,64 @@ def test_controlnet_pipe_loads_once_and_caches(monkeypatch):
     assert b._cn_models["flux-union-pro"] is p1.controlnet
 
 
+def test_controlnet_pipe_blocks_flagged_remote_repo(monkeypatch):
+    # A bare owner/name ControlNet is accepted by resolve_controlnet without the base
+    # trust gate, so the load path must run the Hub malware preflight: a flagged remote
+    # repo must raise BEFORE from_pretrained downloads/deserializes it.
+    import threading
+
+    import utils.security
+    from core.inference.diffusion import DiffusionBackend
+
+    loaded = {"called": False}
+
+    class _TrapModel(_FakeCNModel):
+        @classmethod
+        def from_pretrained(cls, path, torch_dtype = None, token = None):
+            loaded["called"] = True
+            return super().from_pretrained(path, torch_dtype = torch_dtype, token = token)
+
+    mod = _fake_diffusers()
+    mod.FluxControlNetModel = _TrapModel
+    monkeypatch.setitem(sys.modules, "diffusers", mod)
+    monkeypatch.setattr(
+        utils.security,
+        "evaluate_file_security",
+        lambda name, hf_token = None, **kw: types.SimpleNamespace(
+            blocked = True, reason = "Hugging Face security scan flagged unsafe files: evil.bin"
+        ),
+    )
+    b = DiffusionBackend()
+    st = _state()
+    b._state = st
+    resolved = dc.ResolvedControlNet("evil/cn", "evil/cn", is_local = False)
+    with pytest.raises(ValueError, match = "security scan flagged"):
+        b._controlnet_pipe(st, resolved, threading.Event())
+    assert loaded["called"] is False
+
+
+def test_controlnet_pipe_skips_scan_for_local_dir(monkeypatch, tmp_path):
+    # A local dir the user picked has no Hub scan; the preflight must not block it even
+    # if the (unused) scan stub would say blocked.
+    import threading
+
+    import utils.security
+    from core.inference.diffusion import DiffusionBackend
+
+    monkeypatch.setitem(sys.modules, "diffusers", _fake_diffusers())
+    monkeypatch.setattr(
+        utils.security,
+        "evaluate_file_security",
+        lambda name, hf_token = None, **kw: types.SimpleNamespace(blocked = True, reason = "x"),
+    )
+    b = DiffusionBackend()
+    st = _state()
+    b._state = st
+    resolved = dc.ResolvedControlNet("my-cn", str(tmp_path), is_local = True)
+    p = b._controlnet_pipe(st, resolved, threading.Event())
+    assert isinstance(p, _FakeCNPipe)
+
+
 def test_controlnet_pipe_rejects_family_without_classes():
     import threading
 
diff --git a/studio/backend/tests/test_diffusion_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py
index b10dd04f4c..451f594ba1 100644
--- a/studio/backend/tests/test_diffusion_dataset_api.py
+++ b/studio/backend/tests/test_diffusion_dataset_api.py
@@ -177,15 +177,29 @@ def test_delete_image_cleans_sidecar_and_thumb(client, ds_root):
     folder.mkdir()
     _write_png(folder / "x.png")
     (folder / "x.txt").write_text("cap", encoding = "utf-8")
-    # Generate a thumbnail so we can assert it is cleaned up too.
+    # Generate a thumbnail so we can assert it is cleaned up too. Thumbs are keyed on
+    # the full filename (stem + extension) to avoid same-stem collisions across formats.
     client.get("/api/train/diffusion/dataset/d/image/x.png?thumb=32")
-    assert list((folder / ".thumbs").glob("x_*.jpg"))
+    assert list((folder / ".thumbs").glob("x.png_*.jpg"))
 
     r = client.delete("/api/train/diffusion/dataset/d/image/x.png")
     assert r.status_code == 200, r.text
     assert not (folder / "x.png").exists()
     assert not (folder / "x.txt").exists()
-    assert not list((folder / ".thumbs").glob("x_*.jpg"))
+    assert not list((folder / ".thumbs").glob("x.png_*.jpg"))
+
+
+def test_thumb_cache_key_distinguishes_same_stem_extensions(client, ds_root):
+    # sample.png and sample.jpg share a stem; each must get its OWN thumbnail cache
+    # file, so the labeling grid never serves one image's thumbnail for the other.
+    folder = ds_root / "d"
+    folder.mkdir()
+    Image.new("RGB", (8, 8), (10, 20, 30)).save(folder / "sample.png", format = "PNG")
+    Image.new("RGB", (8, 8), (200, 210, 220)).save(folder / "sample.jpg", format = "JPEG")
+    client.get("/api/train/diffusion/dataset/d/image/sample.png?thumb=32")
+    client.get("/api/train/diffusion/dataset/d/image/sample.jpg?thumb=32")
+    thumbs = sorted(p.name for p in (folder / ".thumbs").glob("*.jpg"))
+    assert thumbs == ["sample.jpg_32.jpg", "sample.png_32.jpg"]
 
 
 # ── traversal / validation ───────────────────────────────────────────────────
diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py
index b5acaa7bb9..b138de3980 100644
--- a/studio/backend/tests/test_diffusion_training.py
+++ b/studio/backend/tests/test_diffusion_training.py
@@ -20,7 +20,7 @@ import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 
-from auth.authentication import get_current_subject
+from auth.authentication import authenticated_via_api_key, get_current_subject
 from core.training.diffusion_training_service import DiffusionTrainingService
 from routes.training import router as training_router
 
@@ -275,8 +275,12 @@ def client(monkeypatch):
     app = FastAPI()
     app.include_router(training_router, prefix = "/api/train")
     app.dependency_overrides[get_current_subject] = lambda: "test-user"
+    # Default to session (UI) auth: the API-key inference-in-flight guard is a no-op there,
+    # so the wiring tests below behave as before. The guard test flips this override.
+    app.dependency_overrides[authenticated_via_api_key] = lambda: False
     c = TestClient(app)
     c._fake = fake  # type: ignore[attr-defined]
+    c._app = app  # type: ignore[attr-defined]
     return c
 
 
@@ -365,6 +369,37 @@ def test_route_start_conflict_maps_to_409(client):
     assert r.status_code == 409
 
 
+def test_route_start_over_api_with_inference_in_flight_is_409(client, monkeypatch):
+    # An API-key client must not start diffusion training (which frees VRAM by unloading
+    # chat) while an inference request is streaming; it should 409 instead of killing it.
+    client._app.dependency_overrides[authenticated_via_api_key] = lambda: True
+    monkeypatch.setattr(
+        "core.inference.llama_keepwarm.other_inference_request_count",
+        lambda current_request_counted = False: 1,
+    )
+    freed = {"called": False}
+    import routes.training as tr
+
+    monkeypatch.setattr(
+        tr, "_free_gpu_for_diffusion_training", lambda: freed.__setitem__("called", True)
+    )
+    r = client.post("/api/train/diffusion/start", json = _BODY)
+    assert r.status_code == 409
+    # The guard must run BEFORE any GPU is freed, so the live inference stream survives.
+    assert freed["called"] is False
+
+
+def test_route_start_over_api_without_inference_proceeds(client, monkeypatch):
+    # Same API-key path but no inference in flight: the start proceeds normally.
+    client._app.dependency_overrides[authenticated_via_api_key] = lambda: True
+    monkeypatch.setattr(
+        "core.inference.llama_keepwarm.other_inference_request_count",
+        lambda current_request_counted = False: 0,
+    )
+    r = client.post("/api/train/diffusion/start", json = _BODY)
+    assert r.status_code == 200
+
+
 def test_route_status_and_stop(client):
     client.post("/api/train/diffusion/start", json = _BODY)
     s = client.get("/api/train/diffusion/status")
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
index c380e04244..f096af73c8 100644
--- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
+++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
@@ -2282,6 +2282,14 @@ export function HubModelPicker({
     }
 
     if (section === "recommended") {
+      // Curated safetensors rows render ABOVE the recommended rows (and call
+      // getOptionProps), so their keys must lead here or they fall back to the
+      // duplicate ...-option-missing id and drop out of arrow-key navigation.
+      keys.push(
+        ...curatedSafetensorsRows.map((m) =>
+          makeModelOptionKey("curated-safetensors", m.id),
+        ),
+      );
       keys.push(
         ...recommendedRows.map((r) => makeModelOptionKey("recommended", r.id)),
       );
@@ -2291,6 +2299,7 @@ export function HubModelPicker({
   }, [
     cachedReady,
     chatOnly,
+    curatedSafetensorsRows,
     sortedCustomFolderModels,
     customFoldersCollapsed,
     downloadedCollapsed,

From e80067512845c19f21ef89b92ced7c52307cc148 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 11:40:31 +0000
Subject: [PATCH 8/8] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 studio/backend/core/inference/diffusion.py        | 1 -
 studio/backend/tests/test_diffusion_controlnet.py | 8 ++++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py
index 603e4ca637..b40add9dd2 100644
--- a/studio/backend/core/inference/diffusion.py
+++ b/studio/backend/core/inference/diffusion.py
@@ -1544,7 +1544,6 @@ class DiffusionBackend:
             # local dir the user picked has no Hub scan and is exempt (fail-open there).
             if not getattr(resolved_cn, "is_local", False):
                 from utils.security import evaluate_file_security
-
                 _cn_fs = evaluate_file_security(resolved_cn.path, hf_token = state.hf_token or None)
                 if _cn_fs.blocked:
                     raise ValueError(_cn_fs.reason)
diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py
index b7ae037197..bcb7393457 100644
--- a/studio/backend/tests/test_diffusion_controlnet.py
+++ b/studio/backend/tests/test_diffusion_controlnet.py
@@ -234,7 +234,6 @@ def _state():
 def _allow_cn_security(monkeypatch):
     """Stub the Hub malware preflight to allow the load (hermetic, no network)."""
     import utils.security
-
     monkeypatch.setattr(
         utils.security,
         "evaluate_file_security",
@@ -277,7 +276,12 @@ def test_controlnet_pipe_blocks_flagged_remote_repo(monkeypatch):
 
     class _TrapModel(_FakeCNModel):
         @classmethod
-        def from_pretrained(cls, path, torch_dtype = None, token = None):
+        def from_pretrained(
+            cls,
+            path,
+            torch_dtype = None,
+            token = None,
+        ):
             loaded["called"] = True
             return super().from_pretrained(path, torch_dtype = torch_dtype, token = token)