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", diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index c36d597197..06496bb1a4 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -437,11 +437,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). @@ -450,12 +456,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, @@ -672,7 +682,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, @@ -899,6 +909,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 @@ -1015,6 +1026,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 @@ -1058,7 +1070,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, @@ -1094,7 +1111,7 @@ 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 @@ -1326,6 +1343,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)``. @@ -1373,7 +1391,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 @@ -1381,7 +1399,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, @@ -1402,13 +1422,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 @@ -1555,6 +1576,16 @@ class DiffusionBackend: self._cn_models.clear() self._cn_pipes.clear() clear_gpu_cache() + # 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 @@ -1598,7 +1629,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/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index 065fd4eb11..276dcad3eb 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -507,6 +507,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] = { @@ -536,11 +546,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: diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 02237b86fb..a6dd63f4bc 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -1100,6 +1100,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 61be165992..832d1d38a2 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: @@ -1234,6 +1252,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() @@ -1601,6 +1634,10 @@ async def get_diffusion_dataset_image( thumbs_dir = folder / _THUMBS_DIRNAME thumbs_dir.mkdir(exist_ok = True) + # 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: @@ -1646,12 +1683,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) @@ -1674,6 +1723,9 @@ 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(): + # 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_backend.py b/studio/backend/tests/test_diffusion_backend.py index b0b5eb1807..00bbd7316c 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -2121,3 +2121,34 @@ 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() diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index ced1b752ac..ab5eb71b22 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -232,14 +232,28 @@ 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 + # 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) @@ -250,6 +264,69 @@ 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 @@ -260,3 +337,19 @@ 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_dataset_api.py b/studio/backend/tests/test_diffusion_dataset_api.py index 154b671215..dd8f94f556 100644 --- a/studio/backend/tests/test_diffusion_dataset_api.py +++ b/studio/backend/tests/test_diffusion_dataset_api.py @@ -177,7 +177,8 @@ 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") # Thumb cache key includes the extension (x.png_32.jpg), so png and jpg # siblings can't collide. @@ -190,6 +191,19 @@ def test_delete_image_cleans_sidecar_and_thumb(client, ds_root): 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 ─────────────────────────────────────────────────── def test_dataset_name_traversal_rejected_over_http(client, ds_root): # A name that fails the folder-name validator returns 400, never touches disk. diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index 39a67345cc..0610423d0b 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 @@ -265,11 +265,22 @@ 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" + # 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 @@ -358,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,