Hot Fix: coldstart weight download (version 0.9.3) (#235)

This commit is contained in:
Anthony Wu 2025-07-08 20:19:34 -07:00 committed by GitHub
parent 4fc63f535a
commit 42bf9e3d4d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 14 additions and 41 deletions

View File

@ -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'`

View File

@ -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"]

View File

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

View File

@ -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)