mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 01:04:32 +02:00
refactor(anima): migrate to generic native_loader
Replaces anima_lora.py's bespoke try_load_lora / group_keys / resolve_network_key with thin wrappers binding native_loader's generics to anima's prefix tuples and resolve_targets, mirroring flux2 / zimage / chroma / ernie. The hand-rolled apply_lora_alphas bake-with-balance pass goes away; alpha / scale / dora_scale flow through NetworkWeights.w to NetworkModule.calc_scale at apply time. native_loader gains an optional network_prefix kwarg (str or Callable[[prefix_used], str], default "lora_transformer_") used when constructing network_key. Anima passes a callable picking lora_transformer_ / lora_llm_adapter_ / lora_te_ per matched prefix. Single-component siblings keep the default and are unchanged. network.NetworkModule.apply_weight_decompose grows a dual-path DoRA convention detector. The pre-fix implementation only handled per-input dora_scale (DoRA paper / kohya, shape (1, in)), silently broadcasting per-output LyCORIS / PEFT dora_scale (shape (out, 1)) into an incoherent element-wise rescaling. Detection is structural: (out, 1, ...) routes to per-output; everything else (including the square-weight 1D ambiguity) defaults to per-input for legacy compat. Pre-existing bug surfaced by the LoKR+DoRA LyCORIS files Anima now loads. Behavior changes: - LoHA via the generic try_load_loha (NetworkModuleHada); covers scenery-anima-base and any other LyCORIS .hada_w* export. - Kohya lora_te_ prefix recognized. The legacy resolver only matched BFL text_encoders.qwen3_06b.transformer.model. and silently dropped lora_te_layers_N_* keys (41% of BlueArcStyle's bases were unloaded). - LoKR+DoRA LyCORIS files now apply correctly; the per-output dora_scale is honored instead of silently scrambled. Adds test/test-anima-native-adapters.py: 37 offline tests across all five prefixes (LoRA + LoHA), every COSMOS_2_FLAT_RENAME entry, DoRA threading, marker disambiguation, try_load_chain dispatch, calc_updown sanity, and both DoRA conventions (per-input / per-output / 1D ambiguous). Adapter mock mirrors AnimaLLMAdapter's real module tree.
This commit is contained in:
CalamitousFelicitousness
@@ -56,6 +56,23 @@ KNOWN_PREFIXES_DEFAULT = ("diffusion_model.", "transformer.", "lora_unet_")
|
||||
BARE_DIFFUSERS_PREFIX_USED = "bare_diffusers"
|
||||
|
||||
|
||||
# Default network-key prefix. Single-component arches (flux2, zimage, chroma,
|
||||
# ernie) keep this default; multi-component arches (anima: transformer plus
|
||||
# llm_adapter plus text_encoder) pass a callable that picks per ``prefix_used``.
|
||||
NETWORK_PREFIX_DEFAULT = "lora_transformer_"
|
||||
|
||||
|
||||
def _resolve_prefix(network_prefix, prefix_used):
|
||||
"""Return the network-key prefix for one parsed group.
|
||||
|
||||
``network_prefix`` is either a literal string (single-component arches) or
|
||||
a ``Callable[[str | None], str]`` that picks based on which arch prefix
|
||||
was matched (multi-component arches route to ``lora_te_`` / ``lora_llm_adapter_``
|
||||
/ ``lora_transformer_`` etc.).
|
||||
"""
|
||||
return network_prefix(prefix_used) if callable(network_prefix) else network_prefix
|
||||
|
||||
|
||||
SUFFIX_NORMALIZE = {
|
||||
"lora_A.weight": "lora_down.weight",
|
||||
"lora_B.weight": "lora_up.weight",
|
||||
@@ -345,8 +362,12 @@ read_state_dict = sd_models.read_state_dict
|
||||
#
|
||||
# Each loader takes a per-arch ``resolve_targets`` callable returning a list of
|
||||
# ``(diffusers_path, ChunkSpec | None)`` tuples for each parsed group. The
|
||||
# loader builds the network key as ``"lora_transformer_" + path.replace(".", "_")``
|
||||
# and instantiates the appropriate ``network.NetworkModule*`` subclass.
|
||||
# loader builds the network key as ``network_prefix + path.replace(".", "_")``
|
||||
# where ``network_prefix`` defaults to ``"lora_transformer_"`` and may be
|
||||
# either a literal string (single-component arches) or a callable that picks
|
||||
# per ``prefix_used`` (multi-component arches such as Anima that route to
|
||||
# ``lora_te_`` or ``lora_llm_adapter_`` based on which arch prefix matched).
|
||||
# The loader then instantiates the appropriate ``network.NetworkModule*`` subclass.
|
||||
#
|
||||
# Fused-target handling per family:
|
||||
#
|
||||
@@ -383,6 +404,7 @@ def _slice_lora_chunk(w, chunk: ChunkSpec):
|
||||
def try_load_lora(name, network_on_disk, lora_scale, *,
|
||||
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
|
||||
bare_prefixes=(), bare_diffusers_prefixes=(),
|
||||
network_prefix=NETWORK_PREFIX_DEFAULT,
|
||||
arch_name="generic"):
|
||||
"""Generic LoRA loader (handles DoRA via the universal ``finalize_updown`` hook).
|
||||
|
||||
@@ -408,8 +430,9 @@ def try_load_lora(name, network_on_disk, lora_scale, *,
|
||||
for (prefix, base), w in groups.items():
|
||||
if "lora_down.weight" not in w or "lora_up.weight" not in w:
|
||||
continue
|
||||
arch_prefix = _resolve_prefix(network_prefix, prefix)
|
||||
for diffusers_path, chunk in resolve_targets(prefix, base):
|
||||
network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
|
||||
network_key = arch_prefix + diffusers_path.replace(".", "_")
|
||||
sd_module = mapping.get(network_key)
|
||||
if sd_module is None:
|
||||
unmapped += 1
|
||||
@@ -436,6 +459,7 @@ def try_load_lora(name, network_on_disk, lora_scale, *,
|
||||
def try_load_lokr(name, network_on_disk, lora_scale, *,
|
||||
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
|
||||
bare_prefixes=(), bare_diffusers_prefixes=(),
|
||||
network_prefix=NETWORK_PREFIX_DEFAULT,
|
||||
arch_name="generic"):
|
||||
"""Generic LoKR loader.
|
||||
|
||||
@@ -466,8 +490,9 @@ def try_load_lokr(name, network_on_disk, lora_scale, *,
|
||||
has_2 = "lokr_w2" in w or ("lokr_w2_a" in w and "lokr_w2_b" in w)
|
||||
if not (has_1 and has_2):
|
||||
continue
|
||||
arch_prefix = _resolve_prefix(network_prefix, prefix)
|
||||
for diffusers_path, chunk in resolve_targets(prefix, base):
|
||||
network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
|
||||
network_key = arch_prefix + diffusers_path.replace(".", "_")
|
||||
sd_module = mapping.get(network_key)
|
||||
if sd_module is None:
|
||||
unmapped += 1
|
||||
@@ -486,6 +511,7 @@ def try_load_lokr(name, network_on_disk, lora_scale, *,
|
||||
def try_load_loha(name, network_on_disk, lora_scale, *,
|
||||
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
|
||||
bare_prefixes=(), bare_diffusers_prefixes=(),
|
||||
network_prefix=NETWORK_PREFIX_DEFAULT,
|
||||
arch_name="generic"):
|
||||
"""Generic LoHA (Hadamard product) loader.
|
||||
|
||||
@@ -521,8 +547,9 @@ def try_load_loha(name, network_on_disk, lora_scale, *,
|
||||
log.warning(f'Network load: type=LoHA name="{name}" arch={arch_name} key={base} Tucker fused QKV skipped (unsupported)')
|
||||
skipped += 1
|
||||
continue
|
||||
arch_prefix = _resolve_prefix(network_prefix, prefix)
|
||||
for diffusers_path, chunk in targets:
|
||||
network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
|
||||
network_key = arch_prefix + diffusers_path.replace(".", "_")
|
||||
sd_module = mapping.get(network_key)
|
||||
if sd_module is None:
|
||||
unmapped += 1
|
||||
@@ -542,6 +569,7 @@ def try_load_loha(name, network_on_disk, lora_scale, *,
|
||||
def try_load_oft(name, network_on_disk, lora_scale, *,
|
||||
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
|
||||
bare_prefixes=(), bare_diffusers_prefixes=(),
|
||||
network_prefix=NETWORK_PREFIX_DEFAULT,
|
||||
arch_name="generic"):
|
||||
"""Generic OFT/BOFT loader.
|
||||
|
||||
@@ -582,8 +610,9 @@ def try_load_oft(name, network_on_disk, lora_scale, *,
|
||||
log.warning(f'Network load: type={"BOFT" if is_boft else "OFT"} name="{name}" arch={arch_name} key={base} fused QKV skipped (unsupported)')
|
||||
skipped += 1
|
||||
continue
|
||||
arch_prefix = _resolve_prefix(network_prefix, prefix)
|
||||
for diffusers_path, _ in targets:
|
||||
network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
|
||||
network_key = arch_prefix + diffusers_path.replace(".", "_")
|
||||
sd_module = mapping.get(network_key)
|
||||
if sd_module is None:
|
||||
unmapped += 1
|
||||
@@ -600,6 +629,7 @@ def try_load_oft(name, network_on_disk, lora_scale, *,
|
||||
def try_load_ia3(name, network_on_disk, lora_scale, *,
|
||||
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
|
||||
bare_prefixes=(), bare_diffusers_prefixes=(),
|
||||
network_prefix=NETWORK_PREFIX_DEFAULT,
|
||||
arch_name="generic"):
|
||||
"""Generic IA3 loader.
|
||||
|
||||
@@ -638,8 +668,9 @@ def try_load_ia3(name, network_on_disk, lora_scale, *,
|
||||
log.warning(f'Network load: type=IA3 name="{name}" arch={arch_name} key={base} fused QKV skipped (unsupported)')
|
||||
skipped += 1
|
||||
continue
|
||||
arch_prefix = _resolve_prefix(network_prefix, prefix)
|
||||
for diffusers_path, _ in targets:
|
||||
network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
|
||||
network_key = arch_prefix + diffusers_path.replace(".", "_")
|
||||
sd_module = mapping.get(network_key)
|
||||
if sd_module is None:
|
||||
unmapped += 1
|
||||
@@ -653,6 +684,7 @@ def try_load_ia3(name, network_on_disk, lora_scale, *,
|
||||
def try_load_glora(name, network_on_disk, lora_scale, *,
|
||||
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
|
||||
bare_prefixes=(), bare_diffusers_prefixes=(),
|
||||
network_prefix=NETWORK_PREFIX_DEFAULT,
|
||||
arch_name="generic"):
|
||||
"""Generic GLoRA loader.
|
||||
|
||||
@@ -687,8 +719,9 @@ def try_load_glora(name, network_on_disk, lora_scale, *,
|
||||
log.warning(f'Network load: type=GLoRA name="{name}" arch={arch_name} key={base} fused QKV skipped (unsupported)')
|
||||
skipped += 1
|
||||
continue
|
||||
arch_prefix = _resolve_prefix(network_prefix, prefix)
|
||||
for diffusers_path, _ in targets:
|
||||
network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
|
||||
network_key = arch_prefix + diffusers_path.replace(".", "_")
|
||||
sd_module = mapping.get(network_key)
|
||||
if sd_module is None:
|
||||
unmapped += 1
|
||||
@@ -702,6 +735,7 @@ def try_load_glora(name, network_on_disk, lora_scale, *,
|
||||
def try_load_norm(name, network_on_disk, lora_scale, *,
|
||||
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
|
||||
bare_prefixes=(), bare_diffusers_prefixes=(),
|
||||
network_prefix=NETWORK_PREFIX_DEFAULT,
|
||||
arch_name="generic"):
|
||||
"""Generic Norm (LayerNorm / RMSNorm weight + bias delta) loader.
|
||||
|
||||
@@ -737,10 +771,11 @@ def try_load_norm(name, network_on_disk, lora_scale, *,
|
||||
if not targets:
|
||||
unmapped += 1
|
||||
continue
|
||||
arch_prefix = _resolve_prefix(network_prefix, prefix)
|
||||
for diffusers_path, chunk in targets:
|
||||
if chunk is not None:
|
||||
continue # norm targets are not fused
|
||||
network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
|
||||
network_key = arch_prefix + diffusers_path.replace(".", "_")
|
||||
sd_module = mapping.get(network_key)
|
||||
if sd_module is None:
|
||||
unmapped += 1
|
||||
@@ -756,6 +791,7 @@ def try_load_norm(name, network_on_disk, lora_scale, *,
|
||||
def try_load_full(name, network_on_disk, lora_scale, *,
|
||||
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
|
||||
bare_prefixes=(), bare_diffusers_prefixes=(),
|
||||
network_prefix=NETWORK_PREFIX_DEFAULT,
|
||||
arch_name="generic"):
|
||||
"""Generic Full (full-rank weight delta) loader.
|
||||
|
||||
@@ -789,8 +825,9 @@ def try_load_full(name, network_on_disk, lora_scale, *,
|
||||
log.warning(f'Network load: type=Full name="{name}" arch={arch_name} key={base} fused QKV skipped (unsupported)')
|
||||
skipped += 1
|
||||
continue
|
||||
arch_prefix = _resolve_prefix(network_prefix, prefix)
|
||||
for diffusers_path, _ in targets:
|
||||
network_key = "lora_transformer_" + diffusers_path.replace(".", "_")
|
||||
network_key = arch_prefix + diffusers_path.replace(".", "_")
|
||||
sd_module = mapping.get(network_key)
|
||||
if sd_module is None:
|
||||
unmapped += 1
|
||||
|
||||
+46
-10
@@ -202,17 +202,53 @@ class NetworkModule:
|
||||
updown = updown.to(orig_weight.device)
|
||||
|
||||
merged_scale1 = updown + orig_weight
|
||||
merged_scale1_norm = (
|
||||
merged_scale1.transpose(0, 1)
|
||||
.reshape(merged_scale1.shape[1], -1)
|
||||
.norm(dim=1, keepdim=True)
|
||||
.reshape(merged_scale1.shape[1], *[1] * self.dora_norm_dims)
|
||||
.transpose(0, 1)
|
||||
)
|
||||
|
||||
dora_merged = (
|
||||
merged_scale1 * (dora_scale / merged_scale1_norm)
|
||||
)
|
||||
# DoRA convention detection. Two flavors coexist in the wild:
|
||||
#
|
||||
# - per-input (DoRA paper / kohya): dora_scale stores per-column magnitudes,
|
||||
# shape ``(1, in, ...)`` or ``(in,)``. ``W' = W * (m / ||W||_col)`` rescales
|
||||
# each column to magnitude ``m[i]``.
|
||||
# - per-output (LyCORIS / PEFT / diffusers): dora_scale stores per-row
|
||||
# magnitudes, shape ``(out, 1, ...)`` or ``(out,)``. ``W' = W * (m / ||W||_row)``
|
||||
# rescales each row to magnitude ``m[o]``.
|
||||
#
|
||||
# PyTorch silently broadcasts ``(out, 1) / (1, in)`` into ``(out, in)``, so
|
||||
# mismatched conventions are not a shape error but a semantic one (the
|
||||
# update gets scrambled). Detection is structural rather than numeric:
|
||||
# a 2-D dora_scale with shape ``(out, 1, ...)`` is unambiguously per-output
|
||||
# even when ``out == in`` (square weights like self-attention q/k/v).
|
||||
# 1-D dora_scale falls back to comparing the length against out / in;
|
||||
# when both match (square weight), default to per-input for legacy compat.
|
||||
out_dim = merged_scale1.shape[0]
|
||||
in_dim = merged_scale1.shape[1] if merged_scale1.ndim >= 2 else None
|
||||
per_output = False
|
||||
if dora_scale.ndim >= 2:
|
||||
# ND form: leading dim equals out_dim and every trailing dim is 1.
|
||||
if dora_scale.shape[0] == out_dim and all(d == 1 for d in dora_scale.shape[1:]):
|
||||
per_output = True
|
||||
elif dora_scale.ndim == 1:
|
||||
# 1D vector: per-output only when length unambiguously matches out_dim.
|
||||
if dora_scale.shape[0] == out_dim and dora_scale.shape[0] != in_dim:
|
||||
per_output = True
|
||||
|
||||
if per_output:
|
||||
# Per-output: norm along all non-output axes; result broadcasts as (out, 1, ...).
|
||||
merged_scale1_norm = (
|
||||
merged_scale1.reshape(out_dim, -1)
|
||||
.norm(dim=1, keepdim=True)
|
||||
.reshape(out_dim, *[1] * self.dora_norm_dims)
|
||||
)
|
||||
else:
|
||||
# Per-input: norm along all non-input axes; result broadcasts as (1, in, ...).
|
||||
merged_scale1_norm = (
|
||||
merged_scale1.transpose(0, 1)
|
||||
.reshape(merged_scale1.shape[1], -1)
|
||||
.norm(dim=1, keepdim=True)
|
||||
.reshape(merged_scale1.shape[1], *[1] * self.dora_norm_dims)
|
||||
.transpose(0, 1)
|
||||
)
|
||||
|
||||
dora_merged = merged_scale1 * (dora_scale / merged_scale1_norm)
|
||||
final_updown = dora_merged - orig_weight
|
||||
return final_updown
|
||||
|
||||
|
||||
+206
-218
@@ -1,239 +1,227 @@
|
||||
"""Anima native LoRA loader.
|
||||
"""Anima native adapter loader.
|
||||
|
||||
Handles three trainer formats observed in community Anima LoRAs:
|
||||
- Kohya: lora_unet_blocks_0_self_attn_q_proj.lora_down.weight (+ .alpha)
|
||||
- BFL/AI-toolkit: diffusion_model.blocks.0.self_attn.q_proj.lora_A.weight
|
||||
- Hybrid: BFL key shape with .alpha scalars, optionally with a Qwen3 text
|
||||
encoder branch under text_encoders.qwen3_06b.transformer.model.layers...
|
||||
Anima ships with three trainable surfaces:
|
||||
|
||||
Routes paths to one of three live components (transformer, llm_adapter,
|
||||
text_encoder) via lora_convert.assign_network_names_to_compvis_modules and
|
||||
the resulting network_layer_mapping.
|
||||
- ``transformer`` (Cosmos 2.0-style DiT), stamped as ``lora_transformer_*``
|
||||
- ``llm_adapter`` (custom MLP that projects Qwen3 hidden states into the DiT
|
||||
cross-attention dim), stamped as ``lora_llm_adapter_*``
|
||||
- ``text_encoder`` (Qwen3Model), stamped as ``lora_te_*``
|
||||
|
||||
Path renames mirror the Cosmos 2.0 table from
|
||||
diffusers.loaders.single_file_utils.convert_cosmos_transformer_checkpoint_to_diffusers,
|
||||
transposed into flat (underscore) form so the rewritten path matches
|
||||
network_layer_mapping keys without further conversion.
|
||||
The loader recognizes four ``ss_network_module`` formats in the wild:
|
||||
|
||||
- kohya transformer: ``lora_unet_blocks_0_self_attn_q_proj.lora_down.weight``
|
||||
- kohya text-encoder: ``lora_te_layers_0_self_attn_q_proj.lora_down.weight``
|
||||
- BFL/AI-toolkit: ``diffusion_model.blocks.0.self_attn.q_proj.lora_A.weight``
|
||||
with the adapter under ``diffusion_model.llm_adapter.*`` and the TE under
|
||||
``text_encoders.qwen3_06b.transformer.model.*``
|
||||
- LyCORIS Hadamard (LoHA): same kohya path structure but with
|
||||
``hada_w1_a`` / ``hada_w1_b`` / ``hada_w2_a`` / ``hada_w2_b`` weights
|
||||
|
||||
All transformer paths route through :func:`cosmos_rename_flat` which mirrors
|
||||
``diffusers.loaders.single_file_utils.convert_cosmos_transformer_checkpoint_to_diffusers``'s
|
||||
``TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0``, transposed into flat (underscore)
|
||||
form so str.replace produces a path that matches the diffusers module name
|
||||
already stamped on the network_layer_mapping. Adapter and TE paths bypass the
|
||||
rename and are flattened verbatim.
|
||||
|
||||
Network-key construction (transformer vs llm_adapter vs te) is parameterized
|
||||
in :mod:`modules.lora.native_loader` via the ``network_prefix`` kwarg; this
|
||||
module supplies :func:`network_prefix_for` to pick per ``prefix_used``.
|
||||
Family-specific dispatch (LoRA, LoHA, LoKR, OFT, IA3, GLoRA, Norm, Full) is
|
||||
inherited from native_loader's generics; alpha / scale / DoRA flow through
|
||||
the standard ``NetworkWeights.w`` slots rather than being baked into the
|
||||
factor weights at load time.
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from modules import shared, sd_models
|
||||
from modules.logger import log
|
||||
from modules.lora import network, network_lora, lora_convert
|
||||
from modules.lora import lora_common as l
|
||||
|
||||
from modules.lora import native_loader
|
||||
|
||||
|
||||
KOHYA_PREFIX = 'lora_unet_'
|
||||
BFL_ADAPTER_PREFIX = 'diffusion_model.llm_adapter.'
|
||||
BFL_TRANSFORMER_PREFIX = 'diffusion_model.'
|
||||
BFL_TE_PREFIX = 'text_encoders.qwen3_06b.transformer.model.'
|
||||
# === Arch-specific prefix configuration ===
|
||||
#
|
||||
# Order matters: longer / more-specific prefixes must precede shorter ones,
|
||||
# because :func:`native_loader.parse_key` returns the first match. Both
|
||||
# ``diffusion_model.llm_adapter.`` and ``text_encoders.qwen3_06b.transformer.model.``
|
||||
# start with ``diffusion_model.`` / ``text_encoders.`` so they must be listed first.
|
||||
|
||||
LORA_MARKERS = ('.lora_down.', '.lora_up.', '.lora_A.', '.lora_B.')
|
||||
ANIMA_PREFIXES = (
|
||||
"diffusion_model.llm_adapter.",
|
||||
"text_encoders.qwen3_06b.transformer.model.",
|
||||
"diffusion_model.",
|
||||
"lora_te_",
|
||||
"lora_unet_",
|
||||
)
|
||||
|
||||
|
||||
# === Re-exports for test / back-compat ===
|
||||
# Tests address these through the anima_lora module surface; sibling pipelines
|
||||
# do the same (see flux2_lora / zimage_lora / ernie_lora).
|
||||
|
||||
LORA_SUFFIXES = native_loader.LORA_SUFFIXES
|
||||
LOKR_SUFFIXES = native_loader.LOKR_SUFFIXES
|
||||
LOHA_SUFFIXES = native_loader.LOHA_SUFFIXES
|
||||
OFT_SUFFIXES = native_loader.OFT_SUFFIXES
|
||||
IA3_SUFFIXES = native_loader.IA3_SUFFIXES
|
||||
GLORA_SUFFIXES = native_loader.GLORA_SUFFIXES
|
||||
NORM_SUFFIXES = native_loader.NORM_SUFFIXES
|
||||
FULL_SUFFIXES = native_loader.FULL_SUFFIXES
|
||||
|
||||
LORA_MARKERS = native_loader.LORA_MARKERS
|
||||
LOKR_MARKERS = native_loader.LOKR_MARKERS
|
||||
LOHA_MARKERS = native_loader.LOHA_MARKERS
|
||||
OFT_MARKERS = native_loader.OFT_MARKERS
|
||||
IA3_MARKERS = native_loader.IA3_MARKERS
|
||||
GLORA_MARKERS = native_loader.GLORA_MARKERS
|
||||
NORM_MARKERS = native_loader.NORM_MARKERS
|
||||
FULL_MARKERS = native_loader.FULL_MARKERS
|
||||
|
||||
SUFFIX_NORMALIZE = native_loader.SUFFIX_NORMALIZE
|
||||
BARE_DIFFUSERS_PREFIX_USED = native_loader.BARE_DIFFUSERS_PREFIX_USED
|
||||
has_marker = native_loader.has_marker
|
||||
|
||||
|
||||
def parse_key(key, suffixes):
|
||||
"""Anima-bound :func:`native_loader.parse_key`."""
|
||||
return native_loader.parse_key(key, suffixes, prefixes=ANIMA_PREFIXES)
|
||||
|
||||
|
||||
def group_by_suffixes(state_dict, suffixes):
|
||||
"""Anima-bound :func:`native_loader.group_by_suffixes`."""
|
||||
return native_loader.group_by_suffixes(state_dict, suffixes, prefixes=ANIMA_PREFIXES)
|
||||
|
||||
|
||||
# === Cosmos 2.0 path rename (transformer only) ===
|
||||
#
|
||||
# Applied to underscore-flattened paths. Order mirrors diffusers
|
||||
# ``TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0`` exactly: longer substrings must
|
||||
# precede substrings nested inside them so ``str.replace`` does not consume a
|
||||
# fragment that a later rule still needs to match. Numeric-suffixed entries
|
||||
# (e.g. ``t_embedder_1``) come first because the bare ``t_embedder`` would
|
||||
# otherwise eat the ``_1`` suffix.
|
||||
|
||||
# Cosmos 2.0 path rename, applied to underscore-flattened paths. Order
|
||||
# mirrors diffusers TRANSFORMER_KEYS_RENAME_DICT_COSMOS_2_0: longer
|
||||
# substrings precede substrings nested inside them, so str.replace does
|
||||
# not consume a fragment that a later rule still needs to match.
|
||||
COSMOS_2_FLAT_RENAME = OrderedDict([
|
||||
('t_embedder_1', 'time_embed_t_embedder'),
|
||||
('t_embedding_norm', 'time_embed_norm'),
|
||||
('blocks', 'transformer_blocks'),
|
||||
('adaln_modulation_self_attn_1', 'norm1_linear_1'),
|
||||
('adaln_modulation_self_attn_2', 'norm1_linear_2'),
|
||||
('adaln_modulation_cross_attn_1', 'norm2_linear_1'),
|
||||
('adaln_modulation_cross_attn_2', 'norm2_linear_2'),
|
||||
('adaln_modulation_mlp_1', 'norm3_linear_1'),
|
||||
('adaln_modulation_mlp_2', 'norm3_linear_2'),
|
||||
('self_attn', 'attn1'),
|
||||
('cross_attn', 'attn2'),
|
||||
('q_proj', 'to_q'),
|
||||
('k_proj', 'to_k'),
|
||||
('v_proj', 'to_v'),
|
||||
('output_proj', 'to_out_0'),
|
||||
('q_norm', 'norm_q'),
|
||||
('k_norm', 'norm_k'),
|
||||
('mlp_layer1', 'ff_net_0_proj'),
|
||||
('mlp_layer2', 'ff_net_2'),
|
||||
('x_embedder_proj_1', 'patch_embed_proj'),
|
||||
('final_layer_adaln_modulation_1', 'norm_out_linear_1'),
|
||||
('final_layer_adaln_modulation_2', 'norm_out_linear_2'),
|
||||
('final_layer_linear', 'proj_out'),
|
||||
("t_embedder_1", "time_embed_t_embedder"),
|
||||
("t_embedding_norm", "time_embed_norm"),
|
||||
("blocks", "transformer_blocks"),
|
||||
("adaln_modulation_self_attn_1", "norm1_linear_1"),
|
||||
("adaln_modulation_self_attn_2", "norm1_linear_2"),
|
||||
("adaln_modulation_cross_attn_1", "norm2_linear_1"),
|
||||
("adaln_modulation_cross_attn_2", "norm2_linear_2"),
|
||||
("adaln_modulation_mlp_1", "norm3_linear_1"),
|
||||
("adaln_modulation_mlp_2", "norm3_linear_2"),
|
||||
("self_attn", "attn1"),
|
||||
("cross_attn", "attn2"),
|
||||
("q_proj", "to_q"),
|
||||
("k_proj", "to_k"),
|
||||
("v_proj", "to_v"),
|
||||
("output_proj", "to_out_0"),
|
||||
("q_norm", "norm_q"),
|
||||
("k_norm", "norm_k"),
|
||||
("mlp_layer1", "ff_net_0_proj"),
|
||||
("mlp_layer2", "ff_net_2"),
|
||||
("x_embedder_proj_1", "patch_embed_proj"),
|
||||
("final_layer_adaln_modulation_1", "norm_out_linear_1"),
|
||||
("final_layer_adaln_modulation_2", "norm_out_linear_2"),
|
||||
("final_layer_linear", "proj_out"),
|
||||
])
|
||||
|
||||
|
||||
def try_load(name, network_on_disk, lora_scale):
|
||||
"""Single dispatcher entry point. Anima only supports the LoRA family."""
|
||||
return try_load_lora(name, network_on_disk, lora_scale)
|
||||
|
||||
|
||||
def try_load_lora(name, network_on_disk, lora_scale):
|
||||
"""Try loading an Anima LoRA as native modules.
|
||||
|
||||
Returns a Network with native modules, or None if the file is not a
|
||||
recognized Anima format. Recognition is gated on Anima-specific path
|
||||
fragments so a non-Anima file routed here under a mounted Anima model
|
||||
is rejected rather than force-loaded.
|
||||
"""
|
||||
t0 = time.time()
|
||||
state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
|
||||
has_lora = any(any(m in k for m in LORA_MARKERS) for k in state_dict)
|
||||
if not has_lora:
|
||||
return None
|
||||
is_anima = any(
|
||||
k.startswith(KOHYA_PREFIX + 'blocks_')
|
||||
or k.startswith(BFL_TRANSFORMER_PREFIX + 'blocks.')
|
||||
or k.startswith(BFL_ADAPTER_PREFIX)
|
||||
or k.startswith(BFL_TE_PREFIX)
|
||||
for k in state_dict
|
||||
)
|
||||
if not is_anima:
|
||||
return None
|
||||
state_dict = apply_lora_alphas(state_dict)
|
||||
sd_model = getattr(shared.sd_model, 'pipe', shared.sd_model)
|
||||
lora_convert.assign_network_names_to_compvis_modules(sd_model)
|
||||
net = network.Network(name, network_on_disk)
|
||||
net.mtime = os.path.getmtime(network_on_disk.filename)
|
||||
matched = 0
|
||||
unmatched = 0
|
||||
unmatched_samples = []
|
||||
for base, weights in group_keys(state_dict):
|
||||
network_key = resolve_network_key(base)
|
||||
if network_key is None:
|
||||
unmatched += 1
|
||||
if len(unmatched_samples) < 5:
|
||||
unmatched_samples.append(base)
|
||||
continue
|
||||
sd_module = sd_model.network_layer_mapping.get(network_key)
|
||||
if sd_module is None:
|
||||
unmatched += 1
|
||||
if len(unmatched_samples) < 5:
|
||||
unmatched_samples.append(f'{base} (key={network_key})')
|
||||
continue
|
||||
if not shapes_match(sd_module, weights):
|
||||
log.warning(f'Network load: type=LoRA name="{name}" shape mismatch key={network_key}')
|
||||
continue
|
||||
nw = network.NetworkWeights(network_key=network_key, sd_key=network_key, w=weights, sd_module=sd_module)
|
||||
net.modules[network_key] = network_lora.NetworkModuleLora(net, nw)
|
||||
matched += 1
|
||||
if matched == 0:
|
||||
return None
|
||||
log.debug(f'Network load: type=LoRA name="{name}" method=native modules={matched} unmatched={unmatched} scale={lora_scale}')
|
||||
if unmatched > 0 and l.debug:
|
||||
log.debug(f'Network load: type=LoRA name="{name}" unmatched_samples={unmatched_samples}')
|
||||
l.timer.activate += time.time() - t0
|
||||
return net
|
||||
|
||||
|
||||
def apply_lora_alphas(state_dict):
|
||||
"""Bake .alpha scalars into lora_down / lora_up.
|
||||
|
||||
Native module loading consumes lora_down/lora_up directly; there is no
|
||||
alpha kwarg path. Mirrors flux2_lora.apply_lora_alphas, with kohya and
|
||||
BFL down/up names both supported (the hybrid format uses BFL keys with
|
||||
kohya-style .alpha).
|
||||
"""
|
||||
alpha_keys = [k for k in state_dict if k.endswith('.alpha')]
|
||||
if not alpha_keys:
|
||||
return state_dict
|
||||
for alpha_key in alpha_keys:
|
||||
base = alpha_key[:-len('.alpha')]
|
||||
down_key = next((c for c in (f'{base}.lora_down.weight', f'{base}.lora_A.weight') if c in state_dict), None)
|
||||
if down_key is None:
|
||||
continue
|
||||
rank = state_dict[down_key].shape[0]
|
||||
alpha = state_dict.pop(alpha_key).item()
|
||||
scale = alpha / rank
|
||||
scale_down = scale
|
||||
scale_up = 1.0
|
||||
while scale_down * 2 < scale_up:
|
||||
scale_down *= 2
|
||||
scale_up /= 2
|
||||
state_dict[down_key] = state_dict[down_key] * scale_down
|
||||
up_key = next((c for c in (f'{base}.lora_up.weight', f'{base}.lora_B.weight') if c in state_dict), None)
|
||||
if up_key is not None:
|
||||
state_dict[up_key] = state_dict[up_key] * scale_up
|
||||
remaining = [k for k in state_dict if k.endswith('.alpha')]
|
||||
if remaining:
|
||||
log.debug(f'Network load: type=LoRA stripped {len(remaining)} orphaned alpha keys')
|
||||
for k in remaining:
|
||||
del state_dict[k]
|
||||
return state_dict
|
||||
|
||||
|
||||
def normalize_weight_key(suffix):
|
||||
"""Rewrite PEFT-style suffixes to kohya/native form (lora_A to lora_down, lora_B to lora_up)."""
|
||||
return suffix.replace('lora_A.', 'lora_down.').replace('lora_B.', 'lora_up.')
|
||||
|
||||
|
||||
def group_keys(state_dict):
|
||||
"""Group state-dict keys by base path and yield (base, weights) pairs.
|
||||
|
||||
Bases keep their source key shape (kohya flat or BFL dotted) for
|
||||
component routing in resolve_network_key.
|
||||
"""
|
||||
groups = {}
|
||||
for key, weight in state_dict.items():
|
||||
base = None
|
||||
suffix = None
|
||||
if key.startswith(KOHYA_PREFIX):
|
||||
base, _, suffix = key.partition('.')
|
||||
else:
|
||||
for marker in LORA_MARKERS:
|
||||
pos = key.find(marker)
|
||||
if pos != -1:
|
||||
base = key[:pos]
|
||||
suffix = key[pos + 1:]
|
||||
break
|
||||
if base is None or suffix is None:
|
||||
continue
|
||||
groups.setdefault(base, {})[normalize_weight_key(suffix)] = weight
|
||||
for base, weights in groups.items():
|
||||
if 'lora_down.weight' in weights and 'lora_up.weight' in weights:
|
||||
yield base, weights
|
||||
|
||||
|
||||
def resolve_network_key(base):
|
||||
"""Map a grouped base path to the network_layer_mapping lookup key.
|
||||
|
||||
Adapter and TE prefixes must be checked before the bare BFL transformer
|
||||
prefix because both are extensions of diffusion_model.
|
||||
"""
|
||||
if base.startswith(KOHYA_PREFIX):
|
||||
flat = base[len(KOHYA_PREFIX):]
|
||||
return 'lora_transformer_' + cosmos_rename_flat(flat)
|
||||
if base.startswith(BFL_ADAPTER_PREFIX):
|
||||
rest = base[len(BFL_ADAPTER_PREFIX):]
|
||||
return 'lora_llm_adapter_' + rest.replace('.', '_')
|
||||
if base.startswith(BFL_TE_PREFIX):
|
||||
rest = base[len(BFL_TE_PREFIX):]
|
||||
return 'lora_te_' + rest.replace('.', '_')
|
||||
if base.startswith(BFL_TRANSFORMER_PREFIX):
|
||||
rest = base[len(BFL_TRANSFORMER_PREFIX):]
|
||||
flat = rest.replace('.', '_')
|
||||
return 'lora_transformer_' + cosmos_rename_flat(flat)
|
||||
return None
|
||||
|
||||
|
||||
def cosmos_rename_flat(flat):
|
||||
"""Apply the Cosmos 2.0 path rename to a flattened path."""
|
||||
"""Apply the Cosmos 2.0 path rename to a flattened (underscore) path."""
|
||||
for k, v in COSMOS_2_FLAT_RENAME.items():
|
||||
flat = flat.replace(k, v)
|
||||
return flat
|
||||
|
||||
|
||||
def shapes_match(sd_module, weights):
|
||||
"""Confirm LoRA rank dimensions line up with the live module weight."""
|
||||
if not hasattr(sd_module, 'weight'):
|
||||
return True
|
||||
if hasattr(sd_module, 'sdnq_dequantizer'):
|
||||
mod_shape = sd_module.sdnq_dequantizer.original_shape
|
||||
else:
|
||||
mod_shape = sd_module.weight.shape
|
||||
if len(mod_shape) < 2:
|
||||
return False
|
||||
return (
|
||||
weights['lora_down.weight'].shape[1] == mod_shape[1]
|
||||
and weights['lora_up.weight'].shape[0] == mod_shape[0]
|
||||
# === Target resolution (arch-specific) ===
|
||||
|
||||
|
||||
def resolve_targets(prefix_used, base):
|
||||
"""Return ``[(diffusers_path, ChunkSpec | None), ...]`` for a parsed group.
|
||||
|
||||
Anima has no fused QKV (every kohya / BFL key targets a single diffusers
|
||||
module), so ``ChunkSpec`` is always ``None`` and each call returns at most
|
||||
one target.
|
||||
|
||||
The returned ``diffusers_path`` is already underscore-flattened; the
|
||||
generic loader's ``path.replace('.', '_')`` is a no-op on it. Routing into
|
||||
the right namespace (transformer / llm_adapter / te) is handled separately
|
||||
by :func:`network_prefix_for`.
|
||||
"""
|
||||
if prefix_used == "lora_unet_":
|
||||
return [(cosmos_rename_flat(base), None)]
|
||||
if prefix_used == "diffusion_model.":
|
||||
return [(cosmos_rename_flat(base.replace(".", "_")), None)]
|
||||
if prefix_used == "diffusion_model.llm_adapter.":
|
||||
return [(base.replace(".", "_"), None)]
|
||||
if prefix_used in ("text_encoders.qwen3_06b.transformer.model.", "lora_te_"):
|
||||
return [(base.replace(".", "_"), None)]
|
||||
return []
|
||||
|
||||
|
||||
def network_prefix_for(prefix_used):
|
||||
"""Pick the ``network_layer_mapping`` namespace prefix.
|
||||
|
||||
``lora_convert.assign_network_names_to_compvis_modules`` stamps three
|
||||
namespaces on Anima modules; this picks the right one based on which arch
|
||||
prefix the key matched.
|
||||
"""
|
||||
if prefix_used == "diffusion_model.llm_adapter.":
|
||||
return "lora_llm_adapter_"
|
||||
if prefix_used in ("text_encoders.qwen3_06b.transformer.model.", "lora_te_"):
|
||||
return "lora_te_"
|
||||
return "lora_transformer_"
|
||||
|
||||
|
||||
# === Native loaders (thin wrappers over native_loader generics) ===
|
||||
|
||||
_BIND_KWARGS = dict(
|
||||
resolve_targets=resolve_targets,
|
||||
prefixes=ANIMA_PREFIXES,
|
||||
network_prefix=network_prefix_for,
|
||||
arch_name="anima",
|
||||
)
|
||||
|
||||
|
||||
def try_load_lora(name, network_on_disk, lora_scale):
|
||||
return native_loader.try_load_lora(name, network_on_disk, lora_scale, **_BIND_KWARGS)
|
||||
|
||||
|
||||
def try_load_lokr(name, network_on_disk, lora_scale):
|
||||
return native_loader.try_load_lokr(name, network_on_disk, lora_scale, **_BIND_KWARGS)
|
||||
|
||||
|
||||
def try_load_loha(name, network_on_disk, lora_scale):
|
||||
return native_loader.try_load_loha(name, network_on_disk, lora_scale, **_BIND_KWARGS)
|
||||
|
||||
|
||||
def try_load_oft(name, network_on_disk, lora_scale):
|
||||
return native_loader.try_load_oft(name, network_on_disk, lora_scale, **_BIND_KWARGS)
|
||||
|
||||
|
||||
def try_load_ia3(name, network_on_disk, lora_scale):
|
||||
return native_loader.try_load_ia3(name, network_on_disk, lora_scale, **_BIND_KWARGS)
|
||||
|
||||
|
||||
def try_load_glora(name, network_on_disk, lora_scale):
|
||||
return native_loader.try_load_glora(name, network_on_disk, lora_scale, **_BIND_KWARGS)
|
||||
|
||||
|
||||
def try_load_norm(name, network_on_disk, lora_scale):
|
||||
return native_loader.try_load_norm(name, network_on_disk, lora_scale, **_BIND_KWARGS)
|
||||
|
||||
|
||||
def try_load_full(name, network_on_disk, lora_scale):
|
||||
return native_loader.try_load_full(name, network_on_disk, lora_scale, **_BIND_KWARGS)
|
||||
|
||||
|
||||
def try_load(name, network_on_disk, lora_scale):
|
||||
"""Run every Anima family loader, merge any that match."""
|
||||
return native_loader.try_load_chain(
|
||||
name, network_on_disk, lora_scale,
|
||||
family_loaders=(
|
||||
try_load_lora, try_load_lokr, try_load_loha, try_load_oft,
|
||||
try_load_ia3, try_load_glora, try_load_norm, try_load_full,
|
||||
),
|
||||
)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user