"""Ordinary user-LoRA preparation and request-scoped adapter lifecycle. This module owns backend LoRA policy only. Gradio widgets and event wiring remain in ``app.py``. The internal Full/SFT Stage-2 adapter and IC-Colorizer adapter remain separate domains; this module deliberately does not generalize them into one manager. """ from __future__ import annotations import gc import time import urllib.parse from pathlib import Path from typing import Callable import torch from huggingface_hub import HfApi, hf_hub_download, parse_hf_uri from huggingface_hub.errors import ( GatedRepoError, HfHubHTTPError, LocalEntryNotFoundError, RemoteEntryNotFoundError, RepositoryNotFoundError, RevisionNotFoundError, ) from safetensors import safe_open from safetensors.torch import load_file as load_safetensors_file, save_file as save_safetensors_file from ltx import civitai as civitai_backend from ltx.runtime_utils import resolved_revision_from_hub_path class LoraSelectionError(ValueError): """User-facing ordinary LoRA selection/preparation error.""" def sanitize_defs(items, source: str = "configured") -> list[dict]: """Normalize portable LoRA definitions without persisting transport secrets.""" clean: list[dict] = [] seen: set[str] = set() for raw in items or []: if not isinstance(raw, dict): continue label = str(raw.get("label") or "").strip() weight_name = str(raw.get("weight_name") or "").strip() item_source = str(raw.get("source") or source).strip() or source if not label or not weight_name or label in seen: continue if item_source == "civitai": try: model_id = int(raw.get("civitai_model_id")) version_id = int(raw.get("civitai_version_id")) file_id = int(raw.get("civitai_file_id")) except Exception: continue clean.append( { "label": label, "source": "civitai", "repo_id": "", "weight_name": weight_name, "revision": "", "civitai_model_id": model_id, "civitai_version_id": version_id, "civitai_file_id": file_id, "civitai_sha256": str(raw.get("civitai_sha256") or "").strip().lower(), "civitai_page_url": str(raw.get("civitai_page_url") or "").strip(), "trained_words": [str(x) for x in (raw.get("trained_words") or []) if str(x).strip()], } ) seen.add(label) continue repo_id = str(raw.get("repo_id") or "").strip() revision = str(raw.get("revision") or "").strip() if repo_id: clean.append( { "label": label, "repo_id": repo_id, "weight_name": weight_name, "revision": revision, "source": item_source, } ) seen.add(label) return clean def all_defs(builtin_loras, custom_loras=None) -> list[dict]: builtins = sanitize_defs(builtin_loras, source="built_in") sessions = sanitize_defs(custom_loras, source="session") builtin_labels = {item["label"] for item in builtins} sessions = [item for item in sessions if item["label"] not in builtin_labels] return builtins + sessions def validate_repo_id(repo_id: str) -> str: repo_id = str(repo_id or "").strip() if repo_id.count("/") != 1 or repo_id.startswith("/") or repo_id.endswith("/"): raise LoraSelectionError("HF repo ID must look like owner/repo.") return repo_id def validate_weight_name(weight_name: str) -> str: weight_name = str(weight_name or "").strip().replace("\\", "/") if not weight_name.endswith(".safetensors"): raise LoraSelectionError("LoRA weight filename must end with .safetensors.") if weight_name.startswith("/") or ".." in Path(weight_name).parts: raise LoraSelectionError("LoRA weight filename must be a repository-relative path.") return weight_name def normalize_hf_lora_source(source: str, revision: str = "", weight_name: str = "") -> dict: """Canonicalize a Hub repo ID, model URL, file URL, or hf:// model URI.""" raw = str(source or "").strip() if not raw: raise LoraSelectionError("Enter an HF model repo ID or URL.") revision = str(revision or "").strip() weight_name = str(weight_name or "").strip().replace("\\", "/") url_revision = "" url_weight = "" source_kind = "repo_id" is_hf_uri = raw.startswith("hf://") is_hf_web = "://" in raw or raw.startswith(("huggingface.co/", "www.huggingface.co/", "hf.co/")) if is_hf_uri: source_kind = "hf_uri" try: uri = parse_hf_uri(raw) except Exception as exc: raise LoraSelectionError("Invalid Hugging Face hf:// URI.") from exc if uri.type != "model": raise LoraSelectionError("HF LoRA source must point to a model repository, not a dataset/Space/kernel/bucket.") repo_id = validate_repo_id(uri.id) url_revision = str(uri.revision or "").strip() url_weight = str(uri.path_in_repo or "").strip().replace("\\", "/") elif not is_hf_web: repo_id = validate_repo_id(raw) else: source_kind = "url" web_raw = raw if "://" in raw else f"https://{raw}" # Newer huggingface_hub releases parse supported web URLs directly. Keep # the small fallback below so the product remains tolerant of older # environments while requirements/runtime converge. uri = None try: uri = parse_hf_uri(web_raw) except Exception: uri = None if uri is not None: if uri.type != "model": raise LoraSelectionError("HF LoRA source must point to a model repository, not a dataset/Space/kernel/bucket.") repo_id = validate_repo_id(uri.id) url_revision = str(uri.revision or "").strip() url_weight = str(uri.path_in_repo or "").strip().replace("\\", "/") else: try: parsed = urllib.parse.urlsplit(web_raw) except Exception as exc: raise LoraSelectionError("Invalid Hugging Face URL.") from exc if (parsed.hostname or "").casefold() not in {"huggingface.co", "www.huggingface.co", "hf.co"}: raise LoraSelectionError("HF LoRA URL must point to huggingface.co or hf.co.") raw_parts = [x for x in (parsed.path or "").strip("/").split("/") if x] parts = [urllib.parse.unquote(x) for x in raw_parts] if parts and parts[0] == "models": parts = parts[1:] if parts and parts[0] in {"datasets", "spaces", "kernels", "buckets", "collections"}: raise LoraSelectionError("HF LoRA source must point to a model repository.") if len(parts) < 2: raise LoraSelectionError("HF model URL must include owner/repo.") repo_id = validate_repo_id(f"{parts[0]}/{parts[1]}") if len(parts) > 2: route = parts[2] if route in {"blob", "resolve", "raw"}: rest = parts[3:] if len(rest) < 2: raise LoraSelectionError("HF file URL must include revision and repository-relative file path.") if len(rest) >= 4 and rest[0] == "refs" and rest[1] in {"pr", "convert"}: url_revision = "/".join(rest[:3]) url_weight = "/".join(rest[3:]) else: url_revision = rest[0] url_weight = "/".join(rest[1:]) elif route == "tree": rest = parts[3:] if not rest: raise LoraSelectionError("HF tree URL must include a revision.") if len(rest) >= 3 and rest[0] == "refs" and rest[1] in {"pr", "convert"}: url_revision = "/".join(rest[:3]) else: url_revision = rest[0] else: raise LoraSelectionError("Unsupported Hugging Face model URL route. Use the repo page or a blob/resolve/raw file URL.") if url_revision and revision and url_revision != revision: raise LoraSelectionError(f"HF URL revision {url_revision!r} conflicts with Revision {revision!r}.") resolved_revision = url_revision or revision if url_weight: url_weight = validate_weight_name(url_weight) if weight_name and validate_weight_name(weight_name) != url_weight: raise LoraSelectionError("HF file URL conflicts with the Safetensors file field.") resolved_weight = url_weight else: resolved_weight = validate_weight_name(weight_name) if weight_name else "" return { "repo_id": repo_id, "revision": resolved_revision, "weight_name": resolved_weight, "source_kind": source_kind, } def _download_hf_lora_file(item: dict, hf_token: str | None) -> Path: """Download one exact Hub LoRA file and normalize common Hub failures.""" repo_id = str(item.get("repo_id") or "").strip() weight_name = str(item.get("weight_name") or "").strip() revision = str(item.get("revision") or "").strip() or None try: return Path( hf_hub_download( repo_id=repo_id, filename=weight_name, revision=revision, token=hf_token, ) ) except GatedRepoError as exc: raise LoraSelectionError( f"HF LoRA repo {repo_id} is gated and the current Space token cannot access it. " "Request access for the token owner or configure an authorized HF token." ) from exc except RepositoryNotFoundError as exc: raise LoraSelectionError( f"HF LoRA repo {repo_id} was not found or is private to the current Space token." ) from exc except RevisionNotFoundError as exc: raise LoraSelectionError( f"HF revision {revision!r} was not found in {repo_id}." ) from exc except RemoteEntryNotFoundError as exc: where = f" at revision {revision!r}" if revision else "" raise LoraSelectionError( f"HF LoRA file {weight_name!r} was not found in {repo_id}{where}. " "Inspect the repo and choose an exact .safetensors file." ) from exc except LocalEntryNotFoundError as exc: raise LoraSelectionError( f"HF LoRA file {weight_name!r} is not cached locally and the Hub could not be reached. " "Check network/offline mode and try again." ) from exc except HfHubHTTPError as exc: status = getattr(getattr(exc, "response", None), "status_code", None) if status in {401, 403}: raise LoraSelectionError( f"HF denied access while downloading {repo_id}/{weight_name} (HTTP {status}). " "Check gated/private access and the Space HF token." ) from exc if status == 429: raise LoraSelectionError( "HF Hub rate-limited the LoRA download (HTTP 429). Wait briefly and try again." ) from exc suffix = f" (HTTP {status})" if status else "" raise LoraSelectionError( f"HF LoRA download failed for {repo_id}/{weight_name}{suffix}." ) from exc except OSError as exc: raise LoraSelectionError( f"HF LoRA download could not complete for {repo_id}/{weight_name}: {type(exc).__name__}." ) from exc def resolve_selected(selected, builtin_loras, custom_loras) -> list[dict]: defs = {item["label"]: item for item in all_defs(builtin_loras, custom_loras)} resolved = [] for label in selected or []: if label not in defs: raise LoraSelectionError(f"Selected LoRA is unavailable: {label}") resolved.append(defs[label]) return resolved def inspect_model_repo(repo_id: str, hf_token: str | None, revision: str | None = None) -> dict: """CPU/network-only Hub metadata inspection for the session-LoRA form.""" repo_id = validate_repo_id(repo_id) revision = str(revision or "").strip() or None try: info = HfApi(token=hf_token).model_info(repo_id=repo_id, revision=revision, files_metadata=False) except GatedRepoError as exc: raise LoraSelectionError(f"HF repo {repo_id} is gated and the current Space token cannot access it.") from exc except RepositoryNotFoundError as exc: raise LoraSelectionError(f"HF model repo {repo_id} was not found or is private to the current Space token.") from exc except RevisionNotFoundError as exc: raise LoraSelectionError(f"HF revision {revision!r} was not found in {repo_id}.") from exc except HfHubHTTPError as exc: status = getattr(getattr(exc, "response", None), "status_code", None) suffix = f" (HTTP {status})" if status else "" raise LoraSelectionError(f"HF repo inspection failed for {repo_id}{suffix}.") from exc except Exception as exc: raise LoraSelectionError(f"HF repo inspection failed for {repo_id}: {type(exc).__name__}.") from exc safetensors = sorted( { str(getattr(sibling, "rfilename", "") or "") for sibling in (getattr(info, "siblings", None) or []) if str(getattr(sibling, "rfilename", "") or "").lower().endswith(".safetensors") } ) return { "repo_id": repo_id, "requested_revision": revision, "resolved_revision": str(getattr(info, "sha", "") or "") or None, "private": bool(getattr(info, "private", False)), "gated": getattr(info, "gated", None), "safetensors": safetensors, "tags": [str(x) for x in (getattr(info, "tags", None) or [])], } _LTX2_LORA_PREFIXES = ("diffusion_model.", "text_embedding_projection.", "transformer.", "connectors.") def _checkpoint_key_sample(keys, limit: int = 6) -> list[str]: return [str(key) for key in list(keys)[: max(1, int(limit))]] def _mapped_lora_key(key: str) -> str: return str(key).replace(".lora_down.weight", ".lora_A.weight").replace(".lora_up.weight", ".lora_B.weight") def _inspect_civitai_ltx2_checkpoint(local: Path) -> dict: """Inspect a Civitai safetensors header without loading tensor payloads.""" try: with safe_open(str(local), framework="pt", device="cpu") as handle: keys = list(handle.keys()) except Exception as exc: raise LoraSelectionError(f"Could not inspect Civitai LoRA safetensors header: {type(exc).__name__}: {exc}") from exc if not keys: raise LoraSelectionError("Civitai checkpoint is empty; refusing to request GPU quota.") unsupported_prefixes = [key for key in keys if not str(key).startswith(_LTX2_LORA_PREFIXES)] if unsupported_prefixes: sample = ", ".join(_checkpoint_key_sample(unsupported_prefixes)) raise LoraSelectionError( "Unsupported LTX-2 LoRA checkpoint dialect from Civitai. " f"Unexpected parameter namespace(s): {sample}. Try another file/version. GPU quota was not requested." ) alpha_keys = [key for key in keys if str(key).endswith(".alpha")] down_up_keys = [key for key in keys if ".lora_down.weight" in str(key) or ".lora_up.weight" in str(key)] unsupported_non_lora = [key for key in keys if "lora" not in str(key).lower() and not str(key).endswith(".alpha")] if unsupported_non_lora: sample = ", ".join(_checkpoint_key_sample(unsupported_non_lora)) raise LoraSelectionError( "Civitai file is not a supported LTX-2 LoRA-only checkpoint. " f"Unexpected parameter(s): {sample}. Try another exact .safetensors file/version. GPU quota was not requested." ) if alpha_keys or down_up_keys: dialect = "ltx2_comfy_alpha_or_down_up" needs_normalization = True elif any(str(key).startswith("diffusion_model.") or str(key).startswith("text_embedding_projection.") for key in keys): dialect = "ltx2_comfy_native" needs_normalization = False else: dialect = "ltx2_diffusers_peft" needs_normalization = False return { "dialect": dialect, "key_count": len(keys), "key_sample": _checkpoint_key_sample(keys), "alpha_key_count": len(alpha_keys), "down_up_key_count": len(down_up_keys), "needs_normalization": needs_normalization, } def _normalize_civitai_ltx2_checkpoint(local: Path, cache_root: Path, source_sha256: str) -> tuple[Path, dict]: """Normalize only known LTX-2 Comfy LoRA variants before GPU allocation. Current pinned Diffusers accepts LTX2 Comfy-style ``diffusion_model.`` keys but does not normalize per-module ``.alpha`` or generic ``lora_down/lora_up`` keys. For that narrow, known dialect we rename down/up to A/B and fold alpha/rank into lora_B. Unknown or mixed checkpoint dialects fail closed. """ evidence = _inspect_civitai_ltx2_checkpoint(local) if not evidence["needs_normalization"]: evidence.update(normalized=False, normalized_path=None, alpha_fold_count=0) return local, evidence try: state = load_safetensors_file(str(local), device="cpu") except Exception as exc: raise LoraSelectionError(f"Could not read Civitai LoRA tensors for CPU normalization: {type(exc).__name__}: {exc}") from exc normalized: dict[str, torch.Tensor] = {} alpha_tensors: dict[str, torch.Tensor] = {} for raw_key, tensor in state.items(): key = str(raw_key) if key.endswith(".alpha"): alpha_tensors[key[:-len(".alpha")]] = tensor continue mapped = _mapped_lora_key(key) if mapped in normalized: raise LoraSelectionError(f"Checkpoint normalization produced a duplicate LoRA key: {mapped}") normalized[mapped] = tensor a_keys = [key for key in normalized if key.endswith(".lora_A.weight")] b_keys = [key for key in normalized if key.endswith(".lora_B.weight")] if not a_keys or not b_keys: raise LoraSelectionError( "Unsupported Civitai LoRA checkpoint: no complete lora_A/lora_B pairs were found after known-format normalization. " "Try another file/version. GPU quota was not requested." ) pair_bases = set() for key in a_keys: base = key[:-len(".lora_A.weight")] b_key = base + ".lora_B.weight" if b_key not in normalized: raise LoraSelectionError(f"Incomplete Civitai LoRA pair for {base}: lora_B is missing. GPU quota was not requested.") pair_bases.add(base) for key in b_keys: base = key[:-len(".lora_B.weight")] if base + ".lora_A.weight" not in normalized: raise LoraSelectionError(f"Incomplete Civitai LoRA pair for {base}: lora_A is missing. GPU quota was not requested.") pair_bases.add(base) alpha_fold_count = 0 for base, alpha_tensor in alpha_tensors.items(): a_key = base + ".lora_A.weight" b_key = base + ".lora_B.weight" if a_key not in normalized or b_key not in normalized: raise LoraSelectionError( f"Civitai LoRA alpha has no matching A/B pair for {base}. Refusing unsafe fallback before GPU allocation." ) if alpha_tensor.numel() != 1: raise LoraSelectionError(f"Civitai LoRA alpha for {base} is not scalar; unsupported checkpoint dialect.") a_tensor = normalized[a_key] if a_tensor.ndim < 1 or int(a_tensor.shape[0]) <= 0: raise LoraSelectionError(f"Could not derive LoRA rank for {base}; unsupported checkpoint dialect.") rank = int(a_tensor.shape[0]) alpha = float(alpha_tensor.detach().float().item()) scale = alpha / float(rank) original_b = normalized[b_key] normalized[b_key] = (original_b.detach().float() * scale).to(dtype=original_b.dtype) alpha_fold_count += 1 remaining_bad = [key for key in normalized if "lora" not in key.lower()] if remaining_bad: sample = ", ".join(_checkpoint_key_sample(remaining_bad)) raise LoraSelectionError(f"Unsupported non-LoRA parameter(s) remain after normalization: {sample}") source_sha256 = str(source_sha256 or "").strip().lower() if len(source_sha256) != 64: import hashlib digest = hashlib.sha256() with local.open("rb") as fh: for chunk in iter(lambda: fh.read(1024 * 1024), b""): digest.update(chunk) source_sha256 = digest.hexdigest() out_dir = Path(cache_root) / "normalized_ltx2_lora" out_dir.mkdir(parents=True, exist_ok=True) out_path = out_dir / f"{source_sha256}.diffusers-compatible.safetensors" try: save_safetensors_file(normalized, str(out_path)) except Exception as exc: raise LoraSelectionError(f"Could not write runtime-normalized Civitai LoRA: {type(exc).__name__}: {exc}") from exc evidence.update( normalized=True, normalized_path=str(out_path), normalized_key_count=len(normalized), alpha_fold_count=alpha_fold_count, source_sha256=source_sha256, ) return out_path, evidence def definition_key(item: dict) -> tuple[str, ...]: source = str(item.get("source") or "").strip() if source == "civitai": return ( "civitai", str(item.get("civitai_model_id") or ""), str(item.get("civitai_version_id") or ""), str(item.get("civitai_file_id") or ""), str(item.get("civitai_sha256") or "").lower(), ) return ( "hf", str(item.get("repo_id") or ""), str(item.get("weight_name") or ""), str(item.get("revision") or item.get("requested_revision") or ""), ) def prepare_selected(selected, builtin_loras, custom_loras, hf_token: str | None, civitai_cache_root: Path | None = None, civitai_api_key: str = "") -> tuple[list[dict], float]: """CPU-side source preparation. This function must never run inside a GPU callback.""" resolved = resolve_selected(selected, builtin_loras, custom_loras) if not resolved: return [], 0.0 prepared = [] total_started = time.monotonic() for item in resolved: item_started = time.monotonic() if str(item.get("source") or "") == "civitai": if civitai_cache_root is None: raise LoraSelectionError("Civitai LoRA cache root is unavailable.") try: acquired = civitai_backend.download_exact_file(item, Path(civitai_cache_root), api_key=civitai_api_key) except civitai_backend.CivitaiError as exc: raise LoraSelectionError(str(exc)) from exc local = Path(str(acquired["local_path"])) source_sha256 = str(acquired.get("sha256") or item.get("civitai_sha256") or "").strip().lower() local, checkpoint_compatibility = _normalize_civitai_ltx2_checkpoint( local, Path(civitai_cache_root), source_sha256 ) record = { "label": item["label"], "source": "civitai", "repo_id": None, "weight_name": item["weight_name"], "requested_revision": None, "resolved_revision": None, "civitai_model_id": item.get("civitai_model_id"), "civitai_version_id": item.get("civitai_version_id"), "civitai_file_id": item.get("civitai_file_id"), "civitai_sha256": acquired.get("sha256") or item.get("civitai_sha256"), "local_path": str(local), "source_size_bytes": int(acquired.get("size_bytes") or 0), "size_bytes": int(local.stat().st_size), "checkpoint_compatibility": checkpoint_compatibility, "cache_hit": bool(acquired.get("cache_hit")), "prepare_seconds": time.monotonic() - item_started, } else: local = _download_hf_lora_file(item, hf_token) record = { "label": item["label"], "source": item.get("source"), "repo_id": item["repo_id"], "weight_name": item["weight_name"], "requested_revision": item["revision"] or None, "resolved_revision": resolved_revision_from_hub_path(local), "local_path": str(local), "size_bytes": int(local.stat().st_size), "prepare_seconds": time.monotonic() - item_started, } record["definition_key"] = list(definition_key(item)) prepared.append(record) return prepared, time.monotonic() - total_started def prepared_lookup(prepared_loras) -> dict: lookup = {} for item in prepared_loras or []: if not isinstance(item, dict): continue raw_key = item.get("definition_key") if isinstance(raw_key, (list, tuple)) and raw_key: key = tuple(str(x) for x in raw_key) else: key = definition_key(item) lookup[key] = item return lookup def adapter_state(pipe) -> dict: state = {} for name in ("get_active_adapters", "get_list_adapters"): fn = getattr(pipe, name, None) if callable(fn): try: state[name] = fn() except Exception as exc: state[name] = {"error": f"{type(exc).__name__}: {exc}"} return state def load_request_loras( *, pipe, pipe_i2v, pipe_condition, pipe_ic, selected, builtin_loras, custom_loras, prepared_loras, strength: float, request_id: str, gpu_state: Callable[[], dict], ) -> tuple[list[dict], dict]: metrics = { "requested_labels": [str(x) for x in (selected or [])], "requested_count": len(selected or []), "strength": float(strength), "hub_download_inside_gpu_callback": False, "request_scoped": True, "fused": False, } if not selected: metrics["status"] = "not_requested" return [], metrics resolved = resolve_selected(selected, builtin_loras, custom_loras) prepared = prepared_lookup(prepared_loras) loaded = [] adapter_names = [] metrics["gpu_before_load"] = gpu_state() total_started = time.monotonic() try: per_adapter = [] for idx, item in enumerate(resolved): key = definition_key(item) prep = prepared.get(key) if not prep: raise LoraSelectionError( f"Selected LoRA is not CPU-prepared: {item['label']}. " "Use Prepare selected LoRAs; Generate normally performs this pre-step automatically." ) local = Path(str(prep.get("local_path") or "")) if not local.is_file(): raise LoraSelectionError(f"Prepared LoRA file is no longer available locally: {item['label']}") adapter_name = f"req_{request_id[:8]}_{idx}" load_started = time.monotonic() pipe.load_lora_weights(str(local.parent), weight_name=local.name, adapter_name=adapter_name) load_seconds = time.monotonic() - load_started adapter_names.append(adapter_name) public_record = { "label": item["label"], "source": item.get("source"), "repo_id": item.get("repo_id") or None, "weight_name": item["weight_name"], "requested_revision": item.get("revision") or None, "resolved_revision": prep.get("resolved_revision"), "size_bytes": prep.get("size_bytes"), "adapter_name": adapter_name, "load_seconds": load_seconds, } if item.get("source") == "civitai": public_record.update( civitai_model_id=item.get("civitai_model_id"), civitai_version_id=item.get("civitai_version_id"), civitai_file_id=item.get("civitai_file_id"), civitai_sha256=prep.get("civitai_sha256") or item.get("civitai_sha256"), checkpoint_compatibility=prep.get("checkpoint_compatibility"), ) loaded.append(public_record) per_adapter.append({k: public_record[k] for k in ("label", "adapter_name", "size_bytes", "load_seconds")}) set_started = time.monotonic() pipe.set_adapters(adapter_names, adapter_weights=[float(strength)] * len(adapter_names)) pipe.enable_lora() metrics["set_adapters_seconds"] = time.monotonic() - set_started metrics["load_total_seconds"] = time.monotonic() - total_started metrics["per_adapter"] = per_adapter metrics["adapter_names"] = list(adapter_names) metrics["adapter_state_after_set"] = adapter_state(pipe) metrics["gpu_after_load"] = gpu_state() metrics["shared_component_identity"] = { "pipe_i2v_transformer_shared": bool(pipe_i2v is None or pipe_i2v.transformer is pipe.transformer), "pipe_i2v_connectors_shared": bool(pipe_i2v is None or pipe_i2v.connectors is pipe.connectors), "pipe_condition_transformer_shared": bool(pipe_condition is None or pipe_condition.transformer is pipe.transformer), "pipe_condition_connectors_shared": bool(pipe_condition is None or pipe_condition.connectors is pipe.connectors), "pipe_ic_transformer_shared": bool(pipe_ic is None or pipe_ic.transformer is pipe.transformer), "pipe_ic_connectors_shared": bool(pipe_ic is None or pipe_ic.connectors is pipe.connectors), } metrics["status"] = "loaded_active" return loaded, metrics except Exception: try: pipe.disable_lora() if adapter_names: pipe.delete_adapters(adapter_names) except Exception: pass raise def cleanup_request_loras( *, pipe, loaded, metrics=None, full_sft_profile: bool, gpu_state: Callable[[], dict], ) -> dict: metrics = metrics if isinstance(metrics, dict) else {} if not loaded: metrics.setdefault("cleanup_status", "not_needed") return metrics adapter_names = [str(item.get("adapter_name")) for item in loaded if item.get("adapter_name")] metrics["gpu_before_cleanup"] = gpu_state() metrics["cleanup_adapter_names"] = list(adapter_names) metrics["internal_stage2_adapter_preserved"] = bool(full_sft_profile) started = time.monotonic() try: pipe.disable_lora() if adapter_names: pipe.delete_adapters(adapter_names) metrics["cleanup_status"] = "PASS" finally: gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() metrics["cleanup_seconds"] = time.monotonic() - started metrics["adapter_state_after_cleanup"] = adapter_state(pipe) metrics["gpu_after_cleanup"] = gpu_state() return metrics