diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index aafa6abe8f..188d14016e 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -133,6 +133,16 @@ logger = get_logger(__name__) _MODEL_KINDS = frozenset({"gguf", "single_file", "pipeline"}) +def hub_cache_dir() -> str: + """The cache root every loader call must be pinned to. + + diffusers resolves an unset cache_dir through huggingface_hub's import-time constant, + which a mid-session cache-folder change does not update. The prefetch reads the live + setting, so without this a single load could split across two roots.""" + from utils.hf_cache_settings import active_hf_hub_cache + return active_hf_hub_cache() + + def resolve_model_kind(gguf_filename: Optional[str], model_kind: Optional[str] = None) -> str: """Classify a load request into one of ``_MODEL_KINDS``. @@ -1041,9 +1051,12 @@ class DiffusionBackend: @staticmethod def _hub_cache_repo_dir(repo_id: str) -> Path: - """Local HF hub cache dir for ``repo_id``.""" - from huggingface_hub import constants - return Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}" + """Local HF hub cache dir for ``repo_id``. + + Reads the live setting, not huggingface_hub's import-time constant: changing the + cache folder does not update the constant, so the old one would count bytes in a + root the download no longer writes to (progress stuck at 0 for the whole pull).""" + return Path(hub_cache_dir()) / f"models--{repo_id.replace('/', '--')}" @staticmethod def _cache_bytes(repo_id: str) -> int: @@ -1503,7 +1516,10 @@ class DiffusionBackend: # per-component too (see diffusion_ideogram4.py). pipe = load_ideogram4_pipeline(repo_id, dtype, hf_token = hf_token) else: - pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype} + pipe_kwargs: dict[str, Any] = { + "torch_dtype": dtype, + "cache_dir": hub_cache_dir(), + } if hf_token: pipe_kwargs["token"] = hf_token if fam.name == HIDREAM_FAMILY_NAME: @@ -1541,7 +1557,11 @@ class DiffusionBackend: # A single-file SDXL-style checkpoint is the WHOLE pipeline, so load it # through the pipeline class; ``config`` points at the base repo so diffusers # builds the correct structure around the single-file weights. - sf_pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "config": base} + sf_pipe_kwargs: dict[str, Any] = { + "torch_dtype": dtype, + "config": base, + "cache_dir": hub_cache_dir(), + } if hf_token: sf_pipe_kwargs["token"] = hf_token pipe = pipeline_cls.from_single_file(single_file_path, **sf_pipe_kwargs) @@ -1553,6 +1573,7 @@ class DiffusionBackend: "subfolder": "transformer", # Config is fetched from the (possibly gated) base before auth. "token": hf_token, + "cache_dir": hub_cache_dir(), } if kind == "gguf": # Dequantise the GGUF transformer on-device at the compute dtype. @@ -1585,7 +1606,11 @@ class DiffusionBackend: ).get("text_encoder"), ) else: - pipe_kwargs = {"torch_dtype": dtype, "transformer": transformer} + pipe_kwargs = { + "torch_dtype": dtype, + "transformer": transformer, + "cache_dir": hub_cache_dir(), + } if hf_token: pipe_kwargs["token"] = hf_token if fam.name == HIDREAM_FAMILY_NAME: @@ -2016,7 +2041,11 @@ class DiffusionBackend: "prequant checkpoint unavailable and the dense transformer does not fit resident" ) transformer = transformer_cls.from_pretrained( - base, subfolder = "transformer", torch_dtype = dtype, token = hf_token + base, + subfolder = "transformer", + torch_dtype = dtype, + token = hf_token, + cache_dir = hub_cache_dir(), ) pipe = self._assemble_pipe( pipeline_cls, @@ -2104,7 +2133,11 @@ class DiffusionBackend: ) pipe.to(device) return pipe - pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "transformer": transformer} + pipe_kwargs: dict[str, Any] = { + "torch_dtype": dtype, + "transformer": transformer, + "cache_dir": hub_cache_dir(), + } if hf_token: pipe_kwargs["token"] = hf_token if getattr(fam, "name", None) == HIDREAM_FAMILY_NAME: @@ -2327,7 +2360,7 @@ class DiffusionBackend: cn_dtype = getattr(torch, str(state.dtype).replace("torch.", ""), None) # Force safetensors for an untrusted remote repo: if the Hub scan failed open above, an # embedded pickle would still deserialize on load. A local dir the user chose is exempt. - cn_from_pretrained_kwargs: dict[str, Any] = {} + cn_from_pretrained_kwargs: dict[str, Any] = {"cache_dir": hub_cache_dir()} if remote_cn: cn_from_pretrained_kwargs["use_safetensors"] = True cn_model = getattr(diffusers, model_cls_name).from_pretrained( diff --git a/studio/backend/core/inference/video.py b/studio/backend/core/inference/video.py index 272890fa47..60a0764a1c 100644 --- a/studio/backend/core/inference/video.py +++ b/studio/backend/core/inference/video.py @@ -92,6 +92,8 @@ from .video_families import ( supported_video_family_names, ) from utils.hardware import clear_gpu_cache +# Shared with the image backend so both pin every loader call to the same live cache root. +from core.inference.diffusion import hub_cache_dir logger = get_logger(__name__) @@ -784,9 +786,7 @@ class VideoBackend: try: import os - from huggingface_hub.constants import HF_HUB_CACHE - - folder = Path(HF_HUB_CACHE) / ("models--" + repo_id.strip().replace("/", "--")) + folder = Path(hub_cache_dir()) / ("models--" + repo_id.strip().replace("/", "--")) if not folder.is_dir(): return 0 total = 0 @@ -999,7 +999,9 @@ class VideoBackend: # ── build the pipeline. pipeline_cls = getattr(diffusers, fam.pipeline_class) - pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype} + # cache_dir pins every loader call to the live cache root, so a mid-session + # change can't split one model across the old and new roots. + pipe_kwargs: dict[str, Any] = {"torch_dtype": dtype, "cache_dir": hub_cache_dir()} if getattr(fam, "vae_force_fp32", False): # Wan's VAE must decode in float32. A scalar torch_dtype truncates its fp32 weights # to bf16 (no _keep_in_fp32_modules); a later .to(float32) only widens lossy values @@ -1036,6 +1038,7 @@ class VideoBackend: "config": base, "subfolder": "transformer", "token": hf_token, + "cache_dir": hub_cache_dir(), } if kind == "gguf": sf_kwargs["quantization_config"] = diffusers.GGUFQuantizationConfig( diff --git a/studio/backend/tests/test_diffusion_controlnet.py b/studio/backend/tests/test_diffusion_controlnet.py index 36ce756773..fb11160960 100644 --- a/studio/backend/tests/test_diffusion_controlnet.py +++ b/studio/backend/tests/test_diffusion_controlnet.py @@ -240,10 +240,14 @@ class _FakeCNModel: torch_dtype = None, token = None, use_safetensors = None, + # cache_dir (and any future loader kwarg) rides through: the real call pins the + # live cache root so a load cannot split across two of them. + **kwargs, ): m = cls() m.path = path m.use_safetensors = use_safetensors + m.cache_dir = kwargs.get("cache_dir") return m def to(self, device): diff --git a/studio/backend/tests/test_hf_cache_settings.py b/studio/backend/tests/test_hf_cache_settings.py index 1875d61809..4997c6609f 100644 --- a/studio/backend/tests/test_hf_cache_settings.py +++ b/studio/backend/tests/test_hf_cache_settings.py @@ -288,3 +288,37 @@ def test_inactive_cache_model_loads_from_snapshot_path(tmp_path): assert row.model_id == "org/model" assert row.active_cache is False assert row.load_id == str(snapshot) + + +def test_diffusion_cache_root_follows_a_live_switch(settings_store, tmp_path): + # The image/video backends used huggingface_hub's import-time HF_HUB_CACHE constant, + # which set_hf_cache_home does not update. The download then wrote to the new root + # while progress counted the old one, and a load could split across both. + import core.inference.diffusion as diffusion + + moved = tmp_path / "external-c" / "huggingface" + # Write the setting straight into the store: set_hf_cache_home's folder validation is + # not what is under test, and it rejects the pytest tmp root on macOS. + settings_store[hf_cache_settings.CACHE_HOME_SETTING_KEY] = str(moved) + + assert diffusion.hub_cache_dir() == str(moved / "hub") + assert diffusion.DiffusionBackend._hub_cache_repo_dir("org/model") == ( + moved / "hub" / "models--org--model" + ) + + +def test_diffusion_loader_calls_pin_the_cache_dir(): + # Every from_pretrained / from_single_file must carry cache_dir, else diffusers + # resolves it through the stale constant and half a load lands in the old root. + for rel in ("core/inference/diffusion.py", "core/inference/video.py"): + source = (Path(_BACKEND_DIR) / rel).read_text(encoding = "utf-8") + for call in ("from_pretrained(", "from_single_file("): + for index, line in enumerate(source.splitlines(), start = 1): + if not line.strip().startswith(("pipe = ", "transformer = ", "cn_model = ")): + continue + if call not in line: + continue + window = "\n".join(source.splitlines()[index - 1 : index + 8]) + assert "cache_dir" in window or "kwargs" in window, ( + f"{rel}:{index} calls {call} without a pinned cache_dir" + ) diff --git a/studio/backend/tests/test_video_backend.py b/studio/backend/tests/test_video_backend.py index e936ca7e94..e875102484 100644 --- a/studio/backend/tests/test_video_backend.py +++ b/studio/backend/tests/test_video_backend.py @@ -1150,7 +1150,7 @@ def test_cache_bytes_counts_incomplete_blobs(fake_runtime, tmp_path, monkeypatch # scan_cache_dir skips in-flight *.incomplete blobs, so the old counter froze at the # last completed blob for the whole multi-GB shard pull. The walk must count both, # without double-counting snapshot symlinks. - import huggingface_hub.constants as hub_constants + import core.inference.video as video_mod repo_dir = tmp_path / "models--Wan-AI--Wan2.2-TI2V-5B-Diffusers" blobs = repo_dir / "blobs" @@ -1160,7 +1160,9 @@ def test_cache_bytes_counts_incomplete_blobs(fake_runtime, tmp_path, monkeypatch snap = repo_dir / "snapshots" / "deadbeef" snap.mkdir(parents = True) (snap / "model_index.json").symlink_to(blobs / "aa11") # must not double-count - monkeypatch.setattr(hub_constants, "HF_HUB_CACHE", str(tmp_path)) + # The live cache root, not huggingface_hub's import-time constant: the counter follows + # a mid-session cache-folder change, which the constant does not. + monkeypatch.setattr(video_mod, "hub_cache_dir", lambda: str(tmp_path)) backend = VideoBackend() assert backend._cache_bytes("Wan-AI/Wan2.2-TI2V-5B-Diffusers") == 1500