Small LoRA improvements (#291)

CHANGELOG.md

@@ -5,6 +5,15 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.13.2] - 2025-12-05
+
+### ✨ Improvements
+
+- **Better error messages for multi-file LoRA repos**: When a HuggingFace LoRA repo contains multiple `.safetensors` files, the error message now displays copy-paste ready options instead of a raw list
+- **Z-Image LoRA format support**: Added support for Kohya and ComfyUI LoRA naming conventions, enabling compatibility with more community LoRAs.
+
+---
+
 ## [0.13.1] - 2025-12-03
 
 ### 🐛 Bug Fixes

pyproject.toml

@@ -14,7 +14,7 @@ source-exclude = [
 
 [project]
 name = "mflux"
-version = "0.13.1"
+version = "0.13.2"
 description = "A MLX port of FLUX based on the Huggingface Diffusers implementation."
 readme = "README.md"
 keywords = ["diffusers", "flux", "mlx"]

src/mflux/cli/parser/parsers.py

@@ -309,7 +309,7 @@ class CommandLineParser(argparse.ArgumentParser):
                 try:
                     resolved_path = LoraResolution.resolve(lora_path)
                     resolved_paths.append(resolved_path)
-                except FileNotFoundError as e:  # noqa: PERF203
+                except (FileNotFoundError, ValueError) as e:  # noqa: PERF203
                     self.error(str(e))
             namespace.lora_paths = resolved_paths
 

src/mflux/models/common/resolution/lora_resolution.py

@@ -170,10 +170,10 @@ class LoraResolution:
             raise FileNotFoundError(f"No .safetensors file found in cached repo: {repo_id}")
 
         if len(safetensor_files) > 1:
-            file_names = [f.name for f in safetensor_files]
+            file_names = sorted([f.name for f in safetensor_files])
+            file_list = "\n".join(f"  - {repo_id}:{name}" for name in file_names)
             raise ValueError(
-                f"Multiple .safetensors files found in '{repo_id}': {file_names}. "
+                f"Multiple .safetensors files found in '{repo_id}'. Please specify which file to use:\n{file_list}"
-                f"Please specify which file to use with the collection format: '{repo_id}:<filename>.safetensors'"
             )
 
         return str(safetensor_files[0])
@@ -197,10 +197,10 @@ class LoraResolution:
             raise FileNotFoundError(f"No .safetensors file found in HuggingFace repo: {repo_id}")
 
         if len(safetensor_files) > 1:
-            file_names = [f.name for f in safetensor_files]
+            file_names = sorted([f.name for f in safetensor_files])
+            file_list = "\n".join(f"  - {repo_id}:{name}" for name in file_names)
             raise ValueError(
-                f"Multiple .safetensors files found in '{repo_id}': {file_names}. "
+                f"Multiple .safetensors files found in '{repo_id}'. Please specify which file to use:\n{file_list}"
-                f"Please specify which file to use with the collection format: '{repo_id}:<filename>.safetensors'"
             )
 
         return str(safetensor_files[0])

src/mflux/models/z_image/weights/z_image_lora_mapping.py

@@ -5,11 +5,8 @@ class ZImageLoRAMapping(LoRAMapping):
     @staticmethod
     def get_mapping() -> list[LoRATarget]:
         targets = []
-
-        # Generate mappings for all layer types
         for layer_type in ["layers", "noise_refiner", "context_refiner"]:
             targets.extend(ZImageLoRAMapping._get_layer_targets(layer_type))
-
         return targets
 
     @staticmethod
@@ -19,136 +16,181 @@ class ZImageLoRAMapping(LoRAMapping):
                 model_path=f"{layer_type}.{{block}}.adaLN_modulation.0",
                 possible_up_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.adaLN_modulation.0.lora_B.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.adaLN_modulation.0.lora_up.weight",
                     f"{layer_type}.{{block}}.adaLN_modulation.0.lora_B.weight",
                     f"{layer_type}.{{block}}.adaLN_modulation.0.lora_up.weight",
+                    f"lora_unet_{layer_type}_{{block}}_adaLN_modulation_0.lora_up.weight",
                 ],
                 possible_down_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.adaLN_modulation.0.lora_A.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.adaLN_modulation.0.lora_down.weight",
                     f"{layer_type}.{{block}}.adaLN_modulation.0.lora_A.weight",
                     f"{layer_type}.{{block}}.adaLN_modulation.0.lora_down.weight",
+                    f"lora_unet_{layer_type}_{{block}}_adaLN_modulation_0.lora_down.weight",
                 ],
                 possible_alpha_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.adaLN_modulation.0.alpha",
                     f"{layer_type}.{{block}}.adaLN_modulation.0.alpha",
