Spaces:

jickman125
/

mmh3testing

Running on Zero

fix: classify new catalog LoRAs as kohya, sniff format from keys

bc09615 about 1 month ago

19.8 kB

"""LoRA fetch + load from Hugging Face (or any direct .safetensors path).
Functions:
- `_parse_hf_link(url)`            → `(repo, file | None)`
- `_resolve_repo_file(repo)`       → pick the single .safetensors of a repo link
- `_resolve_repo_revision(repo)`   → Hub tip SHA (immutable cache key)
- `_lora_cache_path(repo, file, revision)` → bucket path, or None if not writable
- `_fetch_lora_file(repo, file, revision)` → local .safetensors path (cached)
- `_load_lora_from_cache(path)`    → load any local .safetensors into _CUSTOM_LoRA
- `_ensure_nsfw_lora_loaded(key)`  → lazy fetch + convert a catalog NSFW LoRA
- `_ensure_custom_lora_loaded(repo, file)` → same for a user-supplied HF LoRA
- `_ensure_custom_lora_loaded_from_path(key, path, label)` → load an already-fetched file
Three layers of `H3_MAX_LORA_BYTES` enforcement: HEAD pre-check, post-download
`os.path.getsize`, and a running byte counter during the bucket-cache copy
(defends against a Hub lying about Content-Length). Concurrent writers use
`_LORA_CACHE_LOCK` + temp file + `os.replace`.
"""
from __future__ import annotations

import json
import os
import re
import struct
import time

import gradio as gr

from config import NSFW_LORA_CONFIGS
from lora.cache import _evict_lora_cache_for_quota
from lora.formats import (
    _convert_lora,
    _detect_wrong_arch,
)
import state as _state
from state import (
    _CUSTOM_LoRA,
    _H3_MAX_LORA_BYTES,
    _LORA_CACHE_DIR,
    _LORA_CACHE_LOCK,
    _NSFW_LoRA,
    _log,
)


def _parse_hf_link(url):
    """Parse a Hugging Face repo or file link into (repo, file). `file` is None for a bare repo link."""
    url = (url or "").strip().split("?")[0].split("#")[0]
    if not url:
        raise gr.Error("Paste a Hugging Face or CivitAI link first.")
    if not url.startswith("http"):
        parts = [p for p in url.split("/") if p]
        if len(parts) >= 2:
            return parts[0] + "/" + parts[1], "/".join(parts[2:]) or None
        raise gr.Error(f"Couldn't parse `{url}` as a Hugging Face repo link.")
    if "huggingface.co" not in url:
        raise gr.Error("Only Hugging Face links are supported (huggingface.co).")
    parts = [p for p in url.split("huggingface.co", 1)[1].strip("/").split("/") if p]
    if parts and parts[0] in ("spaces", "models", "datasets"):
        parts = parts[1:]  # https://huggingface.co/spaces/Owner/Repo → Owner/Repo
    if len(parts) < 2:
        raise gr.Error(f"Couldn't parse `{url}` as a Hugging Face repo link.")
    repo = parts[0] + "/" + parts[1]
    if len(parts) == 2:
        return repo, None
    if parts[2] in ("resolve", "tree", "blob"):
        if len(parts) == 4 and parts[2] in ("tree", "blob"):
            return repo, None  # tree/main or blob/main with no subpath → treat as a repo link
        if len(parts) < 5:
            raise gr.Error(f"Couldn't parse a filename from `{url}`.")
        if parts[3] != "main":
            raise gr.Error("Only `main`-branch links are supported — use a `resolve/main/...` link.")
        return repo, "/".join(parts[4:])
    return repo, "/".join(parts[2:])


def _resolve_repo_file(repo):
    """Pick the single .safetensors file of a repo link; error if there isn't exactly one."""
    from huggingface_hub import list_repo_files
    try:
        files = [f for f in list_repo_files(repo) if f.endswith(".safetensors")]
    except Exception as _e:
        raise gr.Error(f"Couldn't list files in `{repo}`: {_e}")
    if len(files) == 1:
        return files[0]
    if not files:
        raise gr.Error(f"No .safetensors files found in `{repo}`.")
    raise gr.Error(f"`{repo}` has multiple .safetensors files — paste a direct file link instead: "
                   + ", ".join(files[:5]) + ("…" if len(files) > 5 else ""))


def _resolve_repo_revision(repo: str) -> str:
    """Resolve the current tip of *repo* to an immutable commit SHA.
    Falls back to the literal string 'main' only if the Hub is unreachable so
    cache keys remain stable within a process; callers must not treat 'main' as
    immutable across restarts.
    """
    try:
        from huggingface_hub import HfApi
        info = HfApi().repo_info(repo, repo_type="model")
        sha = getattr(info, "sha", None) or getattr(info, "revision", None)
        if sha and isinstance(sha, str) and len(sha) >= 7:
            return sha
    except Exception as _e:
        _log(f"revision resolve failed for {repo}: {_e}", "WARN")
    return "main"


