ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Pin diffusion and video loads to the live HF cache root

Browse source 
Both read huggingface_hub's import-time HF_HUB_CACHE, which changing the cache
folder does not update: progress counted the old root while the download wrote to
the new one, and from_pretrained could split one model across both.
This commit is contained in:
Unsloth 2026-07-26 03:52:45 -07:00
commit c095ccb96a
5 changed files with 91 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

51
studio/backend/core/inference/diffusion.py
View file

@ -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(

11
studio/backend/core/inference/video.py
View file

@ -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(

4
studio/backend/tests/test_diffusion_controlnet.py
View file

@ -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):

34
studio/backend/tests/test_hf_cache_settings.py
View file

@ -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"
)

6
studio/backend/tests/test_video_backend.py
View file

@ -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