+                    f"lora_unet_{layer_type}_{{block}}_adaLN_modulation_0.alpha",
                 ],
             ),
             LoRATarget(
                 model_path=f"{layer_type}.{{block}}.attention.to_q",
                 possible_up_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.attention.to_q.lora_B.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.attention.to_q.lora_up.weight",
                     f"{layer_type}.{{block}}.attention.to_q.lora_B.weight",
                     f"{layer_type}.{{block}}.attention.to_q.lora_up.weight",
+                    f"lora_unet_{layer_type}_{{block}}_attention_to_q.lora_up.weight",
                 ],
                 possible_down_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.attention.to_q.lora_A.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.attention.to_q.lora_down.weight",
                     f"{layer_type}.{{block}}.attention.to_q.lora_A.weight",
                     f"{layer_type}.{{block}}.attention.to_q.lora_down.weight",
+                    f"lora_unet_{layer_type}_{{block}}_attention_to_q.lora_down.weight",
                 ],
                 possible_alpha_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.attention.to_q.alpha",
                     f"{layer_type}.{{block}}.attention.to_q.alpha",
+                    f"lora_unet_{layer_type}_{{block}}_attention_to_q.alpha",
                 ],
             ),
             LoRATarget(
                 model_path=f"{layer_type}.{{block}}.attention.to_k",
                 possible_up_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.attention.to_k.lora_B.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.attention.to_k.lora_up.weight",
                     f"{layer_type}.{{block}}.attention.to_k.lora_B.weight",
                     f"{layer_type}.{{block}}.attention.to_k.lora_up.weight",
+                    f"lora_unet_{layer_type}_{{block}}_attention_to_k.lora_up.weight",
                 ],
                 possible_down_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.attention.to_k.lora_A.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.attention.to_k.lora_down.weight",
                     f"{layer_type}.{{block}}.attention.to_k.lora_A.weight",
                     f"{layer_type}.{{block}}.attention.to_k.lora_down.weight",
+                    f"lora_unet_{layer_type}_{{block}}_attention_to_k.lora_down.weight",
                 ],
                 possible_alpha_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.attention.to_k.alpha",
                     f"{layer_type}.{{block}}.attention.to_k.alpha",
+                    f"lora_unet_{layer_type}_{{block}}_attention_to_k.alpha",
                 ],
             ),
             LoRATarget(
                 model_path=f"{layer_type}.{{block}}.attention.to_v",
                 possible_up_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.attention.to_v.lora_B.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.attention.to_v.lora_up.weight",
                     f"{layer_type}.{{block}}.attention.to_v.lora_B.weight",
                     f"{layer_type}.{{block}}.attention.to_v.lora_up.weight",
+                    f"lora_unet_{layer_type}_{{block}}_attention_to_v.lora_up.weight",
                 ],
                 possible_down_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.attention.to_v.lora_A.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.attention.to_v.lora_down.weight",
                     f"{layer_type}.{{block}}.attention.to_v.lora_A.weight",
                     f"{layer_type}.{{block}}.attention.to_v.lora_down.weight",
+                    f"lora_unet_{layer_type}_{{block}}_attention_to_v.lora_down.weight",
                 ],
                 possible_alpha_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.attention.to_v.alpha",
                     f"{layer_type}.{{block}}.attention.to_v.alpha",
+                    f"lora_unet_{layer_type}_{{block}}_attention_to_v.alpha",
                 ],
             ),
             LoRATarget(
                 model_path=f"{layer_type}.{{block}}.attention.to_out.0",
                 possible_up_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.attention.to_out.0.lora_B.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.attention.to_out.0.lora_up.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.attention.out.lora_B.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.attention.out.lora_up.weight",
                     f"{layer_type}.{{block}}.attention.to_out.0.lora_B.weight",
                     f"{layer_type}.{{block}}.attention.to_out.0.lora_up.weight",
+                    f"lora_unet_{layer_type}_{{block}}_attention_to_out_0.lora_up.weight",
                 ],
                 possible_down_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.attention.to_out.0.lora_A.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.attention.to_out.0.lora_down.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.attention.out.lora_A.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.attention.out.lora_down.weight",
                     f"{layer_type}.{{block}}.attention.to_out.0.lora_A.weight",
                     f"{layer_type}.{{block}}.attention.to_out.0.lora_down.weight",
+                    f"lora_unet_{layer_type}_{{block}}_attention_to_out_0.lora_down.weight",
                 ],
                 possible_alpha_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.attention.to_out.0.alpha",