def _lora_cache_path(repo, file, revision: str | None = None):
    """Bucket path for a (repo, revision, file) LoRA, or None if the bucket isn't usable.
    Cache identity is (repo, revision_sha, file) so a moving `main` cannot serve
    stale weights under the same path.
    """
    try:
        os.makedirs(_LORA_CACHE_DIR, exist_ok=True)
        _probe = os.path.join(_LORA_CACHE_DIR, ".probe")
        with open(_probe, "w") as _f:
            _f.write("ok")
        os.remove(_probe)
    except Exception:
        return None
    rev = (revision or "main")[:40]
    _key = re.sub(r"[^A-Za-z0-9._-]", "_", f"{repo.replace('/', '__')}__rev_{rev}__{file}")
    return os.path.join(_LORA_CACHE_DIR, _key)


def _fetch_lora_file(repo, file, revision: str | None = None, max_bytes: int | None = None):
    """Local path to a LoRA .safetensors, preferring the bucket cache over a fresh download.
    Cache key includes the immutable Hub revision SHA when available.
    Enforces H3_MAX_LORA_BYTES (reject before/during download) and
    H3_MAX_LORA_CACHE_BYTES (evict unreferenced entries or reject).
    Concurrent writers use temp files + os.replace so consumers never see
    partial artifacts.
    ``max_bytes`` overrides the user-LoRA size cap (used for the system turbo
    weights, which are larger than the custom-LoRA quota by design).
    """
    from huggingface_hub import hf_hub_download, hf_hub_url
    import requests

    limit = int(max_bytes) if max_bytes is not None else _H3_MAX_LORA_BYTES
    rev = revision or _resolve_repo_revision(repo)
    with _LORA_CACHE_LOCK:
        _cached = _lora_cache_path(repo, file, revision=rev)
        if _cached and os.path.exists(_cached):
            try:
                with open(_cached, "rb") as _f:
                    _n = struct.unpack("<Q", _f.read(8))[0]
                    _header_keys = list(json.loads(_f.read(_n)).keys())
                _hint = _detect_wrong_arch(_header_keys)
                if _hint:
                    _log(f"LoRA cache evict (incompatible): {repo}/{file}@{rev[:7]} — {_hint}", "WARN")
                    try:
                        os.remove(_cached)
                    except OSError:
                        pass
                    raise gr.Error(f"`{file}` is not a MiniMax-H3 LoRA. {_hint}")
                size = os.path.getsize(_cached)
                if size > limit:
                    try:
                        os.remove(_cached)
                    except OSError:
                        pass
                    raise gr.Error(
                        f"`{file}` exceeds H3_MAX_LORA_BYTES "
                        f"({size / 1e6:.0f} MB > {limit / 1e6:.0f} MB)."
                    )
                _log(f"LoRA cache hit: {repo}/{file}@{rev[:7]} ({size / 1e6:.0f} MB)")
                return _cached
            except gr.Error:
                raise
            except Exception:
                _log(f"LoRA cache corrupt, re-downloading: {repo}/{file}@{rev[:7]}", "WARN")
                try:
                    os.remove(_cached)
                except OSError:
                    pass

        # Pre-check size via HEAD when Content-Length is available.
        try:
            url = hf_hub_url(repo_id=repo, filename=file, revision=rev if rev != "main" else "main")
            head = requests.head(url, allow_redirects=True, timeout=30)
            cl = head.headers.get("Content-Length")
            if cl is not None:
                expected = int(cl)
                if expected > limit:
                    raise gr.Error(
                        f"`{file}` is too large ({expected / 1e6:.0f} MB > "
                        f"{limit / 1e6:.0f} MB limit)."
                    )
                _evict_lora_cache_for_quota(expected)
        except gr.Error:
            raise
        except Exception as _e:
            _log(f"LoRA HEAD pre-check skipped ({_e})", "WARN")

        _t0 = time.time()
        # Download via hub; then enforce size on the resulting file before cache admit.
        _path = hf_hub_download(repo, file, revision=rev if rev != "main" else None)
        size = os.path.getsize(_path)
        _log(f"LoRA downloaded: {repo}/{file}@{rev[:7]} ({size / 1e6:.0f} MB in {time.time() - _t0:.0f}s)")
        if size > limit:
            try:
                os.remove(_path)
            except OSError:
                pass
            raise gr.Error(
                f"`{file}` exceeds H3_MAX_LORA_BYTES "
                f"({size / 1e6:.0f} MB > {limit / 1e6:.0f} MB)."
            )

        # Architecture gate before admitting to persistent cache.
        try:
            with open(_path, "rb") as _f:
                _n = struct.unpack("<Q", _f.read(8))[0]
                _header = json.loads(_f.read(_n))
            _hint = _detect_wrong_arch(list(_header.keys()))
            if _hint:
                _log(f"LoRA rejected (incompatible architecture): {repo}/{file} — {_hint}", "WARN")
                try:
                    os.remove(_path)
                except OSError:
                    pass
                raise gr.Error(f"`{file}` is not a MiniMax-H3 LoRA. {_hint}")
        except gr.Error:
            raise
        except Exception as _e:
            _log(f"LoRA header peek failed ({_e}); proceeding with copy", "WARN")

        if _cached:
            try:
                _evict_lora_cache_for_quota(size)
                _tmp = _cached + ".tmp"
                # Stream with a running byte counter so a lied Content-Length still cannot exceed the ceiling.
                written = 0
                with open(_path, "rb") as _src, open(_tmp, "wb") as _dst:
                    while True:
                        chunk = _src.read(1024 * 1024)
                        if not chunk:
                            break
                        written += len(chunk)
                        if written > limit:
                            _dst.close()
                            try:
                                os.remove(_tmp)
                            except OSError:
                                pass
                            raise gr.Error(
                                f"`{file}` exceeded H3_MAX_LORA_BYTES while copying "
                                f"({written / 1e6:.0f} MB)."
                            )
                        _dst.write(chunk)
                os.replace(_tmp, _cached)
                _log(f"LoRA cached to bucket: {repo}/{file}@{rev[:7]}")
                return _cached
            except gr.Error:
                raise
            except Exception as _e:
                _log(f"LoRA bucket cache write failed ({_e}); using HF cache", "WARN")
                try:
                    if os.path.exists(_cached + ".tmp"):
                        os.remove(_cached + ".tmp")
                except OSError:
                    pass
        return _path


