From a187bbbab0d353e1d770234cda8f412da1a5ba5e Mon Sep 17 00:00:00 2001 From: Anthony Wu <462072+anthonywu@users.noreply.github.com> Date: Sat, 5 Jul 2025 04:55:17 -0700 Subject: [PATCH] =?UTF-8?q?=F0=9F=97=83=EF=B8=8F=20Implement=20Local-first?= =?UTF-8?q?=20behavior=20to=20be=20resilient=20when=20HF/network=20is=20do?= =?UTF-8?q?wn=20(#226)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Anthony Wu --- .../controlnet/weight_handler_controlnet.py | 2 +- .../flux_tools/redux/weight_handler_redux.py | 3 +- src/mflux/tokenizer/tokenizer_handler.py | 2 +- src/mflux/weights/download.py | 41 +++++++++++++++++++ src/mflux/weights/weight_handler.py | 2 +- .../weight_handler_lora_huggingface.py | 3 +- 6 files changed, 46 insertions(+), 7 deletions(-) create mode 100644 src/mflux/weights/download.py diff --git a/src/mflux/controlnet/weight_handler_controlnet.py b/src/mflux/controlnet/weight_handler_controlnet.py index 411aee7..e4ab354 100644 --- a/src/mflux/controlnet/weight_handler_controlnet.py +++ b/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 diff --git a/src/mflux/flux_tools/redux/weight_handler_redux.py b/src/mflux/flux_tools/redux/weight_handler_redux.py index ee3a1b8..6140c0f 100644 --- a/src/mflux/flux_tools/redux/weight_handler_redux.py +++ b/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 diff --git a/src/mflux/tokenizer/tokenizer_handler.py b/src/mflux/tokenizer/tokenizer_handler.py index 68ae3b6..d90accd 100644 --- a/src/mflux/tokenizer/tokenizer_handler.py +++ b/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: diff --git a/src/mflux/weights/download.py b/src/mflux/weights/download.py new file mode 100644 index 0000000..7b4e7d7 --- /dev/null +++ b/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) diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py index b04cbe2..bc4de4d 100644 --- a/src/mflux/weights/weight_handler.py +++ b/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 diff --git a/src/mflux/weights/weight_handler_lora_huggingface.py b/src/mflux/weights/weight_handler_lora_huggingface.py index 012f03e..baa10d7 100644 --- a/src/mflux/weights/weight_handler_lora_huggingface.py +++ b/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: