cp: fix(diffusion): reuse warm HF cache instead of re-downloading models (2747) into r0.5.0

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)

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)
 

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__()

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"

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,

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,

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(

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,

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(