lora_utils.py

import hashlib
import os
from urllib.parse import urlparse


ADAPTER_NAME_PREFIX = "custom"
HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("hf")


def _parse_hf_lora_url(url: str):
    parsed = urlparse(url)
    if "huggingface.co" not in parsed.netloc:
        return None, None, None

    path_parts = [part for part in parsed.path.split("/") if part]
    if len(path_parts) < 2:
        return None, None, None

    repo_id = f"{path_parts[0]}/{path_parts[1]}"
    weight_parts = path_parts[2:]
    revision = None
    if len(weight_parts) >= 2 and weight_parts[0] in {"blob", "resolve"}:
        revision = weight_parts[1]
        weight_parts = weight_parts[2:]
    weight_name = "/".join(weight_parts) if weight_parts else None
    if not weight_name or not weight_name.endswith(".safetensors"):
        return repo_id, None, revision
    return repo_id, weight_name, revision


def _split_lora_spec(spec: str):
    if not spec:
        return None, None, None

    spec = spec.strip()
    if not spec:
        return None, None, None

    if spec.startswith("http://") or spec.startswith("https://"):
        return _parse_hf_lora_url(spec)
    if ":" in spec:
        repo_id, weight_name = spec.split(":", 1)
        return repo_id.strip(), weight_name.strip(), None
    return spec, None, None


def _split_adapter_line_scale(line: str):
    if "@" not in line:
        return line, 1.0

    spec_candidate, scale_candidate = line.rsplit("@", 1)
    try:
        inline_scale = float(scale_candidate.strip())
    except ValueError:
        return line, 1.0
    return spec_candidate.strip(), inline_scale


def parse_adapter_specs(spec_text: str, global_scale: float):
    if not spec_text or not spec_text.strip():
        return []

    requested_entries = []
    seen_keys = set()

    for line_number, raw_line in enumerate(spec_text.splitlines(), start=1):
        line = raw_line.strip()
        if not line:
            continue

        spec, inline_scale = _split_adapter_line_scale(line)
        repo_id, weight_name, revision = _split_lora_spec(spec)
        if not repo_id or not weight_name:
            raise ValueError(
                "Please provide LoRA entries as "
                "'user/repo:weights.safetensors' or direct .safetensors URLs. "
                f"Invalid line {line_number}: {raw_line!r}"
            )

        adapter_key = (repo_id, weight_name, revision)
        if adapter_key in seen_keys:
            raise ValueError(
                f"Duplicate LoRA entry for '{repo_id}:{weight_name}' on line {line_number}."
            )
        seen_keys.add(adapter_key)

        requested_entries.append(
            {
                "key": adapter_key,
                "repo_id": repo_id,
                "weight_name": weight_name,
                "revision": revision,
                "adapter_name": adapter_runtime_name(adapter_key),
                "inline_scale": inline_scale,
                "global_scale": global_scale,
                "scale": inline_scale * global_scale,
            }
        )

    return requested_entries


def adapter_runtime_name(adapter_key):
    key_parts = [part for part in adapter_key if part is not None]
    digest = hashlib.sha1(":".join(str(part) for part in key_parts).encode("utf-8")).hexdigest()[:12]
    return f"{ADAPTER_NAME_PREFIX}_{digest}"


def _iter_named_adapter_hosts(pipe):
    seen = set()
    for host_name, host in (
        (None, pipe),
        ("transformer", getattr(pipe, "transformer", None)),
        ("unconditional_transformer", getattr(pipe, "unconditional_transformer", None)),
    ):
        if host is None or id(host) in seen:
            continue
        seen.add(id(host))
        yield host_name, host


def _iter_adapter_hosts(pipe):
    for _, host in _iter_named_adapter_hosts(pipe):
        yield host


def _flatten_adapter_names(adapter_mapping):
    if isinstance(adapter_mapping, dict):
        names = set()
        for adapters in adapter_mapping.values():
            if isinstance(adapters, str):
                names.add(adapters)
            else:
                names.update(adapters)
        return names
    if isinstance(adapter_mapping, str):
        return {adapter_mapping}
    if adapter_mapping is None:
        return set()
    return set(adapter_mapping)


