diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index f1fbcccfa1..0eef14de04 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -396,6 +396,34 @@ class DiffusionBackend: return hf_hub_download(repo_id, gguf_filename, token = hf_token) + def _dense_quant_prefetch_needed(self, fam: DiffusionFamily, kwargs: dict) -> bool: + """True when ``load_pipeline`` may take the dense transformer-quant path, so + the prefetch should also pull the base repo's ``transformer/`` shards. + + Those shards are excluded from the prefetch by default (the GGUF supplies + the transformer), but ``_load_dense_quant_pipeline`` fetches them with + ``from_pretrained(subfolder = "transformer")`` under the load lock during + "finalizing", after the previous pipeline was already evicted, where + unload/cancellation cannot preempt the download. Mirrors the dense-path + gates in ``load_pipeline``: quant requested and supported for this device, + and no pre-quantized checkpoint that would shortcut the dense build.""" + mode = normalize_transformer_quant(kwargs.get("transformer_quant")) + if mode is None: + return False + try: + target = self._resolve_device_target(fam) + if not dense_transformer_supported(target): + return False + scheme = select_transformer_quant_scheme(target, mode) + if scheme is None: + return False + source = resolve_prequant_source( + fam, scheme, path_override = kwargs.get("transformer_prequant_path") + ) + return source is None + except Exception: # noqa: BLE001 — widening the prefetch is best-effort only + return False + def _prefetch_files( self, repo_id: str, @@ -594,6 +622,14 @@ class DiffusionBackend: kwargs.get("hf_token"), kind = kind, single_file_is_pipeline = bool(fam and fam.single_file_is_pipeline), + # The dense transformer-quant path downloads the base repo's + # transformer/ shards via from_pretrained(subfolder="transformer") + # INSIDE the locked finalize phase, where unload/cancellation cannot + # preempt the multi-GB pull. When that path can actually run, pull the + # shards here in the preemptible prefetch instead. (Pipeline loads + # already include transformer/ via their own filter.) + include_transformer = kind == "gguf" + and self._dense_quant_prefetch_needed(fam, kwargs), ) with self._lock: # Stamp progress only if this load is still current; a superseding @@ -673,6 +709,7 @@ class DiffusionBackend: *, kind: str = "gguf", single_file_is_pipeline: bool = False, + include_transformer: bool = False, ) -> tuple[int, list[str]]: """Total download size for the progress bar, plus the base-repo files to fetch (the prefetch reuses this list, so the base is listed only once). @@ -714,11 +751,13 @@ class DiffusionBackend: total += sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename) # A whole-pipeline single file (SDXL) needs only the base repo's config/tokenizer, # not its (unused, multi-GB) weight files. - base_filter = ( - _base_config_file_downloaded - if (kind == "single_file" and single_file_is_pipeline) - else _base_file_downloaded - ) + if kind == "single_file" and single_file_is_pipeline: + base_filter = _base_config_file_downloaded + else: + + def base_filter(rfilename: str) -> bool: + return _base_file_downloaded(rfilename, include_transformer = include_transformer) + base_info = api.model_info(base_repo, files_metadata = True, token = hf_token) for s in base_info.siblings: if base_filter(s.rfilename): @@ -2105,16 +2144,18 @@ def _offload_controlnet_module(cn_model: Any, device: str, logger: Any) -> bool: return False -def _base_file_downloaded(rfilename: str) -> bool: +def _base_file_downloaded(rfilename: str, *, include_transformer: bool = False) -> bool: """True for base-repo files ``from_pretrained`` actually fetches. The transformer is supplied by the GGUF, and repo docs (``assets/``, the top-level README/PDF/images) are never downloaded — counting them would peg the progress estimate above what lands on disk, so the bar would sit short of 100% for the whole pipeline-load phase instead of advancing to "finalizing". - """ + ``include_transformer`` admits the ``transformer/`` shards for loads where the + dense transformer-quant path will fetch them anyway (see + ``_dense_quant_prefetch_needed``).""" if rfilename.startswith("transformer/"): - return False + return include_transformer if "/" not in rfilename: # top-level: only the pipeline manifest is fetched return rfilename == "model_index.json" return not rfilename.startswith("assets/") diff --git a/studio/backend/core/inference/diffusion_memory.py b/studio/backend/core/inference/diffusion_memory.py index 1735ccbbd6..c56d1592ce 100644 --- a/studio/backend/core/inference/diffusion_memory.py +++ b/studio/backend/core/inference/diffusion_memory.py @@ -378,7 +378,17 @@ def plan_diffusion_memory( policy = OFFLOAD_MODEL reasons.append("companions exceed budget; whole-module offload of every component") - if explicit_offload and policy == OFFLOAD_NONE and can_offload and not device_memory.is_unified: + # The legacy cpu_offload flag only applies when NO memory_mode was supplied: + # the API documents memory_mode as overriding cpu_offload when set, so an + # explicit `fast` request must stay resident even if the caller also left the + # old flag enabled. + if ( + explicit_offload + and normalize_memory_mode(requested_mode) is None + and policy == OFFLOAD_NONE + and can_offload + and not device_memory.is_unified + ): policy = OFFLOAD_MODEL reasons.append("explicit cpu_offload overrides resident placement") diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 0aa9597225..d70088b759 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -1922,6 +1922,52 @@ def test_transformer_quant_unsupported_scheme_skips_dense_download( assert _FakeTransformer.last["path"] # GGUF from_single_file used +def test_base_file_downloaded_include_transformer_flag(): + # Default: transformer/ shards are the GGUF's job, so they are excluded from + # the prefetch list; the dense transformer-quant path opts them back in. + from core.inference.diffusion import _base_file_downloaded + + assert _base_file_downloaded("transformer/diffusion_pytorch_model-00001.safetensors") is False + assert ( + _base_file_downloaded( + "transformer/diffusion_pytorch_model-00001.safetensors", include_transformer = True + ) + is True + ) + # The flag must not admit anything else that is normally excluded. + assert _base_file_downloaded("assets/teaser.png", include_transformer = True) is False + assert _base_file_downloaded("README.md", include_transformer = True) is False + + +def test_dense_quant_prefetch_needed_gates(fake_runtime, monkeypatch): + # The transformer/ prefetch only widens when the dense quant path can really + # run: quant requested + device supported + scheme resolvable + no prequant + # checkpoint shortcutting the dense build. + from core.inference import diffusion as dmod + + backend = DiffusionBackend() + _force_cuda_target(backend, monkeypatch) + fam = detect_family("unsloth/Z-Image-Turbo-GGUF") + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: True) + monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8") + monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None) + + assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is True + # No quant requested -> never widen. + assert backend._dense_quant_prefetch_needed(fam, {}) is False + # A resolvable pre-quantized checkpoint shortcuts the dense download. + monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: object()) + assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False + # Unsupported scheme bails before the dense path (and so must the prefetch). + monkeypatch.setattr(dmod, "resolve_prequant_source", lambda fam, scheme, **kw: None) + monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: None) + assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False + # Device without dense support (e.g. non-CUDA) never widens. + monkeypatch.setattr(dmod, "select_transformer_quant_scheme", lambda target, mode: "fp8") + monkeypatch.setattr(dmod, "dense_transformer_supported", lambda target: False) + assert backend._dense_quant_prefetch_needed(fam, {"transformer_quant": "fp8"}) is False + + def test_companion_cache_bytes_local_dir_excludes_transformer(tmp_path): # A LOCAL diffusers base: sum the on-disk VAE / text-encoder weights so auto memory # planning sees the resident companions, but exclude transformer/ (the GGUF supplies diff --git a/studio/backend/tests/test_diffusion_memory.py b/studio/backend/tests/test_diffusion_memory.py index bc9ba592ce..e6cc054648 100644 --- a/studio/backend/tests/test_diffusion_memory.py +++ b/studio/backend/tests/test_diffusion_memory.py @@ -270,6 +270,21 @@ def test_explicit_cpu_offload_overrides_resident_auto_choice(): assert any("explicit cpu_offload" in r for r in plan.reasons) +def test_explicit_memory_mode_wins_over_legacy_cpu_offload(): + # The API documents memory_mode as overriding cpu_offload when set: fast + + # the legacy flag must stay resident, not silently downgrade to offload. + plan = plan_diffusion_memory( + target = _target(), + device_memory = _discrete(80000), + model_dense_mib = 4000, + runtime_headroom_mib = 2000, + requested_mode = MEMORY_MODE_FAST, + explicit_offload = True, + ) + assert plan.offload_policy == OFFLOAD_NONE + assert not any("explicit cpu_offload" in r for r in plan.reasons) + + def test_explicit_cpu_offload_ignored_on_cpu_target(): plan = plan_diffusion_memory( target = _target(device = "cpu", backend = "cpu", supports_offload = False),