| """Private, content-addressed MOSS transformer LoRAs. No executable model code. |
| |
| Only standard A/B (up/down) linear LoRA on the 42 MOSS DiT blocks is supported. |
| Weights merge from immutable host originals; no merged checkpoint is written. |
| """ |
| from __future__ import annotations |
| import hashlib |
| import json |
| import math |
| import os |
| from pathlib import Path |
| import re |
| import threading |
|
|
| MAX_BYTES = 512 * 1024 * 1024 |
| MAX_TOTAL_BYTES = 5 * 1024 * 1024 * 1024 |
| MAX_MODULES = 4 |
| ID = re.compile(r"^[a-f0-9]{64}$") |
| STORE_LOCK = threading.RLock() |
|
|
|
|
| def owner_hash(token): |
| if not isinstance(token, str) or not re.fullmatch(r"[A-Za-z0-9_-]{43,128}", token): |
| raise ValueError("A private LoRA client token is required") |
| return hashlib.sha256(token.encode()).hexdigest() |
|
|
|
|
| def normalize_config(raw): |
| config = json.loads(raw) |
| if not isinstance(config, dict) or set(config) != {"alpha", "base_model"}: |
| raise ValueError("Invalid LoRA configuration") |
| alpha = config["alpha"] |
| if alpha is not None and (isinstance(alpha, bool) or not isinstance(alpha, (int, float)) or not math.isfinite(alpha) or not 0 < alpha <= 4096): |
| raise ValueError("Network alpha must be positive and at most 4096") |
| model = config["base_model"] |
| if model is not None and (not isinstance(model, str) or len(model) > 200): |
| raise ValueError("Invalid adapter base model") |
| return config |
|
|
|
|
| def target_key(name): |
| name = re.sub(r"^(?:base_model\.model\.|transformer\.|diffusion_model\.)", "", name) |
| name = re.sub(r"^lora_unet_layers_(\d+)_proj_(q|k|v|o|ffn_up_gate|ffn_down)$", r"layers.\1.proj.\2", name) |
| match = re.fullmatch(r"layers\.(\d+)\.proj\.(q|k|v|o|ffn_up_gate|ffn_down)", name) |
| if not match or int(match[1]) >= 42: |
| raise ValueError("Unsupported tensor target: use a MOSS DiT linear LoRA (not SDXL, FLUX, LoKr or DoRA)") |
| return name + ".weight", match[2] |
|
|
|
|
| def inspect_adapter(file, config): |
| from safetensors import SafetensorError |
| try: |
| return _inspect_adapter(file, config) |
| except SafetensorError as error: |
| raise ValueError("Invalid or incomplete safetensors adapter") from error |
|
|
|
|
| def _inspect_adapter(file, config): |
| from safetensors import safe_open |
| pairs = {} |
| with safe_open(str(file), framework="pt", device="cpu") as sf: |
| metadata = sf.metadata() or {} |
| for key in sf.keys(): |
| match = re.fullmatch(r"(.+)\.(lora_A(?:\.default)?\.weight|lora_B(?:\.default)?\.weight|lora_down\.weight|lora_up\.weight|alpha)", key) |
| if not match: |
| raise ValueError("Unsupported adapter tensors; only standard linear LoRA is supported") |
| target, kind = target_key(match[1]) |
| group = pairs.setdefault(target, {"kind": kind}) |
| suffix = match[2] |
| role = "alpha" if suffix == "alpha" else "a" if suffix.startswith(("lora_A", "lora_down")) else "b" |
| if role in group: |
| raise ValueError("Duplicate LoRA tensor target") |
| shape = list(sf.get_slice(key).get_shape()) |
| if role == "alpha": |
| if shape not in ([], [1]): |
| raise ValueError("Invalid per-layer network alpha") |
| alpha = float(sf.get_tensor(key).item()) |
| if not math.isfinite(alpha) or not 0 < alpha <= 4096: |
| raise ValueError("Invalid per-layer network alpha") |
| group[role] = alpha |
| else: |
| if sf.get_slice(key).get_dtype() not in ("F16", "BF16", "F32") or len(shape) != 2: |
| raise ValueError("LoRA matrices must be float16, bfloat16 or float32") |
| group[role] = key |
| group[role + "_shape"] = shape |
| if not pairs: |
| raise ValueError("No LoRA matrices found") |
| for target, group in pairs.items(): |
| if "a" not in group or "b" not in group: |
| raise ValueError("Incomplete LoRA matrix pair") |
| a, b = group["a_shape"], group["b_shape"] |
| expected = [16384, 3072] if group["kind"] == "ffn_up_gate" else [3072, 8192] if group["kind"] == "ffn_down" else [3072, 3072] |
| if a[0] != b[1] or not 1 <= a[0] <= 512 or [b[0], a[1]] != expected: |
| raise ValueError("Adapter tensor dimensions do not match MOSS-Image") |
| alpha = group.get("alpha", config.get("alpha")) |
| if alpha is None and metadata.get("ss_network_alpha") is not None: |
| alpha = float(metadata["ss_network_alpha"]) |
| if alpha is None: |
| alpha = a[0] |
| if not math.isfinite(alpha) or not 0 < alpha <= 4096: |
| raise ValueError("Invalid network alpha") |
| group["factor"] = alpha / a[0] |
| return pairs |
|
|
|
|
| def validate_selection(selection): |
| if not isinstance(selection, list) or len(selection) > MAX_MODULES: |
| raise ValueError("Choose at most four LoRA modules") |
| seen = set() |
| for item in selection: |
| if not isinstance(item, dict) or set(item) != {"id", "strength"} or not ID.fullmatch(str(item.get("id", ""))): |
| raise ValueError("Invalid LoRA module selection") |
| strength = item["strength"] |
| if isinstance(strength, bool) or not isinstance(strength, (int, float)) or not math.isfinite(strength) or not -2 <= strength <= 2: |
| raise ValueError("LoRA strength must be between -2 and 2") |
| if item["id"] in seen: |
| raise ValueError("Duplicate LoRA module") |
| seen.add(item["id"]) |
| return selection |
|
|
|
|
| class AdapterStore: |
| def __init__(self, root, model): |
| self.root = Path(root) |
| self.model = model |
|
|
| def paths(self, owner, identifier): |
| if not ID.fullmatch(owner) or not ID.fullmatch(identifier): |
| raise ValueError("Invalid LoRA identifier") |
| return self.root / owner / (identifier + ".safetensors"), self.root / owner / (identifier + ".json") |
|
|
| def compatible(self, config): |
| declared = config.get("base_model") |
| |
| if declared and declared != self.model and not (declared == "OpenMOSS-Team/MOSS-Image-Base" and self.model == "OpenMOSS-Team/MOSS-Image-SFT"): |
| raise ValueError("Adapter declares a different base model from the serving model") |
|
|
| def install(self, temporary, owner, identifier, config_raw, name): |
| config = normalize_config(config_raw) |
| self.compatible(config) |
| digest = hashlib.sha256() |
| with open(temporary, "rb") as handle: |
| while chunk := handle.read(1024 * 1024): |
| digest.update(chunk) |
| calculated = hashlib.sha256((digest.hexdigest() + "\n" + config_raw).encode()).hexdigest() |
| if calculated != identifier: |
| raise ValueError("LoRA checksum mismatch") |
| pairs = inspect_adapter(temporary, config) |
| file, manifest = self.paths(owner, identifier) |
| record = {"id": identifier, "name": name[:160], "bytes": Path(temporary).stat().st_size, |
| "sha256": digest.hexdigest(), "config": config, "config_raw": config_raw, "targets": len(pairs)} |
| with STORE_LOCK: |
| current = sum(p.stat().st_size for p in self.root.glob("*/*.safetensors")) |
| if not file.exists() and current + record["bytes"] > MAX_TOTAL_BYTES: |
| raise ValueError("Server LoRA cache is full; remove unused server modules") |
| file.parent.mkdir(parents=True, exist_ok=True, mode=0o700) |
| os.replace(temporary, file) |
| file.chmod(0o600) |
| tmp = manifest.with_suffix(".json.tmp") |
| with open(tmp, "w") as handle: |
| json.dump(record, handle) |
| handle.flush() |
| os.fsync(handle.fileno()) |
| tmp.chmod(0o600) |
| os.replace(tmp, manifest) |
| return record |
|
|
| def get(self, owner, identifier): |
| file, manifest = self.paths(owner, identifier) |
| if not file.is_file() or not manifest.is_file(): |
| raise FileNotFoundError("LoRA module not uploaded") |
| record = json.loads(manifest.read_text()) |
| self.compatible(record["config"]) |
| return record |
|
|
| def resolve(self, owner, selection): |
| validate_selection(selection) |
| resolved = [] |
| for item in selection: |
| record = self.get(owner, item["id"]) |
| file, _ = self.paths(owner, item["id"]) |
| resolved.append({**record, "file": str(file), "strength": item["strength"]}) |
| return resolved |
|
|
|
|
| def merge_weights(base, adapters): |
| import torch |
| from safetensors import safe_open |
| merged = dict(base) |
| |
| targets = {} |
| for item in adapters: |
| if item["strength"] == 0: |
| continue |
| digest = hashlib.sha256() |
| with open(item["file"], "rb") as handle: |
| while chunk := handle.read(1024 * 1024): |
| digest.update(chunk) |
| if digest.hexdigest() != item["sha256"]: |
| raise ValueError("Cached LoRA checksum mismatch") |
| for target, pair in inspect_adapter(item["file"], item["config"]).items(): |
| if target not in base: |
| raise ValueError("Adapter target is absent from the loaded model") |
| targets.setdefault(target, []).append((item, pair)) |
| for target, terms in targets.items(): |
| weight = base[target].float().clone() |
| for item, pair in terms: |
| with safe_open(item["file"], framework="pt", device="cpu") as sf: |
| a, b = sf.get_tensor(pair["a"]).float(), sf.get_tensor(pair["b"]).float() |
| if not torch.isfinite(a).all() or not torch.isfinite(b).all(): |
| raise ValueError("LoRA contains non-finite weights") |
| weight.addmm_(b, a, beta=1, alpha=item["strength"] * pair["factor"]) |
| if not torch.isfinite(weight).all(): |
| raise ValueError("LoRA produces non-finite model weights") |
| merged[target] = weight.to(base[target].dtype) |
| return merged |
|
|
|
|
| def apply_adapters(pipe, adapters): |
| fingerprint = tuple((item["id"], item["strength"]) for item in adapters if item["strength"] != 0) |
| if getattr(pipe, "_lora_fingerprint", ()) == fingerprint: |
| return |
| if not hasattr(pipe, "_lora_original_weights"): |
| pipe._lora_original_weights = pipe._load_dit_weights() |
| merged = merge_weights(pipe._lora_original_weights, adapters) |
| if pipe._dit_model is not None: |
| from ttnn_moss.pipeline import _free_ttnn_module |
| _free_ttnn_module(pipe._dit_model) |
| pipe._dit_model = None |
| pipe._dit_weights = merged |
| pipe._lora_fingerprint = fingerprint |
|
|