def _sorted_lora_entries(entries):
    return sorted(entries, key=lambda entry: entry["adapter_name"])


def _download_lora_weight(repo_id: str, weight_name: str, revision=None, token=HF_TOKEN):
    from huggingface_hub import hf_hub_download

    kwargs = {}
    if token:
        kwargs["token"] = token
    if revision:
        kwargs["revision"] = revision
    return hf_hub_download(repo_id, filename=weight_name, **kwargs)


def _load_adapter_state_dict(local_path: str):
    if local_path.endswith(".safetensors"):
        from safetensors.torch import load_file as safetensors_load_file

        return safetensors_load_file(local_path)

    import torch

    return torch.load(local_path, map_location="cpu")


# ---------------------------------------------------------------------------
# LoKR (LyCORIS Kronecker product) support
# ---------------------------------------------------------------------------

_ACTIVE_LOKR_MERGES = {}


def _is_lokr_state_dict(state_dict):
    return any(key.endswith(".lokr_w1") or key.endswith(".lokr_w2") for key in state_dict)


def _strip_diffusion_model_prefix(state_dict):
    if not any(key.startswith("diffusion_model.") for key in state_dict):
        return state_dict
    return {
        key.replace("diffusion_model.", "", 1): value
        for key, value in state_dict.items()
    }


def _is_lokr_key(key):
    suffix = key.rsplit(".", 1)[-1]
    return suffix == "alpha" or suffix.startswith("lokr_")


def _collect_lokr_groups(state_dict):
    groups = {}
    for key, value in state_dict.items():
        if not _is_lokr_key(key):
            continue
        prefix, suffix = key.rsplit(".", 1)
        groups.setdefault(prefix, {})[suffix] = value
    if not groups:
        raise ValueError("Checkpoint does not contain LoKr tensors.")
    return groups


def _materialize_lokr_factor(group, factor_name):
    direct = group.get(factor_name)
    if direct is not None:
        return direct.float()
    part_a = group.get(f"{factor_name}_a")
    part_b = group.get(f"{factor_name}_b")
    if part_a is None and part_b is None:
        return None
    if part_a is None or part_b is None:
        raise ValueError(f"Incomplete LoKr factor '{factor_name}'.")
    return part_a.float() @ part_b.float()


def _infer_lokr_factor_shape(group, factor_name):
    direct = group.get(factor_name)
    if direct is not None:
        return tuple(direct.shape)
    part_a = group.get(f"{factor_name}_a")
    part_b = group.get(f"{factor_name}_b")
    if part_a is None and part_b is None:
        return None
    if part_a is None or part_b is None:
        raise ValueError(f"Incomplete LoKr factor '{factor_name}'.")
    return (part_a.shape[0], part_b.shape[1])


def _lokr_group_scale_multiplier(group):
    alpha = group.get("alpha")
    if alpha is None:
        return 1.0
    rank_tensor = group.get("lokr_w1_b")
    if rank_tensor is None:
        rank_tensor = group.get("lokr_w2_b")
    if rank_tensor is None:
        return 1.0
    if alpha.numel() != 1:
        raise ValueError("Expected scalar alpha for LoKr module.")
    return float(alpha.item()) / float(rank_tensor.shape[0])


def _rebuild_lokr_delta(group):
    import torch

    if group.get("lokr_t1") is not None or group.get("lokr_t2") is not None:
        raise ValueError("Convolutional LoKr tensors are not supported.")
    w1 = _materialize_lokr_factor(group, "lokr_w1")
    w2 = _materialize_lokr_factor(group, "lokr_w2")
    if w1 is None or w2 is None:
        raise ValueError("LoKr checkpoint is missing required w1/w2 factors.")
    return torch.kron(w1.contiguous(), w2.contiguous())


def _get_lokr_linear_module(host, module_path):
    module = host.get_submodule(module_path)
    base_layer = getattr(module, "base_layer", None)
    if base_layer is not None:
        weight = getattr(base_layer, "weight", None)
        if weight is not None and weight.ndim == 2:
            return base_layer
    weight = getattr(module, "weight", None)
    if weight is None:
        raise ValueError(f"Target module '{module_path}' does not expose a weight.")
    if weight.ndim != 2:
        raise ValueError(f"Target module '{module_path}' is not a linear weight (ndim={weight.ndim}).")
    return module


def _merge_lokr_into_host(host, state_dict, scale):
    import torch

    state_dict = _strip_diffusion_model_prefix(state_dict)
    state_dict = _strip_known_peft_prefixes(state_dict)
    groups = _collect_lokr_groups(state_dict)

    applied = []
    try:
        for prefix in sorted(groups.keys()):
            group = groups[prefix]
            delta = _rebuild_lokr_delta(group)
            scale_mul = _lokr_group_scale_multiplier(group)
            effective_scale = scale * scale_mul

            module = _get_lokr_linear_module(host, prefix)
            weight = module.weight

            if delta.shape != weight.shape:
                raise ValueError(
                    f"LoKr delta for '{prefix}' has shape {tuple(delta.shape)}, "
                    f"expected {tuple(weight.shape)}."
                )

            chunk = delta.to(device=weight.device, dtype=weight.dtype)
            with torch.no_grad():
                weight.add_(chunk, alpha=effective_scale)
            applied.append((weight, chunk, effective_scale))
            del delta
    except Exception:
        for weight, chunk, eff_scale in reversed(applied):
            with torch.no_grad():
                weight.add_(chunk, alpha=-eff_scale)
        raise


def _unload_lokr_merges(pipe):
    if not _ACTIVE_LOKR_MERGES:
        return
    transformer = getattr(pipe, "transformer", None)
    if transformer is None:
        _ACTIVE_LOKR_MERGES.clear()
        return
    for key in list(_ACTIVE_LOKR_MERGES.keys()):
        merge_info = _ACTIVE_LOKR_MERGES.pop(key)
        try:
            sd = _load_adapter_state_dict(merge_info["local_path"])
            _merge_lokr_into_host(transformer, sd, -merge_info["scale"])
        except Exception as e:
            print(f"[lokr] Warning: failed to unmerge {key}: {e}")


# ---------------------------------------------------------------------------


def _ensure_pipeline_lora_prefix(state_dict):
    if any(key.startswith("transformer.") for key in state_dict.keys()):
        return state_dict

    if all(
        key.startswith("single_transformer_blocks.")
        or key.startswith("transformer_blocks.")
        for key in state_dict.keys()
    ):
        return {f"transformer.{key}": value for key, value in state_dict.items()}

    return state_dict


def _has_lora_tensors(state_dict):
    return any(
        ".lora_A." in key or ".lora_B." in key
        or ".lora_down." in key or ".lora_up." in key
        or ".lora_linear_layer." in key
        for key in state_dict.keys()
    )


def _strip_state_dict_prefix(state_dict, prefix):
    if not prefix:
        return state_dict
    return {
        key[len(prefix) :] if key.startswith(prefix) else key: value
        for key, value in state_dict.items()
    }


def _strip_known_peft_prefixes(state_dict):
    stripped = dict(state_dict)
    for prefix in ("base_model.model.", "model."):
        if any(key.startswith(prefix) for key in stripped.keys()):
            stripped = _strip_state_dict_prefix(stripped, prefix)
    return stripped


def _state_dict_for_model_host(state_dict, host_name):
    state_dict = _strip_known_peft_prefixes(state_dict)
    if not host_name:
        return state_dict

    own_prefix = f"{host_name}."
    own_state_dict = {
        key[len(own_prefix) :]: value
        for key, value in state_dict.items()
        if key.startswith(own_prefix)
    }
    if _has_lora_tensors(own_state_dict):
        return own_state_dict

    transformer_prefix = "transformer."
    transformer_state_dict = {
        key[len(transformer_prefix) :]: value
        for key, value in state_dict.items()
        if key.startswith(transformer_prefix)
    }
    if _has_lora_tensors(transformer_state_dict):
        return transformer_state_dict

    if not any(
        key.startswith(("transformer.", "unconditional_transformer."))
        for key in state_dict.keys()
        if ".lora_" in key or key.endswith(".alpha")
    ):
        return state_dict

    return own_state_dict


