Add ability to download LoRAs automatically from Huggingface

src/mflux/flux/flux_initializer.py

@@ -10,6 +10,7 @@ from mflux.tokenizer.t5_tokenizer import TokenizerT5
 from mflux.tokenizer.tokenizer_handler import TokenizerHandler
 from mflux.weights.weight_handler import WeightHandler
 from mflux.weights.weight_handler_lora import WeightHandlerLoRA
+from mflux.weights.weight_handler_lora_huggingface import WeightHandlerLoRAHuggingFace
 from mflux.weights.weight_util import WeightUtil
 
 
@@ -20,10 +21,13 @@ class FluxInitializer:
         model_config: ModelConfig,
         quantize: int | None,
         local_path: str | None,
-        lora_paths: list[str] | None,
-        lora_scales: list[float] | None,
+        lora_paths: list[str] | None = None,
+        lora_scales: list[float] | None = None,
+        lora_names: list[str] | None = None,
+        lora_repo_id: str | None = None,
     ) -> None:
         # 0. Set paths, configs and prompt_cache for later
+        lora_paths = lora_paths or []
         flux_model.prompt_cache = {}
         flux_model.lora_paths = lora_paths
         flux_model.lora_scales = lora_scales
@@ -70,9 +74,13 @@ class FluxInitializer:
         )
 
         # 5. Set LoRA weights
+        hf_lora_paths = WeightHandlerLoRAHuggingFace.download_loras(
+            lora_names=lora_names,
+            repo_id=lora_repo_id,
+        )
         lora_weights = WeightHandlerLoRA.load_lora_weights(
             transformer=flux_model.transformer,
-            lora_files=lora_paths,
+            lora_files=lora_paths + hf_lora_paths,
             lora_scales=lora_scales,
         )
         WeightHandlerLoRA.set_lora_weights(
@@ -86,8 +94,10 @@ class FluxInitializer:
         model_config: ModelConfig,
         quantize: int | None,
         local_path: str | None,
-        lora_paths: list[str] | None,
-        lora_scales: list[float] | None,
+        lora_paths: list[str] | None = None,
+        lora_scales: list[float] | None = None,
+        lora_names: list[str] | None = None,
+        lora_repo_id: str | None = None,
     ) -> None:
         # 1. Start with same init as regular Flux
         FluxInitializer.init(
@@ -97,6 +107,8 @@ class FluxInitializer:
             local_path=local_path,
             lora_paths=lora_paths,
             lora_scales=lora_scales,
+            lora_names=lora_names,
+            lora_repo_id=lora_repo_id,
         )
 
         # 2. Apply ControlNet-specific initialization

src/mflux/ui/cli/parsers.py

@@ -5,6 +5,7 @@ import time
 import typing as t
 from pathlib import Path
 
+from mflux.community.in_context_lora.in_context_loras import LORA_NAME_MAP, LORA_REPO_ID
 from mflux.ui import defaults as ui_defaults
 
 
@@ -47,10 +48,15 @@ class CommandLineParser(argparse.ArgumentParser):
         self.add_argument("--base-model", type=str, required=False, choices=ui_defaults.MODEL_CHOICES, help="When using a third-party huggingface model, explicitly specify whether the base model is dev or schnell")
         self.add_argument("--quantize",  "-q", type=int, choices=ui_defaults.QUANTIZE_CHOICES, default=None, help=f"Quantize the model ({' or '.join(map(str, ui_defaults.QUANTIZE_CHOICES))}, Default is None)")
 
-    def add_lora_arguments(self) -> None:
+    def add_lora_arguments(self) -> None:  # fmt: off
         self.supports_lora = True
+        lora_group = self.add_argument_group("LoRA configuration")
+        lora_group.add_argument("--lora-style", type=str, choices=sorted(LORA_NAME_MAP.keys()), help="Style of the LoRA to use (e.g., 'storyboard' for film storyboard style)")
         self.add_argument("--lora-paths", type=str, nargs="*", default=None, help="Local safetensors for applying LORA from disk")
         self.add_argument("--lora-scales", type=float, nargs="*", default=None, help="Scaling factor to adjust the impact of LoRA weights on the model. A value of 1.0 applies the LoRA weights as they are.")
+        lora_group.add_argument("--lora-name", type=str, help="Name of the LoRA to download from Hugging Face")
+        lora_group.add_argument("--lora-repo-id", type=str, default=LORA_REPO_ID, help=f"Hugging Face repository ID for LoRAs (default: {LORA_REPO_ID})")
+    # fmt: on
 
     def _add_image_generator_common_arguments(self) -> None:
         self.supports_image_generation = True

src/mflux/weights/weight_handler_lora_huggingface.py

@@ -0,0 +1,74 @@
+import os
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+
+class WeightHandlerLoRAHuggingFace:
+    @staticmethod
+    def download_loras(
+        lora_names: list[str] = None,
+        repo_id: str = None,
+        cache_dir: str = None,
+    ) -> list[str]:
+        if repo_id is None:
+            return []
+
+        lora_paths = []
+        if lora_names:
+            for lora_name in lora_names:
+                lora_path = WeightHandlerLoRAHuggingFace._download_lora(
+                    repo_id=repo_id,
+                    lora_name=lora_name,
+                    cache_dir=cache_dir,
+                )
+                lora_paths.append(lora_path)
+
+        return lora_paths
+
+    @staticmethod
+    def _download_lora(
+        repo_id: str,
+        lora_name: str,
+        cache_dir: str = None,
+    ) -> str:
+        # Create cache directory if it doesn't exist
+        if cache_dir is None:
+            cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "mflux", "loras")
+
+        os.makedirs(cache_dir, 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):
+            print(f"Using cached LoRA: {cached_file_path}")
+            return cached_file_path
+
+        # Download the LoRA from Hugging Face
+        print(f"Downloading LoRA '{lora_name}' from {repo_id}...")
+        download_path = Path(
+            snapshot_download(
+                repo_id=repo_id,
+                allow_patterns=[f"*{lora_name}*"],
+                cache_dir=cache_dir,
+            )
+        )
+
+        # Find the downloaded file
+        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):
+                    # Create a symlink or copy the file
+                    try:
+                        os.symlink(file, target_path)
+                    except (OSError, AttributeError):
+                        import shutil
+
+                        shutil.copy2(file, target_path)
+
+                print(f"LoRA downloaded to: {target_path}")
+                return target_path
+
+        raise FileNotFoundError(f"Could not find LoRA file '{lora_name}' in the downloaded repository.")