+                    f"diffusion_model.{layer_type}.{{block}}.attention.out.alpha",
                     f"{layer_type}.{{block}}.attention.to_out.0.alpha",
+                    f"lora_unet_{layer_type}_{{block}}_attention_to_out_0.alpha",
                 ],
             ),
             LoRATarget(
                 model_path=f"{layer_type}.{{block}}.feed_forward.w1",
                 possible_up_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.feed_forward.w1.lora_B.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.feed_forward.w1.lora_up.weight",
                     f"{layer_type}.{{block}}.feed_forward.w1.lora_B.weight",
                     f"{layer_type}.{{block}}.feed_forward.w1.lora_up.weight",
+                    f"lora_unet_{layer_type}_{{block}}_feed_forward_w1.lora_up.weight",
                 ],
                 possible_down_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.feed_forward.w1.lora_A.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.feed_forward.w1.lora_down.weight",
                     f"{layer_type}.{{block}}.feed_forward.w1.lora_A.weight",
                     f"{layer_type}.{{block}}.feed_forward.w1.lora_down.weight",
+                    f"lora_unet_{layer_type}_{{block}}_feed_forward_w1.lora_down.weight",
                 ],
                 possible_alpha_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.feed_forward.w1.alpha",
                     f"{layer_type}.{{block}}.feed_forward.w1.alpha",
+                    f"lora_unet_{layer_type}_{{block}}_feed_forward_w1.alpha",
                 ],
             ),
             LoRATarget(
                 model_path=f"{layer_type}.{{block}}.feed_forward.w2",
                 possible_up_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.feed_forward.w2.lora_B.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.feed_forward.w2.lora_up.weight",
                     f"{layer_type}.{{block}}.feed_forward.w2.lora_B.weight",
                     f"{layer_type}.{{block}}.feed_forward.w2.lora_up.weight",
+                    f"lora_unet_{layer_type}_{{block}}_feed_forward_w2.lora_up.weight",
                 ],
                 possible_down_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.feed_forward.w2.lora_A.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.feed_forward.w2.lora_down.weight",
                     f"{layer_type}.{{block}}.feed_forward.w2.lora_A.weight",
                     f"{layer_type}.{{block}}.feed_forward.w2.lora_down.weight",
+                    f"lora_unet_{layer_type}_{{block}}_feed_forward_w2.lora_down.weight",
                 ],
                 possible_alpha_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.feed_forward.w2.alpha",
                     f"{layer_type}.{{block}}.feed_forward.w2.alpha",
+                    f"lora_unet_{layer_type}_{{block}}_feed_forward_w2.alpha",
                 ],
             ),
             LoRATarget(
                 model_path=f"{layer_type}.{{block}}.feed_forward.w3",
                 possible_up_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.feed_forward.w3.lora_B.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.feed_forward.w3.lora_up.weight",
                     f"{layer_type}.{{block}}.feed_forward.w3.lora_B.weight",
                     f"{layer_type}.{{block}}.feed_forward.w3.lora_up.weight",
+                    f"lora_unet_{layer_type}_{{block}}_feed_forward_w3.lora_up.weight",
                 ],
                 possible_down_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.feed_forward.w3.lora_A.weight",
+                    f"diffusion_model.{layer_type}.{{block}}.feed_forward.w3.lora_down.weight",
                     f"{layer_type}.{{block}}.feed_forward.w3.lora_A.weight",
                     f"{layer_type}.{{block}}.feed_forward.w3.lora_down.weight",
+                    f"lora_unet_{layer_type}_{{block}}_feed_forward_w3.lora_down.weight",
                 ],
                 possible_alpha_patterns=[
                     f"diffusion_model.{layer_type}.{{block}}.feed_forward.w3.alpha",
                     f"{layer_type}.{{block}}.feed_forward.w3.alpha",
+                    f"lora_unet_{layer_type}_{{block}}_feed_forward_w3.alpha",
                 ],
             ),
         ]

tests/resolution/test_lora_resolution.py

@@ -95,9 +95,8 @@ class TestLoraResolutionHuggingFace:
 
         error_msg = str(exc_info.value)
         assert "Multiple .safetensors files found" in error_msg
-        assert "lora_v1.safetensors" in error_msg or "lora_v2.safetensors" in error_msg
+        assert "org/multi-lora-repo:lora_v1.safetensors" in error_msg
-        assert "collection format" in error_msg
+        assert "org/multi-lora-repo:lora_v2.safetensors" in error_msg
-        assert "org/multi-lora-repo:" in error_msg
 
     @pytest.mark.fast
     @patch("mflux.models.common.resolution.lora_resolution.snapshot_download")