From e6a123541c2fd4c6e262d2bfb191c52db095de4d Mon Sep 17 00:00:00 2001 From: christopher5106 Date: Fri, 10 Jul 2026 19:32:10 +0200 Subject: [PATCH] [LoRA] add LoKr adapter support (Z-Image, Flux2/Klein) Adds loading of LoKr (LyCORIS Kronecker product) adapters: - `load_lora_adapter` detects `lokr_` keys and injects a peft `LoKrConfig`, inferred from the tensor shapes via `_create_lokr_config` (decompose factor, per-module rank/alpha patterns). - State dict conversions for the formats in the wild: ai-toolkit Z-Image (dotted diffusers paths under `diffusion_model.`), ai-toolkit BFL Flux2 (fused qkv), LyCORIS underscore format, and bare dotted diffusers paths. - BFL fused-QKV LoKr cannot be split exactly into separate Q/K/V Kronecker factors, so `Flux2LoraLoaderMixin.load_lora_weights` fuses the model's QKV projections and maps the adapter 1:1 (exact). - Alpha follows the LyCORIS convention: scaling applies only to rank-decomposed factors and is baked into the weights at conversion. Fixes #13221 --- .../loaders/lora_conversion_utils.py | 160 +++++++++++++++ src/diffusers/loaders/lora_pipeline.py | 192 +++++++++++++----- src/diffusers/loaders/peft.py | 112 +++++----- src/diffusers/utils/peft_utils.py | 79 ++++++- tests/lora/test_lora_layers_flux2.py | 113 +++++++++++ tests/lora/test_lora_layers_z_image.py | 76 +++++++ 6 files changed, 624 insertions(+), 108 deletions(-) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index e1bc530eb29d..7dcc53bcd6fa 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -2741,6 +2741,166 @@ def _convert_to_ai_toolkit_cat(sds_sd, ait_sd, sds_key, ait_keys, dims=None): return ait_sd +def _bake_lokr_alpha(state_dict): + """ + Consume `.alpha` keys by baking the LyCORIS `alpha / rank` scaling into the left Kronecker factor. The scaling only + applies when a factor is rank-decomposed (`lokr_w1_a/b` or `lokr_w2_a/b`); when both factors are stored as full + matrices, LoKr applies no alpha scaling and the alpha key is simply dropped. + """ + for alpha_key in [k for k in state_dict if k.endswith(".alpha")]: + alpha = state_dict.pop(alpha_key).item() + module = alpha_key.removesuffix(".alpha") + w1_b = state_dict.get(f"{module}.lokr_w1_b") + w2_b = state_dict.get(f"{module}.lokr_w2_b") + rank = w2_b.shape[0] if w2_b is not None else w1_b.shape[0] if w1_b is not None else None + if rank is None: + continue + w1_key = f"{module}.lokr_w1" if f"{module}.lokr_w1" in state_dict else f"{module}.lokr_w1_a" + state_dict[w1_key] = state_dict[w1_key] * (alpha / rank) + + +def _convert_non_diffusers_lokr_to_diffusers(state_dict): + """ + Convert a non-diffusers LoKr state dict whose module paths already match the diffusers model (e.g. ai-toolkit + Z-Image checkpoints with keys like `diffusion_model.layers.0.attention.to_q.lokr_w1`) to the peft-loadable format: + the `diffusion_model.` prefix is replaced with `transformer.` and the `.alpha` keys are consumed. + """ + state_dict = {k.removeprefix("diffusion_model."): v for k, v in state_dict.items()} + _bake_lokr_alpha(state_dict) + + non_lokr_keys = [k for k in state_dict if ".lokr_" not in k] + if non_lokr_keys: + raise ValueError(f"`state_dict` contains unexpected non-LoKr keys: {non_lokr_keys}.") + + return {f"transformer.{k}": v for k, v in state_dict.items()} + + +_LOKR_SUFFIXES = ("lokr_w1", "lokr_w1_a", "lokr_w1_b", "lokr_w2", "lokr_w2_a", "lokr_w2_b") + + +def _convert_non_diffusers_flux2_lokr_to_diffusers(state_dict): + """ + Convert a BFL-format Flux2 LoKr state dict (e.g. trained with ai-toolkit, keys like + `diffusion_model.double_blocks.0.img_attn.qkv.lokr_w1`) to the peft-loadable diffusers format. + + BFL checkpoints apply LoKr to the fused QKV projections of the double blocks. Unlike a LoRA delta, a Kronecker + product delta over the fused projection cannot be split exactly into separate Q/K/V factors, so these are mapped to + the model's fused `to_qkv`/`to_added_qkv` projections instead; `Flux2LoraLoaderMixin.load_lora_weights` fuses the + model's projections before injecting such an adapter. + """ + original_state_dict = {k.removeprefix("diffusion_model."): v for k, v in state_dict.items()} + _bake_lokr_alpha(original_state_dict) + + converted_state_dict = {} + + # Some Flux2 LoKr checkpoints already store expanded diffusers block names; accept those as-is. + for key in list(original_state_dict.keys()): + if key.startswith(("single_transformer_blocks.", "transformer_blocks.")): + converted_state_dict[key] = original_state_dict.pop(key) + + num_double_layers = 0 + num_single_layers = 0 + for key in original_state_dict.keys(): + if key.startswith("single_blocks."): + num_single_layers = max(num_single_layers, int(key.split(".")[1]) + 1) + elif key.startswith("double_blocks."): + num_double_layers = max(num_double_layers, int(key.split(".")[1]) + 1) + + def _remap(bfl_path, diffusers_path): + for suffix in _LOKR_SUFFIXES: + weight = original_state_dict.pop(f"{bfl_path}.{suffix}", None) + if weight is not None: + converted_state_dict[f"{diffusers_path}.{suffix}"] = weight + + for sl in range(num_single_layers): + _remap(f"single_blocks.{sl}.linear1", f"single_transformer_blocks.{sl}.attn.to_qkv_mlp_proj") + _remap(f"single_blocks.{sl}.linear2", f"single_transformer_blocks.{sl}.attn.to_out") + + for dl in range(num_double_layers): + tb = f"transformer_blocks.{dl}" + db = f"double_blocks.{dl}" + + _remap(f"{db}.img_attn.qkv", f"{tb}.attn.to_qkv") + _remap(f"{db}.txt_attn.qkv", f"{tb}.attn.to_added_qkv") + + _remap(f"{db}.img_attn.proj", f"{tb}.attn.to_out.0") + _remap(f"{db}.txt_attn.proj", f"{tb}.attn.to_add_out") + + _remap(f"{db}.img_mlp.0", f"{tb}.ff.linear_in") + _remap(f"{db}.img_mlp.2", f"{tb}.ff.linear_out") + _remap(f"{db}.txt_mlp.0", f"{tb}.ff_context.linear_in") + _remap(f"{db}.txt_mlp.2", f"{tb}.ff_context.linear_out") + + extra_mappings = { + "img_in": "x_embedder", + "txt_in": "context_embedder", + "time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1", + "time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2", + "guidance_in.in_layer": "time_guidance_embed.guidance_embedder.linear_1", + "guidance_in.out_layer": "time_guidance_embed.guidance_embedder.linear_2", + "final_layer.linear": "proj_out", + "final_layer.adaLN_modulation.1": "norm_out.linear", + "single_stream_modulation.lin": "single_stream_modulation.linear", + "double_stream_modulation_img.lin": "double_stream_modulation_img.linear", + "double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear", + } + for bfl_key, diffusers_key in extra_mappings.items(): + _remap(bfl_key, diffusers_key) + + if len(original_state_dict) > 0: + raise ValueError(f"`original_state_dict` should be empty at this point but has {original_state_dict.keys()=}.") + + return {f"transformer.{k}": v for k, v in converted_state_dict.items()} + + +# Mapping from LyCORIS underscore-encoded sub-paths to dotted Flux2 module paths. +_LYCORIS_FLUX2_SUBPATH_MAP = { + "attn_to_q": "attn.to_q", + "attn_to_k": "attn.to_k", + "attn_to_v": "attn.to_v", + "attn_to_out_0": "attn.to_out.0", + "attn_to_add_out": "attn.to_add_out", + "attn_add_q_proj": "attn.add_q_proj", + "attn_add_k_proj": "attn.add_k_proj", + "attn_add_v_proj": "attn.add_v_proj", + "attn_to_qkv_mlp_proj": "attn.to_qkv_mlp_proj", + "attn_to_out": "attn.to_out", + "ff_linear_in": "ff.linear_in", + "ff_linear_out": "ff.linear_out", + "ff_context_linear_in": "ff_context.linear_in", + "ff_context_linear_out": "ff_context.linear_out", +} + + +def _convert_lycoris_flux2_lokr_to_diffusers(state_dict): + """ + Convert a LyCORIS-format Flux2 LoKr state dict (keys like `lycoris_transformer_blocks_0_attn_to_q.lokr_w1`) to the + peft-loadable diffusers format. LyCORIS wraps the diffusers model directly and encodes each module path with + underscores, which are decoded through a lookup of the known block sub-paths. + """ + state_dict = dict(state_dict) + _bake_lokr_alpha(state_dict) + + lycoris_key_pattern = re.compile(r"^lycoris_((?:single_)?transformer_blocks)_(\d+)_(.+)\.(.+)$") + + converted_state_dict = {} + for key in list(state_dict.keys()): + match = lycoris_key_pattern.match(key) + if match is None: + continue + container, block_idx, sub_path, suffix = match.groups() + diffusers_sub_path = _LYCORIS_FLUX2_SUBPATH_MAP.get(sub_path) + if diffusers_sub_path is None: + continue + diffusers_key = f"transformer.{container}.{block_idx}.{diffusers_sub_path}.{suffix}" + converted_state_dict[diffusers_key] = state_dict.pop(key) + + if len(state_dict) > 0: + raise ValueError(f"`state_dict` should be empty at this point but has {state_dict.keys()=}.") + + return converted_state_dict + + def _convert_non_diffusers_z_image_lora_to_diffusers(state_dict): """ Convert non-diffusers ZImage LoRA state dict to diffusers format. diff --git a/src/diffusers/loaders/lora_pipeline.py b/src/diffusers/loaders/lora_pipeline.py index 81ebd2f81102..370ab6badb42 100644 --- a/src/diffusers/loaders/lora_pipeline.py +++ b/src/diffusers/loaders/lora_pipeline.py @@ -45,12 +45,15 @@ _convert_hunyuan_video_lora_to_diffusers, _convert_kohya_flux2_lora_to_diffusers, _convert_kohya_flux_lora_to_diffusers, + _convert_lycoris_flux2_lokr_to_diffusers, _convert_musubi_wan_lora_to_diffusers, _convert_non_diffusers_anima_lora_to_diffusers, + _convert_non_diffusers_flux2_lokr_to_diffusers, _convert_non_diffusers_flux2_lora_to_diffusers, _convert_non_diffusers_hidream_lora_to_diffusers, _convert_non_diffusers_ideogram4_lora_to_diffusers, _convert_non_diffusers_krea2_lora_to_diffusers, + _convert_non_diffusers_lokr_to_diffusers, _convert_non_diffusers_lora_to_diffusers, _convert_non_diffusers_ltx2_lora_to_diffusers, _convert_non_diffusers_ltxv_lora_to_diffusers, @@ -216,9 +219,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, network_alphas, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_unet( state_dict, @@ -643,9 +648,11 @@ def load_lora_weights( **kwargs, ) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_unet( state_dict, @@ -1083,9 +1090,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -1379,9 +1388,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -2508,9 +2519,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -2705,9 +2718,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -2908,9 +2923,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -3117,9 +3134,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) transformer_peft_state_dict = { k: v for k, v in state_dict.items() if k.startswith(f"{self.transformer_name}.") @@ -3335,9 +3354,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -3536,9 +3557,11 @@ def load_lora_weights( # First, ensure that the checkpoint is a compatible one and can be successfully loaded. kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -3739,9 +3762,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -3943,9 +3968,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -4143,9 +4170,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -4397,9 +4426,11 @@ def load_lora_weights( transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer, state_dict=state_dict, ) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) load_into_transformer_2 = kwargs.pop("load_into_transformer_2", False) if load_into_transformer_2: @@ -4674,9 +4705,11 @@ def load_lora_weights( transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer, state_dict=state_dict, ) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) load_into_transformer_2 = kwargs.pop("load_into_transformer_2", False) if load_into_transformer_2: @@ -4894,9 +4927,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -5097,9 +5132,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -5303,9 +5340,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -5508,9 +5547,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -5680,7 +5721,12 @@ def lora_state_dict( has_lora_unet = any(k.startswith("lora_unet_") for k in state_dict) has_diffusion_model = any(k.startswith("diffusion_model.") for k in state_dict) has_default = any("default." in k for k in state_dict) - if has_alphas_in_sd or has_lora_unet or has_diffusion_model or has_default: + is_lokr = any(".lokr_" in k for k in state_dict) + if is_lokr: + # ai-toolkit Z-Image LoKr checkpoints store module paths that already match the diffusers model. + if has_diffusion_model or has_alphas_in_sd: + state_dict = _convert_non_diffusers_lokr_to_diffusers(state_dict) + elif has_alphas_in_sd or has_lora_unet or has_diffusion_model or has_default: state_dict = _convert_non_diffusers_z_image_lora_to_diffusers(state_dict) out = (state_dict, metadata) if return_lora_metadata else state_dict @@ -5714,9 +5760,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -5909,9 +5957,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) transformer_state_dict = {k: v for k, v in state_dict.items() if k.startswith(f"{self.transformer_name}.")} text_conditioner_state_dict = { @@ -6081,14 +6131,22 @@ def lora_state_dict( if is_peft_format: state_dict = {k.replace("base_model.model.", "diffusion_model."): v for k, v in state_dict.items()} + is_lokr = any(".lokr_" in k for k in state_dict) is_ai_toolkit = any(k.startswith("diffusion_model.") for k in state_dict) - if is_ai_toolkit: + if is_lokr: + if any(k.startswith("lycoris_") for k in state_dict): + state_dict = _convert_lycoris_flux2_lokr_to_diffusers(state_dict) + elif is_ai_toolkit: + state_dict = _convert_non_diffusers_flux2_lokr_to_diffusers(state_dict) + elif not any(k.startswith("transformer.") for k in state_dict): + # Bare dotted diffusers module paths (e.g. SimpleTuner exports), possibly with alpha keys. + state_dict = _convert_non_diffusers_lokr_to_diffusers(state_dict) + elif is_ai_toolkit: state_dict = _convert_non_diffusers_flux2_lora_to_diffusers(state_dict) out = (state_dict, metadata) if return_lora_metadata else state_dict return out - # Copied from diffusers.loaders.lora_pipeline.CogVideoXLoraLoaderMixin.load_lora_weights def load_lora_weights( self, pretrained_model_name_or_path_or_dict: str | dict[str, torch.Tensor], @@ -6116,13 +6174,33 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) + + transformer = getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer + + # BFL-format LoKr checkpoints apply LoKr to the fused QKV projections, whose Kronecker product delta cannot + # be split exactly into separate Q/K/V factors. Fuse the model's projections so the adapter maps 1:1. + needs_fused_qkv = any(".attn.to_qkv." in k or ".attn.to_added_qkv." in k for k in state_dict) + if needs_fused_qkv: + if getattr(transformer, "is_quantized", False): + raise ValueError( + "This LoKr checkpoint targets fused QKV projections, which requires fusing the transformer's " + "QKV projections, and that is not supported on quantized models. Please load the transformer " + "without quantization." + ) + logger.info( + "The LoKr checkpoint targets fused QKV projections; calling `fuse_qkv_projections()` on the " + "transformer before loading the adapter." + ) + transformer.fuse_qkv_projections() self.load_lora_into_transformer( state_dict, - transformer=getattr(self, self.transformer_name) if not hasattr(self, "transformer") else self.transformer, + transformer=transformer, adapter_name=adapter_name, metadata=metadata, _pipeline=self, @@ -6323,9 +6401,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -6534,9 +6614,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, @@ -6735,9 +6817,11 @@ def load_lora_weights( kwargs["return_lora_metadata"] = True state_dict, metadata = self.lora_state_dict(pretrained_model_name_or_path_or_dict, **kwargs) - is_correct_format = all("lora" in key for key in state_dict.keys()) + is_correct_format = all("lora" in key or "lokr" in key for key in state_dict.keys()) if not is_correct_format: - raise ValueError("Invalid LoRA checkpoint. Make sure all LoRA param names contain `'lora'` substring.") + raise ValueError( + "Invalid LoRA checkpoint. Make sure all LoRA param names contain the `'lora'` or `'lokr'` substring." + ) self.load_lora_into_transformer( state_dict, diff --git a/src/diffusers/loaders/peft.py b/src/diffusers/loaders/peft.py index daa078bc25d5..2511fb5fb1ee 100644 --- a/src/diffusers/loaders/peft.py +++ b/src/diffusers/loaders/peft.py @@ -38,7 +38,7 @@ set_adapter_layers, set_weights_and_activate_adapters, ) -from ..utils.peft_utils import _create_lora_config, _maybe_warn_for_unhandled_keys +from ..utils.peft_utils import _create_lokr_config, _create_lora_config, _maybe_warn_for_unhandled_keys from .lora_base import _fetch_state_dict, _func_optionally_disable_offloading from .unet_loader_utils import _maybe_expand_lora_scales @@ -46,7 +46,7 @@ logger = logging.get_logger(__name__) _SET_ADAPTER_SCALE_FN_MAPPING = defaultdict( - lambda: (lambda model_cls, weights: weights), + lambda: lambda model_cls, weights: weights, { "UNet2DConditionModel": _maybe_expand_lora_scales, "UNetMotionModel": _maybe_expand_lora_scales, @@ -213,56 +213,64 @@ def load_lora_adapter( "Please choose an existing adapter name or set `hotswap=False` to prevent hotswapping." ) - # check with first key if is not in peft format - first_key = next(iter(state_dict.keys())) - if "lora_A" not in first_key: - state_dict = convert_unet_state_dict_to_peft(state_dict) - - # Control LoRA from SAI is different from BFL Control LoRA - # https://huggingface.co/stabilityai/control-lora - # https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors - is_sai_sd_control_lora = "lora_controlnet" in state_dict - if is_sai_sd_control_lora: - state_dict = convert_sai_sd_control_lora_state_dict_to_peft(state_dict) - - rank = {} - for key, val in state_dict.items(): - # Cannot figure out rank from lora layers that don't have at least 2 dimensions. - # Bias layers in LoRA only have a single dimension - if "lora_B" in key and val.ndim > 1: - # Check out https://github.com/huggingface/peft/pull/2419 for the `^` symbol. - # We may run into some ambiguous configuration values when a model has module - # names, sharing a common prefix (`proj_out.weight` and `blocks.transformer.proj_out.weight`, - # for example) and they have different LoRA ranks. - rank[f"^{key}"] = val.shape[1] - - if network_alphas is not None and len(network_alphas) >= 1: - alpha_keys = [k for k in network_alphas.keys() if k.startswith(f"{prefix}.")] - network_alphas = { - k.removeprefix(f"{prefix}."): v for k, v in network_alphas.items() if k in alpha_keys - } - - # adapter_name if adapter_name is None: adapter_name = get_adapter_name(self) - # create LoraConfig - lora_config = _create_lora_config( - state_dict, - network_alphas, - metadata, - rank, - model_state_dict=self.state_dict(), - adapter_name=adapter_name, - ) + # LoKr adapters (Kronecker product factors) use a different peft config and state dict layout + # (`{module}.lokr_w1` etc.) than LoRA; detect them before any LoRA-specific key handling. + is_lokr = any(".lokr_" in k for k in state_dict) + + if is_lokr: + if hotswap: + raise ValueError("Hotswapping LoKr adapters is not supported.") + adapter_config = _create_lokr_config(state_dict, metadata) + else: + # check with first key if is not in peft format + first_key = next(iter(state_dict.keys())) + if "lora_A" not in first_key: + state_dict = convert_unet_state_dict_to_peft(state_dict) + + # Control LoRA from SAI is different from BFL Control LoRA + # https://huggingface.co/stabilityai/control-lora + # https://huggingface.co/comfyanonymous/ControlNet-v1-1_fp16_safetensors + is_sai_sd_control_lora = "lora_controlnet" in state_dict + if is_sai_sd_control_lora: + state_dict = convert_sai_sd_control_lora_state_dict_to_peft(state_dict) + + rank = {} + for key, val in state_dict.items(): + # Cannot figure out rank from lora layers that don't have at least 2 dimensions. + # Bias layers in LoRA only have a single dimension + if "lora_B" in key and val.ndim > 1: + # Check out https://github.com/huggingface/peft/pull/2419 for the `^` symbol. + # We may run into some ambiguous configuration values when a model has module + # names, sharing a common prefix (`proj_out.weight` and `blocks.transformer.proj_out.weight`, + # for example) and they have different LoRA ranks. + rank[f"^{key}"] = val.shape[1] + + if network_alphas is not None and len(network_alphas) >= 1: + alpha_keys = [k for k in network_alphas.keys() if k.startswith(f"{prefix}.")] + network_alphas = { + k.removeprefix(f"{prefix}."): v for k, v in network_alphas.items() if k in alpha_keys + } + + # create LoraConfig + adapter_config = _create_lora_config( + state_dict, + network_alphas, + metadata, + rank, + model_state_dict=self.state_dict(), + adapter_name=adapter_name, + ) - # Adjust LoRA config for Control LoRA - if is_sai_sd_control_lora: - lora_config.lora_alpha = lora_config.r - lora_config.alpha_pattern = lora_config.rank_pattern - lora_config.bias = "all" - lora_config.modules_to_save = lora_config.exclude_modules - lora_config.exclude_modules = None + # Adjust LoRA config for Control LoRA + if is_sai_sd_control_lora: + adapter_config.lora_alpha = adapter_config.r + adapter_config.alpha_pattern = adapter_config.rank_pattern + adapter_config.bias = "all" + adapter_config.modules_to_save = adapter_config.exclude_modules + adapter_config.exclude_modules = None # None: ) +def _create_lokr_config(state_dict, metadata): + """ + Create a `LoKrConfig` from a peft-format LoKr state dict (keys like `{module}.lokr_w1`). + + Without metadata, the config is inferred from the tensor shapes. The checkpoint alpha is expected to be already + baked into the weights by the state dict conversion, so `alpha` is set equal to the rank (runtime scaling 1.0). + peft re-derives each module's Kronecker factorization from `decompose_factor` and only creates rank-decomposed + factors when the rank is small compared to the factorized dimensions, so for modules whose checkpoint factors are + full matrices the rank is set to `max(lokr_w2.shape)` to make peft create full matrices as well. + """ + from peft import LoKrConfig + from peft.tuners.lokr.layer import factorization + + if metadata is not None: + try: + return LoKrConfig(**metadata) + except TypeError as e: + raise TypeError("`LoKrConfig` class could not be instantiated.") from e + + modules = sorted({k.rpartition(".lokr_")[0] for k in state_dict if ".lokr_" in k}) + + # Reconstruct each module's factorized dimensions, (out_l, out_k) x (in_m, in_n), from the checkpoint. A factor + # is either stored as a full matrix (`lokr_w1`) or rank-decomposed into `lokr_w1_a @ lokr_w1_b` (same for w2). + factorizations = {} + rank_dict = {} + for module in modules: + w1, w1_a, w1_b = (state_dict.get(f"{module}.lokr_w1{s}") for s in ("", "_a", "_b")) + w2, w2_a, w2_b = (state_dict.get(f"{module}.lokr_w2{s}") for s in ("", "_a", "_b")) + out_l, in_m = w1.shape if w1 is not None else (w1_a.shape[0], w1_b.shape[1]) + out_k, in_n = w2.shape if w2 is not None else (w2_a.shape[0], w2_b.shape[1]) + factorizations[module] = ((out_l, out_k), (in_m, in_n)) + if w2_a is not None: + rank_dict[module] = w2_a.shape[1] + elif w1_a is not None: + rank_dict[module] = w1_a.shape[1] + else: + rank_dict[module] = max(w2.shape) + + # Find the `decompose_factor` under which peft reproduces the checkpoint factorizations. A fixed factor shows up + # as the left dimension of the modules it divides (modules it does not divide fall back to a near-square + # factorization, like with factor -1), so every observed left dimension is a candidate. + left_dims = {dims[0][0] for dims in factorizations.values()} + decompose_factor = None + for candidate in sorted(left_dims) + [-1]: + if all( + factorization(out_l * out_k, candidate) == (out_l, out_k) + and factorization(in_m * in_n, candidate) == (in_m, in_n) + for (out_l, out_k), (in_m, in_n) in factorizations.values() + ): + decompose_factor = candidate + break + if decompose_factor is None: + raise ValueError( + "Could not infer a `decompose_factor` that reproduces the Kronecker factorizations of this LoKr " + "state dict. Please open an issue: https://github.com/huggingface/diffusers/issues/new" + ) + + r = collections.Counter(rank_dict.values()).most_common(1)[0][0] + rank_pattern = {k: v for k, v in rank_dict.items() if v != r} + + lokr_config_kwargs = { + "r": r, + "alpha": r, + "rank_pattern": rank_pattern, + "alpha_pattern": dict(rank_pattern), + "target_modules": modules, + "decompose_both": any(".lokr_w1_a" in k for k in state_dict), + "decompose_factor": decompose_factor, + } + try: + return LoKrConfig(**lokr_config_kwargs) + except TypeError as e: + raise TypeError("`LoKrConfig` class could not be instantiated.") from e + + def _create_lora_config( state_dict, network_alphas, metadata, rank_pattern_dict, is_unet=True, model_state_dict=None, adapter_name=None ): @@ -404,7 +479,7 @@ def _maybe_warn_for_unhandled_keys(incompatible_keys, adapter_name): # Check only for unexpected keys. unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) if unexpected_keys: - lora_unexpected_keys = [k for k in unexpected_keys if "lora_" in k and adapter_name in k] + lora_unexpected_keys = [k for k in unexpected_keys if ("lora_" in k or "lokr_" in k) and adapter_name in k] if lora_unexpected_keys: warn_msg = ( f"Loading adapter weights from state_dict led to unexpected keys found in the model:" @@ -414,7 +489,7 @@ def _maybe_warn_for_unhandled_keys(incompatible_keys, adapter_name): # Filter missing keys specific to the current adapter. missing_keys = getattr(incompatible_keys, "missing_keys", None) if missing_keys: - lora_missing_keys = [k for k in missing_keys if "lora_" in k and adapter_name in k] + lora_missing_keys = [k for k in missing_keys if ("lora_" in k or "lokr_" in k) and adapter_name in k] if lora_missing_keys: warn_msg += ( f"Loading adapter weights from state_dict led to missing keys in the model:" diff --git a/tests/lora/test_lora_layers_flux2.py b/tests/lora/test_lora_layers_flux2.py index 25451fa51957..fa5f2118fe51 100644 --- a/tests/lora/test_lora_layers_flux2.py +++ b/tests/lora/test_lora_layers_flux2.py @@ -148,3 +148,116 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se @unittest.skip("Not supported in Flux2.") def test_modify_padding_mode(self): pass + + +@require_peft_backend +class Flux2LoKrTests(unittest.TestCase): + factor = 4 + + def get_transformer(self): + return Flux2Transformer2DModel( + patch_size=1, + in_channels=4, + num_layers=1, + num_single_layers=1, + attention_head_dim=16, + num_attention_heads=2, + joint_attention_dim=16, + timestep_guidance_channels=256, + axes_dims_rope=[4, 4, 4, 4], + ) + + def make_factors(self, out_features, in_features): + from peft.tuners.lokr.layer import factorization + + (out_l, out_k) = factorization(out_features, self.factor) + (in_m, in_n) = factorization(in_features, self.factor) + return torch.randn(out_l, in_m), torch.randn(out_k, in_n) + + def test_lokr_bfl_checkpoint_fuses_qkv_and_loads(self): + # ai-toolkit Flux2 LoKr checkpoints use BFL layer names with LoKr applied to the fused QKV projections + # of the double blocks; loading fuses the model's QKV projections and maps the adapter 1:1. + from peft.tuners.lokr.layer import LoKrLayer + + torch.manual_seed(0) + transformer = self.get_transformer() + named = dict(transformer.named_modules()) + state_dict = {} + expected_deltas = {} + + for bfl_path, diffusers_path, target in [ + ("single_blocks.0.linear1", "single_transformer_blocks.0.attn.to_qkv_mlp_proj", None), + ("single_blocks.0.linear2", "single_transformer_blocks.0.attn.to_out", None), + ("double_blocks.0.img_attn.proj", "transformer_blocks.0.attn.to_out.0", None), + ("double_blocks.0.img_mlp.0", "transformer_blocks.0.ff.linear_in", None), + ("double_blocks.0.img_attn.qkv", "transformer_blocks.0.attn.to_qkv", "transformer_blocks.0.attn.to_q"), + ( + "double_blocks.0.txt_attn.qkv", + "transformer_blocks.0.attn.to_added_qkv", + "transformer_blocks.0.attn.add_q_proj", + ), + ]: + # fused QKV modules do not exist before fusion; derive their dimensions from the query projection + linear = named[target if target is not None else diffusers_path] + out_features = linear.out_features if target is None else 3 * linear.out_features + w1, w2 = self.make_factors(out_features, linear.in_features) + state_dict[f"diffusion_model.{bfl_path}.lokr_w1"] = w1 + state_dict[f"diffusion_model.{bfl_path}.lokr_w2"] = w2 + state_dict[f"diffusion_model.{bfl_path}.alpha"] = torch.tensor(9999220736.0) + expected_deltas[diffusers_path] = torch.kron(w1, w2) + + pipe = Flux2Pipeline( + scheduler=FlowMatchEulerDiscreteScheduler(), + vae=None, + text_encoder=None, + tokenizer=None, + transformer=transformer, + ) + pipe.load_lora_weights(state_dict, adapter_name="default") + + self.assertTrue(transformer.transformer_blocks[0].attn.fused_projections) + named = dict(transformer.named_modules()) + for module, expected in expected_deltas.items(): + layer = named[module] + self.assertIsInstance(layer, LoKrLayer, module) + delta = layer.get_delta_weight("default") + self.assertLess((delta - expected).abs().max().item(), 1e-5, module) + + def test_lokr_lycoris_checkpoint(self): + # LyCORIS wraps the diffusers model directly and encodes module paths with underscores under a + # `lycoris_` prefix. + from peft.tuners.lokr.layer import LoKrLayer + + from diffusers.loaders.lora_pipeline import Flux2LoraLoaderMixin + + torch.manual_seed(0) + transformer = self.get_transformer() + named = dict(transformer.named_modules()) + state_dict = {} + expected_deltas = {} + + for lycoris_path, diffusers_path in [ + ( + "lycoris_single_transformer_blocks_0_attn_to_qkv_mlp_proj", + "single_transformer_blocks.0.attn.to_qkv_mlp_proj", + ), + ("lycoris_transformer_blocks_0_attn_to_q", "transformer_blocks.0.attn.to_q"), + ("lycoris_transformer_blocks_0_attn_to_out_0", "transformer_blocks.0.attn.to_out.0"), + ("lycoris_transformer_blocks_0_ff_linear_in", "transformer_blocks.0.ff.linear_in"), + ]: + linear = named[diffusers_path] + w1, w2 = self.make_factors(linear.out_features, linear.in_features) + state_dict[f"{lycoris_path}.lokr_w1"] = w1 + state_dict[f"{lycoris_path}.lokr_w2"] = w2 + state_dict[f"{lycoris_path}.alpha"] = torch.tensor(16.0) + expected_deltas[diffusers_path] = torch.kron(w1, w2) + + converted = Flux2LoraLoaderMixin.lora_state_dict(state_dict) + transformer.load_lora_adapter(converted, prefix="transformer", adapter_name="default") + + named = dict(transformer.named_modules()) + for module, expected in expected_deltas.items(): + layer = named[module] + self.assertIsInstance(layer, LoKrLayer, module) + delta = layer.get_delta_weight("default") + self.assertLess((delta - expected).abs().max().item(), 1e-5, module) diff --git a/tests/lora/test_lora_layers_z_image.py b/tests/lora/test_lora_layers_z_image.py index 5bdaf9147667..9f8c6790b16f 100644 --- a/tests/lora/test_lora_layers_z_image.py +++ b/tests/lora/test_lora_layers_z_image.py @@ -265,3 +265,79 @@ def test_simple_inference_with_text_denoiser_block_scale_for_all_dict_options(se @unittest.skip("Not supported in ZImage.") def test_modify_padding_mode(self): pass + + +@require_peft_backend +class ZImageLoKrTests(unittest.TestCase): + def get_transformer(self): + return ZImageTransformer2DModel( + all_patch_size=(2,), + all_f_patch_size=(1,), + in_channels=16, + dim=32, + n_layers=2, + n_refiner_layers=1, + n_heads=2, + n_kv_heads=2, + norm_eps=1e-5, + qk_norm=True, + cap_feat_dim=16, + rope_theta=256.0, + t_scale=1000.0, + axes_dims=[8, 4, 4], + axes_lens=[256, 32, 32], + ) + + def test_lokr_ai_toolkit_checkpoint(self): + # ai-toolkit Z-Image LoKr checkpoints store dotted diffusers module paths under a `diffusion_model.` + # prefix, with full Kronecker factors and a placeholder alpha, plus optionally rank-decomposed factors + # with a meaningful alpha. + from peft.tuners.lokr.layer import LoKrLayer, factorization + + from diffusers.loaders.lora_pipeline import ZImageLoraLoaderMixin + + torch.manual_seed(0) + transformer = self.get_transformer() + named = dict(transformer.named_modules()) + factor = 4 + state_dict = {} + expected_deltas = {} + + for module in ["layers.0.attention.to_q", "layers.0.feed_forward.w1", "layers.0.adaLN_modulation.0"]: + linear = named[module] + (out_l, out_k) = factorization(linear.out_features, factor) + (in_m, in_n) = factorization(linear.in_features, factor) + w1 = torch.randn(out_l, in_m) + w2 = torch.randn(out_k, in_n) + state_dict[f"diffusion_model.{module}.lokr_w1"] = w1 + state_dict[f"diffusion_model.{module}.lokr_w2"] = w2 + # alpha applies no scaling when both factors are full matrices; ai-toolkit stores a placeholder + state_dict[f"diffusion_model.{module}.alpha"] = torch.tensor(9999220736.0) + expected_deltas[module] = torch.kron(w1, w2) + + # module with a rank-decomposed right factor and a meaningful alpha + rank, alpha = 2, 1.0 + module = "layers.1.attention.to_v" + linear = named[module] + (out_l, out_k) = factorization(linear.out_features, factor) + (in_m, in_n) = factorization(linear.in_features, factor) + w1 = torch.randn(out_l, in_m) + w2_a = torch.randn(out_k, rank) + w2_b = torch.randn(rank, in_n) + state_dict[f"diffusion_model.{module}.lokr_w1"] = w1 + state_dict[f"diffusion_model.{module}.lokr_w2_a"] = w2_a + state_dict[f"diffusion_model.{module}.lokr_w2_b"] = w2_b + state_dict[f"diffusion_model.{module}.alpha"] = torch.tensor(alpha) + expected_deltas[module] = (alpha / rank) * torch.kron(w1, w2_a @ w2_b) + + converted = ZImageLoraLoaderMixin.lora_state_dict(state_dict) + self.assertTrue(all(k.startswith("transformer.") and ".lokr_" in k for k in converted)) + + transformer.load_lora_adapter(converted, prefix="transformer", adapter_name="default") + + named = dict(transformer.named_modules()) + wrapped = {name for name, module in transformer.named_modules() if isinstance(module, LoKrLayer)} + self.assertEqual(wrapped, set(expected_deltas)) + for module, expected in expected_deltas.items(): + delta = named[module].get_delta_weight("default") + self.assertLess((delta - expected).abs().max().item(), 1e-5, module)