diff --git a/README.md b/README.md index 14b8ba2..cfa43ee 100644 --- a/README.md +++ b/README.md @@ -163,13 +163,9 @@ This is useful for integrating MFLUX into shell scripts or dynamically generatin ⚠️ *If the specific model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell or Dev model respectively). See the [quantization](#%EF%B8%8F-quantization) section for running compressed versions of the model.* ⚠️ -*By default, model files are downloaded to the `.cache` folder within your home directory. For example, in my setup, the path looks like this:* +*By default, mflux caches files in `~/Library/Caches/mflux/`. The Hugging Face model files themselves are cached separately in the Hugging Face cache directory (e.g., `~/.cache/huggingface/`).* -``` -/Users/filipstrand/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-dev -``` - -*To change this default behavior, you can do so by modifying the `HF_HOME` environment variable. For more details on how to adjust this setting, please refer to the [Hugging Face documentation](https://huggingface.co/docs/huggingface_hub/en/package_reference/environment_variables)*. +*To change the mflux cache location, set the `MFLUX_CACHE_DIR` environment variable. To change the Hugging Face cache location, you can modify the `HF_HOME` environment variable. For more details on Hugging Face cache settings, please refer to the [Hugging Face documentation](https://huggingface.co/docs/huggingface_hub/en/package_reference/environment_variables)*. 🔒 [FLUX.1-dev currently requires granted access to its Huggingface repo. For troubleshooting, see the issue tracker](https://github.com/filipstrand/mflux/issues/14) 🔒 diff --git a/pyproject.toml b/pyproject.toml index 187232d..115fb67 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,6 +21,7 @@ dependencies = [ "piexif>=1.1.3,<2.0", "pillow>=10.4.0,<11.0; python_version<'3.13'", "pillow>=11.0,<12.0; python_version>='3.13'", + "platformdirs>=4.0,<5.0", "safetensors>=0.4.4,<1.0", # python 3.13 workaround for now: # use temporary community build of py13 wheel, use until official project build diff --git a/src/mflux/models/depth_pro/weight_handler_depth_pro.py b/src/mflux/models/depth_pro/weight_handler_depth_pro.py index d932889..ddabbf4 100644 --- a/src/mflux/models/depth_pro/weight_handler_depth_pro.py +++ b/src/mflux/models/depth_pro/weight_handler_depth_pro.py @@ -1,13 +1,12 @@ import logging -import os import urllib.error import urllib.request -from pathlib import Path import mlx.core as mx import torch from mlx.utils import tree_unflatten +from mflux.ui.defaults import MFLUX_CACHE_DIR from mflux.weights.weight_handler import MetaData from mflux.weights.weight_util import WeightUtil @@ -47,7 +46,8 @@ class WeightHandlerDepthPro: APPLE_MODEL_URL = "https://ml-site.cdn-apple.com/models/depth-pro/depth_pro.pt" # 1. Create cache directory for the model - cache_dir = Path(os.path.expanduser("~/.cache/mflux/depth_pro")) + cache_dir = MFLUX_CACHE_DIR / "depth_pro" + cache_dir.mkdir(parents=True, exist_ok=True) model_path = cache_dir / "depth_pro.pt" diff --git a/src/mflux/ui/defaults.py b/src/mflux/ui/defaults.py index 512a4db..4441d56 100644 --- a/src/mflux/ui/defaults.py +++ b/src/mflux/ui/defaults.py @@ -1,3 +1,12 @@ +import logging +import os +import shutil +from pathlib import Path + +import platformdirs + +logger = logging.getLogger(__name__) + BATTERY_PERCENTAGE_STOP_LIMIT = 5 CONTROLNET_STRENGTH = 0.4 DEFAULT_DEV_FILL_GUIDANCE = 30 @@ -14,3 +23,60 @@ MODEL_INFERENCE_STEPS = { "schnell": 4, } QUANTIZE_CHOICES = [3, 4, 6, 8] + + +def _migrate_legacy_cache(new_cache_dir: Path) -> None: + """Migrate legacy ~/.cache/mflux to new location if needed.""" + legacy_cache = Path.home() / ".cache" / "mflux" + + # Skip if legacy path doesn't exist or is already a symlink + if not legacy_cache.exists() or legacy_cache.is_symlink(): + return + + # Skip if we're already using the legacy path + if new_cache_dir == legacy_cache: + return + + try: + logger.warning(f"Migrating cache from {legacy_cache} to {new_cache_dir}") + + # Create new directory + new_cache_dir.mkdir(parents=True, exist_ok=True) + + # Move all contents from old to new location + for item in legacy_cache.iterdir(): + src = legacy_cache / item.name + dst = new_cache_dir / item.name + + if dst.exists(): + logger.warning(f" Skipping {item.name} (already exists in destination)") + continue + + logger.warning(f" Moving {item.name}") + shutil.move(str(src), str(dst)) + + # Remove the now-empty old directory + legacy_cache.rmdir() + + # Create symlink from old location to new location for backward compatibility + legacy_cache.parent.mkdir(parents=True, exist_ok=True) + legacy_cache.symlink_to(new_cache_dir) + logger.info(f"Created symlink: {legacy_cache} -> {new_cache_dir}") + + except (OSError, IOError, shutil.Error) as e: + logger.warning(f"Cache migration failed: {e}") + logger.info("Continuing with existing location") + + +# Determine cache directory +if os.environ.get("MFLUX_CACHE_DIR"): + # User specified cache directory (e.g. external storage) + MFLUX_CACHE_DIR = Path(os.environ["MFLUX_CACHE_DIR"]).resolve() +else: + # macOS-idiomatic cache directory @ /Users/username/Library/Caches/mflux + MFLUX_CACHE_DIR = Path(platformdirs.user_cache_dir(appname="mflux")) + +# Perform one-time migration if needed +_migrate_legacy_cache(MFLUX_CACHE_DIR) + +MFLUX_LORA_CACHE_DIR = MFLUX_CACHE_DIR / "loras" diff --git a/src/mflux/weights/weight_handler_lora_huggingface.py b/src/mflux/weights/weight_handler_lora_huggingface.py index cbf641a..8df7c23 100644 --- a/src/mflux/weights/weight_handler_lora_huggingface.py +++ b/src/mflux/weights/weight_handler_lora_huggingface.py @@ -1,15 +1,17 @@ -import os +import shutil from pathlib import Path from huggingface_hub import snapshot_download +from mflux.ui.defaults import MFLUX_LORA_CACHE_DIR + class WeightHandlerLoRAHuggingFace: @staticmethod def download_loras( lora_names: list[str] | None = None, repo_id: str | None = None, - cache_dir: str | None = None, + cache_dir: Path | str | None = None, ) -> list[str]: if repo_id is None: return [] @@ -30,19 +32,21 @@ class WeightHandlerLoRAHuggingFace: def download_lora( repo_id: str, lora_name: str, - cache_dir: str | None = None, + cache_dir: Path | str | None = None, ) -> str: - # Create cache directory if it doesn't exist + # Ensure cache_dir is a Path object if cache_dir is None: - cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "mflux", "loras") + cache_path = MFLUX_LORA_CACHE_DIR + else: + cache_path = Path(cache_dir) - os.makedirs(cache_dir, exist_ok=True) + cache_path.mkdir(parents=True, exist_ok=True) # Check if the file already exists in the cache - cached_file_path = os.path.join(cache_dir, lora_name) - if os.path.exists(cached_file_path): + cached_file_path = cache_path / lora_name + if cached_file_path.exists(): print(f"Using cached LoRA: {cached_file_path}") - return cached_file_path + return str(cached_file_path) # Download the LoRA from Hugging Face print(f"Downloading LoRA '{lora_name}' from {repo_id}...") @@ -50,7 +54,7 @@ class WeightHandlerLoRAHuggingFace: snapshot_download( repo_id=repo_id, allow_patterns=[f"*{lora_name}*"], - cache_dir=cache_dir, + cache_dir=str(cache_path), ) ) @@ -58,17 +62,15 @@ class WeightHandlerLoRAHuggingFace: for file in download_path.glob(f"**/*{lora_name}*"): if file.is_file() and file.suffix in [".safetensors", ".bin"]: # Copy or link the file to the cache directory with the expected name - target_path = os.path.join(cache_dir, lora_name) - if not os.path.exists(target_path): + target_path = cache_path / lora_name + if not target_path.exists(): # Create a symlink or copy the file try: - os.symlink(file, target_path) + target_path.symlink_to(file) except (OSError, AttributeError): - import shutil - shutil.copy2(file, target_path) print(f"LoRA downloaded to: {target_path}") - return target_path + return str(target_path) raise FileNotFoundError(f"Could not find LoRA file '{lora_name}' in the downloaded repository.")