diff --git a/src/mflux/flux/flux_initializer.py b/src/mflux/flux/flux_initializer.py index 1997486..0ecab05 100644 --- a/src/mflux/flux/flux_initializer.py +++ b/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 diff --git a/src/mflux/ui/cli/parsers.py b/src/mflux/ui/cli/parsers.py index 59e930f..b62cd88 100644 --- a/src/mflux/ui/cli/parsers.py +++ b/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 diff --git a/src/mflux/weights/weight_handler_lora_huggingface.py b/src/mflux/weights/weight_handler_lora_huggingface.py new file mode 100644 index 0000000..fe270c5 --- /dev/null +++ b/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.")