def _lora_module_name_from_key(key):
    for marker in (".lora_A.", ".lora_B."):
        if marker in key:
            return key.split(marker, 1)[0]
    return None


def _module_name_from_alpha_key(key):
    if key.endswith(".alpha"):
        return key[: -len(".alpha")]
    return None


def _scalar_to_float(value):
    if hasattr(value, "detach"):
        return float(value.detach().cpu().reshape(-1)[0].item())
    if hasattr(value, "item"):
        return float(value.item())
    return float(value)


def _build_lora_config(state_dict):
    from peft import LoraConfig

    rank_pattern = {}
    alpha_pattern = {}
    for key, value in state_dict.items():
        module_name = _lora_module_name_from_key(key)
        if module_name is None:
            continue
        if ".lora_A." in key and hasattr(value, "shape") and value.shape:
            rank_pattern[module_name] = int(value.shape[0])

    for key, value in state_dict.items():
        module_name = _module_name_from_alpha_key(key)
        if module_name is not None:
            alpha_pattern[module_name] = _scalar_to_float(value)

    if not rank_pattern:
        return LoraConfig()

    default_rank = max(rank_pattern.values())
    for module_name, rank in rank_pattern.items():
        alpha_pattern.setdefault(module_name, rank)
    return LoraConfig(
        r=default_rank,
        lora_alpha=default_rank,
        rank_pattern=rank_pattern,
        alpha_pattern=alpha_pattern,
    )


def _peft_load_state_dict(state_dict):
    return {
        key: value
        for key, value in state_dict.items()
        if not key.endswith(".alpha")
    }


def _load_lora_with_peft(host, state_dict, adapter_name):
    from peft import inject_adapter_in_model
    from peft.utils import set_peft_model_state_dict

    state_dict = _strip_known_peft_prefixes(state_dict)
    config = _build_lora_config(state_dict)
    inject_adapter_in_model(config, host, adapter_name=adapter_name, state_dict=state_dict)
    result = set_peft_model_state_dict(host, _peft_load_state_dict(state_dict), adapter_name=adapter_name)
    unexpected_keys = [
        key
        for key in getattr(result, "unexpected_keys", [])
        if ".lora_" in key
    ]
    if unexpected_keys:
        raise ValueError(f"Unexpected LoRA keys while loading adapter: {unexpected_keys[:5]}")
    missing_keys = [
        key
        for key in getattr(result, "missing_keys", [])
        if ".lora_" in key and f".{adapter_name}." in key
    ]
    if missing_keys:
        raise ValueError(f"Missing LoRA keys while loading adapter: {missing_keys[:5]}")
    return result


def _iter_host_modules(host):
    if not hasattr(host, "modules"):
        return []
    try:
        return list(host.modules())
    except Exception:
        return []


def _iter_host_and_modules(host):
    seen = set()
    for target in (host, *_iter_host_modules(host)):
        if id(target) in seen:
            continue
        seen.add(id(target))
        yield target


def _peft_adapter_names_on_host(host):
    adapter_names = set()
    peft_config = getattr(host, "peft_config", None)
    if isinstance(peft_config, dict):
        adapter_names.update(peft_config.keys())

    for module in _iter_host_modules(host):
        for attr_name in ("lora_A", "lora_B", "scaling"):
            adapters = getattr(module, attr_name, None)
            if hasattr(adapters, "keys"):
                try:
                    adapter_names.update(adapters.keys())
                except Exception:
                    pass
    return adapter_names


