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