diff --git a/nemo_automodel/_diffusers/_hf_cache.py b/nemo_automodel/_diffusers/_hf_cache.py new file mode 100644 index 0000000000..f8225a6560 --- /dev/null +++ b/nemo_automodel/_diffusers/_hf_cache.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""HF cache resolution helpers for the diffusers bridge. + +Diffusers' ``from_pretrained`` resolves a bare repo id in *online* mode by +default: it issues per-file network requests to the Hub to revalidate ETags +before deciding whether to reuse the cache. Even a warm ``HF_HOME`` is then +re-validated over the network, and any ETag drift or partial cache turns into a +fresh download. The transformers bridge avoids this by pre-resolving the repo +to a local snapshot directory (see ``_resolve_model_dir`` in +``nemo_automodel/_transformers/model_init.py``) and handing that directory to +HF, which then does zero network I/O. This module ports the same discipline to +the diffusion path. +""" + +import os + +from nemo_automodel.shared.import_utils import safe_import + +HF_HUB_AVAILABLE, _ = safe_import("huggingface_hub") + +if HF_HUB_AVAILABLE: + from huggingface_hub import snapshot_download +else: + snapshot_download = None + + +def resolve_diffusion_model_dir(model_id: str) -> str: + """Resolve a HF repo id to a local snapshot directory. + + Mirrors the transformers bridge so a warm ``HF_HOME`` is never re-validated + over the network. Local paths are returned unchanged. For repo ids, the + snapshot is downloaded once only when the cache is cold and the process is + online (``HF_HUB_OFFLINE`` unset); the returned directory is then resolved + with ``local_files_only=True`` so the subsequent ``from_pretrained`` call + performs no network I/O. + + Args: + model_id: A HuggingFace repo id or a local filesystem path. + + Returns: + A local directory path containing the model snapshot. When + ``huggingface_hub`` is unavailable, ``model_id`` is returned unchanged + so resolution falls back to HF's own handling. + """ + if os.path.isdir(model_id) or not HF_HUB_AVAILABLE: + return model_id + + if os.environ.get("HF_HUB_OFFLINE", "0") != "1": + # Cold cache + online: fetch the snapshot once. + snapshot_download(model_id) + # Resolve (and require) the local snapshot without revalidating over the network. + return snapshot_download(model_id, local_files_only=True) diff --git a/nemo_automodel/_diffusers/auto_diffusion_pipeline.py b/nemo_automodel/_diffusers/auto_diffusion_pipeline.py index 7ef77ef835..62997ed20f 100644 --- a/nemo_automodel/_diffusers/auto_diffusion_pipeline.py +++ b/nemo_automodel/_diffusers/auto_diffusion_pipeline.py @@ -48,6 +48,7 @@ import torch import torch.nn as nn +from nemo_automodel._diffusers._hf_cache import resolve_diffusion_model_dir from nemo_automodel.components.distributed import DistributedSetup, ParallelismSizes, parallelizer from nemo_automodel.components.distributed.config import DDPConfig, FSDP2Config from nemo_automodel.components.distributed.ddp import DDPManager @@ -611,9 +612,13 @@ def from_pretrained( logger.info("[INFO] Loading pipeline from pretrained: %s", pretrained_model_name_or_path) + # Resolve to a local snapshot dir so a warm HF cache is not re-validated + # (and potentially re-downloaded) over the network on every run. + model_dir = resolve_diffusion_model_dir(pretrained_model_name_or_path) + # Use DiffusionPipeline.from_pretrained for auto-detection pipe: DiffusionPipeline = DiffusionPipeline.from_pretrained( - pretrained_model_name_or_path, + model_dir, *model_args, torch_dtype=torch_dtype, **kwargs, @@ -797,6 +802,10 @@ def from_config( logger.info("[INFO] Model ID: %s", model_id) logger.info("[INFO] Transformer class: %s", spec.transformer_cls) + # Resolve to a local snapshot dir so config/pipeline loads reuse the + # warm HF cache instead of re-validating over the network. + model_id = resolve_diffusion_model_dir(model_id) + # Dynamically import transformer class from diffusers TransformerCls = _import_diffusers_class(spec.transformer_cls) diff --git a/tests/unit_tests/_diffusers/test_auto_diffusion_pipeline.py b/tests/unit_tests/_diffusers/test_auto_diffusion_pipeline.py index adfc6cee36..4b0a4f313f 100644 --- a/tests/unit_tests/_diffusers/test_auto_diffusion_pipeline.py +++ b/tests/unit_tests/_diffusers/test_auto_diffusion_pipeline.py @@ -35,6 +35,19 @@ MODULE_PATH = "nemo_automodel._diffusers.auto_diffusion_pipeline" +@pytest.fixture(autouse=True) +def _no_hf_cache_resolution(): + """Keep unit tests offline. + + ``from_pretrained``/``from_config`` resolve the repo id to a local HF + snapshot dir via ``resolve_diffusion_model_dir``. Tests pass fake repo ids + (e.g. ``"dummy"``), so leaving this unmocked makes the helper hit the Hub. + Pass the id through unchanged instead. + """ + with patch(f"{MODULE_PATH}.resolve_diffusion_model_dir", side_effect=lambda model_id: model_id): + yield + + class DummyModule(torch.nn.Module): def __init__(self): super().__init__() diff --git a/tests/unit_tests/_diffusers/test_hf_cache.py b/tests/unit_tests/_diffusers/test_hf_cache.py new file mode 100644 index 0000000000..a68f1d3fa5 --- /dev/null +++ b/tests/unit_tests/_diffusers/test_hf_cache.py @@ -0,0 +1,59 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import patch + +from nemo_automodel._diffusers import _hf_cache +from nemo_automodel._diffusers._hf_cache import resolve_diffusion_model_dir + +MODULE = "nemo_automodel._diffusers._hf_cache" + + +def test_resolve_returns_local_path_unchanged(tmp_path): + # A directory that already exists should be returned verbatim without any + # Hub interaction. + with patch(f"{MODULE}.snapshot_download") as mock_sd: + assert resolve_diffusion_model_dir(str(tmp_path)) == str(tmp_path) + mock_sd.assert_not_called() + + +def test_resolve_offline_uses_cache_only(monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setattr(_hf_cache, "HF_HUB_AVAILABLE", True) + + with patch(f"{MODULE}.snapshot_download", return_value="/cache/snapshot") as mock_sd: + resolved = resolve_diffusion_model_dir("some/repo-id") + + # Offline: no cold-cache download, single cache-only resolution. + assert resolved == "/cache/snapshot" + mock_sd.assert_called_once_with("some/repo-id", local_files_only=True) + + +def test_resolve_online_downloads_then_resolves_locally(monkeypatch): + monkeypatch.delenv("HF_HUB_OFFLINE", raising=False) + monkeypatch.setattr(_hf_cache, "HF_HUB_AVAILABLE", True) + + with patch(f"{MODULE}.snapshot_download", return_value="/cache/snapshot") as mock_sd: + resolved = resolve_diffusion_model_dir("some/repo-id") + + assert resolved == "/cache/snapshot" + # Online: fetch once (cold cache), then resolve the local dir without revalidation. + assert mock_sd.call_count == 2 + assert mock_sd.call_args_list[0].args == ("some/repo-id",) + assert mock_sd.call_args_list[1].kwargs == {"local_files_only": True} + + +def test_resolve_passthrough_when_hub_unavailable(monkeypatch): + monkeypatch.setattr(_hf_cache, "HF_HUB_AVAILABLE", False) + assert resolve_diffusion_model_dir("some/repo-id") == "some/repo-id" diff --git a/tools/diffusion/processors/flux.py b/tools/diffusion/processors/flux.py index 16c267ca45..a54b0ee71a 100644 --- a/tools/diffusion/processors/flux.py +++ b/tools/diffusion/processors/flux.py @@ -70,11 +70,17 @@ def load_models(self, model_name: str, device: str) -> Dict[str, Any]: """ from diffusers import FluxPipeline + from nemo_automodel._diffusers._hf_cache import resolve_diffusion_model_dir + logger.info("[FLUX] Loading models from %s via FluxPipeline...", model_name) # Patch T5 layer norm so it can run in bf16 (apex FusedRMSNorm doesn't support it) patch_t5_layer_norm() + # Resolve to a local snapshot dir so a warm HF cache is not re-validated + # (and potentially re-downloaded) over the network on every run. + model_name = resolve_diffusion_model_dir(model_name) + # Load pipeline without transformer (not needed for preprocessing) pipeline = FluxPipeline.from_pretrained( model_name, diff --git a/tools/diffusion/processors/flux2.py b/tools/diffusion/processors/flux2.py index 707d9cee56..6d8d604cb4 100644 --- a/tools/diffusion/processors/flux2.py +++ b/tools/diffusion/processors/flux2.py @@ -73,8 +73,14 @@ def load_models(self, model_name: str, device: str) -> Dict[str, Any]: """ from diffusers import Flux2Pipeline + from nemo_automodel._diffusers._hf_cache import resolve_diffusion_model_dir + logger.info("[FLUX.2] Loading models from %s via Flux2Pipeline...", model_name) + # Resolve to a local snapshot dir so a warm HF cache is not re-validated + # (and potentially re-downloaded) over the network on every run. + model_name = resolve_diffusion_model_dir(model_name) + # Load without transformer (not needed for preprocessing) pipeline = Flux2Pipeline.from_pretrained( model_name, diff --git a/tools/diffusion/processors/hunyuan.py b/tools/diffusion/processors/hunyuan.py index e2ef6b4b5e..a46bec3dca 100644 --- a/tools/diffusion/processors/hunyuan.py +++ b/tools/diffusion/processors/hunyuan.py @@ -98,8 +98,14 @@ def load_models(self, model_name: str, device: str) -> Dict[str, Any]: # loading so that the ByT5 text encoder uses a native implementation. patch_t5_layer_norm() + from nemo_automodel._diffusers._hf_cache import resolve_diffusion_model_dir + logger.info("[HunyuanVideo] Loading pipeline from %s...", model_name) + # Resolve to a local snapshot dir so a warm HF cache is not re-validated + # (and potentially re-downloaded) over the network on every run. + model_name = resolve_diffusion_model_dir(model_name) + # Load pipeline without transformer to save memory # cpu_offload=True helps manage VRAM pipeline = HunyuanVideo15ImageToVideoPipeline.from_pretrained( diff --git a/tools/diffusion/processors/qwen_image.py b/tools/diffusion/processors/qwen_image.py index 892a6f3071..26e9c6c268 100644 --- a/tools/diffusion/processors/qwen_image.py +++ b/tools/diffusion/processors/qwen_image.py @@ -65,8 +65,14 @@ def load_models(self, model_name: str, device: str) -> Dict[str, Any]: """ from diffusers import QwenImagePipeline + from nemo_automodel._diffusers._hf_cache import resolve_diffusion_model_dir + logger.info("[Qwen-Image] Loading models from %s...", model_name) + # Resolve to a local snapshot dir so a warm HF cache is not re-validated + # (and potentially re-downloaded) over the network on every run. + model_name = resolve_diffusion_model_dir(model_name) + # Load pipeline without transformer (not needed for preprocessing) pipeline = QwenImagePipeline.from_pretrained( model_name, diff --git a/tools/diffusion/processors/wan.py b/tools/diffusion/processors/wan.py index b6baf77df9..698fd6a1fa 100644 --- a/tools/diffusion/processors/wan.py +++ b/tools/diffusion/processors/wan.py @@ -119,8 +119,15 @@ def load_models(self, model_name: str, device: str) -> Dict[str, Any]: # UMT5 requires bfloat16 (float16 causes overflow/zeros in attention and layer norm) text_encoder_dtype = torch.bfloat16 if "cuda" in device else torch.float32 + from nemo_automodel._diffusers._hf_cache import resolve_diffusion_model_dir + logger.info("[Wan] Loading models from %s...", model_name) + # Resolve to a local snapshot dir once so the per-component subfolder + # loads below reuse the warm HF cache instead of re-validating over the + # network on every run. + model_name = resolve_diffusion_model_dir(model_name) + # Load text encoder logger.info(" Loading UMT5 text encoder...") text_encoder = UMT5EncoderModel.from_pretrained(