def _adapter_names_on_host(host):
    adapter_names = set()
    if hasattr(host, "get_list_adapters"):
        try:
            adapter_names.update(_flatten_adapter_names(host.get_list_adapters()))
        except Exception:
            pass
    adapter_names.update(_peft_adapter_names_on_host(host))
    return adapter_names


def _delete_peft_adapter_on_host(host, adapter_name):
    deleted = False
    for target in _iter_host_and_modules(host):
        if not hasattr(target, "delete_adapter"):
            continue
        try:
            target.delete_adapter(adapter_name)
            deleted = True
        except Exception:
            pass
    peft_config = getattr(host, "peft_config", None)
    if isinstance(peft_config, dict) and adapter_name in peft_config:
        peft_config.pop(adapter_name, None)
        deleted = True
    return deleted


def _set_peft_adapters_on_host(host, adapter_names, adapter_weights):
    changed = False
    if not adapter_names:
        for target in _iter_host_and_modules(host):
            if hasattr(target, "enable_adapters"):
                try:
                    target.enable_adapters(False)
                    changed = True
                except Exception:
                    pass
        return changed

    for target in _iter_host_and_modules(host):
        if hasattr(target, "set_adapter"):
            try:
                target.set_adapter(adapter_names)
                changed = True
            except TypeError:
                if len(adapter_names) == 1:
                    try:
                        target.set_adapter(adapter_names[0])
                        changed = True
                    except Exception:
                        pass
            except Exception:
                pass
        if hasattr(target, "enable_adapters"):
            try:
                target.enable_adapters(True)
                changed = True
            except Exception:
                pass
        if hasattr(target, "set_scale"):
            for adapter_name, adapter_weight in zip(adapter_names, adapter_weights):
                try:
                    target.set_scale(adapter_name, adapter_weight)
                    changed = True
                except Exception:
                    pass
    return changed


def _is_model_adapter_host(host):
    return hasattr(host, "named_modules") and hasattr(host, "modules")


def _describe_adapter_hosts(pipe):
    descriptions = []
    for host_name, host in _iter_named_adapter_hosts(pipe):
        methods = [
            method_name
            for method_name in (
                "load_lora_weights",
                "load_lora_adapter",
                "set_adapters",
                "set_adapter",
                "delete_adapters",
                "delete_adapter",
            )
            if hasattr(host, method_name)
        ]
        label = host_name or "pipeline"
        method_text = ", ".join(methods) if methods else "no adapter methods"
        descriptions.append(f"{label}={host.__class__.__name__} ({method_text})")
    return "; ".join(descriptions)


def safe_unload_lora_adapters(pipe):
    _unload_lokr_merges(pipe)
    deleted = False
    for host in _iter_adapter_hosts(pipe):
        if hasattr(host, "delete_adapters"):
            try:
                adapter_names = sorted(_flatten_adapter_names(host.get_list_adapters()))
            except Exception:
                adapter_names = []
            for adapter_name in adapter_names:
                try:
                    host.delete_adapters(adapter_name)
                    deleted = True
                except Exception:
                    pass
        for adapter_name in sorted(_peft_adapter_names_on_host(host)):
            if _delete_peft_adapter_on_host(host, adapter_name):
                deleted = True
    if deleted:
        return

    if hasattr(pipe, "unload_lora_weights"):
        try:
            pipe.unload_lora_weights()
            return
        except Exception:
            pass

    for host in _iter_adapter_hosts(pipe):
        if hasattr(host, "set_adapters"):
            try:
                host.set_adapters([])
            except Exception:
                pass
        if hasattr(host, "disable_adapters"):
            try:
                host.disable_adapters()
            except Exception:
                pass
        if hasattr(host, "disable_lora"):
            try:
                host.disable_lora()
            except Exception:
                pass
        _set_peft_adapters_on_host(host, [], [])


def _set_adapters_on_host(host, adapter_names, adapter_weights):
    if not hasattr(host, "set_adapters"):
        return False

    if not adapter_names:
        try:
            host.set_adapters([])
            return True
        except Exception:
            return False

    for kwargs in (
        {"adapter_weights": adapter_weights},
        {"weights": adapter_weights},
    ):
        try:
            host.set_adapters(adapter_names, **kwargs)
            return True
        except TypeError:
            continue
        except Exception:
            return False
    return False


