🗂️ Cache Refactor - use platformdirs helper (#201)

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

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

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

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"
 

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"

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.")