🗃️ Implement Local-first behavior to be resilient when HF/network is down (#226)

Co-authored-by: Anthony Wu <pls-file-gh-issue@users.noreply.github.com>

src/mflux/controlnet/weight_handler_controlnet.py

@@ -2,9 +2,9 @@ import json
 from pathlib import Path
 
 import mlx.core as mx
-from huggingface_hub import snapshot_download
 from mlx.utils import tree_unflatten
 
+from mflux.weights.download import snapshot_download
 from mflux.weights.weight_handler import MetaData
 from mflux.weights.weight_util import WeightUtil
 

src/mflux/flux_tools/redux/weight_handler_redux.py

@@ -1,8 +1,7 @@
 from pathlib import Path
 
-from huggingface_hub import snapshot_download
-
 from mflux import ModelConfig
+from mflux.weights.download import snapshot_download
 from mflux.weights.weight_handler import MetaData, WeightHandler
 
 

src/mflux/tokenizer/tokenizer_handler.py

@@ -1,9 +1,9 @@
 from pathlib import Path
 
 import transformers
-from huggingface_hub import snapshot_download
 
 from mflux.tokenizer.clip_tokenizer import TokenizerCLIP
+from mflux.weights.download import snapshot_download
 
 
 class TokenizerHandler:

src/mflux/weights/download.py

@@ -0,0 +1,41 @@
+"""Download utilities for mflux weights with cache-first behavior."""
+
+from huggingface_hub import snapshot_download as hf_snapshot_download
+from huggingface_hub.utils import LocalEntryNotFoundError
+
+
+def snapshot_download(repo_id: str, **kwargs) -> str:
+    """Download repo files with cache-first behavior.
+
+    This wrapper function attempts to use cached files first before downloading.
+    If local_files_only is explicitly set to True, it will not fall back to downloading.
+    Otherwise, it will try local cache first, then download if cache is not found.
+
+    Args:
+        repo_id: A user or an organization name and a repo name separated by a `/`.
+        **kwargs: Additional arguments passed to huggingface_hub.snapshot_download
+
+    Returns:
+        str: Path to the downloaded/cached repository
+
+    Docs: https://huggingface.co/docs/huggingface_hub/v0.33.1/en/package_reference/file_download#huggingface_hub.snapshot_download
+    """  # fmt: off
+    # Extract relevant kwargs for our logic
+    force_download = kwargs.get("force_download", False)
+    local_files_only = kwargs.get("local_files_only", False)
+
+    # If force_download or local_files_only is explicitly set, use original behavior
+    if force_download or local_files_only:
+        return hf_snapshot_download(repo_id=repo_id, **kwargs)
+
+    # Try to use cached files first
+    try:
+        # Override local_files_only to True for cache check
+        cache_kwargs = kwargs.copy()
+        cache_kwargs["local_files_only"] = True
+        return hf_snapshot_download(repo_id=repo_id, **cache_kwargs)
+    except LocalEntryNotFoundError:
+        # Cache doesn't exist, download from Hugging Face
+        download_kwargs = kwargs.copy()
+        download_kwargs["local_files_only"] = False
+        return hf_snapshot_download(repo_id=repo_id, **download_kwargs)

src/mflux/weights/weight_handler.py

@@ -2,9 +2,9 @@ from dataclasses import dataclass
 from pathlib import Path
 
 import mlx.core as mx
-from huggingface_hub import snapshot_download
 from mlx.utils import tree_unflatten
 
+from mflux.weights.download import snapshot_download
 from mflux.weights.lora_converter import LoRAConverter
 from mflux.weights.weight_util import WeightUtil
 

src/mflux/weights/weight_handler_lora_huggingface.py

@@ -1,9 +1,8 @@
 import shutil
 from pathlib import Path
 
-from huggingface_hub import snapshot_download
-
 from mflux.ui.defaults import MFLUX_LORA_CACHE_DIR
+from mflux.weights.download import snapshot_download
 
 
 class WeightHandlerLoRAHuggingFace: