From 42bf9e3d4d3a39ad4e2fdbf9fa01103145891ad6 Mon Sep 17 00:00:00 2001 From: Anthony Wu <462072+anthonywu@users.noreply.github.com> Date: Tue, 8 Jul 2025 20:19:34 -0700 Subject: [PATCH] Hot Fix: coldstart weight download (version 0.9.3) (#235) --- README.md | 1 + pyproject.toml | 2 +- src/mflux/flux/flux_initializer.py | 16 ++++++------- src/mflux/weights/download.py | 36 ++++-------------------------- 4 files changed, 14 insertions(+), 41 deletions(-) diff --git a/README.md b/README.md index 4d1f657..bd61ca5 100644 --- a/README.md +++ b/README.md @@ -1854,6 +1854,7 @@ See `uv run tools/rename_images.py --help` for full CLI usage help. ### 💡Workflow Tips - To hide the model fetching status progress bars, `export HF_HUB_DISABLE_PROGRESS_BARS=1` +- To opt-out of [HuggingFace telemetry](https://huggingface.co/docs/huggingface_hub/main/en/package_reference/environment_variables#hfhubdisabletelemetry), `export HF_HUB_DISABLE_TELEMETRY=1` - Use config files to save complex job parameters in a file instead of passing many `--args` - Set up shell aliases for required args examples: - shortcut for dev model: `alias mflux-dev='mflux-generate --model dev'` diff --git a/pyproject.toml b/pyproject.toml index bbd5eb6..bed57f8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ source-exclude = [ [project] name = "mflux" -version = "0.9.2" +version = "0.9.3" description = "A MLX port of FLUX based on the Huggingface Diffusers implementation." readme = "README.md" keywords = ["diffusers", "flux", "mlx"] diff --git a/src/mflux/flux/flux_initializer.py b/src/mflux/flux/flux_initializer.py index dd6e601..1ae87f5 100644 --- a/src/mflux/flux/flux_initializer.py +++ b/src/mflux/flux/flux_initializer.py @@ -36,7 +36,14 @@ class FluxInitializer: flux_model.prompt_cache = {} flux_model.model_config = model_config - # 1. Initialize tokenizers + # 1. Load the regular weights + weights = WeightHandler.load_regular_weights( + repo_id=model_config.model_name, + local_path=local_path, + transformer_repo_id=model_config.custom_transformer_model, + ) + + # 2. Initialize tokenizers tokenizers = TokenizerHandler( repo_id=model_config.model_name, max_t5_length=model_config.max_sequence_length, @@ -50,13 +57,6 @@ class FluxInitializer: tokenizer=tokenizers.clip, ) - # 2. Load the regular weights - weights = WeightHandler.load_regular_weights( - repo_id=model_config.model_name, - local_path=local_path, - transformer_repo_id=model_config.custom_transformer_model, - ) - # 3. Initialize all models flux_model.vae = VAE() flux_model.t5_text_encoder = T5Encoder() diff --git a/src/mflux/weights/download.py b/src/mflux/weights/download.py index 7b4e7d7..96744de 100644 --- a/src/mflux/weights/download.py +++ b/src/mflux/weights/download.py @@ -1,41 +1,13 @@ """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. + This wrapper function is a wrapper for upstream hf_snapshot_download - 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) + 2025-07-08 HOT FIX: just pass the args through for now, see #235 discussion + """ + return hf_snapshot_download(repo_id, **kwargs)