def apply_lora_adapters(pipe, lora_entries):
    if not lora_entries:
        safe_unload_lora_adapters(pipe)
        return

    sorted_entries = _sorted_lora_entries(lora_entries)
    adapter_names = [entry["adapter_name"] for entry in sorted_entries]
    adapter_weights = [entry["scale"] for entry in sorted_entries]

    activated = False
    missing_on_hosts = []
    for host_name, host in _iter_named_adapter_hosts(pipe):
        host_adapter_names = _adapter_names_on_host(host)
        if not host_adapter_names:
            continue
        missing = set(adapter_names) - host_adapter_names
        if missing:
            missing_on_hosts.append(f"{host_name or 'pipeline'} missing {sorted(missing)}")
            continue
        if not (
            _set_adapters_on_host(host, adapter_names, adapter_weights)
            or _set_peft_adapters_on_host(host, adapter_names, adapter_weights)
        ):
            raise ValueError(f"Could not activate LoRA adapters on {host_name or 'pipeline'}.")
        activated = True

    if missing_on_hosts:
        raise ValueError("Partial LoRA adapter state: " + "; ".join(missing_on_hosts))

    if activated:
        return

    for host in _iter_adapter_hosts(pipe):
        if _set_adapters_on_host(host, adapter_names, adapter_weights):
            activated = True

    if activated:
        return

    if len(adapter_names) == 1 and hasattr(pipe, "set_lora_scale"):
        pipe.set_lora_scale(adapter_weights[0])
        return

    raise ValueError("This runtime does not support activating multiple LoRA adapters.")


def _load_lora_adapter_on_host(host, state_dict, adapter_name):
    try:
        host.load_lora_adapter(dict(state_dict), adapter_name=adapter_name, prefix=None)
        return
    except TypeError:
        host.load_lora_adapter(dict(state_dict), adapter_name=adapter_name)


def _pipeline_load_kwargs(entry, token):
    base_kwargs = {
        "weight_name": entry["weight_name"],
        "adapter_name": entry["adapter_name"],
    }
    if entry.get("revision"):
        base_kwargs["revision"] = entry["revision"]
    if token:
        base_kwargs["token"] = token

    variants = [base_kwargs]
    if "token" in base_kwargs:
        without_token = dict(base_kwargs)
        without_token.pop("token", None)
        variants.append(without_token)
    if "revision" in base_kwargs:
        without_revision = dict(base_kwargs)
        without_revision.pop("revision", None)
        variants.append(without_revision)
        without_token_revision = dict(without_revision)
        without_token_revision.pop("token", None)
        variants.append(without_token_revision)

    unique_variants = []
    seen = set()
    for kwargs in variants:
        key = tuple(sorted(kwargs.items()))
        if key not in seen:
            seen.add(key)
            unique_variants.append(kwargs)
    return unique_variants


