Small LoRA improvements (#291)

This commit is contained in:
Filip Strand 2025-12-05 03:18:27 +01:00 committed by GitHub
parent 5369d5adc6
commit f45082f880
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 1249 additions and 1314 deletions

View File

@ -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

View File

@ -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"]

View File

@ -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

View File

@ -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"Please specify which file to use with the collection format: '{repo_id}:<filename>.safetensors'"
f"Multiple .safetensors files found in '{repo_id}'. Please specify which file to use:\n{file_list}"
)
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"Please specify which file to use with the collection format: '{repo_id}:<filename>.safetensors'"
f"Multiple .safetensors files found in '{repo_id}'. Please specify which file to use:\n{file_list}"
)
return str(safetensor_files[0])

View File

@ -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",
],
),
]

View File

@ -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 "collection format" in error_msg
assert "org/multi-lora-repo:" in error_msg
assert "org/multi-lora-repo:lora_v1.safetensors" in error_msg
assert "org/multi-lora-repo:lora_v2.safetensors" in error_msg
@pytest.mark.fast
@patch("mflux.models.common.resolution.lora_resolution.snapshot_download")

2485
uv.lock generated

File diff suppressed because it is too large Load Diff