vladi/ideogram4
1
1import hashlib2import os3from urllib.parse import urlparse4 5 6ADAPTER_NAME_PREFIX = "custom"7HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("hf")8 9 10def _parse_hf_lora_url(url: str):11 parsed = urlparse(url)12 if "huggingface.co" not in parsed.netloc:13 return None, None, None14 15 path_parts = [part for part in parsed.path.split("/") if part]16 if len(path_parts) < 2:17 return None, None, None18 19 repo_id = f"{path_parts[0]}/{path_parts[1]}"20 weight_parts = path_parts[2:]21 revision = None22 if len(weight_parts) >= 2 and weight_parts[0] in {"blob", "resolve"}:23 revision = weight_parts[1]24 weight_parts = weight_parts[2:]25 weight_name = "/".join(weight_parts) if weight_parts else None26 if not weight_name or not weight_name.endswith(".safetensors"):27 return repo_id, None, revision28 return repo_id, weight_name, revision29 30 31def _split_lora_spec(spec: str):32 if not spec:33 return None, None, None34 35 spec = spec.strip()36 if not spec:37 return None, None, None38 39 if spec.startswith("http://") or spec.startswith("https://"):40 return _parse_hf_lora_url(spec)41 if ":" in spec:42 repo_id, weight_name = spec.split(":", 1)43 return repo_id.strip(), weight_name.strip(), None44 return spec, None, None45 46 47def _split_adapter_line_scale(line: str):48 if "@" not in line:49 return line, 1.050 51 spec_candidate, scale_candidate = line.rsplit("@", 1)52 try:53 inline_scale = float(scale_candidate.strip())54 except ValueError:55 return line, 1.056 return spec_candidate.strip(), inline_scale57 58 59def parse_adapter_specs(spec_text: str, global_scale: float):60 if not spec_text or not spec_text.strip():61 return []62 63 requested_entries = []64 seen_keys = set()65 66 for line_number, raw_line in enumerate(spec_text.splitlines(), start=1):67 line = raw_line.strip()68 if not line:69 continue70 71 spec, inline_scale = _split_adapter_line_scale(line)72 repo_id, weight_name, revision = _split_lora_spec(spec)73 if not repo_id or not weight_name:74 raise ValueError(75 "Please provide LoRA entries as "76 "'user/repo:weights.safetensors' or direct .safetensors URLs. "77 f"Invalid line {line_number}: {raw_line!r}"78 )79 80 adapter_key = (repo_id, weight_name, revision)81 if adapter_key in seen_keys:82 raise ValueError(83 f"Duplicate LoRA entry for '{repo_id}:{weight_name}' on line {line_number}."84 )85 seen_keys.add(adapter_key)86 87 requested_entries.append(88 {89 "key": adapter_key,90 "repo_id": repo_id,91 "weight_name": weight_name,92 "revision": revision,93 "adapter_name": adapter_runtime_name(adapter_key),94 "inline_scale": inline_scale,95 "global_scale": global_scale,96 "scale": inline_scale * global_scale,97 }98 )99 100 return requested_entries101 102 103def adapter_runtime_name(adapter_key):104 key_parts = [part for part in adapter_key if part is not None]105 digest = hashlib.sha1(":".join(str(part) for part in key_parts).encode("utf-8")).hexdigest()[:12]106 return f"{ADAPTER_NAME_PREFIX}_{digest}"107 108 109def _iter_named_adapter_hosts(pipe):110 seen = set()111 for host_name, host in (112 (None, pipe),113 ("transformer", getattr(pipe, "transformer", None)),114 ("unconditional_transformer", getattr(pipe, "unconditional_transformer", None)),115 ):116 if host is None or id(host) in seen:117 continue118 seen.add(id(host))119 yield host_name, host120 121 122def _iter_adapter_hosts(pipe):123 for _, host in _iter_named_adapter_hosts(pipe):124 yield host125 126 127def _flatten_adapter_names(adapter_mapping):128 if isinstance(adapter_mapping, dict):129 names = set()130 for adapters in adapter_mapping.values():131 if isinstance(adapters, str):132 names.add(adapters)133 else:134 names.update(adapters)135 return names136 if isinstance(adapter_mapping, str):137 return {adapter_mapping}138 if adapter_mapping is None:139 return set()140 return set(adapter_mapping)141 142 143def _sorted_lora_entries(entries):144 return sorted(entries, key=lambda entry: entry["adapter_name"])145 146 147def _download_lora_weight(repo_id: str, weight_name: str, revision=None, token=HF_TOKEN):148 from huggingface_hub import hf_hub_download149 150 kwargs = {}151 if token:152 kwargs["token"] = token153 if revision:154 kwargs["revision"] = revision155 return hf_hub_download(repo_id, filename=weight_name, **kwargs)156 157 158def _load_adapter_state_dict(local_path: str):159 if local_path.endswith(".safetensors"):160 from safetensors.torch import load_file as safetensors_load_file161 162 return safetensors_load_file(local_path)163 164 import torch165 166 return torch.load(local_path, map_location="cpu")167 168 169# ---------------------------------------------------------------------------170# LoKR (LyCORIS Kronecker product) support171# ---------------------------------------------------------------------------172 173_ACTIVE_LOKR_MERGES = {}174 175 176def _is_lokr_state_dict(state_dict):177 return any(key.endswith(".lokr_w1") or key.endswith(".lokr_w2") for key in state_dict)178 179 180def _strip_diffusion_model_prefix(state_dict):181 if not any(key.startswith("diffusion_model.") for key in state_dict):182 return state_dict183 return {184 key.replace("diffusion_model.", "", 1): value185 for key, value in state_dict.items()186 }187 188 189def _is_lokr_key(key):190 suffix = key.rsplit(".", 1)[-1]191 return suffix == "alpha" or suffix.startswith("lokr_")192 193 194def _collect_lokr_groups(state_dict):195 groups = {}196 for key, value in state_dict.items():197 if not _is_lokr_key(key):198 continue199 prefix, suffix = key.rsplit(".", 1)200 groups.setdefault(prefix, {})[suffix] = value201 if not groups:202 raise ValueError("Checkpoint does not contain LoKr tensors.")203 return groups204 205 206def _materialize_lokr_factor(group, factor_name):207 direct = group.get(factor_name)208 if direct is not None:209 return direct.float()210 part_a = group.get(f"{factor_name}_a")211 part_b = group.get(f"{factor_name}_b")212 if part_a is None and part_b is None:213 return None214 if part_a is None or part_b is None:215 raise ValueError(f"Incomplete LoKr factor '{factor_name}'.")216 return part_a.float() @ part_b.float()217 218 219def _infer_lokr_factor_shape(group, factor_name):220 direct = group.get(factor_name)221 if direct is not None:222 return tuple(direct.shape)223 part_a = group.get(f"{factor_name}_a")224 part_b = group.get(f"{factor_name}_b")225 if part_a is None and part_b is None:226 return None227 if part_a is None or part_b is None:228 raise ValueError(f"Incomplete LoKr factor '{factor_name}'.")229 return (part_a.shape[0], part_b.shape[1])230 231 232def _lokr_group_scale_multiplier(group):233 alpha = group.get("alpha")234 if alpha is None:235 return 1.0236 rank_tensor = group.get("lokr_w1_b")237 if rank_tensor is None:238 rank_tensor = group.get("lokr_w2_b")239 if rank_tensor is None:240 return 1.0241 if alpha.numel() != 1:242 raise ValueError("Expected scalar alpha for LoKr module.")243 return float(alpha.item()) / float(rank_tensor.shape[0])244 245 246def _rebuild_lokr_delta(group):247 import torch248 249 if group.get("lokr_t1") is not None or group.get("lokr_t2") is not None:250 raise ValueError("Convolutional LoKr tensors are not supported.")251 w1 = _materialize_lokr_factor(group, "lokr_w1")252 w2 = _materialize_lokr_factor(group, "lokr_w2")253 if w1 is None or w2 is None:254 raise ValueError("LoKr checkpoint is missing required w1/w2 factors.")255 return torch.kron(w1.contiguous(), w2.contiguous())256 257 258def _get_lokr_linear_module(host, module_path):259 module = host.get_submodule(module_path)260 base_layer = getattr(module, "base_layer", None)261 if base_layer is not None:262 weight = getattr(base_layer, "weight", None)263 if weight is not None and weight.ndim == 2:264 return base_layer265 weight = getattr(module, "weight", None)266 if weight is None:267 raise ValueError(f"Target module '{module_path}' does not expose a weight.")268 if weight.ndim != 2:269 raise ValueError(f"Target module '{module_path}' is not a linear weight (ndim={weight.ndim}).")270 return module271 272 273def _merge_lokr_into_host(host, state_dict, scale):274 import torch275 276 state_dict = _strip_diffusion_model_prefix(state_dict)277 state_dict = _strip_known_peft_prefixes(state_dict)278 groups = _collect_lokr_groups(state_dict)279 280 applied = []281 try:282 for prefix in sorted(groups.keys()):283 group = groups[prefix]284 delta = _rebuild_lokr_delta(group)285 scale_mul = _lokr_group_scale_multiplier(group)286 effective_scale = scale * scale_mul287 288 module = _get_lokr_linear_module(host, prefix)289 weight = module.weight290 291 if delta.shape != weight.shape:292 raise ValueError(293 f"LoKr delta for '{prefix}' has shape {tuple(delta.shape)}, "294 f"expected {tuple(weight.shape)}."295 )296 297 chunk = delta.to(device=weight.device, dtype=weight.dtype)298 with torch.no_grad():299 weight.add_(chunk, alpha=effective_scale)300 applied.append((weight, chunk, effective_scale))301 del delta302 except Exception:303 for weight, chunk, eff_scale in reversed(applied):304 with torch.no_grad():305 weight.add_(chunk, alpha=-eff_scale)306 raise307 308 309def _unload_lokr_merges(pipe):310 if not _ACTIVE_LOKR_MERGES:311 return312 transformer = getattr(pipe, "transformer", None)313 if transformer is None:314 _ACTIVE_LOKR_MERGES.clear()315 return316 for key in list(_ACTIVE_LOKR_MERGES.keys()):317 merge_info = _ACTIVE_LOKR_MERGES.pop(key)318 try:319 sd = _load_adapter_state_dict(merge_info["local_path"])320 _merge_lokr_into_host(transformer, sd, -merge_info["scale"])321 except Exception as e:322 print(f"[lokr] Warning: failed to unmerge {key}: {e}")323 324 325# ---------------------------------------------------------------------------326 327 328def _ensure_pipeline_lora_prefix(state_dict):329 if any(key.startswith("transformer.") for key in state_dict.keys()):330 return state_dict331 332 if all(333 key.startswith("single_transformer_blocks.")334 or key.startswith("transformer_blocks.")335 for key in state_dict.keys()336 ):337 return {f"transformer.{key}": value for key, value in state_dict.items()}338 339 return state_dict340 341 342def _has_lora_tensors(state_dict):343 return any(344 ".lora_A." in key or ".lora_B." in key345 or ".lora_down." in key or ".lora_up." in key346 or ".lora_linear_layer." in key347 for key in state_dict.keys()348 )349 350 351def _strip_state_dict_prefix(state_dict, prefix):352 if not prefix:353 return state_dict354 return {355 key[len(prefix) :] if key.startswith(prefix) else key: value356 for key, value in state_dict.items()357 }358 359 360def _strip_known_peft_prefixes(state_dict):361 stripped = dict(state_dict)362 for prefix in ("base_model.model.", "model."):363 if any(key.startswith(prefix) for key in stripped.keys()):364 stripped = _strip_state_dict_prefix(stripped, prefix)365 return stripped366 367 368def _state_dict_for_model_host(state_dict, host_name):369 state_dict = _strip_known_peft_prefixes(state_dict)370 if not host_name:371 return state_dict372 373 own_prefix = f"{host_name}."374 own_state_dict = {375 key[len(own_prefix) :]: value376 for key, value in state_dict.items()377 if key.startswith(own_prefix)378 }379 if _has_lora_tensors(own_state_dict):380 return own_state_dict381 382 transformer_prefix = "transformer."383 transformer_state_dict = {384 key[len(transformer_prefix) :]: value385 for key, value in state_dict.items()386 if key.startswith(transformer_prefix)387 }388 if _has_lora_tensors(transformer_state_dict):389 return transformer_state_dict390 391 if not any(392 key.startswith(("transformer.", "unconditional_transformer."))393 for key in state_dict.keys()394 if ".lora_" in key or key.endswith(".alpha")395 ):396 return state_dict397 398 return own_state_dict399 400 401def _lora_module_name_from_key(key):402 for marker in (".lora_A.", ".lora_B."):403 if marker in key:404 return key.split(marker, 1)[0]405 return None406 407 408def _module_name_from_alpha_key(key):409 if key.endswith(".alpha"):410 return key[: -len(".alpha")]411 return None412 413 414def _scalar_to_float(value):415 if hasattr(value, "detach"):416 return float(value.detach().cpu().reshape(-1)[0].item())417 if hasattr(value, "item"):418 return float(value.item())419 return float(value)420 421 422def _build_lora_config(state_dict):423 from peft import LoraConfig424 425 rank_pattern = {}426 alpha_pattern = {}427 for key, value in state_dict.items():428 module_name = _lora_module_name_from_key(key)429 if module_name is None:430 continue431 if ".lora_A." in key and hasattr(value, "shape") and value.shape:432 rank_pattern[module_name] = int(value.shape[0])433 434 for key, value in state_dict.items():435 module_name = _module_name_from_alpha_key(key)436 if module_name is not None:437 alpha_pattern[module_name] = _scalar_to_float(value)438 439 if not rank_pattern:440 return LoraConfig()441 442 default_rank = max(rank_pattern.values())443 for module_name, rank in rank_pattern.items():444 alpha_pattern.setdefault(module_name, rank)445 return LoraConfig(446 r=default_rank,447 lora_alpha=default_rank,448 rank_pattern=rank_pattern,449 alpha_pattern=alpha_pattern,450 )451 452 453def _peft_load_state_dict(state_dict):454 return {455 key: value456 for key, value in state_dict.items()457 if not key.endswith(".alpha")458 }459 460 461def _load_lora_with_peft(host, state_dict, adapter_name):462 from peft import inject_adapter_in_model463 from peft.utils import set_peft_model_state_dict464 465 state_dict = _strip_known_peft_prefixes(state_dict)466 config = _build_lora_config(state_dict)467 inject_adapter_in_model(config, host, adapter_name=adapter_name, state_dict=state_dict)468 result = set_peft_model_state_dict(host, _peft_load_state_dict(state_dict), adapter_name=adapter_name)469 unexpected_keys = [470 key471 for key in getattr(result, "unexpected_keys", [])472 if ".lora_" in key473 ]474 if unexpected_keys:475 raise ValueError(f"Unexpected LoRA keys while loading adapter: {unexpected_keys[:5]}")476 missing_keys = [477 key478 for key in getattr(result, "missing_keys", [])479 if ".lora_" in key and f".{adapter_name}." in key480 ]481 if missing_keys:482 raise ValueError(f"Missing LoRA keys while loading adapter: {missing_keys[:5]}")483 return result484 485 486def _iter_host_modules(host):487 if not hasattr(host, "modules"):488 return []489 try:490 return list(host.modules())491 except Exception:492 return []493 494 495def _iter_host_and_modules(host):496 seen = set()497 for target in (host, *_iter_host_modules(host)):498 if id(target) in seen:499 continue500 seen.add(id(target))501 yield target502 503 504def _peft_adapter_names_on_host(host):505 adapter_names = set()506 peft_config = getattr(host, "peft_config", None)507 if isinstance(peft_config, dict):508 adapter_names.update(peft_config.keys())509 510 for module in _iter_host_modules(host):511 for attr_name in ("lora_A", "lora_B", "scaling"):512 adapters = getattr(module, attr_name, None)513 if hasattr(adapters, "keys"):514 try:515 adapter_names.update(adapters.keys())516 except Exception:517 pass518 return adapter_names519 520 521def _adapter_names_on_host(host):522 adapter_names = set()523 if hasattr(host, "get_list_adapters"):524 try:525 adapter_names.update(_flatten_adapter_names(host.get_list_adapters()))526 except Exception:527 pass528 adapter_names.update(_peft_adapter_names_on_host(host))529 return adapter_names530 531 532def _delete_peft_adapter_on_host(host, adapter_name):533 deleted = False534 for target in _iter_host_and_modules(host):535 if not hasattr(target, "delete_adapter"):536 continue537 try:538 target.delete_adapter(adapter_name)539 deleted = True540 except Exception:541 pass542 peft_config = getattr(host, "peft_config", None)543 if isinstance(peft_config, dict) and adapter_name in peft_config:544 peft_config.pop(adapter_name, None)545 deleted = True546 return deleted547 548 549def _set_peft_adapters_on_host(host, adapter_names, adapter_weights):550 changed = False551 if not adapter_names:552 for target in _iter_host_and_modules(host):553 if hasattr(target, "enable_adapters"):554 try:555 target.enable_adapters(False)556 changed = True557 except Exception:558 pass559 return changed560 561 for target in _iter_host_and_modules(host):562 if hasattr(target, "set_adapter"):563 try:564 target.set_adapter(adapter_names)565 changed = True566 except TypeError:567 if len(adapter_names) == 1:568 try:569 target.set_adapter(adapter_names[0])570 changed = True571 except Exception:572 pass573 except Exception:574 pass575 if hasattr(target, "enable_adapters"):576 try:577 target.enable_adapters(True)578 changed = True579 except Exception:580 pass581 if hasattr(target, "set_scale"):582 for adapter_name, adapter_weight in zip(adapter_names, adapter_weights):583 try:584 target.set_scale(adapter_name, adapter_weight)585 changed = True586 except Exception:587 pass588 return changed589 590 591def _is_model_adapter_host(host):592 return hasattr(host, "named_modules") and hasattr(host, "modules")593 594 595def _describe_adapter_hosts(pipe):596 descriptions = []597 for host_name, host in _iter_named_adapter_hosts(pipe):598 methods = [599 method_name600 for method_name in (601 "load_lora_weights",602 "load_lora_adapter",603 "set_adapters",604 "set_adapter",605 "delete_adapters",606 "delete_adapter",607 )608 if hasattr(host, method_name)609 ]610 label = host_name or "pipeline"611 method_text = ", ".join(methods) if methods else "no adapter methods"612 descriptions.append(f"{label}={host.__class__.__name__} ({method_text})")613 return "; ".join(descriptions)614 615 616def safe_unload_lora_adapters(pipe):617 _unload_lokr_merges(pipe)618 deleted = False619 for host in _iter_adapter_hosts(pipe):620 if hasattr(host, "delete_adapters"):621 try:622 adapter_names = sorted(_flatten_adapter_names(host.get_list_adapters()))623 except Exception:624 adapter_names = []625 for adapter_name in adapter_names:626 try:627 host.delete_adapters(adapter_name)628 deleted = True629 except Exception:630 pass631 for adapter_name in sorted(_peft_adapter_names_on_host(host)):632 if _delete_peft_adapter_on_host(host, adapter_name):633 deleted = True634 if deleted:635 return636 637 if hasattr(pipe, "unload_lora_weights"):638 try:639 pipe.unload_lora_weights()640 return641 except Exception:642 pass643 644 for host in _iter_adapter_hosts(pipe):645 if hasattr(host, "set_adapters"):646 try:647 host.set_adapters([])648 except Exception:649 pass650 if hasattr(host, "disable_adapters"):651 try:652 host.disable_adapters()653 except Exception:654 pass655 if hasattr(host, "disable_lora"):656 try:657 host.disable_lora()658 except Exception:659 pass660 _set_peft_adapters_on_host(host, [], [])661 662 663def _set_adapters_on_host(host, adapter_names, adapter_weights):664 if not hasattr(host, "set_adapters"):665 return False666 667 if not adapter_names:668 try:669 host.set_adapters([])670 return True671 except Exception:672 return False673 674 for kwargs in (675 {"adapter_weights": adapter_weights},676 {"weights": adapter_weights},677 ):678 try:679 host.set_adapters(adapter_names, **kwargs)680 return True681 except TypeError:682 continue683 except Exception:684 return False685 return False686 687 688def apply_lora_adapters(pipe, lora_entries):689 if not lora_entries:690 safe_unload_lora_adapters(pipe)691 return692 693 sorted_entries = _sorted_lora_entries(lora_entries)694 adapter_names = [entry["adapter_name"] for entry in sorted_entries]695 adapter_weights = [entry["scale"] for entry in sorted_entries]696 697 activated = False698 missing_on_hosts = []699 for host_name, host in _iter_named_adapter_hosts(pipe):700 host_adapter_names = _adapter_names_on_host(host)701 if not host_adapter_names:702 continue703 missing = set(adapter_names) - host_adapter_names704 if missing:705 missing_on_hosts.append(f"{host_name or 'pipeline'} missing {sorted(missing)}")706 continue707 if not (708 _set_adapters_on_host(host, adapter_names, adapter_weights)709 or _set_peft_adapters_on_host(host, adapter_names, adapter_weights)710 ):711 raise ValueError(f"Could not activate LoRA adapters on {host_name or 'pipeline'}.")712 activated = True713 714 if missing_on_hosts:715 raise ValueError("Partial LoRA adapter state: " + "; ".join(missing_on_hosts))716 717 if activated:718 return719 720 for host in _iter_adapter_hosts(pipe):721 if _set_adapters_on_host(host, adapter_names, adapter_weights):722 activated = True723 724 if activated:725 return726 727 if len(adapter_names) == 1 and hasattr(pipe, "set_lora_scale"):728 pipe.set_lora_scale(adapter_weights[0])729 return730 731 raise ValueError("This runtime does not support activating multiple LoRA adapters.")732 733 734def _load_lora_adapter_on_host(host, state_dict, adapter_name):735 try:736 host.load_lora_adapter(dict(state_dict), adapter_name=adapter_name, prefix=None)737 return738 except TypeError:739 host.load_lora_adapter(dict(state_dict), adapter_name=adapter_name)740 741 742def _pipeline_load_kwargs(entry, token):743 base_kwargs = {744 "weight_name": entry["weight_name"],745 "adapter_name": entry["adapter_name"],746 }747 if entry.get("revision"):748 base_kwargs["revision"] = entry["revision"]749 if token:750 base_kwargs["token"] = token751 752 variants = [base_kwargs]753 if "token" in base_kwargs:754 without_token = dict(base_kwargs)755 without_token.pop("token", None)756 variants.append(without_token)757 if "revision" in base_kwargs:758 without_revision = dict(base_kwargs)759 without_revision.pop("revision", None)760 variants.append(without_revision)761 without_token_revision = dict(without_revision)762 without_token_revision.pop("token", None)763 variants.append(without_token_revision)764 765 unique_variants = []766 seen = set()767 for kwargs in variants:768 key = tuple(sorted(kwargs.items()))769 if key not in seen:770 seen.add(key)771 unique_variants.append(kwargs)772 return unique_variants773 774 775def load_lora_adapter(pipe, entry, token=HF_TOKEN):776 native_error = None777 if hasattr(pipe, "load_lora_weights"):778 for load_kwargs in _pipeline_load_kwargs(entry, token):779 try:780 pipe.load_lora_weights(entry["repo_id"], **load_kwargs)781 return782 except TypeError as exc:783 native_error = exc784 except Exception as exc:785 native_error = exc786 break787 788 local_path = _download_lora_weight(789 entry["repo_id"],790 entry["weight_name"],791 revision=entry.get("revision"),792 token=token,793 )794 state_dict = _load_adapter_state_dict(local_path)795 796 # LoKR (LyCORIS Kronecker) — merge directly into weights797 if _is_lokr_state_dict(state_dict):798 transformer = getattr(pipe, "transformer", None)799 if transformer is None:800 raise ValueError("Pipeline has no transformer for LoKr merging.")801 _merge_lokr_into_host(transformer, state_dict, entry["scale"])802 _ACTIVE_LOKR_MERGES[entry["key"]] = {803 "local_path": local_path,804 "scale": entry["scale"],805 }806 return807 808 native_hosts = [809 (host_name, host)810 for host_name, host in _iter_named_adapter_hosts(pipe)811 if hasattr(host, "load_lora_adapter")812 ]813 if native_hosts:814 loaded_hosts = []815 try:816 for host_name, host in native_hosts:817 host_state_dict = _state_dict_for_model_host(state_dict, host_name)818 if not _has_lora_tensors(host_state_dict):819 continue820 _load_lora_adapter_on_host(host, host_state_dict, entry["adapter_name"])821 loaded_hosts.append(host)822 if loaded_hosts:823 return824 except Exception as exc:825 for host in loaded_hosts:826 _delete_peft_adapter_on_host(host, entry["adapter_name"])827 if hasattr(host, "delete_adapters"):828 try:829 host.delete_adapters(entry["adapter_name"])830 except Exception:831 pass832 native_error = exc833 834 peft_hosts = [835 (host_name, host)836 for host_name, host in _iter_named_adapter_hosts(pipe)837 if host_name is not None and _is_model_adapter_host(host)838 ]839 if not peft_hosts and _is_model_adapter_host(pipe):840 peft_hosts = [(None, pipe)]841 842 if peft_hosts:843 loaded_hosts = []844 try:845 for host_name, host in peft_hosts:846 host_state_dict = _state_dict_for_model_host(state_dict, host_name)847 if not _has_lora_tensors(host_state_dict):848 continue849 _load_lora_with_peft(host, host_state_dict, entry["adapter_name"])850 loaded_hosts.append(host)851 if loaded_hosts:852 return853 except Exception as exc:854 for host in loaded_hosts:855 _delete_peft_adapter_on_host(host, entry["adapter_name"])856 native_error = exc857 858 if hasattr(pipe, "load_lora_weights"):859 try:860 fallback_kwargs = {"adapter_name": entry["adapter_name"]}861 pipe.load_lora_weights(_ensure_pipeline_lora_prefix(state_dict), **fallback_kwargs)862 return863 except Exception as exc:864 if native_error is not None:865 raise ValueError(f"{native_error}; fallback failed with {exc}") from exc866 raise867 868 details = _describe_adapter_hosts(pipe)869 sample_keys = list(state_dict.keys())[:8]870 if native_error is not None:871 raise ValueError(f"Could not load LoRA adapter with native or PEFT fallback: {native_error}. Hosts: {details}. Sample keys: {sample_keys}") from native_error872 raise ValueError(f"This pipeline does not expose a usable LoRA loader. Hosts: {details}. Sample keys: {sample_keys}")873 874 875def ensure_loras_loaded(pipe, spec_text: str, global_scale: float, active_by_key: dict, token=HF_TOKEN):876 desired_entries = parse_adapter_specs(spec_text, global_scale)877 desired_by_key = {entry["key"]: entry for entry in desired_entries}878 879 if not desired_entries:880 if active_by_key:881 safe_unload_lora_adapters(pipe)882 active_by_key.clear()883 return []884 885 if set(active_by_key.keys()) != set(desired_by_key.keys()):886 try:887 safe_unload_lora_adapters(pipe)888 loaded_entries = []889 for entry in _sorted_lora_entries(desired_entries):890 load_lora_adapter(pipe, entry, token=token)891 loaded_entries.append(entry)892 # LoKR entries are already merged into weights — only apply PEFT/LoRA entries893 peft_entries = [e for e in loaded_entries if e["key"] not in _ACTIVE_LOKR_MERGES]894 if peft_entries:895 apply_lora_adapters(pipe, peft_entries)896 except Exception:897 safe_unload_lora_adapters(pipe)898 active_by_key.clear()899 raise900 else:901 peft_entries = [e for e in desired_entries if e["key"] not in _ACTIVE_LOKR_MERGES]902 if peft_entries:903 apply_lora_adapters(pipe, peft_entries)904 905 active_by_key.clear()906 active_by_key.update(desired_by_key)907 return desired_entries