def _ensure_nsfw_lora_loaded(key):
    """Download + convert + store a lazy NSFW LoRA's A/B factors on CPU. No-op if already loaded."""
    if key in _NSFW_LoRA:
        return
    cfg = NSFW_LORA_CONFIGS[key]
    import torch
    from safetensors.torch import load_file as _load_lora
    _t0 = time.time()
    _log(f"loading NSFW LoRA {key} from {cfg['repo']} ...")
    # Catalog NSFW weights can exceed the user custom-LoRA quota (default 512 MB);
    # allow up to 1.5 GB for official catalog entries only.
    _path = _fetch_lora_file(cfg["repo"], cfg["file"], max_bytes=1536 * 1024 * 1024)
    _raw = _load_lora(_path)
    # The catalog's `format` is a hint, not a contract — a mislabeled entry converts under whatever
    # format its keys actually are, and only a file no converter recognizes becomes a user error.
    _sd, _fmt = _convert_lora(_raw, cfg.get("format", ""))
    if not _sd:
        raise gr.Error(
            f"`{cfg['file']}` converted to no usable modules — it isn't a MiniMax-H3 LoRA "
            f"(or its `format` is mislabeled). {_detect_wrong_arch(_raw)}"
        )
    _log(f"NSFW LoRA {key}: {len(_raw)} keys as {_fmt}")
    _pairs: dict[str, dict[str, torch.Tensor]] = {}
    for _k, _v in _sd.items():
        if ".lora_A." in _k:
            _pairs.setdefault(_k.split(".lora_A.")[0], {})["A"] = _v
        elif ".lora_B." in _k:
            _pairs.setdefault(_k.split(".lora_B.")[0], {})["B"] = _v
    _loaded = []
    for _base, _pair in _pairs.items():
        if "A" not in _pair or "B" not in _pair:
            continue
        _parts = _base.split(".")
        _obj = _state.PIPE.transformer
        for _p in _parts:
            _obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p)
        _loaded.append((_obj, _pair["A"].to(_obj.weight.dtype), _pair["B"].to(_obj.weight.dtype)))
    if not _loaded:
        raise gr.Error(
            f"`{cfg['file']}` has no complete LoRA module pairs — it isn't a MiniMax-H3 LoRA "
            f"(or its `format` is mislabeled). {_detect_wrong_arch(_raw)}"
        )
    del _sd, _pairs, _raw
    import gc as _gc; _gc.collect()
    _NSFW_LoRA[key] = _loaded
    _log(f"NSFW LoRA {key} parsed ({len(_loaded)} modules in {time.time() - _t0:.1f}s)")


def _ensure_custom_lora_loaded(repo, file):
    """Download + convert + store a user-supplied Hugging Face LoRA's A/B factors on CPU. No-op if already loaded."""
    key = (repo, file)
    if key in _CUSTOM_LoRA:
        return
    import torch
    from safetensors.torch import load_file as _load_lora
    _t0 = time.time()
    _log(f"loading custom LoRA {repo}/{file} ...")
    try:
        _path = _fetch_lora_file(repo, file)
        _raw = _load_lora(_path)
    except Exception as _e:
        raise gr.Error(f"Couldn't fetch custom LoRA `{repo}/{file}`: {_e}")
    _sd, _fmt = _convert_lora(_raw)
    if not _sd:
        raise gr.Error(f"`{repo}/{file}` doesn't look like a MiniMax-H3 LoRA (no lora_A/lora_down/lokr keys found).")
    _log(f"custom LoRA {repo}/{file}: {len(_raw)} keys as {_fmt}")
    _pairs: dict[str, dict[str, torch.Tensor]] = {}
    for _k, _v in _sd.items():
        if ".lora_A." in _k:
            _pairs.setdefault(_k.split(".lora_A.")[0], {})["A"] = _v
        elif ".lora_B." in _k:
            _pairs.setdefault(_k.split(".lora_B.")[0], {})["B"] = _v
    _loaded = []
    for _base, _pair in _pairs.items():
        if "A" not in _pair or "B" not in _pair:
            continue
        _parts = _base.split(".")
        _obj = _state.PIPE.transformer
        for _p in _parts:
            _obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p)
        _loaded.append((_obj, _pair["A"].to(_obj.weight.dtype), _pair["B"].to(_obj.weight.dtype)))
    if not _loaded:
        raise gr.Error(f"`{repo}/{file}` converted to no usable modules — is it a MiniMax-H3 LoRA? {_detect_wrong_arch(_raw)}")
    del _sd, _pairs, _raw
    import gc as _gc; _gc.collect()
    _CUSTOM_LoRA[key] = _loaded
    _log(f"custom LoRA {repo}/{file} parsed ({len(_loaded)} modules in {time.time() - _t0:.1f}s)")


