diff --git a/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx b/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx index 98fb6b2067f..872b533803b 100644 --- a/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx +++ b/docs/src/content/docs/Users Guide/03.Models/03.model-format-support.mdx @@ -223,6 +223,12 @@ starter models do, needs no name. TAEF2, TAEQI2.1 and TAEW2.2 are recognized by FLUX.1 and FLUX.2 also accept ComfyUI/OneTrainer PEFT LoRAs whose module names start with `diffusion_model.`. Per-layer `.alpha` values in these `diffusion_model.` files set the LoRA scale. +SimpleTuner's LyCORIS training writes its layers with the model's own module names flattened behind a `lycoris_` +prefix (for example `lycoris_single_transformer_blocks_0_attn_to_k.lokr_w1`). Files in this layout are recognized +on FLUX.1 (both `lycoris_` and `lycoris__` spellings), FLUX.2, SDXL, Z-Image and Wan, as they already were on +SD 3 / 3.5; FLUX.3's own flattened names are a different layout and are still not recognized, as noted on its +family page. + For SD 1.x / 2.x and FLUX.2, LyCORIS LoKr files may store either Kronecker operand whole or as a factorized `_a` / `_b` pair. FLUX.2 also accepts LoHA files in the same BFL and kohya layouts. diff --git a/invokeai/backend/model_manager/configs/lora.py b/invokeai/backend/model_manager/configs/lora.py index 62d2031f1df..db19e3ac4d3 100644 --- a/invokeai/backend/model_manager/configs/lora.py +++ b/invokeai/backend/model_manager/configs/lora.py @@ -74,6 +74,7 @@ has_non_ltx2_architecture_keys, has_unsupported_ltx2_lora_variant_keys, ) +from invokeai.backend.patches.lora_conversions.lycoris_flat_lora_constants import lycoris_flat_family from invokeai.backend.patches.lora_conversions.minimax_h3_lora_constants import ( find_minimax_h3_lora_defect, has_minimax_h3_lora_keys, @@ -706,6 +707,8 @@ def _validate_looks_like_lora(cls, mod: ModelOnDisk) -> None: "lora_te1_", "lora_te2_", "lora_transformer_", + # SimpleTuner's LyCORIS flattens the model's own module names behind this prefix (#10174). + "lycoris_", }, ) @@ -746,6 +749,15 @@ def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType: if is_sd3_lora(state_dict): raise NotAMatchError("model looks like an SD 3 LoRA, not an SD 1/2/XL LoRA") + # SimpleTuner's LyCORIS flattens the model's own module names behind `lycoris_` (#10174). FLUX.1 and SDXL + # are the families this base detection serves; FLUX.2, Z-Image and Wan have their own probes, and names + # that settle nothing fall through to the Stable Diffusion heuristics below. + flat_family = lycoris_flat_family(state_dict) + if flat_family == "flux1": + return BaseModelType.Flux + if flat_family == "sdxl": + return BaseModelType.StableDiffusionXL + # If we've gotten here, we assume that the model is a Stable Diffusion model token_vector_length = lora_token_vector_length(state_dict) if token_vector_length == 768: @@ -800,6 +812,10 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType: if _get_flux_lora_format(mod) and _is_flux2_lora(mod): return BaseModelType.Flux2 + # SimpleTuner's LyCORIS flat keys: FLUX.2 Klein's own module names (`to_qkv_mlp_proj`, `ff.linear_in`) say + # so, whatever the tensor shapes can (#10174). + if lycoris_flat_family(mod.load_state_dict()) == "flux2": + return BaseModelType.Flux2 raise NotAMatchError("model is not a FLUX.2 LoRA") @@ -839,6 +855,11 @@ def _validate_looks_like_lora(cls, mod: ModelOnDisk) -> None: if is_state_dict_likely_z_image_kohya_lora(state_dict): return + # SimpleTuner's LyCORIS flattens the module names behind `lycoris_` (`lycoris_layers_0_attention_to_k`); + # the family detection resolves them and keeps Ideogram 4's same-named blocks out by width (#10174). + if lycoris_flat_family(state_dict) == "z_image": + return + # Check for Z-Image specific LoRA patterns (dot-notation formats) has_z_image_lora_keys = state_dict_has_any_keys_starting_with( state_dict, @@ -886,6 +907,9 @@ def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType: # Check for Kohya format if is_state_dict_likely_z_image_kohya_lora(state_dict): return BaseModelType.ZImage + # SimpleTuner's LyCORIS flat names (#10174) + if lycoris_flat_family(state_dict) == "z_image": + return BaseModelType.ZImage # Check for Z-Image transformer layer patterns (dot-notation formats) # Z-Image uses diffusion_model.layers.X structure (unlike Flux which uses double_blocks/single_blocks) @@ -1538,7 +1562,9 @@ def _validate_looks_like_lora(cls, mod: ModelOnDisk) -> None: state_dict = mod.load_state_dict() str_keys = [k for k in state_dict.keys() if isinstance(k, str)] - has_wan_keys = has_wan_kohya_keys(str_keys) or has_wan_peft_keys(str_keys) + has_wan_keys = ( + has_wan_kohya_keys(str_keys) or has_wan_peft_keys(str_keys) or lycoris_flat_family(state_dict) == "wan" + ) has_lora_suffix = state_dict_has_any_keys_ending_with( state_dict, { @@ -1565,9 +1591,10 @@ def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType: state_dict = mod.load_state_dict() str_keys = [k for k in state_dict.keys() if isinstance(k, str)] - if (has_wan_kohya_keys(str_keys) or has_wan_peft_keys(str_keys)) and not has_non_wan_architecture_keys( - str_keys - ): + has_wan_keys = ( + has_wan_kohya_keys(str_keys) or has_wan_peft_keys(str_keys) or lycoris_flat_family(state_dict) == "wan" + ) + if has_wan_keys and not has_non_wan_architecture_keys(str_keys): return BaseModelType.Wan raise NotAMatchError("model does not look like a Wan LoRA") diff --git a/invokeai/backend/model_manager/load/model_loaders/lora.py b/invokeai/backend/model_manager/load/model_loaders/lora.py index 49fbc299004..0c491288f44 100644 --- a/invokeai/backend/model_manager/load/model_loaders/lora.py +++ b/invokeai/backend/model_manager/load/model_loaders/lora.py @@ -69,6 +69,10 @@ lora_model_from_krea2_state_dict, ) from invokeai.backend.patches.lora_conversions.ltx2_lora_conversion_utils import lora_model_from_ltx2_state_dict +from invokeai.backend.patches.lora_conversions.lycoris_flat_lora_constants import lycoris_flat_family +from invokeai.backend.patches.lora_conversions.lycoris_flat_lora_conversion_utils import ( + lora_model_from_lycoris_flat_state_dict, +) from invokeai.backend.patches.lora_conversions.minimax_h3_lora_conversion_utils import ( lora_model_from_minimax_h3_state_dict, ) @@ -158,8 +162,13 @@ def _load_model( factorized_dora_layers = _count_factorized_layers_with_dora_magnitude(state_dict) + # SimpleTuner's LyCORIS flattens the model's own module names behind `lycoris_` (#10174); map them back and + # hand each family its own converter. State dicts this layout does not claim (SD 3's flat spelling, which + # `lora_model_from_sd3_state_dict` resolves already, and every other format) fall through to the dispatch. + if lycoris_flat_family(state_dict) is not None: + model = lora_model_from_lycoris_flat_state_dict(state_dict, base=self._model_base) # Apply state_dict key conversions, if necessary. - if self._model_base == BaseModelType.StableDiffusionXL: + elif self._model_base == BaseModelType.StableDiffusionXL: state_dict = convert_sdxl_keys_to_diffusers_format(state_dict) model = lora_model_from_sd_state_dict(state_dict=state_dict) elif self._model_base in (BaseModelType.Flux, BaseModelType.Flux2): diff --git a/invokeai/backend/patches/lora_conversions/lycoris_flat_lora_constants.py b/invokeai/backend/patches/lora_conversions/lycoris_flat_lora_constants.py new file mode 100644 index 00000000000..37f5ebca765 --- /dev/null +++ b/invokeai/backend/patches/lora_conversions/lycoris_flat_lora_constants.py @@ -0,0 +1,334 @@ +# SimpleTuner's LyCORIS flat-key LoRA naming (`lycoris_` prefix): key parsing, per-family module names and the +# family detection shared by the probes (`configs/lora.py`) and the conversion. +# +# SimpleTuner's LyCORIS output names each trained module by flattening the model's own diffusers module path with +# underscores behind a `lycoris_` prefix -- `lycoris_single_transformer_blocks_0_attn_to_k.lokr_w1` -- or behind +# `lycoris__` (the wrapper's separator plus the name), as in `lycoris__context_embedder.lokr_w1`. Module names +# contain underscores of their own (`to_k`, `add_k_proj`, `feed_forward.w1`), so a flattened name is resolvable +# only against the whole module list of a family, the way `sd3_lora_constants` already resolves SD 3's spelling. +# +# Family detection rests on names only this family's module list holds (`single_transformer_blocks.N.attn.to_q` +# for FLUX.1, `to_qkv_mlp_proj` for FLUX.2 Klein, the UNet's `down_blocks` for SDXL, `layers.N.attention` for +# Z-Image, `blocks.N.attn1` for Wan), plus, for SDXL, the transformer-block depth that separates it from SD 1/2 +# (`configs/lora.py`'s `_state_dict_looks_like_sdxl_unet_lora` documents the same rule), and, for Z-Image, the +# width that separates it from Ideogram 4. A file with no such name claims nothing: the flattening alone is no +# evidence of any base. SD 3's own flat spelling (`sd3_lokr_simpletuner_format`) is left to `sd3_lora_constants`, +# whose module names a family here must not resolve -- checked per name at detection time. +# +# This file imports nothing from the config package, which keeps it free of cycles. + +import re +from collections.abc import Mapping +from typing import Any, Literal + +import torch + +from invokeai.backend.patches.lora_conversions.ideogram4_lora_constants import IDEOGRAM4_WIDTHS, layer_delta_shape +from invokeai.backend.patches.lora_conversions.sd3_lora_constants import resolve_sd3_transformer_module + +LycorisFlatFamily = Literal["flux1", "flux2", "sdxl", "z_image", "wan"] + +LYCORIS_FLAT_PREFIX = "lycoris_" + +# The parameters a layer can carry, as `any_lora_layer_from_state_dict` reads them (plus the DoRA magnitudes it +# dispatches on and the `diff` of LyCORIS's full algorithm). A key whose tail is none of these is not a layer +# parameter and keeps a file out of this layout entirely. +LYCORIS_LORA_PARAM_NAMES = ( + "lora_A.weight", + "lora_B.weight", + "lora_down.weight", + "lora_up.weight", + "lora_mid.weight", + "alpha", + "lokr_w1", + "lokr_w1_a", + "lokr_w1_b", + "lokr_w2", + "lokr_w2_a", + "lokr_w2_b", + "lokr_t2", + "hada_w1_a", + "hada_w1_b", + "hada_w2_a", + "hada_w2_b", + "hada_t1", + "hada_t2", + "dora_scale", + "dora_magnitude", + "diff", +) + + +def split_lycoris_flat_key(key: str) -> tuple[str, str] | None: + """``(flattened module name, parameter name)`` for a `lycoris_` key, or ``None``. + + The `lycoris__` spelling (the wrapper's separator in front of the name) drops the second underscore with the + first: `lycoris__context_embedder` names the module `context_embedder`, like `lycoris_context_embedder` would. + """ + if not key.startswith(LYCORIS_FLAT_PREFIX): + return None + remainder = key[len(LYCORIS_FLAT_PREFIX) :] + for param_name in LYCORIS_LORA_PARAM_NAMES: + if remainder.endswith("." + param_name): + flat_name = remainder[: -len(param_name) - 1] + return flat_name.removeprefix("_"), param_name + return None + + +def group_lycoris_flat_layers(state_dict: Mapping[str | int, Any]) -> dict[str, dict[str, torch.Tensor]] | None: + """The tensors of a SimpleTuner LyCORIS file grouped by flattened module name, or ``None``. + + ``None`` means the state dict is not this layout: some key is not a `lycoris_` layer parameter. Callers fall + back to the other formats, so nothing that was identified before this layout was recognized changes. + """ + grouped: dict[str, dict[str, torch.Tensor]] = {} + for key, value in state_dict.items(): + if not isinstance(key, str): + return None + split = split_lycoris_flat_key(key) + if split is None: + return None + grouped.setdefault(split[0], {})[split[1]] = value + return grouped or None + + +# Every module of each family's transformer a LoRA can patch, as diffusers names them, `{i}`/`{j}`/`{k}` standing +# for block indices. The flattened spelling of each template is what the file holds; a name is resolved by matching +# it against the whole template rather than by guessing where the dots were. + +FLUX1_TRANSFORMER_MODULE_TEMPLATES = ( + "context_embedder", + "x_embedder", + "proj_out", + "norm_out.linear", + "time_text_embed.timestep_embedder.linear_1", + "time_text_embed.timestep_embedder.linear_2", + "time_text_embed.text_embedder.linear_1", + "time_text_embed.text_embedder.linear_2", + "time_text_embed.guidance_embedder.linear_1", + "time_text_embed.guidance_embedder.linear_2", + "transformer_blocks.{i}.attn.to_q", + "transformer_blocks.{i}.attn.to_k", + "transformer_blocks.{i}.attn.to_v", + "transformer_blocks.{i}.attn.to_out.0", + "transformer_blocks.{i}.attn.add_q_proj", + "transformer_blocks.{i}.attn.add_k_proj", + "transformer_blocks.{i}.attn.add_v_proj", + "transformer_blocks.{i}.attn.to_add_out", + "transformer_blocks.{i}.ff.net.0.proj", + "transformer_blocks.{i}.ff.net.2", + "transformer_blocks.{i}.ff_context.net.0.proj", + "transformer_blocks.{i}.ff_context.net.2", + "transformer_blocks.{i}.norm1.linear", + "transformer_blocks.{i}.norm1_context.linear", + "single_transformer_blocks.{i}.attn.to_q", + "single_transformer_blocks.{i}.attn.to_k", + "single_transformer_blocks.{i}.attn.to_v", + "single_transformer_blocks.{i}.proj_mlp", + "single_transformer_blocks.{i}.proj_out", + "single_transformer_blocks.{i}.norm.linear", +) + +FLUX2_TRANSFORMER_MODULE_TEMPLATES = ( + "transformer_blocks.{i}.attn.to_q", + "transformer_blocks.{i}.attn.to_k", + "transformer_blocks.{i}.attn.to_v", + "transformer_blocks.{i}.attn.to_out.0", + "transformer_blocks.{i}.attn.add_q_proj", + "transformer_blocks.{i}.attn.add_k_proj", + "transformer_blocks.{i}.attn.add_v_proj", + "transformer_blocks.{i}.attn.to_add_out", + "transformer_blocks.{i}.ff.linear_in", + "transformer_blocks.{i}.ff.linear_out", + "transformer_blocks.{i}.ff_context.linear_in", + "transformer_blocks.{i}.ff_context.linear_out", + "single_transformer_blocks.{i}.attn.to_qkv_mlp_proj", + "single_transformer_blocks.{i}.attn.to_out", +) + +# Z-Image's S3-DiT keeps its blocks under `layers`, `context_refiner` and `noise_refiner` (`z_image_lora_conversion_utils`). +_Z_IMAGE_BLOCK_GROUPS = ("layers", "context_refiner", "noise_refiner") +Z_IMAGE_TRANSFORMER_MODULE_TEMPLATES = tuple( + f"{group}.{{i}}.{tail}" + for group in _Z_IMAGE_BLOCK_GROUPS + for tail in ( + "attention.to_q", + "attention.to_k", + "attention.to_v", + "attention.to_out.0", + "feed_forward.w1", + "feed_forward.w2", + "feed_forward.w3", + "adaLN_modulation.0", + ) +) + +WAN_TRANSFORMER_MODULE_TEMPLATES = ( + "blocks.{i}.attn1.to_q", + "blocks.{i}.attn1.to_k", + "blocks.{i}.attn1.to_v", + "blocks.{i}.attn1.to_out.0", + "blocks.{i}.attn1.norm_q", + "blocks.{i}.attn1.norm_k", + "blocks.{i}.attn2.to_q", + "blocks.{i}.attn2.to_k", + "blocks.{i}.attn2.to_v", + "blocks.{i}.attn2.to_out.0", + "blocks.{i}.attn2.norm_q", + "blocks.{i}.attn2.norm_k", + "blocks.{i}.ffn.net.0.proj", + "blocks.{i}.ffn.net.2", +) + +# The SDXL UNet's module names. The first down block and the last up block carry no attention (`DownBlock2D`, +# `UpBlock2D` in `stabilityai/stable-diffusion-xl-base-1.0`'s unet config), which is why the `attentions` templates +# sit on blocks `{i}` = 1, 2 only by data, not by rule: the indices are free and depth is what identifies the family. +_SDXL_ATTENTION_TAILS = ( + "attn1.to_q", + "attn1.to_k", + "attn1.to_v", + "attn1.to_out.0", + "attn2.to_q", + "attn2.to_k", + "attn2.to_v", + "attn2.to_out.0", + "ff.net.0.proj", + "ff.net.2", + "norm1", + "norm2", +) +_SDXL_RESNET_TAILS = ("norm1", "conv1", "norm2", "conv2", "time_emb_proj", "conv_shortcut") +SDXL_UNET_MODULE_TEMPLATES = ( + tuple( + f"{stage}_blocks.{{i}}.attentions.{{j}}.transformer_blocks.{{k}}.{tail}" + for stage in ("down", "up") + for tail in _SDXL_ATTENTION_TAILS + ) + + tuple(f"mid_block.attentions.{{j}}.transformer_blocks.{{k}}.{tail}" for tail in _SDXL_ATTENTION_TAILS) + + tuple(f"{stage}_blocks.{{i}}.resnets.{{j}}.{tail}" for stage in ("down", "up") for tail in _SDXL_RESNET_TAILS) + + tuple(f"mid_block.resnets.{{j}}.{tail}" for tail in _SDXL_RESNET_TAILS) + + ( + "down_blocks.{i}.downsamplers.{j}.conv", + "up_blocks.{i}.upsamplers.{j}.conv", + "up_blocks.{i}.upsamplers.{j}.conv_shortcut", + "conv_in", + "conv_norm_out", + "conv_out", + "time_embedding.linear_{i}", + "add_embedding.linear_{i}", + ) +) + +_FAMILY_MODULE_TEMPLATES: dict[LycorisFlatFamily, tuple[str, ...]] = { + "flux1": FLUX1_TRANSFORMER_MODULE_TEMPLATES, + "flux2": FLUX2_TRANSFORMER_MODULE_TEMPLATES, + "sdxl": SDXL_UNET_MODULE_TEMPLATES, + "z_image": Z_IMAGE_TRANSFORMER_MODULE_TEMPLATES, + "wan": WAN_TRANSFORMER_MODULE_TEMPLATES, +} + +_INDEX_PLACEHOLDERS = ("{i}", "{j}", "{k}") + + +def _template_regexes(templates: tuple[str, ...]) -> dict[str, re.Pattern[str]]: + """One regex per template, matching the flattened spelling with the indices captured in order.""" + compiled: dict[str, re.Pattern[str]] = {} + for template in templates: + flat = template.replace(".", "_") + pattern = re.escape(flat) + for placeholder in _INDEX_PLACEHOLDERS: + pattern = pattern.replace(re.escape(placeholder), r"(\d+)") + compiled[template] = re.compile(f"^{pattern}$") + return compiled + + +_FLAT_TEMPLATE_RE: dict[LycorisFlatFamily, dict[str, re.Pattern[str]]] = { + family: _template_regexes(templates) for family, templates in _FAMILY_MODULE_TEMPLATES.items() +} + +# A transformer block past the single stack of SD 1.x/SD 2.x UNet attentions: the depth that identifies SDXL, the +# same rule `_state_dict_looks_like_sdxl_unet_lora` applies to kohya's `lora_unet_` spelling. +_SDXL_DEEP_ATTENTION_RE = re.compile( + r"^(?:(?:down|up)_blocks_\d+_attentions_\d+|mid_block_attentions_\d+)_transformer_blocks_(\d+)_" +) + + +def resolve_lycoris_flat_module(family: LycorisFlatFamily, flat_name: str) -> str | None: + """The diffusers module path a flattened name names in ``family``, or ``None``.""" + for template, pattern in _FLAT_TEMPLATE_RE[family].items(): + if match := pattern.match(flat_name): + resolved = template + for group in match.groups(): + for placeholder in _INDEX_PLACEHOLDERS: + if placeholder in resolved: + resolved = resolved.replace(placeholder, group, 1) + break + return resolved + return None + + +def _flattened_name_resolves_only_in(family: LycorisFlatFamily, flat_name: str) -> bool: + """True for names only ``family``'s module list holds — the evidence a file is of that family. + + SD 3's flat spelling counts as a neighbour here: FLUX.1's double blocks share SD 3's module names, so those + names identify neither (`sd3_lokr_simpletuner_format` is an SD 3 file holding exactly them). + """ + if resolve_lycoris_flat_module(family, flat_name) is None: + return False + if resolve_sd3_transformer_module(LYCORIS_FLAT_PREFIX + flat_name) is not None: + return False + return all( + other is family or resolve_lycoris_flat_module(other, flat_name) is None for other in _FAMILY_MODULE_TEMPLATES + ) + + +def _candidate_family(layers: dict[str, dict[str, torch.Tensor]]) -> LycorisFlatFamily | None: + matches: list[LycorisFlatFamily] = [ + family + for family in _FAMILY_MODULE_TEMPLATES + if any(_flattened_name_resolves_only_in(family, flat_name) for flat_name in layers) + ] + if len(matches) != 1: + return None + family = matches[0] + + if family == "sdxl" and not any(_SDXL_DEEP_ATTENTION_RE.match(flat_name + "_") for flat_name in layers): + # A LoRA confined to the shallow attention blocks is indistinguishable from SD 1/2 by structure alone + # (the limitation `_state_dict_looks_like_sdxl_unet_lora` documents). + return None + if family == "z_image" and any( + (shape := layer_delta_shape(values)) is not None + and (shape[0] in IDEOGRAM4_WIDTHS or shape[1] in IDEOGRAM4_WIDTHS) + for values in layers.values() + ): + # Ideogram 4 names its blocks exactly like Z-Image; its widths (4608) are what tell the two apart, the + # same check `LoRA_LyCORIS_ZImage_Config` applies to the other spellings. + return None + return family + + +def lycoris_flat_candidate_family(state_dict: Mapping[str | int, Any]) -> LycorisFlatFamily | None: + """The family a state dict's exclusive names point to, whatever else it holds. + + The conversion resolves every name against this family and names any it cannot place, so a refused file says + which layer is the problem. The probes use `lycoris_flat_family`, which additionally requires every name to + resolve: a file with a key of no family is not identified at all. + """ + layers = group_lycoris_flat_layers(state_dict) + return None if layers is None else _candidate_family(layers) + + +def lycoris_flat_family(state_dict: Mapping[str | int, Any]) -> LycorisFlatFamily | None: + """Which family's transformer a SimpleTuner LyCORIS flat-key state dict was trained for, or ``None``. + + ``None`` covers everything this layout does not settle: state dicts of other layouts, files whose names no + family holds, files with a name of no family, and files whose names several families share without one + only-this-family name to break the tie. + """ + layers = group_lycoris_flat_layers(state_dict) + if layers is None: + return None + family = _candidate_family(layers) + if family is None or not all(resolve_lycoris_flat_module(family, flat_name) is not None for flat_name in layers): + return None + return family diff --git a/invokeai/backend/patches/lora_conversions/lycoris_flat_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/lycoris_flat_lora_conversion_utils.py new file mode 100644 index 00000000000..b50136f62d7 --- /dev/null +++ b/invokeai/backend/patches/lora_conversions/lycoris_flat_lora_conversion_utils.py @@ -0,0 +1,128 @@ +"""Convert SimpleTuner's LyCORIS flat-key LoRAs (`lycoris_` prefix) into a ``ModelPatchRaw`` per family. + +The layout is resolved back to each family's diffusers module names (`lycoris_flat_lora_constants`) and handed to +the same paths the dotted spellings take: FLUX.1 through `lora_layers_from_flux_diffusers_grouped_state_dict`, +whose diffusers-to-BFL mapping (fused QKV, swapped final modulation) it reuses in full; Wan through the kohya +converter, whose `lora_unet_blocks__` names the flattened names already are; SDXL through the SD +converter's flattened `lora_unet_` keys. FLUX.2 Klein and Z-Image take diffusers module paths as their patch keys +directly, which is what their converters produce from the dotted spellings. + +A file is refused rather than half-applied: any key that names no module of the claimed family raises, so a layer +can never be silently dropped and the LoRA quietly do less than it should. +""" + +from collections.abc import Mapping + +import torch + +from invokeai.backend.model_manager.taxonomy import BaseModelType +from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict +from invokeai.backend.patches.lora_conversions.flux_diffusers_lora_conversion_utils import ( + lora_layers_from_flux_diffusers_grouped_state_dict, +) +from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX +from invokeai.backend.patches.lora_conversions.lycoris_flat_lora_constants import ( + LycorisFlatFamily, + group_lycoris_flat_layers, + lycoris_flat_candidate_family, + resolve_lycoris_flat_module, +) +from invokeai.backend.patches.lora_conversions.sd_lora_conversion_utils import lora_model_from_sd_state_dict +from invokeai.backend.patches.lora_conversions.wan_lora_conversion_utils import lora_model_from_wan_state_dict +from invokeai.backend.patches.lora_conversions.z_image_lora_constants import Z_IMAGE_LORA_TRANSFORMER_PREFIX +from invokeai.backend.patches.model_patch_raw import ModelPatchRaw + +_FAMILY_BASES: dict[LycorisFlatFamily, BaseModelType] = { + "flux1": BaseModelType.Flux, + "flux2": BaseModelType.Flux2, + "sdxl": BaseModelType.StableDiffusionXL, + "z_image": BaseModelType.ZImage, + "wan": BaseModelType.Wan, +} + + +def lora_model_from_lycoris_flat_state_dict( + state_dict: Mapping[str | int, torch.Tensor], base: BaseModelType | None = None +) -> ModelPatchRaw: + """Convert a SimpleTuner LyCORIS flat-key state dict into a ``ModelPatchRaw``. + + ``base`` is the base the model was registered with; a disagreement with the keys raises, like the FLUX.3 + re-identification guard, instead of mapping layers onto the wrong model. + """ + family = lycoris_flat_candidate_family(state_dict) + if family is None: + raise ValueError("LoRA model is not in a recognized SimpleTuner LyCORIS flat-key layout") + if base is not None and base is not _FAMILY_BASES[family]: + raise ValueError( + f"LoRA model's keys are those of a {_FAMILY_BASES[family].value} LoRA, but it is registered as a " + f"{base.value} LoRA by an earlier version of InvokeAI. Use Re-identify model in the Model Manager." + ) + + layers = group_lycoris_flat_layers(state_dict) + assert layers is not None # `lycoris_flat_family` grouped it + resolved: dict[str, dict[str, torch.Tensor]] = {} + for flat_name, values in layers.items(): + module = resolve_lycoris_flat_module(family, flat_name) + if module is None: + raise ValueError( + f"LoRA layer 'lycoris_{flat_name}' targets no {_FAMILY_BASES[family].value} module; " + "the file is truncated, edited, or not this family's." + ) + resolved[module] = values + + return _FAMILY_BUILDERS[family](resolved) + + +def _flux1_builder(resolved: dict[str, dict[str, torch.Tensor]]) -> ModelPatchRaw: + # The grouped state dict keys are bare diffusers module paths; the converter's prefix handling is a no-op on + # them and its diffusers-to-BFL mapping applies as it does to dotted files. + layers = lora_layers_from_flux_diffusers_grouped_state_dict(resolved, alpha=None) + return ModelPatchRaw(layers=layers) + + +def _flux2_builder(resolved: dict[str, dict[str, torch.Tensor]]) -> ModelPatchRaw: + # What `lora_model_from_flux2_diffusers_state_dict` does to dotted keys: the module path is the patch key. + layers = { + f"{FLUX_LORA_TRANSFORMER_PREFIX}{module}": any_lora_layer_from_state_dict(values) + for module, values in resolved.items() + } + return ModelPatchRaw(layers=layers) + + +def _z_image_builder(resolved: dict[str, dict[str, torch.Tensor]]) -> ModelPatchRaw: + layers = { + f"{Z_IMAGE_LORA_TRANSFORMER_PREFIX}{module}": any_lora_layer_from_state_dict(values) + for module, values in resolved.items() + } + return ModelPatchRaw(layers=layers) + + +def _wan_builder(resolved: dict[str, dict[str, torch.Tensor]]) -> ModelPatchRaw: + # The flattened names are kohya's `lora_unet_blocks__` spelling without the prefix; let the kohya + # converter map them onto diffusers' `blocks.X.attn1.to_q` paths as it does for kohya files. + state_dict: dict[str, torch.Tensor] = { + f"lora_unet_{module.replace('.', '_')}.{param}": value + for module, values in resolved.items() + for param, value in values.items() + } + return lora_model_from_wan_state_dict(state_dict, alpha=None) + + +def _sdxl_builder(resolved: dict[str, dict[str, torch.Tensor]]) -> ModelPatchRaw: + # SD/SDXL patches keep flattened keys behind `lora_unet_` (see `LayerPatcher.apply_smart_model_patch`); the + # kohya spelling of these very names is what `convert_sdxl_keys_to_diffusers_format` leaves behind. + state_dict: dict[str, torch.Tensor] = { + f"lora_unet_{module.replace('.', '_')}.{param}": value + for module, values in resolved.items() + for param, value in values.items() + } + return lora_model_from_sd_state_dict(state_dict=state_dict) + + +_FAMILY_BUILDERS = { + "flux1": _flux1_builder, + "flux2": _flux2_builder, + "sdxl": _sdxl_builder, + "z_image": _z_image_builder, + "wan": _wan_builder, +} diff --git a/invokeai/backend/patches/lora_conversions/wan_lora_constants.py b/invokeai/backend/patches/lora_conversions/wan_lora_constants.py index b420c1df291..27d89beed74 100644 --- a/invokeai/backend/patches/lora_conversions/wan_lora_constants.py +++ b/invokeai/backend/patches/lora_conversions/wan_lora_constants.py @@ -158,9 +158,45 @@ def detect_wan_lora_variant(state_dict: dict) -> WanLoRAVariantType | None: # Any other inner_dim is uncharted — bail rather than guess. return None + # A LoKR attention layer states the same dimension as the Kronecker product of its two operands, whole or + # factorized — SimpleTuner's LyCORIS files hold no low-rank pair at all (#10174). + for key, _tensor in state_dict.items(): + if not isinstance(key, str) or not (key.endswith("attn1_to_q.lokr_w1") or key.endswith("attn1_to_q.lokr_w1_a")): + continue + prefix = key[: -len("lokr_w1_a") if key.endswith("lokr_w1_a") else -len("lokr_w1")] + operand_dims = [] + for operand in ("w1", "w2"): + dims = _lokr_operand_shape(state_dict, prefix, operand) + if dims is None: + break + operand_dims.append(dims) + if len(operand_dims) != 2: + continue + inner_dim = max(operand_dims[0][0] * operand_dims[1][0], operand_dims[0][1] * operand_dims[1][1]) + if inner_dim == 5120: + return WanLoRAVariantType.A14B + if inner_dim == 3072: + return WanLoRAVariantType.Wan5B + return None + return None +def _lokr_operand_shape(state_dict: dict, key_prefix: str, operand: str) -> tuple[int, int] | None: + """A LoKR operand's (rows, cols), stored whole (`lokr_w1`) or as a low-rank pair (`_a` @ `_b`).""" + whole = state_dict.get(f"{key_prefix}lokr_{operand}") + if whole is not None: + shape = getattr(whole, "shape", None) + return (int(shape[0]), int(shape[1])) if shape is not None and len(shape) >= 2 else None + a, b = state_dict.get(f"{key_prefix}lokr_{operand}_a"), state_dict.get(f"{key_prefix}lokr_{operand}_b") + if a is None or b is None: + return None + a_shape, b_shape = getattr(a, "shape", None), getattr(b, "shape", None) + if a_shape is None or b_shape is None or len(a_shape) < 2 or len(b_shape) < 2: + return None + return int(a_shape[0]), int(b_shape[1]) + + def has_non_wan_architecture_keys(str_keys: list[str]) -> bool: """True if any key indicates a non-Wan architecture (Anima, Qwen, Flux, Z-Image, MiniMax H3). diff --git a/tests/backend/model_manager/configs/test_lycoris_flat_lora_config.py b/tests/backend/model_manager/configs/test_lycoris_flat_lora_config.py new file mode 100644 index 00000000000..2b96428e8ed --- /dev/null +++ b/tests/backend/model_manager/configs/test_lycoris_flat_lora_config.py @@ -0,0 +1,140 @@ +"""Identification of SimpleTuner's LyCORIS flat-key LoRAs (`lycoris_` prefix) — issue #10174. + +SimpleTuner's LyCORIS output names each trained module by flattening the model's own diffusers module +path with underscores behind a `lycoris_` prefix (`lycoris_single_transformer_blocks_0_attn_to_k`), or +behind `lycoris__` (the wrapper's separator plus the name). The probes look for container prefixes +(`transformer.`, `lora_unet_`) or for the dotted diffusers layout, so none of this matches and such +files install as unknown on every base. Each family's probe now resolves the flattened names against +that family's module list — the same move the SD 3 probe already made for its own flat spelling +(`sd3_lokr_simpletuner_format`). + +The fixtures are the real key layouts of the Hugging Face files named in issue #10174, captured from +their safetensors headers. Which config claims a file must not depend on probe order (`Config_Base. +CONFIG_CLASSES` is a set), so every layout is checked against every LoRA config. +""" + +import atexit +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import MagicMock + +import pytest +import torch + +from invokeai.backend.model_manager.configs.base import Config_Base +from invokeai.backend.model_manager.configs.lora import LoRA_LyCORIS_Wan_Config +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType, WanLoRAVariantType +from tests.backend.patches.lora_conversions.lora_state_dicts import ( + flux1_lycoris_flat_lokr_double_underscore_format, + flux1_lycoris_flat_lokr_format, + flux2_lycoris_flat_lokr_format, + sd3_lokr_simpletuner_format, + sdxl_lycoris_flat_lokr_format, + wan_lycoris_flat_lokr_format, + z_image_lycoris_flat_lokr_format, +) + +_TMP_DIR = TemporaryDirectory() +atexit.register(_TMP_DIR.cleanup) + + +def _shape(value) -> list[int]: + """Fixtures record either a bare shape or `(shape, dtype)`.""" + return list(value[0]) if isinstance(value, tuple) else list(value) + + +def _state_dict(keys: dict) -> dict[str, torch.Tensor]: + """The layout with a single element of storage per tensor: probes read names and shapes, never values.""" + return {key: torch.zeros((), dtype=torch.float32).expand(_shape(value)) for key, value in keys.items()} + + +def _mod(sd: dict[str, torch.Tensor], name: str) -> tuple[MagicMock, dict]: + path = Path(_TMP_DIR.name) / f"{name}.safetensors" + path.touch() + mod = MagicMock() + mod.path = path + mod.load_state_dict.return_value = sd + overrides = { + "hash": "test-hash", + "path": str(path), + "file_size": 0, + "name": name, + "source": str(path), + "source_type": "path", + } + return mod, overrides + + +def _accepting_lora_configs(sd: dict[str, torch.Tensor], name: str) -> list[str]: + """Every LoRA config class that claims ``sd`` as a single file.""" + accepting: list[str] = [] + for cls in list(Config_Base.CONFIG_CLASSES): + if getattr(cls.model_fields.get("type", None), "default", None) != ModelType.LoRA: + continue + mod, overrides = _mod(sd, name) + try: + cls.from_model_on_disk(mod, overrides) + except Exception: + # Identification records any failure as "no match"; only a returned config claims the file. + continue + accepting.append(cls.__name__) + return accepting + + +_FLAT_LAYOUTS = { + "flux1-lokr": (flux1_lycoris_flat_lokr_format, "LoRA_LyCORIS_FLUX_Config", BaseModelType.Flux), + "flux1-lokr-double-underscore": ( + flux1_lycoris_flat_lokr_double_underscore_format, + "LoRA_LyCORIS_FLUX_Config", + BaseModelType.Flux, + ), + "flux2-klein-lokr": (flux2_lycoris_flat_lokr_format, "LoRA_LyCORIS_Flux2_Config", BaseModelType.Flux2), + "sdxl-lokr": (sdxl_lycoris_flat_lokr_format, "LoRA_LyCORIS_SDXL_Config", BaseModelType.StableDiffusionXL), + "z-image-lokr": (z_image_lycoris_flat_lokr_format, "LoRA_LyCORIS_ZImage_Config", BaseModelType.ZImage), + "wan-lokr": (wan_lycoris_flat_lokr_format, "LoRA_LyCORIS_Wan_Config", BaseModelType.Wan), +} + + +@pytest.mark.parametrize("layout", list(_FLAT_LAYOUTS)) +def test_a_flat_lycoris_lora_is_claimed_by_its_family_alone(layout: str) -> None: + fixture, expected_config, _base = _FLAT_LAYOUTS[layout] + sd = _state_dict(fixture.state_dict_keys) + + assert _accepting_lora_configs(sd, layout) == [expected_config] + + +@pytest.mark.parametrize("layout", list(_FLAT_LAYOUTS)) +def test_a_flat_lycoris_lora_records_its_base(layout: str) -> None: + fixture, expected_config, base = _FLAT_LAYOUTS[layout] + mod, overrides = _mod(_state_dict(fixture.state_dict_keys), layout) + + for cls in list(Config_Base.CONFIG_CLASSES): + if cls.__name__ == expected_config: + break + else: + raise AssertionError(f"{expected_config} is not registered") + assert cls.from_model_on_disk(mod, overrides).base is base + + +def test_the_wan_flat_layout_records_its_variant_from_lokr_widths() -> None: + """The inner dimension LoKR factors multiply to tells the Wan family apart the way a low-rank pair does.""" + mod, overrides = _mod(_state_dict(wan_lycoris_flat_lokr_format.state_dict_keys), "wan-lokr") + + assert LoRA_LyCORIS_Wan_Config.from_model_on_disk(mod, overrides).variant is WanLoRAVariantType.A14B + + +def test_the_sd3_flat_layout_is_still_sd3() -> None: + """The one flat `lycoris_` layout identified before this change must not move.""" + sd = _state_dict(sd3_lokr_simpletuner_format.state_dict_keys) + + assert _accepting_lora_configs(sd, "sd3-simpletuner-lokr") == ["LoRA_LyCORIS_SD3_Config"] + + +def test_a_flat_lycoris_file_that_resolves_in_no_family_stays_unidentified() -> None: + """The flattening alone is no evidence of any base: names no family's module list holds claim nothing.""" + keys = { + "lycoris_totally_unseen_module.alpha": ([], "BF16"), + "lycoris_totally_unseen_module.lokr_w1": ([8, 8], "BF16"), + "lycoris_totally_unseen_module.lokr_w2": ([384, 384], "BF16"), + } + assert _accepting_lora_configs(_state_dict(keys), "unseen") == [] diff --git a/tests/backend/patches/lora_conversions/lora_state_dicts/flux1_lycoris_flat_lokr_double_underscore_format.py b/tests/backend/patches/lora_conversions/lora_state_dicts/flux1_lycoris_flat_lokr_double_underscore_format.py new file mode 100644 index 00000000000..9b7a3d97985 --- /dev/null +++ b/tests/backend/patches/lora_conversions/lora_state_dicts/flux1_lycoris_flat_lokr_double_underscore_format.py @@ -0,0 +1,106 @@ +"""Key layout of SimpleTuner's LyCORIS LoKR over FLUX.1, the `lycoris__` spelling +(`jimmycarter/flux-training-losercity-next-lycoris5`). + +Captured on 2026-11-11 with a header-only read of +`jimmycarter/flux-training-losercity-next-lycoris5/pytorch_lora_weights.safetensors`, the issue +#10174 variant file. Subsetting rule: every top-level module the file trains, plus single block 0 +and double block 0 in full. Values are `(shape, dtype)`. + +Same flattened diffusers module names as `flux1_lycoris_flat_lokr_format`, but behind a +`lycoris__` prefix (double underscore): the wrapper's separator plus the flattened name, e.g. +`lycoris__single_transformer_blocks_0_attn_to_k`. 1512 keys in the file, 504 layers. This file +also trains the top-level modules (`context_embedder`, `x_embedder`, `proj_out`, `norm_out.linear`, +the `time_text_embed` embedders), which the toxicwind file does not.""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "lycoris__context_embedder.alpha": ([], "BF16"), + "lycoris__context_embedder.lokr_w1": ([12, 8], "BF16"), + "lycoris__context_embedder.lokr_w2": ([256, 512], "BF16"), + "lycoris__norm_out_linear.alpha": ([], "BF16"), + "lycoris__norm_out_linear.lokr_w1": ([12, 12], "BF16"), + "lycoris__norm_out_linear.lokr_w2": ([512, 256], "BF16"), + "lycoris__proj_out.alpha": ([], "BF16"), + "lycoris__proj_out.lokr_w1": ([8, 12], "BF16"), + "lycoris__proj_out.lokr_w2": ([8, 256], "BF16"), + "lycoris__single_transformer_blocks_0_attn_to_k.alpha": ([], "BF16"), + "lycoris__single_transformer_blocks_0_attn_to_k.lokr_w1": ([12, 12], "BF16"), + "lycoris__single_transformer_blocks_0_attn_to_k.lokr_w2": ([256, 256], "BF16"), + "lycoris__single_transformer_blocks_0_attn_to_q.alpha": ([], "BF16"), + "lycoris__single_transformer_blocks_0_attn_to_q.lokr_w1": ([12, 12], "BF16"), + "lycoris__single_transformer_blocks_0_attn_to_q.lokr_w2": ([256, 256], "BF16"), + "lycoris__single_transformer_blocks_0_attn_to_v.alpha": ([], "BF16"), + "lycoris__single_transformer_blocks_0_attn_to_v.lokr_w1": ([12, 12], "BF16"), + "lycoris__single_transformer_blocks_0_attn_to_v.lokr_w2": ([256, 256], "BF16"), + "lycoris__single_transformer_blocks_0_norm_linear.alpha": ([], "BF16"), + "lycoris__single_transformer_blocks_0_norm_linear.lokr_w1": ([12, 12], "BF16"), + "lycoris__single_transformer_blocks_0_norm_linear.lokr_w2": ([768, 256], "BF16"), + "lycoris__single_transformer_blocks_0_proj_mlp.alpha": ([], "BF16"), + "lycoris__single_transformer_blocks_0_proj_mlp.lokr_w1": ([12, 12], "BF16"), + "lycoris__single_transformer_blocks_0_proj_mlp.lokr_w2": ([1024, 256], "BF16"), + "lycoris__single_transformer_blocks_0_proj_out.alpha": ([], "BF16"), + "lycoris__single_transformer_blocks_0_proj_out.lokr_w1": ([12, 12], "BF16"), + "lycoris__single_transformer_blocks_0_proj_out.lokr_w2": ([256, 1280], "BF16"), + "lycoris__time_text_embed_guidance_embedder_linear_1.alpha": ([], "BF16"), + "lycoris__time_text_embed_guidance_embedder_linear_1.lokr_w1": ([12, 8], "BF16"), + "lycoris__time_text_embed_guidance_embedder_linear_1.lokr_w2": ([256, 32], "BF16"), + "lycoris__time_text_embed_guidance_embedder_linear_2.alpha": ([], "BF16"), + "lycoris__time_text_embed_guidance_embedder_linear_2.lokr_w1": ([12, 12], "BF16"), + "lycoris__time_text_embed_guidance_embedder_linear_2.lokr_w2": ([256, 256], "BF16"), + "lycoris__time_text_embed_text_embedder_linear_1.alpha": ([], "BF16"), + "lycoris__time_text_embed_text_embedder_linear_1.lokr_w1": ([12, 12], "BF16"), + "lycoris__time_text_embed_text_embedder_linear_1.lokr_w2": ([256, 64], "BF16"), + "lycoris__time_text_embed_text_embedder_linear_2.alpha": ([], "BF16"), + "lycoris__time_text_embed_text_embedder_linear_2.lokr_w1": ([12, 12], "BF16"), + "lycoris__time_text_embed_text_embedder_linear_2.lokr_w2": ([256, 256], "BF16"), + "lycoris__time_text_embed_timestep_embedder_linear_1.alpha": ([], "BF16"), + "lycoris__time_text_embed_timestep_embedder_linear_1.lokr_w1": ([12, 8], "BF16"), + "lycoris__time_text_embed_timestep_embedder_linear_1.lokr_w2": ([256, 32], "BF16"), + "lycoris__time_text_embed_timestep_embedder_linear_2.alpha": ([], "BF16"), + "lycoris__time_text_embed_timestep_embedder_linear_2.lokr_w1": ([12, 12], "BF16"), + "lycoris__time_text_embed_timestep_embedder_linear_2.lokr_w2": ([256, 256], "BF16"), + "lycoris__transformer_blocks_0_attn_add_k_proj.alpha": ([], "BF16"), + "lycoris__transformer_blocks_0_attn_add_k_proj.lokr_w1": ([12, 12], "BF16"), + "lycoris__transformer_blocks_0_attn_add_k_proj.lokr_w2": ([256, 256], "BF16"), + "lycoris__transformer_blocks_0_attn_add_q_proj.alpha": ([], "BF16"), + "lycoris__transformer_blocks_0_attn_add_q_proj.lokr_w1": ([12, 12], "BF16"), + "lycoris__transformer_blocks_0_attn_add_q_proj.lokr_w2": ([256, 256], "BF16"), + "lycoris__transformer_blocks_0_attn_add_v_proj.alpha": ([], "BF16"), + "lycoris__transformer_blocks_0_attn_add_v_proj.lokr_w1": ([12, 12], "BF16"), + "lycoris__transformer_blocks_0_attn_add_v_proj.lokr_w2": ([256, 256], "BF16"), + "lycoris__transformer_blocks_0_attn_to_add_out.alpha": ([], "BF16"), + "lycoris__transformer_blocks_0_attn_to_add_out.lokr_w1": ([12, 12], "BF16"), + "lycoris__transformer_blocks_0_attn_to_add_out.lokr_w2": ([256, 256], "BF16"), + "lycoris__transformer_blocks_0_attn_to_k.alpha": ([], "BF16"), + "lycoris__transformer_blocks_0_attn_to_k.lokr_w1": ([12, 12], "BF16"), + "lycoris__transformer_blocks_0_attn_to_k.lokr_w2": ([256, 256], "BF16"), + "lycoris__transformer_blocks_0_attn_to_out_0.alpha": ([], "BF16"), + "lycoris__transformer_blocks_0_attn_to_out_0.lokr_w1": ([12, 12], "BF16"), + "lycoris__transformer_blocks_0_attn_to_out_0.lokr_w2": ([256, 256], "BF16"), + "lycoris__transformer_blocks_0_attn_to_q.alpha": ([], "BF16"), + "lycoris__transformer_blocks_0_attn_to_q.lokr_w1": ([12, 12], "BF16"), + "lycoris__transformer_blocks_0_attn_to_q.lokr_w2": ([256, 256], "BF16"), + "lycoris__transformer_blocks_0_attn_to_v.alpha": ([], "BF16"), + "lycoris__transformer_blocks_0_attn_to_v.lokr_w1": ([12, 12], "BF16"), + "lycoris__transformer_blocks_0_attn_to_v.lokr_w2": ([256, 256], "BF16"), + "lycoris__transformer_blocks_0_ff_context_net_0_proj.alpha": ([], "BF16"), + "lycoris__transformer_blocks_0_ff_context_net_0_proj.lokr_w1": ([6, 6], "BF16"), + "lycoris__transformer_blocks_0_ff_context_net_0_proj.lokr_w2": ([2048, 512], "BF16"), + "lycoris__transformer_blocks_0_ff_context_net_2.alpha": ([], "BF16"), + "lycoris__transformer_blocks_0_ff_context_net_2.lokr_w1": ([6, 6], "BF16"), + "lycoris__transformer_blocks_0_ff_context_net_2.lokr_w2": ([512, 2048], "BF16"), + "lycoris__transformer_blocks_0_ff_net_0_proj.alpha": ([], "BF16"), + "lycoris__transformer_blocks_0_ff_net_0_proj.lokr_w1": ([6, 6], "BF16"), + "lycoris__transformer_blocks_0_ff_net_0_proj.lokr_w2": ([2048, 512], "BF16"), + "lycoris__transformer_blocks_0_ff_net_2.alpha": ([], "BF16"), + "lycoris__transformer_blocks_0_ff_net_2.lokr_w1": ([6, 6], "BF16"), + "lycoris__transformer_blocks_0_ff_net_2.lokr_w2": ([512, 2048], "BF16"), + "lycoris__transformer_blocks_0_norm1_context_linear.alpha": ([], "BF16"), + "lycoris__transformer_blocks_0_norm1_context_linear.lokr_w1": ([12, 12], "BF16"), + "lycoris__transformer_blocks_0_norm1_context_linear.lokr_w2": ([1536, 256], "BF16"), + "lycoris__transformer_blocks_0_norm1_linear.alpha": ([], "BF16"), + "lycoris__transformer_blocks_0_norm1_linear.lokr_w1": ([12, 12], "BF16"), + "lycoris__transformer_blocks_0_norm1_linear.lokr_w2": ([1536, 256], "BF16"), + "lycoris__x_embedder.alpha": ([], "BF16"), + "lycoris__x_embedder.lokr_w1": ([12, 8], "BF16"), + "lycoris__x_embedder.lokr_w2": ([256, 8], "BF16"), +} diff --git a/tests/backend/patches/lora_conversions/lora_state_dicts/flux1_lycoris_flat_lokr_format.py b/tests/backend/patches/lora_conversions/lora_state_dicts/flux1_lycoris_flat_lokr_format.py new file mode 100644 index 00000000000..f0381dd9cf9 --- /dev/null +++ b/tests/backend/patches/lora_conversions/lora_state_dicts/flux1_lycoris_flat_lokr_format.py @@ -0,0 +1,60 @@ +"""Key layout of SimpleTuner's LyCORIS LoKR over FLUX.1 (`toxicwind/glassobject-lokr-flux`). + +Captured on 2026-11-11 with a header-only read (HTTP range request) of +`toxicwind/glassobject-lokr-flux/pytorch_lora_weights.safetensors`, the issue #10174 repro file. +Subsetting rule: single block 0's attention and double block 0 in full -- one block of each kind, +every module kind the file trains (Attention + FeedForward presets). Values are `(shape, dtype)`. + +diffusers' `FluxTransformer2DModel` module names flattened behind `lycoris_` with every `.` turned +into `_` (`lycoris_single_transformer_blocks_0_attn_to_k`). Module names contain underscores of +their own (`to_k`, `add_k_proj`), so a flattened name is only resolvable against the whole module +list. LoKR with `alpha`, a full `lokr_w1` and a full `lokr_w2`; the patch weight's shape is the +Kronecker product. 1026 keys in the file, 342 layers.""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "lycoris_single_transformer_blocks_0_attn_to_k.alpha": ([], "BF16"), + "lycoris_single_transformer_blocks_0_attn_to_k.lokr_w1": ([16, 16], "BF16"), + "lycoris_single_transformer_blocks_0_attn_to_k.lokr_w2": ([192, 192], "BF16"), + "lycoris_single_transformer_blocks_0_attn_to_q.alpha": ([], "BF16"), + "lycoris_single_transformer_blocks_0_attn_to_q.lokr_w1": ([16, 16], "BF16"), + "lycoris_single_transformer_blocks_0_attn_to_q.lokr_w2": ([192, 192], "BF16"), + "lycoris_single_transformer_blocks_0_attn_to_v.alpha": ([], "BF16"), + "lycoris_single_transformer_blocks_0_attn_to_v.lokr_w1": ([16, 16], "BF16"), + "lycoris_single_transformer_blocks_0_attn_to_v.lokr_w2": ([192, 192], "BF16"), + "lycoris_transformer_blocks_0_attn_add_k_proj.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_add_k_proj.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_add_k_proj.lokr_w2": ([192, 192], "BF16"), + "lycoris_transformer_blocks_0_attn_add_q_proj.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_add_q_proj.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_add_q_proj.lokr_w2": ([192, 192], "BF16"), + "lycoris_transformer_blocks_0_attn_add_v_proj.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_add_v_proj.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_add_v_proj.lokr_w2": ([192, 192], "BF16"), + "lycoris_transformer_blocks_0_attn_to_add_out.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_to_add_out.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_to_add_out.lokr_w2": ([192, 192], "BF16"), + "lycoris_transformer_blocks_0_attn_to_k.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_to_k.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_to_k.lokr_w2": ([192, 192], "BF16"), + "lycoris_transformer_blocks_0_attn_to_out_0.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_to_out_0.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_to_out_0.lokr_w2": ([192, 192], "BF16"), + "lycoris_transformer_blocks_0_attn_to_q.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_to_q.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_to_q.lokr_w2": ([192, 192], "BF16"), + "lycoris_transformer_blocks_0_attn_to_v.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_to_v.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_to_v.lokr_w2": ([192, 192], "BF16"), + "lycoris_transformer_blocks_0_ff_context_net_0_proj.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_ff_context_net_0_proj.lokr_w1": ([8, 8], "BF16"), + "lycoris_transformer_blocks_0_ff_context_net_0_proj.lokr_w2": ([1536, 384], "BF16"), + "lycoris_transformer_blocks_0_ff_context_net_2.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_ff_context_net_2.lokr_w1": ([8, 8], "BF16"), + "lycoris_transformer_blocks_0_ff_context_net_2.lokr_w2": ([384, 1536], "BF16"), + "lycoris_transformer_blocks_0_ff_net_0_proj.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_ff_net_0_proj.lokr_w1": ([8, 8], "BF16"), + "lycoris_transformer_blocks_0_ff_net_0_proj.lokr_w2": ([1536, 384], "BF16"), + "lycoris_transformer_blocks_0_ff_net_2.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_ff_net_2.lokr_w1": ([8, 8], "BF16"), + "lycoris_transformer_blocks_0_ff_net_2.lokr_w2": ([384, 1536], "BF16"), +} diff --git a/tests/backend/patches/lora_conversions/lora_state_dicts/flux2_lycoris_flat_lokr_format.py b/tests/backend/patches/lora_conversions/lora_state_dicts/flux2_lycoris_flat_lokr_format.py new file mode 100644 index 00000000000..af3ed4e182d --- /dev/null +++ b/tests/backend/patches/lora_conversions/lora_state_dicts/flux2_lycoris_flat_lokr_format.py @@ -0,0 +1,56 @@ +"""Key layout of SimpleTuner's LyCORIS LoKR over FLUX.2 Klein 9B +(`gattaplayer/besch-flux2-klein-9b-lokr-lion-5e-6-bs2-ga2-v01`). + +Captured on 2026-11-11 with a header-only read of the repo's `pytorch_lora_weights.safetensors`, +the issue #10174 FLUX.2 Klein row. Subsetting rule: single block 0 and double block 0 in full -- +every module kind of the file's presets (Flux2Attention, Flux2FeedForward, Flux2ParallelSelfAttention). +Values are `(shape, dtype)`. + +Flux2Transformer2DModel's own module names flattened behind `lycoris_`: the fused +`single_transformer_blocks.N.attn.to_qkv_mlp_proj` and `attn.to_out`, and the double blocks' +`ff.linear_in`/`ff.linear_out` (not FLUX.1's `ff.net.0.proj`). 432 keys in the file, 144 layers.""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "lycoris_single_transformer_blocks_0_attn_to_out.alpha": ([], "BF16"), + "lycoris_single_transformer_blocks_0_attn_to_out.lokr_w1": ([16, 16], "BF16"), + "lycoris_single_transformer_blocks_0_attn_to_out.lokr_w2": ([256, 1024], "BF16"), + "lycoris_single_transformer_blocks_0_attn_to_qkv_mlp_proj.alpha": ([], "BF16"), + "lycoris_single_transformer_blocks_0_attn_to_qkv_mlp_proj.lokr_w1": ([16, 16], "BF16"), + "lycoris_single_transformer_blocks_0_attn_to_qkv_mlp_proj.lokr_w2": ([2304, 256], "BF16"), + "lycoris_transformer_blocks_0_attn_add_k_proj.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_add_k_proj.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_add_k_proj.lokr_w2": ([256, 256], "BF16"), + "lycoris_transformer_blocks_0_attn_add_q_proj.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_add_q_proj.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_add_q_proj.lokr_w2": ([256, 256], "BF16"), + "lycoris_transformer_blocks_0_attn_add_v_proj.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_add_v_proj.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_add_v_proj.lokr_w2": ([256, 256], "BF16"), + "lycoris_transformer_blocks_0_attn_to_add_out.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_to_add_out.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_to_add_out.lokr_w2": ([256, 256], "BF16"), + "lycoris_transformer_blocks_0_attn_to_k.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_to_k.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_to_k.lokr_w2": ([256, 256], "BF16"), + "lycoris_transformer_blocks_0_attn_to_out_0.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_to_out_0.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_to_out_0.lokr_w2": ([256, 256], "BF16"), + "lycoris_transformer_blocks_0_attn_to_q.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_to_q.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_to_q.lokr_w2": ([256, 256], "BF16"), + "lycoris_transformer_blocks_0_attn_to_v.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_attn_to_v.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_attn_to_v.lokr_w2": ([256, 256], "BF16"), + "lycoris_transformer_blocks_0_ff_context_linear_in.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_ff_context_linear_in.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_ff_context_linear_in.lokr_w2": ([1536, 256], "BF16"), + "lycoris_transformer_blocks_0_ff_context_linear_out.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_ff_context_linear_out.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_ff_context_linear_out.lokr_w2": ([256, 768], "BF16"), + "lycoris_transformer_blocks_0_ff_linear_in.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_ff_linear_in.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_ff_linear_in.lokr_w2": ([1536, 256], "BF16"), + "lycoris_transformer_blocks_0_ff_linear_out.alpha": ([], "BF16"), + "lycoris_transformer_blocks_0_ff_linear_out.lokr_w1": ([16, 16], "BF16"), + "lycoris_transformer_blocks_0_ff_linear_out.lokr_w2": ([256, 768], "BF16"), +} diff --git a/tests/backend/patches/lora_conversions/lora_state_dicts/sdxl_lycoris_flat_lokr_format.py b/tests/backend/patches/lora_conversions/lora_state_dicts/sdxl_lycoris_flat_lokr_format.py new file mode 100644 index 00000000000..7911cd357e0 --- /dev/null +++ b/tests/backend/patches/lora_conversions/lora_state_dicts/sdxl_lycoris_flat_lokr_format.py @@ -0,0 +1,134 @@ +"""Key layout of SimpleTuner's LyCORIS LoKR over SDXL (`davidrd123/Hiroshige-SDXL-LoKr`). + +Captured on 2026-11-11 with a header-only read of the repo's `pytorch_lora_weights.safetensors`, +the issue #10174 SDXL row. Subsetting rule: one attention block from each UNet stage, the deep +`down_blocks.2` / `up_blocks.0` ones at a transformer_blocks index past SD1.x/SD2.x's single block +(index 3 and 5 here) so the depth rule can tell SDXL apart. Values are `(shape, dtype)`. + +diffusers' SDXL UNet module names flattened behind `lycoris_` +(`lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn1_to_k`). 2100 keys in the file, +700 layers (Attention + FeedForward presets). Note `attn1`/`attn2` here are the UNet's attention +names, not the DiT families'.""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn1_to_k.alpha": ([], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn1_to_k.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn1_to_k.lokr_w2": ([40, 40], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn1_to_out_0.alpha": ([], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn1_to_out_0.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn1_to_out_0.lokr_w2": ([40, 40], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn1_to_q.alpha": ([], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn1_to_q.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn1_to_q.lokr_w2": ([40, 40], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn1_to_v.alpha": ([], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn1_to_v.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn1_to_v.lokr_w2": ([40, 40], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn2_to_k.alpha": ([], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn2_to_k.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn2_to_k.lokr_w2": ([40, 128], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn2_to_out_0.alpha": ([], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn2_to_out_0.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn2_to_out_0.lokr_w2": ([40, 40], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn2_to_q.alpha": ([], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn2_to_q.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn2_to_q.lokr_w2": ([40, 40], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn2_to_v.alpha": ([], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn2_to_v.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_attn2_to_v.lokr_w2": ([40, 128], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_ff_net_0_proj.alpha": ([], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_ff_net_0_proj.lokr_w1": ([8, 8], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_ff_net_0_proj.lokr_w2": ([640, 80], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_ff_net_2.alpha": ([], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_ff_net_2.lokr_w1": ([8, 8], "BF16"), + "lycoris_down_blocks_1_attentions_0_transformer_blocks_0_ff_net_2.lokr_w2": ([80, 320], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn1_to_k.alpha": ([], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn1_to_k.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn1_to_k.lokr_w2": ([80, 80], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn1_to_out_0.alpha": ([], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn1_to_out_0.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn1_to_out_0.lokr_w2": ([80, 80], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn1_to_q.alpha": ([], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn1_to_q.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn1_to_q.lokr_w2": ([80, 80], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn1_to_v.alpha": ([], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn1_to_v.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn1_to_v.lokr_w2": ([80, 80], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn2_to_k.alpha": ([], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn2_to_k.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn2_to_k.lokr_w2": ([80, 128], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn2_to_out_0.alpha": ([], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn2_to_out_0.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn2_to_out_0.lokr_w2": ([80, 80], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn2_to_q.alpha": ([], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn2_to_q.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn2_to_q.lokr_w2": ([80, 80], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn2_to_v.alpha": ([], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn2_to_v.lokr_w1": ([16, 16], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_attn2_to_v.lokr_w2": ([80, 128], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_ff_net_0_proj.alpha": ([], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_ff_net_0_proj.lokr_w1": ([8, 8], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_ff_net_0_proj.lokr_w2": ([1280, 160], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_ff_net_2.alpha": ([], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_ff_net_2.lokr_w1": ([8, 8], "BF16"), + "lycoris_down_blocks_2_attentions_0_transformer_blocks_3_ff_net_2.lokr_w2": ([160, 640], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn1_to_k.alpha": ([], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn1_to_k.lokr_w1": ([16, 16], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn1_to_k.lokr_w2": ([80, 80], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn1_to_out_0.alpha": ([], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn1_to_out_0.lokr_w1": ([16, 16], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn1_to_out_0.lokr_w2": ([80, 80], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn1_to_q.alpha": ([], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn1_to_q.lokr_w1": ([16, 16], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn1_to_q.lokr_w2": ([80, 80], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn1_to_v.alpha": ([], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn1_to_v.lokr_w1": ([16, 16], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn1_to_v.lokr_w2": ([80, 80], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn2_to_k.alpha": ([], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn2_to_k.lokr_w1": ([16, 16], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn2_to_k.lokr_w2": ([80, 128], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn2_to_out_0.alpha": ([], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn2_to_out_0.lokr_w1": ([16, 16], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn2_to_out_0.lokr_w2": ([80, 80], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn2_to_q.alpha": ([], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn2_to_q.lokr_w1": ([16, 16], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn2_to_q.lokr_w2": ([80, 80], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn2_to_v.alpha": ([], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn2_to_v.lokr_w1": ([16, 16], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_attn2_to_v.lokr_w2": ([80, 128], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_ff_net_0_proj.alpha": ([], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_ff_net_0_proj.lokr_w1": ([8, 8], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_ff_net_0_proj.lokr_w2": ([1280, 160], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_ff_net_2.alpha": ([], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_ff_net_2.lokr_w1": ([8, 8], "BF16"), + "lycoris_mid_block_attentions_0_transformer_blocks_0_ff_net_2.lokr_w2": ([160, 640], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn1_to_k.alpha": ([], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn1_to_k.lokr_w1": ([16, 16], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn1_to_k.lokr_w2": ([80, 80], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn1_to_out_0.alpha": ([], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn1_to_out_0.lokr_w1": ([16, 16], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn1_to_out_0.lokr_w2": ([80, 80], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn1_to_q.alpha": ([], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn1_to_q.lokr_w1": ([16, 16], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn1_to_q.lokr_w2": ([80, 80], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn1_to_v.alpha": ([], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn1_to_v.lokr_w1": ([16, 16], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn1_to_v.lokr_w2": ([80, 80], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn2_to_k.alpha": ([], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn2_to_k.lokr_w1": ([16, 16], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn2_to_k.lokr_w2": ([80, 128], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn2_to_out_0.alpha": ([], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn2_to_out_0.lokr_w1": ([16, 16], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn2_to_out_0.lokr_w2": ([80, 80], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn2_to_q.alpha": ([], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn2_to_q.lokr_w1": ([16, 16], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn2_to_q.lokr_w2": ([80, 80], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn2_to_v.alpha": ([], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn2_to_v.lokr_w1": ([16, 16], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_attn2_to_v.lokr_w2": ([80, 128], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_ff_net_0_proj.alpha": ([], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_ff_net_0_proj.lokr_w1": ([8, 8], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_ff_net_0_proj.lokr_w2": ([1280, 160], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_ff_net_2.alpha": ([], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_ff_net_2.lokr_w1": ([8, 8], "BF16"), + "lycoris_up_blocks_0_attentions_0_transformer_blocks_5_ff_net_2.lokr_w2": ([160, 640], "BF16"), +} diff --git a/tests/backend/patches/lora_conversions/lora_state_dicts/wan_lycoris_flat_lokr_format.py b/tests/backend/patches/lora_conversions/lora_state_dicts/wan_lycoris_flat_lokr_format.py new file mode 100644 index 00000000000..ef071c2064d --- /dev/null +++ b/tests/backend/patches/lora_conversions/lora_state_dicts/wan_lycoris_flat_lokr_format.py @@ -0,0 +1,36 @@ +"""Key layout of SimpleTuner's LyCORIS LoKR over Wan (`uxoah/wan2.1_lokr_simpletuner`). + +Captured on 2026-11-11 with a header-only read of `erkcga_lokr_simpletuner.safetensors` of the +repo, the issue #10174 Wan 2.1 row. Subsetting rule: block 0 in full -- the Attention preset's +eight projections, every module kind of this file. Values are `(shape, dtype)`. + +The Wan transformer's own module names flattened behind `lycoris_` +(`lycoris_blocks_0_attn1_to_q`, `lycoris_blocks_0_attn2_to_out_0`) -- the same tails kohya's +`lora_unet_blocks_` naming uses. 960 keys in the file, 320 layers (40 of 40 blocks).""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "lycoris_blocks_0_attn1_to_k.alpha": ([], "BF16"), + "lycoris_blocks_0_attn1_to_k.lokr_w1": ([16, 16], "BF16"), + "lycoris_blocks_0_attn1_to_k.lokr_w2": ([320, 320], "BF16"), + "lycoris_blocks_0_attn1_to_out_0.alpha": ([], "BF16"), + "lycoris_blocks_0_attn1_to_out_0.lokr_w1": ([16, 16], "BF16"), + "lycoris_blocks_0_attn1_to_out_0.lokr_w2": ([320, 320], "BF16"), + "lycoris_blocks_0_attn1_to_q.alpha": ([], "BF16"), + "lycoris_blocks_0_attn1_to_q.lokr_w1": ([16, 16], "BF16"), + "lycoris_blocks_0_attn1_to_q.lokr_w2": ([320, 320], "BF16"), + "lycoris_blocks_0_attn1_to_v.alpha": ([], "BF16"), + "lycoris_blocks_0_attn1_to_v.lokr_w1": ([16, 16], "BF16"), + "lycoris_blocks_0_attn1_to_v.lokr_w2": ([320, 320], "BF16"), + "lycoris_blocks_0_attn2_to_k.alpha": ([], "BF16"), + "lycoris_blocks_0_attn2_to_k.lokr_w1": ([16, 16], "BF16"), + "lycoris_blocks_0_attn2_to_k.lokr_w2": ([320, 320], "BF16"), + "lycoris_blocks_0_attn2_to_out_0.alpha": ([], "BF16"), + "lycoris_blocks_0_attn2_to_out_0.lokr_w1": ([16, 16], "BF16"), + "lycoris_blocks_0_attn2_to_out_0.lokr_w2": ([320, 320], "BF16"), + "lycoris_blocks_0_attn2_to_q.alpha": ([], "BF16"), + "lycoris_blocks_0_attn2_to_q.lokr_w1": ([16, 16], "BF16"), + "lycoris_blocks_0_attn2_to_q.lokr_w2": ([320, 320], "BF16"), + "lycoris_blocks_0_attn2_to_v.alpha": ([], "BF16"), + "lycoris_blocks_0_attn2_to_v.lokr_w1": ([16, 16], "BF16"), + "lycoris_blocks_0_attn2_to_v.lokr_w2": ([320, 320], "BF16"), +} diff --git a/tests/backend/patches/lora_conversions/lora_state_dicts/z_image_lycoris_flat_lokr_format.py b/tests/backend/patches/lora_conversions/lora_state_dicts/z_image_lycoris_flat_lokr_format.py new file mode 100644 index 00000000000..54dd8b89285 --- /dev/null +++ b/tests/backend/patches/lora_conversions/lora_state_dicts/z_image_lycoris_flat_lokr_format.py @@ -0,0 +1,39 @@ +"""Key layout of SimpleTuner's LyCORIS LoKR over Z-Image +(`noodlecake/Retro_Anime_Painting_lokr-R8-LR2-HIGHNOISEBIAS`). + +Captured on 2026-11-11 with a header-only read of `retanimpaint_lkr_8_v2_2500.safetensors`, the +issue #10174 Z-Image row (ai-toolkit 0.7.9's LyCORIS integration writes the same spelling). +Subsetting rule: layer 0 in full -- every module kind of the file. Values are `(shape, dtype)`. + +Z-Image's transformer module names flattened behind `lycoris_` +(`lycoris_layers_0_attention_to_k`, `lycoris_layers_0_feed_forward_w1`, +`lycoris_layers_0_adaLN_modulation_0`). 720 keys in the file, 240 layers (30 of 36 layers). +The attention width 3840 (LoKR product) is what tells Z-Image from Ideogram 4 (4608), which uses +the same module names.""" + +state_dict_keys: dict[str, tuple[list[int], str]] = { + "lycoris_layers_0_adaLN_modulation_0.alpha": ([], "BF16"), + "lycoris_layers_0_adaLN_modulation_0.lokr_w1": ([8, 8], "BF16"), + "lycoris_layers_0_adaLN_modulation_0.lokr_w2": ([1920, 32], "BF16"), + "lycoris_layers_0_attention_to_k.alpha": ([], "BF16"), + "lycoris_layers_0_attention_to_k.lokr_w1": ([8, 8], "BF16"), + "lycoris_layers_0_attention_to_k.lokr_w2": ([480, 480], "BF16"), + "lycoris_layers_0_attention_to_out_0.alpha": ([], "BF16"), + "lycoris_layers_0_attention_to_out_0.lokr_w1": ([8, 8], "BF16"), + "lycoris_layers_0_attention_to_out_0.lokr_w2": ([480, 480], "BF16"), + "lycoris_layers_0_attention_to_q.alpha": ([], "BF16"), + "lycoris_layers_0_attention_to_q.lokr_w1": ([8, 8], "BF16"), + "lycoris_layers_0_attention_to_q.lokr_w2": ([480, 480], "BF16"), + "lycoris_layers_0_attention_to_v.alpha": ([], "BF16"), + "lycoris_layers_0_attention_to_v.lokr_w1": ([8, 8], "BF16"), + "lycoris_layers_0_attention_to_v.lokr_w2": ([480, 480], "BF16"), + "lycoris_layers_0_feed_forward_w1.alpha": ([], "BF16"), + "lycoris_layers_0_feed_forward_w1.lokr_w1": ([8, 8], "BF16"), + "lycoris_layers_0_feed_forward_w1.lokr_w2": ([1280, 480], "BF16"), + "lycoris_layers_0_feed_forward_w2.alpha": ([], "BF16"), + "lycoris_layers_0_feed_forward_w2.lokr_w1": ([8, 8], "BF16"), + "lycoris_layers_0_feed_forward_w2.lokr_w2": ([480, 1280], "BF16"), + "lycoris_layers_0_feed_forward_w3.alpha": ([], "BF16"), + "lycoris_layers_0_feed_forward_w3.lokr_w1": ([8, 8], "BF16"), + "lycoris_layers_0_feed_forward_w3.lokr_w2": ([1280, 480], "BF16"), +} diff --git a/tests/backend/patches/lora_conversions/test_lycoris_flat_lora_conversion_utils.py b/tests/backend/patches/lora_conversions/test_lycoris_flat_lora_conversion_utils.py new file mode 100644 index 00000000000..9266cd4ce16 --- /dev/null +++ b/tests/backend/patches/lora_conversions/test_lycoris_flat_lora_conversion_utils.py @@ -0,0 +1,322 @@ +"""Conversion of SimpleTuner's LyCORIS flat-key LoRAs (issue #10174) onto each family's modules. + +A layer key that resolves to nothing is not an error at runtime: the patcher logs it and skips it, and the LoRA +quietly does less than it should. So every converted layer is resolved here against the model it will patch — +Invoke's BFL FLUX class for FLUX.1, and diffusers' models built from tiny configs whose module tree matches the +released ones for FLUX.2 Klein, Z-Image, Wan and SDXL (the SDXL UNet below carries the real +`stabilityai/stable-diffusion-xl-base-1.0` block structure at test sizes). The fixtures are the real key layouts +of the issue #10174 files, so each layer's weight delta is also checked against the (out, in) shape of the module +the real file patches, computed from its own LoKR factors. +""" + +import torch + +from invokeai.backend.flux.model import Flux, FluxParams +from invokeai.backend.patches.layers.merged_layer_patch import MergedLayerPatch +from invokeai.backend.patches.layers.swapped_halves_layer_patch import SwappedHalvesLayerPatch +from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX +from invokeai.backend.patches.lora_conversions.lycoris_flat_lora_conversion_utils import ( + lora_model_from_lycoris_flat_state_dict, +) +from invokeai.backend.patches.lora_conversions.wan_lora_constants import WAN_LORA_TRANSFORMER_PREFIX +from invokeai.backend.patches.lora_conversions.z_image_lora_constants import Z_IMAGE_LORA_TRANSFORMER_PREFIX +from tests.backend.patches.lora_conversions.lora_state_dicts import ( + flux1_lycoris_flat_lokr_double_underscore_format, + flux1_lycoris_flat_lokr_format, + flux2_lycoris_flat_lokr_format, + sdxl_lycoris_flat_lokr_format, + wan_lycoris_flat_lokr_format, + z_image_lycoris_flat_lokr_format, +) + + +def _shape(value) -> list[int]: + return list(value[0]) if isinstance(value, tuple) else list(value) + + +def _state_dict(keys: dict) -> dict[str, torch.Tensor]: + """The fixture's layout at its real shapes. LoKR's math runs on these, so they carry real storage.""" + return {key: torch.zeros(_shape(value)) for key, value in keys.items()} + + +def _layer_out_in(keys: dict, layer: str) -> tuple[int, int]: + """The (out, in) shape of the weight a fixture layer patches, from its own LoKR factors.""" + w1, w2 = _shape(keys[f"{layer}.lokr_w1"]), _shape(keys[f"{layer}.lokr_w2"]) + return w1[0] * w2[0], w1[1] * w2[1] + + +def test_the_flux1_flat_layout_lands_on_the_bfl_modules() -> None: + """toxicwind's file: the flattened diffusers names resolve to the same BFL keys the diffusers converter maps + them to, with the Q/K/V projections merged into BFL's fused `qkv`.""" + keys = flux1_lycoris_flat_lokr_format.state_dict_keys + patch = lora_model_from_lycoris_flat_state_dict(_state_dict(keys)) + prefix = FLUX_LORA_TRANSFORMER_PREFIX + + assert set(patch.layers) == { + f"{prefix}double_blocks.0.img_attn.qkv", + f"{prefix}double_blocks.0.img_attn.proj", + f"{prefix}double_blocks.0.img_mlp.0", + f"{prefix}double_blocks.0.img_mlp.2", + f"{prefix}double_blocks.0.txt_attn.qkv", + f"{prefix}double_blocks.0.txt_attn.proj", + f"{prefix}double_blocks.0.txt_mlp.0", + f"{prefix}double_blocks.0.txt_mlp.2", + f"{prefix}single_blocks.0.linear1", + } + + with torch.device("meta"): + model = Flux( + FluxParams( + in_channels=3072, + vec_in_dim=768, + context_in_dim=4096, + hidden_size=64, + mlp_ratio=4.0, + num_heads=2, + depth=1, + depth_single_blocks=1, + axes_dim=[8, 12, 12], + theta=10000, + qkv_bias=True, + guidance_embed=True, + out_channels=64, + ) + ) + for key in patch.layers: + assert model.get_submodule(key.removeprefix(prefix)) is not None, key + + # Every layer's delta is the shape of the module the real file patches (FLUX.1: hidden 3072, MLP 12288). + expected = { + "double_blocks.0.img_attn.proj": (3072, 3072), + "double_blocks.0.img_mlp.0": (12288, 3072), + "double_blocks.0.img_mlp.2": (3072, 12288), + "double_blocks.0.txt_attn.proj": (3072, 3072), + "double_blocks.0.txt_mlp.0": (12288, 3072), + "double_blocks.0.txt_mlp.2": (3072, 12288), + } + for name, shape in expected.items(): + assert patch.layers[f"{prefix}{name}"].get_weight(torch.empty(0)).shape == shape, name + + # The merged img qkv holds to_q, to_k and to_v, each 3072 wide; single_blocks.linear1 the same three. + for name in ("double_blocks.0.img_attn.qkv", "double_blocks.0.txt_attn.qkv", "single_blocks.0.linear1"): + merged = patch.layers[f"{prefix}{name}"] + assert isinstance(merged, MergedLayerPatch) and len(merged.lora_layers) == 3, name + for _sub, sub_range in zip(merged.lora_layers, merged.ranges, strict=True): + assert sub_range.end - sub_range.start == 3072, name + + +def test_the_flux1_double_underscore_layout_also_lands_on_the_bfl_modules() -> None: + """jimmycarter's `lycoris__` file: same mapping, plus the top-level modules and the swapped final modulation.""" + keys = flux1_lycoris_flat_lokr_double_underscore_format.state_dict_keys + patch = lora_model_from_lycoris_flat_state_dict(_state_dict(keys)) + prefix = FLUX_LORA_TRANSFORMER_PREFIX + + assert set(patch.layers) == { + f"{prefix}txt_in", + f"{prefix}img_in", + f"{prefix}final_layer.linear", + f"{prefix}final_layer.adaLN_modulation.1", + f"{prefix}time_in.in_layer", + f"{prefix}time_in.out_layer", + f"{prefix}vector_in.in_layer", + f"{prefix}vector_in.out_layer", + f"{prefix}guidance_in.in_layer", + f"{prefix}guidance_in.out_layer", + f"{prefix}single_blocks.0.linear1", + f"{prefix}single_blocks.0.linear2", + f"{prefix}single_blocks.0.modulation.lin", + f"{prefix}double_blocks.0.img_attn.qkv", + f"{prefix}double_blocks.0.img_attn.proj", + f"{prefix}double_blocks.0.img_mlp.0", + f"{prefix}double_blocks.0.img_mlp.2", + f"{prefix}double_blocks.0.txt_attn.qkv", + f"{prefix}double_blocks.0.txt_attn.proj", + f"{prefix}double_blocks.0.txt_mlp.0", + f"{prefix}double_blocks.0.txt_mlp.2", + f"{prefix}double_blocks.0.img_mod.lin", + f"{prefix}double_blocks.0.txt_mod.lin", + } + assert isinstance(patch.layers[f"{prefix}final_layer.adaLN_modulation.1"], SwappedHalvesLayerPatch) + + # Every layer's delta is the shape of the module the real file patches. + expected = { + "txt_in": (3072, 4096), + "img_in": (3072, 64), + "final_layer.linear": (64, 3072), + "time_in.in_layer": (3072, 256), + "time_in.out_layer": (3072, 3072), + "vector_in.in_layer": (3072, 768), + "vector_in.out_layer": (3072, 3072), + "guidance_in.in_layer": (3072, 256), + "guidance_in.out_layer": (3072, 3072), + # BFL's linear2 consumes the single block's concatenated attention+MLP hidden (h + 4h). + "single_blocks.0.linear2": (3072, 15360), + "single_blocks.0.modulation.lin": (9216, 3072), + "double_blocks.0.img_attn.proj": (3072, 3072), + "double_blocks.0.img_mlp.0": (12288, 3072), + "double_blocks.0.img_mlp.2": (3072, 12288), + "double_blocks.0.txt_attn.proj": (3072, 3072), + "double_blocks.0.txt_mlp.0": (12288, 3072), + "double_blocks.0.txt_mlp.2": (3072, 12288), + "double_blocks.0.img_mod.lin": (18432, 3072), + "double_blocks.0.txt_mod.lin": (18432, 3072), + } + for name, shape in expected.items(): + assert patch.layers[f"{prefix}{name}"].get_weight(torch.empty(0)).shape == shape, name + + +_TINY_FLUX2_CONFIG = { + "patch_size": 1, + "in_channels": 128, + "num_layers": 1, + "num_single_layers": 1, + "attention_head_dim": 8, + "num_attention_heads": 2, + "joint_attention_dim": 8, + "timestep_guidance_channels": 16, + "axes_dims_rope": (2, 2, 2, 2), +} + + +def test_the_flux2_flat_layout_lands_on_the_flux2_transformer() -> None: + from diffusers import Flux2Transformer2DModel + + keys = flux2_lycoris_flat_lokr_format.state_dict_keys + patch = lora_model_from_lycoris_flat_state_dict(_state_dict(keys)) + prefix = FLUX_LORA_TRANSFORMER_PREFIX + + assert set(patch.layers) == { + f"{prefix}single_transformer_blocks.0.attn.to_qkv_mlp_proj", + f"{prefix}single_transformer_blocks.0.attn.to_out", + f"{prefix}transformer_blocks.0.attn.to_q", + f"{prefix}transformer_blocks.0.attn.to_k", + f"{prefix}transformer_blocks.0.attn.to_v", + f"{prefix}transformer_blocks.0.attn.to_out.0", + f"{prefix}transformer_blocks.0.attn.add_q_proj", + f"{prefix}transformer_blocks.0.attn.add_k_proj", + f"{prefix}transformer_blocks.0.attn.add_v_proj", + f"{prefix}transformer_blocks.0.attn.to_add_out", + f"{prefix}transformer_blocks.0.ff.linear_in", + f"{prefix}transformer_blocks.0.ff.linear_out", + f"{prefix}transformer_blocks.0.ff_context.linear_in", + f"{prefix}transformer_blocks.0.ff_context.linear_out", + } + + with torch.device("meta"): + model = Flux2Transformer2DModel(**_TINY_FLUX2_CONFIG) + for key in patch.layers: + assert isinstance(model.get_submodule(key.removeprefix(prefix)), torch.nn.Linear), key + + +def test_the_z_image_flat_layout_lands_on_the_z_image_transformer() -> None: + from diffusers import ZImageTransformer2DModel + + keys = z_image_lycoris_flat_lokr_format.state_dict_keys + patch = lora_model_from_lycoris_flat_state_dict(_state_dict(keys)) + prefix = Z_IMAGE_LORA_TRANSFORMER_PREFIX + + assert set(patch.layers) == { + f"{prefix}layers.0.attention.to_q", + f"{prefix}layers.0.attention.to_k", + f"{prefix}layers.0.attention.to_v", + f"{prefix}layers.0.attention.to_out.0", + f"{prefix}layers.0.feed_forward.w1", + f"{prefix}layers.0.feed_forward.w2", + f"{prefix}layers.0.feed_forward.w3", + f"{prefix}layers.0.adaLN_modulation.0", + } + + with torch.device("meta"): + model = ZImageTransformer2DModel( + in_channels=4, + dim=64, + n_layers=1, + n_refiner_layers=1, + n_heads=2, + n_kv_heads=2, + cap_feat_dim=8, + axes_dims=[8, 12, 12], + axes_lens=[64, 32, 32], + ) + for key in patch.layers: + assert isinstance(model.get_submodule(key.removeprefix(prefix)), torch.nn.Linear), key + + +def test_the_wan_flat_layout_lands_on_the_wan_transformer() -> None: + from diffusers import WanTransformer3DModel + + keys = wan_lycoris_flat_lokr_format.state_dict_keys + patch = lora_model_from_lycoris_flat_state_dict(_state_dict(keys)) + prefix = WAN_LORA_TRANSFORMER_PREFIX + + assert set(patch.layers) == { + f"{prefix}blocks.0.attn1.to_q", + f"{prefix}blocks.0.attn1.to_k", + f"{prefix}blocks.0.attn1.to_v", + f"{prefix}blocks.0.attn1.to_out.0", + f"{prefix}blocks.0.attn2.to_q", + f"{prefix}blocks.0.attn2.to_k", + f"{prefix}blocks.0.attn2.to_v", + f"{prefix}blocks.0.attn2.to_out.0", + } + + with torch.device("meta"): + model = WanTransformer3DModel( + num_layers=1, num_attention_heads=2, attention_head_dim=8, ffn_dim=16, text_dim=16, in_channels=4 + ) + for key in patch.layers: + assert isinstance(model.get_submodule(key.removeprefix(prefix)), torch.nn.Linear), key + + +# The SDXL UNet's block structure, from `stabilityai/stable-diffusion-xl-base-1.0`'s unet config: the first down +# block and the last up block carry no attention, the deep attention blocks hold 10 transformer blocks. At test +# sizes, so only the names matter. +_SDXL_UNET_CONFIG = { + "sample_size": 32, + "in_channels": 4, + "out_channels": 4, + "down_block_types": ("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D"), + "up_block_types": ("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D"), + "block_out_channels": (32, 64, 64), + "layers_per_block": 2, + "transformer_layers_per_block": (1, 2, 10), + "attention_head_dim": (2, 2, 2), + "cross_attention_dim": 16, + "use_linear_projection": True, + "norm_num_groups": 8, +} + + +def test_the_sdxl_flat_layout_lands_on_the_unet_by_flattened_name() -> None: + """SD/SDXL patches keep kohya's flattened keys (`lora_unet_` + the name with `.` as `_`); the patcher resolves + them against the UNet by the same flattening.""" + from diffusers import UNet2DConditionModel + + from invokeai.backend.patches.layer_patcher import LayerPatcher + + keys = sdxl_lycoris_flat_lokr_format.state_dict_keys + patch = lora_model_from_lycoris_flat_state_dict(_state_dict(keys)) + + expected = { + "lora_unet_" + key.rsplit(".", 1)[0].removeprefix("lycoris_") for key in keys if key.endswith(".lokr_w1") + } + assert set(patch.layers) == expected + + with torch.device("meta"): + unet = UNet2DConditionModel(**_SDXL_UNET_CONFIG) + for key in patch.layers: + _, module = LayerPatcher._get_submodule(unet, key.removeprefix("lora_unet_"), layer_key_is_flattened=True) + assert isinstance(module, torch.nn.Linear), key + + +def test_a_flat_file_with_a_key_of_no_family_is_refused() -> None: + keys = dict(flux1_lycoris_flat_lokr_format.state_dict_keys) + keys["lycoris_some_module_nobody_has.lokr_w1"] = ([8, 8], "BF16") + keys["lycoris_some_module_nobody_has.lokr_w2"] = ([384, 384], "BF16") + + try: + lora_model_from_lycoris_flat_state_dict(_state_dict(keys)) + except ValueError as error: + assert "some_module_nobody_has" in str(error) + else: + raise AssertionError("a layer that matches no module must be refused, not dropped")