def load_lora_adapter(pipe, entry, token=HF_TOKEN):
    native_error = None
    if hasattr(pipe, "load_lora_weights"):
        for load_kwargs in _pipeline_load_kwargs(entry, token):
            try:
                pipe.load_lora_weights(entry["repo_id"], **load_kwargs)
                return
            except TypeError as exc:
                native_error = exc
            except Exception as exc:
                native_error = exc
                break

    local_path = _download_lora_weight(
        entry["repo_id"],
        entry["weight_name"],
        revision=entry.get("revision"),
        token=token,
    )
    state_dict = _load_adapter_state_dict(local_path)

    # LoKR (LyCORIS Kronecker) — merge directly into weights
    if _is_lokr_state_dict(state_dict):
        transformer = getattr(pipe, "transformer", None)
        if transformer is None:
            raise ValueError("Pipeline has no transformer for LoKr merging.")
        _merge_lokr_into_host(transformer, state_dict, entry["scale"])
        _ACTIVE_LOKR_MERGES[entry["key"]] = {
            "local_path": local_path,
            "scale": entry["scale"],
        }
        return

    native_hosts = [
        (host_name, host)
        for host_name, host in _iter_named_adapter_hosts(pipe)
        if hasattr(host, "load_lora_adapter")
    ]
    if native_hosts:
        loaded_hosts = []
        try:
            for host_name, host in native_hosts:
                host_state_dict = _state_dict_for_model_host(state_dict, host_name)
                if not _has_lora_tensors(host_state_dict):
                    continue
                _load_lora_adapter_on_host(host, host_state_dict, entry["adapter_name"])
                loaded_hosts.append(host)
            if loaded_hosts:
                return
        except Exception as exc:
            for host in loaded_hosts:
                _delete_peft_adapter_on_host(host, entry["adapter_name"])
                if hasattr(host, "delete_adapters"):
                    try:
                        host.delete_adapters(entry["adapter_name"])
                    except Exception:
                        pass
            native_error = exc

    peft_hosts = [
        (host_name, host)
        for host_name, host in _iter_named_adapter_hosts(pipe)
        if host_name is not None and _is_model_adapter_host(host)
    ]
    if not peft_hosts and _is_model_adapter_host(pipe):
        peft_hosts = [(None, pipe)]

    if peft_hosts:
        loaded_hosts = []
        try:
            for host_name, host in peft_hosts:
                host_state_dict = _state_dict_for_model_host(state_dict, host_name)
                if not _has_lora_tensors(host_state_dict):
                    continue
                _load_lora_with_peft(host, host_state_dict, entry["adapter_name"])
                loaded_hosts.append(host)
            if loaded_hosts:
                return
        except Exception as exc:
            for host in loaded_hosts:
                _delete_peft_adapter_on_host(host, entry["adapter_name"])
            native_error = exc

    if hasattr(pipe, "load_lora_weights"):
        try:
            fallback_kwargs = {"adapter_name": entry["adapter_name"]}
            pipe.load_lora_weights(_ensure_pipeline_lora_prefix(state_dict), **fallback_kwargs)
            return
        except Exception as exc:
            if native_error is not None:
                raise ValueError(f"{native_error}; fallback failed with {exc}") from exc
            raise

    details = _describe_adapter_hosts(pipe)
    sample_keys = list(state_dict.keys())[:8]
    if native_error is not None:
        raise ValueError(f"Could not load LoRA adapter with native or PEFT fallback: {native_error}. Hosts: {details}. Sample keys: {sample_keys}") from native_error
    raise ValueError(f"This pipeline does not expose a usable LoRA loader. Hosts: {details}. Sample keys: {sample_keys}")


def ensure_loras_loaded(pipe, spec_text: str, global_scale: float, active_by_key: dict, token=HF_TOKEN):
    desired_entries = parse_adapter_specs(spec_text, global_scale)
    desired_by_key = {entry["key"]: entry for entry in desired_entries}

    if not desired_entries:
        if active_by_key:
            safe_unload_lora_adapters(pipe)
            active_by_key.clear()
        return []

    if set(active_by_key.keys()) != set(desired_by_key.keys()):
        try:
            safe_unload_lora_adapters(pipe)
            loaded_entries = []
            for entry in _sorted_lora_entries(desired_entries):
                load_lora_adapter(pipe, entry, token=token)
                loaded_entries.append(entry)
            # LoKR entries are already merged into weights — only apply PEFT/LoRA entries
            peft_entries = [e for e in loaded_entries if e["key"] not in _ACTIVE_LOKR_MERGES]
            if peft_entries:
                apply_lora_adapters(pipe, peft_entries)
        except Exception:
            safe_unload_lora_adapters(pipe)
            active_by_key.clear()
            raise
    else:
        peft_entries = [e for e in desired_entries if e["key"] not in _ACTIVE_LOKR_MERGES]
        if peft_entries:
            apply_lora_adapters(pipe, peft_entries)

    active_by_key.clear()
    active_by_key.update(desired_by_key)
    return desired_entries