def _load_lora_from_cache(path):
    """Load a cached .safetensors LoRA directly by file path.  Returns (repo, file) key
    and populates ``_CUSTOM_LoRA`` so it can be merged on generate."""
    import torch
    from safetensors.torch import load_file as _load_sd
    _t0 = time.time()
    fname = os.path.basename(path)
    # Use the filename as a synthetic (repo, file) key.
    key = ("__cache__", fname)
    if key in _CUSTOM_LoRA:
        return key
    # Architecture gate: peek at header keys before loading full tensor data.
    try:
        with open(path, "rb") as _f:
            _n = struct.unpack("<Q", _f.read(8))[0]
            _header_keys = list(json.loads(_f.read(_n)).keys())
        _hint = _detect_wrong_arch(_header_keys)
        if _hint:
            _log(f"cached LoRA rejected (incompatible): {fname} — {_hint}", "WARN")
            try:
                os.remove(path)
            except OSError:
                pass
            raise gr.Error(f"`{fname}` is not a MiniMax-H3 LoRA. {_hint}")
    except gr.Error:
        raise
    except Exception:
        pass  # header peek failed — fall through to full load
    try:
        _raw = _load_sd(path)
    except Exception as _e:
        raise gr.Error(f"Couldn't load cached LoRA `{fname}`: {_e}")
    # Detect format and convert.
    _sd, _fmt = _convert_lora(_raw)
    if not _sd:
        raise gr.Error(f"`{fname}` doesn't look like a MiniMax-H3 LoRA.")
    _log(f"cached LoRA {fname}: {len(_raw)} keys as {_fmt}")
    _pairs: dict[str, dict[str, torch.Tensor]] = {}
    for _k, _v in _sd.items():
        if ".lora_A." in _k:
            _pairs.setdefault(_k.split(".lora_A.")[0], {})["A"] = _v
        elif ".lora_B." in _k:
            _pairs.setdefault(_k.split(".lora_B.")[0], {})["B"] = _v
    _loaded = []
    for _base, _pair in _pairs.items():
        if "A" not in _pair or "B" not in _pair:
            continue
        _parts = _base.split(".")
        _obj = _state.PIPE.transformer
        for _p in _parts:
            _obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p)
        _loaded.append((_obj, _pair["A"].to(_obj.weight.dtype), _pair["B"].to(_obj.weight.dtype)))
    if not _loaded:
        raise gr.Error(f"`{fname}` converted to no usable modules — is it a MiniMax-H3 LoRA? {_detect_wrong_arch(_raw)}")
    del _sd, _pairs, _raw
    import gc as _gc; _gc.collect()
    _CUSTOM_LoRA[key] = _loaded
    _log(f"cached LoRA {fname} loaded ({len(_loaded)} modules in {time.time() - _t0:.1f}s)")
    return key


