"""Bounded, safetensors-only Hugging Face adapters for MiniMax-H3. Hub access and tensor validation run before reserving a GPU. Only model repositories on huggingface.co are accepted; arbitrary URLs, pickle weights and remote code are never loaded. Downloads are temporary; a bounded CPU LRU avoids repeat downloads. """ from __future__ import annotations from collections import OrderedDict import hashlib import json import math import os from pathlib import Path import re import shutil import tempfile import threading from urllib.parse import unquote, urlparse import torch from h3_lora import REFERENCE_MODULES MAX_ADAPTERS = 4 MAX_FILE_BYTES = int(os.environ.get("H3_CUSTOM_LORA_MAX_FILE_BYTES", str(2 * 1024**3))) MAX_STACK_BYTES = int(os.environ.get("H3_CUSTOM_LORA_MAX_STACK_BYTES", str(4 * 1024**3))) MAX_CACHE_BYTES = int(os.environ.get("H3_CUSTOM_LORA_CACHE_BYTES", str(2 * 1024**3))) DISK_RESERVE_BYTES = 256 * 1024**2 _CACHE: OrderedDict = OrderedDict() _CACHE_LOCK = threading.RLock() _PAIR = re.compile(r"^(.*)\.lora_([AB])(?:\.default)?\.weight$") _DOWN_UP_PAIR = re.compile(r"^(.*)\.lora_(down|up)\.weight$") _DIFFUSERS_DOWN_UP_PAIR = re.compile(r"^(.*)\.lora\.(down|up)\.weight$") _ALPHA = re.compile(r"^(.*)\.(?:alpha|lora_alpha)$") # Kohya flattens module paths with underscores, which are also part of real H3 identifiers. Decode only this # explicit H3 grammar; replacing every underscore with a dot corrupts out_proj, qkv_proj, token_refiner, etc. _REFERENCE_SUFFIXES = { "attn_qkv_proj": "attn.qkv_proj", "attn_out_proj": "attn.out_proj", "mlp_fc1": "mlp.fc1", "mlp_fc2": "mlp.fc2", "adaln_proj_linear": "adaln_proj.linear", } _DIFFUSERS_SUFFIXES = { "attn_to_q": "attn.to_q", "attn_to_k": "attn.to_k", "attn_to_v": "attn.to_v", "attn_to_out_0": "attn.to_out.0", "ff_net_0_proj": "ff.net.0.proj", "ff_net_2": "ff.net.2", "adaln_proj_linear": "adaln_proj.linear", } _FLAT_BLOCK = re.compile( r"^(token_refiner_refiner_blocks|token_refiner_blocks|transformer_blocks|blocks)_(0|[1-9][0-9]*)_(.+)$" ) _FLAT_DIRECT_MODULES = { name.replace(".", "_"): name for name in (*REFERENCE_MODULES, *REFERENCE_MODULES.values()) } def _safe_path(value: str, label: str) -> str: if not value or value.startswith(("/", "\\")) or "\\" in value or any(p in ("", ".", "..") for p in value.split("/")): raise ValueError(f"Invalid {label}: use a relative Hugging Face path") if any(ord(c) < 32 for c in value) or ":" in value: raise ValueError(f"Invalid {label}") return value def parse_source(source: str, filename: str = "", revision: str | None = None) -> dict: """Parse owner/repo or a Hub repo/blob/resolve URL without fetching it.""" source = str(source).strip() chosen_file = str(filename or "").strip() chosen_revision = str(revision or "").strip() if "://" in source: url = urlparse(source) if url.scheme != "https" or url.hostname != "huggingface.co" or url.username or url.password or url.port: raise ValueError("LoRA source must be an https://huggingface.co model URL or owner/repo") parts = [unquote(part) for part in url.path.strip("/").split("/")] if len(parts) < 2 or parts[0] in {"spaces", "datasets"}: raise ValueError("Choose a Hugging Face model repository, not a Space or dataset") repo = "/".join(parts[:2]) if len(parts) > 2: if len(parts) < 4 or parts[2] not in {"blob", "resolve", "tree"}: raise ValueError("Unsupported Hugging Face URL; use a repository or blob/resolve file link") link_revision = parts[3] if chosen_revision and chosen_revision != link_revision: raise ValueError("Revision conflicts with the Hugging Face URL") chosen_revision = link_revision if parts[2] != "tree": linked_file = "/".join(parts[4:]) if chosen_file and chosen_file != linked_file: raise ValueError("Filename conflicts with the Hugging Face URL") chosen_file = linked_file elif len(parts) > 4: raise ValueError("Use a repository root or a full safetensors file URL") else: repo = source if not re.fullmatch(r"[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+", repo) or ".." in repo: raise ValueError("LoRA repository must have the form owner/repo") chosen_revision = _safe_path(chosen_revision or "main", "revision") if chosen_file: chosen_file = _safe_path(chosen_file, "filename") if not chosen_file.lower().endswith(".safetensors"): raise ValueError("Only .safetensors LoRA files are supported") return {"repo_id": repo, "revision": chosen_revision, "filename": chosen_file} def _inspect(parsed: dict) -> tuple[dict, object]: from huggingface_hub import HfApi info = HfApi().model_info(parsed["repo_id"], revision=parsed["revision"], files_metadata=True) files = sorted(s.rfilename for s in info.siblings if s.rfilename.lower().endswith(".safetensors")) if not files: raise ValueError("This model repository has no .safetensors files") chosen = parsed["filename"] if chosen and chosen not in files: raise ValueError(f"File `{chosen}` is not in this repository. Available files: {', '.join(files[:30])}") if not chosen and len(files) == 1: chosen = files[0] return {**parsed, "filename": chosen, "files": files, "commit": info.sha}, info def inspect_source(source: str, filename: str = "", revision: str | None = None) -> dict: """Return repo, immutable commit, and candidate filenames; never download model tensors.""" return _inspect(parse_source(source, filename, revision))[0] def _number(value, label, *, low=None, high=None): if isinstance(value, bool): raise ValueError(f"{label} must be a finite number") try: result = float(value) except (TypeError, ValueError) as exc: raise ValueError(f"{label} must be a finite number") from exc if not math.isfinite(result) or (low is not None and result < low) or (high is not None and result > high): raise ValueError(f"{label} must be finite" + (f" and between {low} and {high}" if low is not None else "")) return result def _normalize_name(name): for prefix in ("base_model.model.", "model.diffusion_model.", "diffusion_model.", "transformer."): if name.startswith(prefix): name = name[len(prefix):] for prefix in ("lora_unet_", "lora_transformer_"): if not name.startswith(prefix): continue flat = name[len(prefix):] if flat.startswith("diffusion_model_"): flat = flat[len("diffusion_model_"):] match = _FLAT_BLOCK.fullmatch(flat) if match: container, index, suffix = match.groups() reference = container in ("blocks", "token_refiner_blocks") suffixes = _REFERENCE_SUFFIXES if reference else _DIFFUSERS_SUFFIXES # Token refiners have no timestep AdaLN projection. if suffix in suffixes and not (container.startswith("token_refiner_") and suffix == "adaln_proj_linear"): containers = { "blocks": "blocks", "token_refiner_blocks": "token_refiner.blocks", "transformer_blocks": "transformer_blocks", "token_refiner_refiner_blocks": "token_refiner.refiner_blocks", } return f"{containers[container]}.{index}.{suffixes[suffix]}" if flat in _FLAT_DIRECT_MODULES: return _FLAT_DIRECT_MODULES[flat] raise ValueError(f"Unsupported flattened H3 LoRA module `{name}`; its module path cannot be mapped unambiguously") return name def decode_adapter(tensors: dict, metadata: dict, adapter_config: dict, inner_dim: int) -> dict: """Decode H3 reference/diffusers A/B or down/up pairs; reject unknown or ambiguous tensors.""" from h3_lora import _larry_targets if adapter_config.get("use_dora") or adapter_config.get("use_rslora"): raise ValueError("DoRA and rank-stabilized LoRA formats are not supported") if adapter_config.get("modules_to_save") or adapter_config.get("bias", "none") != "none": raise ValueError("LoRA adapters containing saved modules or bias updates are not supported") pairs, alphas, origins, formats = {}, {}, {}, {} for key, value in tensors.items(): match = _PAIR.fullmatch(key) down_up_match = _DOWN_UP_PAIR.fullmatch(key) diffusers_match = _DIFFUSERS_DOWN_UP_PAIR.fullmatch(key) alpha_match = _ALPHA.fullmatch(key) pair_match = match or down_up_match or diffusers_match supported = pair_match or alpha_match if supported: name = _normalize_name(supported[1]) original = origins.setdefault(name, supported[1]) if original != supported[1]: raise ValueError(f"Ambiguous LoRA module aliases `{original}` and `{supported[1]}` map to `{name}`") if pair_match: part = match[2] if match else {"down": "A", "up": "B"}[pair_match[2]] pair_format = "A/B" if match else "Kohya down/up" if down_up_match else "Diffusers down/up" if formats.setdefault(name, pair_format) != pair_format: raise ValueError(f"Mixed LoRA pair formats for `{name}`: {formats[name]} and {pair_format}") if part in pairs.setdefault(name, {}): raise ValueError(f"Duplicate LoRA {part} tensor for `{name}`") pairs[name][part] = value elif alpha_match and value.numel() == 1: if name in alphas: raise ValueError(f"Duplicate LoRA alpha for `{name}`") alphas[name] = _number(value.item(), "LoRA alpha") else: raise ValueError( f"Unsupported LoRA tensor `{key}`; expected H3 reference/diffusers lora_A/lora_B " "or Kohya/diffusers lora_down/lora_up or lora.down/lora.up pairs" ) if not pairs: raise ValueError("LoRA file contains no supported A/B weight pairs") if alphas.keys() - pairs.keys(): raise ValueError("LoRA file contains alpha tensors without matching A/B pairs") global_alpha = adapter_config.get("lora_alpha", metadata.get("lora_alpha", metadata.get("ss_network_alpha"))) patterns = adapter_config.get("alpha_pattern", {}) if not isinstance(patterns, dict): raise ValueError("adapter_config alpha_pattern must be an object") entries, target_scales = [], {} for name, pair in sorted(pairs.items()): if set(pair) != {"A", "B"}: raise ValueError(f"Missing LoRA A/B partner for `{name}`") a, b = pair["A"], pair["B"] if a.ndim != 2 or b.ndim != 2 or a.shape[0] <= 0 or a.shape[0] != b.shape[1]: raise ValueError(f"Invalid LoRA matrix shapes for `{name}`") if not a.is_floating_point() or not b.is_floating_point() or not torch.isfinite(a).all() or not torch.isfinite(b).all(): raise ValueError(f"LoRA weights for `{name}` must be finite floating-point tensors") rank = a.shape[0] matches = [value for key, value in patterns.items() if name == key or name.endswith("." + key)] if len(matches) > 1: raise ValueError(f"Ambiguous alpha_pattern for `{name}`") alpha = alphas.get(name, matches[0] if matches else global_alpha) scale = 1.0 if alpha is None else _number(alpha, "LoRA alpha") / rank is_reference = name in REFERENCE_MODULES or name.startswith(("blocks.", "token_refiner.blocks.", "final_layer.")) if is_reference: if name.endswith(".attn.qkv_proj") and b.shape[0] != 3 * inner_dim: raise ValueError(f"Reference QKV LoRA `{name}` must contain exactly three attention projections") if name.endswith(".mlp.fc1") and b.shape[0] % 2: raise ValueError(f"Reference SwiGLU LoRA `{name}` must have an even row count") targets = _larry_targets(name, b, inner_dim) else: targets = [(name + ".weight", b)] for key, b_part in targets: if key in target_scales: raise ValueError(f"Multiple tensors map to the same LoRA target `{key}`") entries.append((key, a, b_part)) target_scales[key] = scale return {"entries": entries, "target_scales": target_scales} def _tensor_bytes(spec): # Reference QKV projections share A; count each storage only once. storages = {tensor.untyped_storage().data_ptr(): tensor.untyped_storage().nbytes() for _, a, b in spec["entries"] for tensor in (a, b)} return sum(storages.values()) def _load(source: dict, info, inner_dim: int) -> dict: from huggingface_hub import hf_hub_download from safetensors import safe_open siblings = {s.rfilename: s for s in info.siblings} file_info = siblings[source["filename"]] size = file_info.size if size is None or size <= 0 or size > MAX_FILE_BYTES: raise ValueError(f"LoRA file size must be known and at most {MAX_FILE_BYTES / 1024**3:g} GiB") identity = (source["repo_id"], source["commit"], source["filename"], inner_dim) with _CACHE_LOCK: if identity in _CACHE: _CACHE.move_to_end(identity) return _CACHE[identity] # Serialize downloads so concurrent requests cannot bypass the transient file budget. free = shutil.disk_usage(tempfile.gettempdir()).free required = size + DISK_RESERVE_BYTES if free < required: raise ValueError( f"Not enough temporary disk space for this LoRA: need {required / 1024**3:.2f} GiB " f"including a 256 MiB reserve, but only {free / 1024**3:.2f} GiB is free" ) with tempfile.TemporaryDirectory(prefix="h3-custom-lora-") as directory: config = {} config_name = str(Path(source["filename"]).parent / "adapter_config.json").replace("\\", "/") if config_name not in siblings and "adapter_config.json" in siblings: config_name = "adapter_config.json" config_file = siblings.get(config_name) if config_file: if config_file.size is None or not 0 < config_file.size <= 1024**2: raise ValueError("adapter_config.json exceeds the 1 MiB configuration limit") config_path = hf_hub_download(source["repo_id"], config_name, revision=source["commit"], local_dir=directory) if Path(config_path).stat().st_size > 1024**2: raise ValueError("adapter_config.json exceeds the 1 MiB configuration limit") config = json.loads(Path(config_path).read_text(encoding="utf-8")) if not isinstance(config, dict): raise ValueError("adapter_config.json must be a JSON object") path = hf_hub_download(source["repo_id"], source["filename"], revision=source["commit"], local_dir=directory) if Path(path).stat().st_size > MAX_FILE_BYTES: raise ValueError("Downloaded LoRA exceeds the configured file limit") with safe_open(path, framework="pt", device="cpu") as handle: metadata = handle.metadata() or {} # Own CPU storage permits removing the temporary download on Windows too. tensors = {key: handle.get_tensor(key).clone() for key in handle.keys()} spec = decode_adapter(tensors, metadata, config, inner_dim) spec["bytes"] = _tensor_bytes(spec) spec["id"] = hashlib.sha256(repr(identity).encode()).hexdigest() spec["label"] = f"{source['repo_id']}/{source['filename']}@{source['commit'][:12]}" spec["display_name"] = source["filename"].rsplit("/", 1)[-1] if spec["bytes"] <= MAX_CACHE_BYTES: while _CACHE and sum(item["bytes"] for item in _CACHE.values()) + spec["bytes"] > MAX_CACHE_BYTES: _CACHE.popitem(last=False) _CACHE[identity] = spec return spec def prepare_adapters(config, transformer) -> list[dict]: """Resolve and validate a JSON/list ordinary-LoRA stack without mutating model weights.""" from h3_lora import _prepare_sets if config is None or config == "": return [] if isinstance(config, str): if len(config) > 32768: raise ValueError("LoRA configuration is too large") try: config = json.loads(config) except json.JSONDecodeError as exc: raise ValueError("LoRA configuration must be a JSON list") from exc if not isinstance(config, list) or len(config) > MAX_ADAPTERS: raise ValueError(f"Provide a list with at most {MAX_ADAPTERS} ordinary LoRAs") if not config: return [] inner_dim = transformer.config.num_attention_heads * transformer.config.attention_head_dim specs, seen, total_bytes = [], set(), 0 for item in config: if not isinstance(item, dict): raise ValueError("Each LoRA must be an object containing source and strength") if not isinstance(item.get("enabled", True), bool): raise ValueError("LoRA enabled must be true or false") strength = _number(item.get("strength", 1.0), "LoRA strength", low=-2, high=2) if not item.get("enabled", True) or strength == 0: continue source, info = _inspect(parse_source(item.get("source", ""), item.get("filename", ""), item.get("revision"))) if not source["filename"]: raise ValueError("This repository contains multiple .safetensors files; choose a filename: " + ", ".join(source["files"][:30])) if (source["repo_id"], source["commit"], source["filename"]) in seen: raise ValueError("The same LoRA file was selected more than once; adjust its strength instead") seen.add((source["repo_id"], source["commit"], source["filename"])) size = next(s.size for s in info.siblings if s.rfilename == source["filename"]) if size is None or total_bytes + size > MAX_STACK_BYTES: raise ValueError(f"Combined ordinary LoRAs exceed the {MAX_STACK_BYTES / 1024**3:g} GiB limit") decoded = _load(source, info, inner_dim) spec = {**decoded, "scale": strength} total_bytes += spec["bytes"] if total_bytes > MAX_STACK_BYTES: raise ValueError(f"Combined ordinary LoRAs exceed the {MAX_STACK_BYTES / 1024**3:g} GiB tensor limit") specs.append(spec) _prepare_sets({f"custom:{i}": spec for i, spec in enumerate(specs)}, dict(transformer.named_parameters())) return specs