def _ensure_custom_lora_loaded_from_path(key, path: str, label: str = ""):
    """Load + convert a local .safetensors LoRA into ``_CUSTOM_LoRA[key]``. No-op if present.
    Used by the CivitAI / CivArchive flows after they've fetched the file.
    """
    if key in _CUSTOM_LoRA:
        return
    import torch
    from safetensors.torch import load_file as _load_lora

    _t0 = time.time()
    label = label or f"{key[0]}/{key[1]}"
    _log(f"loading custom LoRA {label} from local path ...")
    try:
        _raw = _load_lora(path)
    except Exception as _e:
        raise gr.Error(f"Couldn't read LoRA `{label}`: {_e}")
    _sd, _fmt = _convert_lora(_raw)
    if not _sd:
        raise gr.Error(f"`{label}` doesn't look like a MiniMax-H3 LoRA (no lora_A/lora_down/lokr keys found).")
    _log(f"custom LoRA {label}: {len(_raw)} keys as {_fmt}")
    _pairs: dict[str, dict[str, torch.Tensor]] = {}
    for _k, _v in _sd.items():
        if ".lora_A." in _k:
            _pairs.setdefault(_k.split(".lora_A.")[0], {})["A"] = _v
        elif ".lora_B." in _k:
            _pairs.setdefault(_k.split(".lora_B.")[0], {})["B"] = _v
    _loaded = []
    for _base, _pair in _pairs.items():
        if "A" not in _pair or "B" not in _pair:
            continue
        _parts = _base.split(".")
        _obj = _state.PIPE.transformer
        for _p in _parts:
            _obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p)
        _loaded.append((_obj, _pair["A"].to(_obj.weight.dtype), _pair["B"].to(_obj.weight.dtype)))
    if not _loaded:
        raise gr.Error(f"`{label}` converted to no usable modules — is it a MiniMax-H3 LoRA? {_detect_wrong_arch(_raw)}")
    del _sd, _pairs, _raw
    import gc as _gc
    _gc.collect()
    _CUSTOM_LoRA[key] = _loaded
    _log(f"custom LoRA {label} parsed ({len(_loaded)} modules in {time.time() - _t0:.1f}s)")

	"""LoRA fetch + load from Hugging Face (or any direct .safetensors path).

	Functions:
	- `_parse_hf_link(url)` → `(repo, file \| None)`
	- `_resolve_repo_file(repo)` → pick the single .safetensors of a repo link
	- `_resolve_repo_revision(repo)` → Hub tip SHA (immutable cache key)
	- `_lora_cache_path(repo, file, revision)` → bucket path, or None if not writable
	- `_fetch_lora_file(repo, file, revision)` → local .safetensors path (cached)
	- `_load_lora_from_cache(path)` → load any local .safetensors into _CUSTOM_LoRA
	- `_ensure_nsfw_lora_loaded(key)` → lazy fetch + convert a catalog NSFW LoRA
	- `_ensure_custom_lora_loaded(repo, file)` → same for a user-supplied HF LoRA
	- `_ensure_custom_lora_loaded_from_path(key, path, label)` → load an already-fetched file

	Three layers of `H3_MAX_LORA_BYTES` enforcement: HEAD pre-check, post-download
	`os.path.getsize`, and a running byte counter during the bucket-cache copy
	(defends against a Hub lying about Content-Length). Concurrent writers use
	`_LORA_CACHE_LOCK` + temp file + `os.replace`.
	"""
	from __future__ import annotations

	import json
	import os
	import re
	import struct
	import time

	import gradio as gr

	from config import NSFW_LORA_CONFIGS
	from lora.cache import _evict_lora_cache_for_quota
	from lora.formats import (
	_convert_lora,
	_detect_wrong_arch,
	)
	import state as _state
	from state import (
	_CUSTOM_LoRA,
	_H3_MAX_LORA_BYTES,
	_LORA_CACHE_DIR,
	_LORA_CACHE_LOCK,
	_NSFW_LoRA,
	_log,
	)


	def _parse_hf_link(url):
	"""Parse a Hugging Face repo or file link into (repo, file). `file` is None for a bare repo link."""
	url = (url or "").strip().split("?")[0].split("#")[0]
	if not url:
	raise gr.Error("Paste a Hugging Face or CivitAI link first.")
	if not url.startswith("http"):
	parts = [p for p in url.split("/") if p]
	if len(parts) >= 2:
	return parts[0] + "/" + parts[1], "/".join(parts[2:]) or None
	raise gr.Error(f"Couldn't parse `{url}` as a Hugging Face repo link.")
	if "huggingface.co" not in url:
	raise gr.Error("Only Hugging Face links are supported (huggingface.co).")
	parts = [p for p in url.split("huggingface.co", 1)[1].strip("/").split("/") if p]
	if parts and parts[0] in ("spaces", "models", "datasets"):
	parts = parts[1:] # https://huggingface.co/spaces/Owner/Repo → Owner/Repo
	if len(parts) < 2:
	raise gr.Error(f"Couldn't parse `{url}` as a Hugging Face repo link.")
	repo = parts[0] + "/" + parts[1]
	if len(parts) == 2:
	return repo, None
	if parts[2] in ("resolve", "tree", "blob"):
	if len(parts) == 4 and parts[2] in ("tree", "blob"):
	return repo, None # tree/main or blob/main with no subpath → treat as a repo link
	if len(parts) < 5:
	raise gr.Error(f"Couldn't parse a filename from `{url}`.")
	if parts[3] != "main":
	raise gr.Error("Only `main`-branch links are supported — use a `resolve/main/...` link.")
	return repo, "/".join(parts[4:])
	return repo, "/".join(parts[2:])


	def _resolve_repo_file(repo):
	"""Pick the single .safetensors file of a repo link; error if there isn't exactly one."""
	from huggingface_hub import list_repo_files
	try:
	files = [f for f in list_repo_files(repo) if f.endswith(".safetensors")]
	except Exception as _e:
	raise gr.Error(f"Couldn't list files in `{repo}`: {_e}")
	if len(files) == 1:
	return files[0]
	if not files:
	raise gr.Error(f"No .safetensors files found in `{repo}`.")
	raise gr.Error(f"`{repo}` has multiple .safetensors files — paste a direct file link instead: "
	+ ", ".join(files[:5]) + ("…" if len(files) > 5 else ""))


	def _resolve_repo_revision(repo: str) -> str:
	"""Resolve the current tip of repo to an immutable commit SHA.

	Falls back to the literal string 'main' only if the Hub is unreachable so
	cache keys remain stable within a process; callers must not treat 'main' as
	immutable across restarts.
	"""
	try:
	from huggingface_hub import HfApi
	info = HfApi().repo_info(repo, repo_type="model")
	sha = getattr(info, "sha", None) or getattr(info, "revision", None)
	if sha and isinstance(sha, str) and len(sha) >= 7:
	return sha
	except Exception as _e:
	_log(f"revision resolve failed for {repo}: {_e}", "WARN")
	return "main"


	def _lora_cache_path(repo, file, revision: str \| None = None):
	"""Bucket path for a (repo, revision, file) LoRA, or None if the bucket isn't usable.

	Cache identity is (repo, revision_sha, file) so a moving `main` cannot serve
	stale weights under the same path.
	"""
	try:
	os.makedirs(_LORA_CACHE_DIR, exist_ok=True)
	_probe = os.path.join(_LORA_CACHE_DIR, ".probe")
	with open(_probe, "w") as _f:
	_f.write("ok")
	os.remove(_probe)
	except Exception:
	return None
	rev = (revision or "main")[:40]
	_key = re.sub(r"[^A-Za-z0-9._-]", "_", f"{repo.replace('/', '__')}__rev_{rev}__{file}")
	return os.path.join(_LORA_CACHE_DIR, _key)


	def _fetch_lora_file(repo, file, revision: str \| None = None, max_bytes: int \| None = None):
	"""Local path to a LoRA .safetensors, preferring the bucket cache over a fresh download.

	Cache key includes the immutable Hub revision SHA when available.
	Enforces H3_MAX_LORA_BYTES (reject before/during download) and
	H3_MAX_LORA_CACHE_BYTES (evict unreferenced entries or reject).
	Concurrent writers use temp files + os.replace so consumers never see
	partial artifacts.

	``max_bytes`` overrides the user-LoRA size cap (used for the system turbo
	weights, which are larger than the custom-LoRA quota by design).
	"""
	from huggingface_hub import hf_hub_download, hf_hub_url
	import requests

	limit = int(max_bytes) if max_bytes is not None else _H3_MAX_LORA_BYTES
	rev = revision or _resolve_repo_revision(repo)
	with _LORA_CACHE_LOCK:
	_cached = _lora_cache_path(repo, file, revision=rev)
	if _cached and os.path.exists(_cached):
	try:
	with open(_cached, "rb") as _f:
	_n = struct.unpack("<Q", _f.read(8))[0]
	_header_keys = list(json.loads(_f.read(_n)).keys())
	_hint = _detect_wrong_arch(_header_keys)
	if _hint:
	_log(f"LoRA cache evict (incompatible): {repo}/{file}@{rev[:7]} — {_hint}", "WARN")
	try:
	os.remove(_cached)
	except OSError:
	pass
	raise gr.Error(f"`{file}` is not a MiniMax-H3 LoRA. {_hint}")
	size = os.path.getsize(_cached)
	if size > limit:
	try:
	os.remove(_cached)
	except OSError:
	pass
	raise gr.Error(
	f"`{file}` exceeds H3_MAX_LORA_BYTES "
	f"({size / 1e6:.0f} MB > {limit / 1e6:.0f} MB)."
	)
	_log(f"LoRA cache hit: {repo}/{file}@{rev[:7]} ({size / 1e6:.0f} MB)")
	return _cached
	except gr.Error:
	raise
	except Exception:
	_log(f"LoRA cache corrupt, re-downloading: {repo}/{file}@{rev[:7]}", "WARN")
	try:
	os.remove(_cached)
	except OSError:
	pass

	# Pre-check size via HEAD when Content-Length is available.
	try:
	url = hf_hub_url(repo_id=repo, filename=file, revision=rev if rev != "main" else "main")
	head = requests.head(url, allow_redirects=True, timeout=30)
	cl = head.headers.get("Content-Length")
	if cl is not None:
	expected = int(cl)
	if expected > limit:
	raise gr.Error(
	f"`{file}` is too large ({expected / 1e6:.0f} MB > "
	f"{limit / 1e6:.0f} MB limit)."
	)
	_evict_lora_cache_for_quota(expected)
	except gr.Error:
	raise
	except Exception as _e:
	_log(f"LoRA HEAD pre-check skipped ({_e})", "WARN")

	_t0 = time.time()
	# Download via hub; then enforce size on the resulting file before cache admit.
	_path = hf_hub_download(repo, file, revision=rev if rev != "main" else None)
	size = os.path.getsize(_path)
	_log(f"LoRA downloaded: {repo}/{file}@{rev[:7]} ({size / 1e6:.0f} MB in {time.time() - _t0:.0f}s)")
	if size > limit:
	try:
	os.remove(_path)
	except OSError:
	pass
	raise gr.Error(
	f"`{file}` exceeds H3_MAX_LORA_BYTES "
	f"({size / 1e6:.0f} MB > {limit / 1e6:.0f} MB)."
	)

	# Architecture gate before admitting to persistent cache.
	try:
	with open(_path, "rb") as _f:
	_n = struct.unpack("<Q", _f.read(8))[0]
	_header = json.loads(_f.read(_n))
	_hint = _detect_wrong_arch(list(_header.keys()))
	if _hint:
	_log(f"LoRA rejected (incompatible architecture): {repo}/{file} — {_hint}", "WARN")
	try:
	os.remove(_path)
	except OSError:
	pass
	raise gr.Error(f"`{file}` is not a MiniMax-H3 LoRA. {_hint}")
	except gr.Error:
	raise
	except Exception as _e:
	_log(f"LoRA header peek failed ({_e}); proceeding with copy", "WARN")

	if _cached:
	try:
	_evict_lora_cache_for_quota(size)
	_tmp = _cached + ".tmp"
	# Stream with a running byte counter so a lied Content-Length still cannot exceed the ceiling.
	written = 0
	with open(_path, "rb") as _src, open(_tmp, "wb") as _dst:
	while True:
	chunk = _src.read(1024 * 1024)
	if not chunk:
	break
	written += len(chunk)
	if written > limit:
	_dst.close()
	try:
	os.remove(_tmp)
	except OSError:
	pass
	raise gr.Error(
	f"`{file}` exceeded H3_MAX_LORA_BYTES while copying "
	f"({written / 1e6:.0f} MB)."
	)
	_dst.write(chunk)
	os.replace(_tmp, _cached)
	_log(f"LoRA cached to bucket: {repo}/{file}@{rev[:7]}")
	return _cached
	except gr.Error:
	raise
	except Exception as _e:
	_log(f"LoRA bucket cache write failed ({_e}); using HF cache", "WARN")
	try:
	if os.path.exists(_cached + ".tmp"):
	os.remove(_cached + ".tmp")
	except OSError:
	pass
	return _path


	def _ensure_nsfw_lora_loaded(key):
	"""Download + convert + store a lazy NSFW LoRA's A/B factors on CPU. No-op if already loaded."""
	if key in _NSFW_LoRA:
	return
	cfg = NSFW_LORA_CONFIGS[key]
	import torch
	from safetensors.torch import load_file as _load_lora
	_t0 = time.time()
	_log(f"loading NSFW LoRA {key} from {cfg['repo']} ...")
	# Catalog NSFW weights can exceed the user custom-LoRA quota (default 512 MB);
	# allow up to 1.5 GB for official catalog entries only.
	_path = _fetch_lora_file(cfg["repo"], cfg["file"], max_bytes=1536 * 1024 * 1024)
	_raw = _load_lora(_path)
	# The catalog's `format` is a hint, not a contract — a mislabeled entry converts under whatever
	# format its keys actually are, and only a file no converter recognizes becomes a user error.
	_sd, _fmt = _convert_lora(_raw, cfg.get("format", ""))
	if not _sd:
	raise gr.Error(
	f"`{cfg['file']}` converted to no usable modules — it isn't a MiniMax-H3 LoRA "
	f"(or its `format` is mislabeled). {_detect_wrong_arch(_raw)}"
	)
	_log(f"NSFW LoRA {key}: {len(_raw)} keys as {_fmt}")
	_pairs: dict[str, dict[str, torch.Tensor]] = {}
	for _k, _v in _sd.items():
	if ".lora_A." in _k:
	_pairs.setdefault(_k.split(".lora_A.")[0], {})["A"] = _v
	elif ".lora_B." in _k:
	_pairs.setdefault(_k.split(".lora_B.")[0], {})["B"] = _v
	_loaded = []
	for _base, _pair in _pairs.items():
	if "A" not in _pair or "B" not in _pair:
	continue
	_parts = _base.split(".")
	_obj = _state.PIPE.transformer
	for _p in _parts:
	_obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p)
	_loaded.append((_obj, _pair["A"].to(_obj.weight.dtype), _pair["B"].to(_obj.weight.dtype)))
	if not _loaded:
	raise gr.Error(
	f"`{cfg['file']}` has no complete LoRA module pairs — it isn't a MiniMax-H3 LoRA "
	f"(or its `format` is mislabeled). {_detect_wrong_arch(_raw)}"
	)
	del _sd, _pairs, _raw
	import gc as _gc; _gc.collect()
	_NSFW_LoRA[key] = _loaded
	_log(f"NSFW LoRA {key} parsed ({len(_loaded)} modules in {time.time() - _t0:.1f}s)")


	def _ensure_custom_lora_loaded(repo, file):
	"""Download + convert + store a user-supplied Hugging Face LoRA's A/B factors on CPU. No-op if already loaded."""
	key = (repo, file)
	if key in _CUSTOM_LoRA:
	return
	import torch
	from safetensors.torch import load_file as _load_lora
	_t0 = time.time()
	_log(f"loading custom LoRA {repo}/{file} ...")
	try:
	_path = _fetch_lora_file(repo, file)
	_raw = _load_lora(_path)
	except Exception as _e:
	raise gr.Error(f"Couldn't fetch custom LoRA `{repo}/{file}`: {_e}")
	_sd, _fmt = _convert_lora(_raw)
	if not _sd:
	raise gr.Error(f"`{repo}/{file}` doesn't look like a MiniMax-H3 LoRA (no lora_A/lora_down/lokr keys found).")
	_log(f"custom LoRA {repo}/{file}: {len(_raw)} keys as {_fmt}")
	_pairs: dict[str, dict[str, torch.Tensor]] = {}
	for _k, _v in _sd.items():
	if ".lora_A." in _k:
	_pairs.setdefault(_k.split(".lora_A.")[0], {})["A"] = _v
	elif ".lora_B." in _k:
	_pairs.setdefault(_k.split(".lora_B.")[0], {})["B"] = _v
	_loaded = []
	for _base, _pair in _pairs.items():
	if "A" not in _pair or "B" not in _pair:
	continue
	_parts = _base.split(".")
	_obj = _state.PIPE.transformer
	for _p in _parts:
	_obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p)
	_loaded.append((_obj, _pair["A"].to(_obj.weight.dtype), _pair["B"].to(_obj.weight.dtype)))
	if not _loaded:
	raise gr.Error(f"`{repo}/{file}` converted to no usable modules — is it a MiniMax-H3 LoRA? {_detect_wrong_arch(_raw)}")
	del _sd, _pairs, _raw
	import gc as _gc; _gc.collect()
	_CUSTOM_LoRA[key] = _loaded
	_log(f"custom LoRA {repo}/{file} parsed ({len(_loaded)} modules in {time.time() - _t0:.1f}s)")


	def _load_lora_from_cache(path):
	"""Load a cached .safetensors LoRA directly by file path. Returns (repo, file) key
	and populates ``_CUSTOM_LoRA`` so it can be merged on generate."""
	import torch
	from safetensors.torch import load_file as _load_sd
	_t0 = time.time()
	fname = os.path.basename(path)
	# Use the filename as a synthetic (repo, file) key.
	key = ("__cache__", fname)
	if key in _CUSTOM_LoRA:
	return key
	# Architecture gate: peek at header keys before loading full tensor data.
	try:
	with open(path, "rb") as _f:
	_n = struct.unpack("<Q", _f.read(8))[0]
	_header_keys = list(json.loads(_f.read(_n)).keys())
	_hint = _detect_wrong_arch(_header_keys)
	if _hint:
	_log(f"cached LoRA rejected (incompatible): {fname} — {_hint}", "WARN")
	try:
	os.remove(path)
	except OSError:
	pass
	raise gr.Error(f"`{fname}` is not a MiniMax-H3 LoRA. {_hint}")
	except gr.Error:
	raise
	except Exception:
	pass # header peek failed — fall through to full load
	try:
	_raw = _load_sd(path)
	except Exception as _e:
	raise gr.Error(f"Couldn't load cached LoRA `{fname}`: {_e}")
	# Detect format and convert.
	_sd, _fmt = _convert_lora(_raw)
	if not _sd:
	raise gr.Error(f"`{fname}` doesn't look like a MiniMax-H3 LoRA.")
	_log(f"cached LoRA {fname}: {len(_raw)} keys as {_fmt}")
	_pairs: dict[str, dict[str, torch.Tensor]] = {}
	for _k, _v in _sd.items():
	if ".lora_A." in _k:
	_pairs.setdefault(_k.split(".lora_A.")[0], {})["A"] = _v
	elif ".lora_B." in _k:
	_pairs.setdefault(_k.split(".lora_B.")[0], {})["B"] = _v
	_loaded = []
	for _base, _pair in _pairs.items():
	if "A" not in _pair or "B" not in _pair:
	continue
	_parts = _base.split(".")
	_obj = _state.PIPE.transformer
	for _p in _parts:
	_obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p)
	_loaded.append((_obj, _pair["A"].to(_obj.weight.dtype), _pair["B"].to(_obj.weight.dtype)))
	if not _loaded:
	raise gr.Error(f"`{fname}` converted to no usable modules — is it a MiniMax-H3 LoRA? {_detect_wrong_arch(_raw)}")
	del _sd, _pairs, _raw
	import gc as _gc; _gc.collect()
	_CUSTOM_LoRA[key] = _loaded
	_log(f"cached LoRA {fname} loaded ({len(_loaded)} modules in {time.time() - _t0:.1f}s)")
	return key


	def _ensure_custom_lora_loaded_from_path(key, path: str, label: str = ""):
	"""Load + convert a local .safetensors LoRA into ``_CUSTOM_LoRA[key]``. No-op if present.

	Used by the CivitAI / CivArchive flows after they've fetched the file.
	"""
	if key in _CUSTOM_LoRA:
	return
	import torch
	from safetensors.torch import load_file as _load_lora

	_t0 = time.time()
	label = label or f"{key[0]}/{key[1]}"
	_log(f"loading custom LoRA {label} from local path ...")
	try:
	_raw = _load_lora(path)
	except Exception as _e:
	raise gr.Error(f"Couldn't read LoRA `{label}`: {_e}")
	_sd, _fmt = _convert_lora(_raw)
	if not _sd:
	raise gr.Error(f"`{label}` doesn't look like a MiniMax-H3 LoRA (no lora_A/lora_down/lokr keys found).")
	_log(f"custom LoRA {label}: {len(_raw)} keys as {_fmt}")
	_pairs: dict[str, dict[str, torch.Tensor]] = {}
	for _k, _v in _sd.items():
	if ".lora_A." in _k:
	_pairs.setdefault(_k.split(".lora_A.")[0], {})["A"] = _v
	elif ".lora_B." in _k:
	_pairs.setdefault(_k.split(".lora_B.")[0], {})["B"] = _v
	_loaded = []
	for _base, _pair in _pairs.items():
	if "A" not in _pair or "B" not in _pair:
	continue
	_parts = _base.split(".")
	_obj = _state.PIPE.transformer
	for _p in _parts:
	_obj = _obj[int(_p)] if _p.isdigit() else getattr(_obj, _p)
	_loaded.append((_obj, _pair["A"].to(_obj.weight.dtype), _pair["B"].to(_obj.weight.dtype)))
	if not _loaded:
	raise gr.Error(f"`{label}` converted to no usable modules — is it a MiniMax-H3 LoRA? {_detect_wrong_arch(_raw)}")
	del _sd, _pairs, _raw
	import gc as _gc
	_gc.collect()
	_CUSTOM_LoRA[key] = _loaded
	_log(f"custom LoRA {label} parsed ({len(_loaded)} modules in {time.time() - _t0:.1f}s)")