"""Model sources: naming heuristics, HF metadata, source objects.""" from pathlib import Path import json import re import struct import urllib.error import urllib.request from urllib.parse import urlparse from vlib import ctx from vlib.ui import _fail, _note, _step from vlib.net import _fetch_bytes_progress, _hf_auth_headers, http_get_json, http_range from vlib.tensors import _gguf_raw_to_f32, _gguf_tensor_f32, _is_gguf_file, _parse_gguf_header, decode_to_f32 OUTPUT_ALIASES = ( "lm_head.weight", "output.weight", "transformer.output_layer.weight", "output_layer.weight", "embed_out.weight", ) EMBEDDING_ALIASES = ( "model.embed_tokens.weight", "embed_tokens.weight", "transformer.embedding.word_embeddings.weight", "wte.weight", "word_embeddings.weight", # GGUF tied checkpoints store the shared head under this name (e.g. Gemma # quants ship only token_embd.weight). Kept here (embed role), not in # OUTPUT_ALIASES, so tie detection reads structure, not aliases. "token_embd.weight", ) _EMBED_SUFFIXES = ("embed_tokens.weight", "word_embeddings.weight", "wte.weight", "token_embd.weight") _HEAD_PATTERNS = ("lm_head", "output_layer", "embed_out") def role_of(name): """Canonical role of one tensor name: 'head', 'embed', or 'body'. Single home for the head/embed duality; cast's BODY machinery (_ROLE_ALIASES, _same_role) is untouched and out of scope.""" if not name: return "body" if name in OUTPUT_ALIASES or name.endswith(tuple(OUTPUT_ALIASES)): return "head" if name in EMBEDDING_ALIASES or name.endswith(_EMBED_SUFFIXES): return "embed" if name.endswith("weight") and any(a in name for a in _HEAD_PATTERNS): return "head" return "body" def is_tied(names): """Structural tie check (no config, no downloads): an embed tensor is present and no head tensor exists. The embed then serves as both input and output projection.""" if not names: return False names = list(names) return any(role_of(n) == "embed" for n in names) and not any(role_of(n) == "head" for n in names) def find_head(names): """Head-role tensor in `names`: exact alias first, pattern second.""" for a in OUTPUT_ALIASES: if a in names: return a for n in names: if role_of(n) == "head": return n return None def find_embed(names): """Embed-role tensor in `names`: exact alias first, suffix second.""" for a in EMBEDDING_ALIASES: if a in names: return a for n in names: if role_of(n) == "embed": return n return None def pair_role(vname, vshape, b_names, bshape_of): """Pair one voice tensor by (role, shape). Returns the base name or None. Same-role alias wins; a head with no base head falls back to the tied base embed. Shapes must match exactly; cross-name pairings announce themselves, exact matches stay silent. None means: fail loud upstream.""" from vlib.ui import _fail, _step if role_of(vname) not in ("head", "embed"): return None cands = [n for n in b_names if role_of(n) == role_of(vname)] tied = False if not cands and role_of(vname) == "head" and is_tied(b_names): cands = [n for n in b_names if role_of(n) == "embed"] tied = bool(cands) if len(cands) != 1: return None b = cands[0] if tuple(vshape) != tuple(bshape_of(b)): _fail(f" ✗ Voice tensor '{vname}' {list(vshape)} doesn't fit base '{b}' {list(tuple(bshape_of(b)))}.") if b == vname: return b if tied: _step("Base model ties its head — using its embedding as the match.") else: _step(f"Using base '{b}' for voice '{vname}'.") return b def _family(arch): a = (arch or "").lower() if "glm" in a or "chatglm" in a: return "GLM" if "gemma4" in a: return "Gemma 4" if "gemma" in a: return "Gemma" if "qwen" in a: return "Qwen" if "llama" in a: return "Llama" if "gpt2" in a or "gptneo" in a: return "GPT" if "opt" in a: return "OPT" return arch or "unknown" def _nested(cfg, *keys): for k in keys: if isinstance(cfg.get(k), (int, str, list)): return cfg[k] tc = cfg.get("text_config") or {} for k in keys: if isinstance(tc.get(k), (int, str, list)): return tc[k] return None def detect_output_tensor(names, config): if not names: return None for alias in OUTPUT_ALIASES: if alias in names: return alias for n in names: if any(a in n for a in ("lm_head", "output_layer", "embed_out")) and n.endswith("weight"): return n tie = config.get("tie_word_embeddings") if config else None if tie is True: for alias in EMBEDDING_ALIASES: if alias in names: return alias # any embedding-like tensor (tied layouts, incl. transformer.wte.weight # and GGUF token_embd.weight) for n in names: if n.endswith(("embed_tokens.weight", "word_embeddings.weight", "wte.weight", "token_embd.weight")): return n return None def largest_2d(names, shapes): best, best_elems = None, 0 for n in names: s = shapes.get(n) if s is not None and len(s) == 2: elems = int(s[0]) * int(s[1]) if elems > best_elems: best, best_elems = n, elems return best def _hf_config(model_id): try: return http_get_json(f"https://huggingface.co/{model_id}/resolve/main/config.json", timeout=30) except Exception: return None def _hf_license(model_id): for url in (f"https://huggingface.co/{model_id}/resolve/main/LICENSE", f"https://huggingface.co/{model_id}/raw/main/LICENSE"): try: req = urllib.request.Request(url, headers=_hf_auth_headers(url)) with urllib.request.urlopen(req, timeout=15) as resp: text = resp.read(2000).decode("utf-8", "replace").strip() if len(text) > 10: return text[:200] except Exception: continue return None def _hf_model_index(model_id, rev="main"): try: return http_get_json(f"https://huggingface.co/{model_id}/resolve/{rev}/model.safetensors.index.json", timeout=20) except Exception: return None _GGUF_SPLIT_RE = re.compile(r"^(?P.+)-(?P\d+)-of-(?P\d+)\.gguf$", re.IGNORECASE) def _gguf_split_key(basename): """Split-shard grouping key, or None for single files. 'model-00001-of-00002.gguf' -> ('model', 2). Total==1 counts as single.""" m = _GGUF_SPLIT_RE.match(basename) if not m or int(m.group("total")) <= 1: return None return (m.group("stem").lower(), int(m.group("total"))) def _gguf_dropped(basename): """Certain non-models, by name alone. Everything else faces the tensor census.""" n = basename.lower() return n.startswith("mmproj") or n.endswith("-imatrix.gguf") def _gguf_name_rank(path): """Fidelity-hint rank for probe ordering (cost only, never correctness). Higher wins. Suspicious-but-possible names demote slightly; the tensor census still gets the final vote.""" n = path.lower() score = 0 for i, key in enumerate(("fp32", "f32", "bf16", "fp16", "f16", "q8_0", "q6_k", "q5_k", "q5_0", "q5_1", "q4_k", "q4_0", "q4_1", "q3_k")): if key in n: score = 100 - i break else: if "iq" in n: score = 84 for bad in ("mtp", "vision", "encoder", "lora", "draft", "instruct-"): if bad in n: score -= 3 return score def _triage_gguf_files(files): """Split a repo .gguf listing into (groups, singles, dropped). Pure function. groups: shard path-lists of one model each (shard order). singles: ranked best-first. Incomplete shard groups are kept; the census/fetch stages fail loud on truly missing tensors.""" dropped = sorted(f for f in files if _gguf_dropped(f.split("/")[-1])) dropped_set = set(dropped) buckets, singles = {}, [] for f in sorted(set(files) - dropped_set): parts = f.split("/") key = _gguf_split_key(parts[-1]) if key is None: singles.append(f) else: buckets.setdefault((tuple(parts[:-1]),) + key, []).append(f) groups = [sorted(g) for g in buckets.values()] groups.sort(key=lambda g: _gguf_name_rank(g[0]), reverse=True) singles.sort(key=_gguf_name_rank, reverse=True) return groups, singles, dropped _GGUF_DTYPE_RANK = { "F32": 7, "FP32": 7, "BF16": 6, "BF16E8M": 6, "F16": 5, "FP16": 5, "Q8_0": 4, "Q6_K": 3, "Q5_0": 2, "Q5_1": 2, "Q5_K": 2, "Q4_0": 1, "Q4_1": 1, "Q4_K": 1, } def _gguf_dtype_rank(dtype_name): """Fidelity rank for a tensor dtype name. Unknown/exotic quants rank last.""" if not dtype_name: return -1 n = str(dtype_name).upper() if n in _GGUF_DTYPE_RANK: return _GGUF_DTYPE_RANK[n] if n.startswith(("Q3", "Q2", "IQ")): return 0 return -1 def _gguf_unit_census(names): """Tensor census for one candidate unit (single file or shard union). Returns (head_name_or_None, has_body). A unit is a base model only when it has BOTH a fetchable head (true head, or tied embed carrier) and transformer block markers — projector, imatrix, draft-head and adapter files fail one side or the other.""" head = find_head(names) if head is None and is_tied(names): head = find_embed(names) body = any(n.startswith("blk.") for n in names) return head, body def _normalize_hf_repo(s): if not s or not s.strip(): return None s = s.strip().rstrip("/") if s.startswith(("http://", "https://")): try: parts = urlparse(s).path.strip("/").split("/") if len(parts) >= 2 and parts[0] and parts[1]: return f"{parts[0]}/{parts[1]}" except Exception: return None return None s = s.split("@")[0].split(":")[0].strip() parts = s.split("/") if len(parts) >= 2 and parts[0] and parts[1]: return f"{parts[0]}/{parts[1]}" return None def _base_paths(base_id): norm = _normalize_hf_repo(base_id) if norm is None: return None org, model = (x.lower() for x in norm.split("/")) file_path = ctx.VOICES_DIR / "base" / org / f"{model}.safetensors" return file_path, ctx.VOICES_DIR / "base" / org / f"{model}.json" def _hf_api_base_model(source_id): try: data = http_get_json(f"https://huggingface.co/api/models/{source_id}", timeout=15) card = data.get("cardData") or {} bm = card.get("base_model") or data.get("base_model") if bm is None: return None if isinstance(bm, str): return [bm] if isinstance(bm, list): return [x for x in bm if isinstance(x, str) and x.strip()] except Exception: return None return None def _hf_mergekit_base(source_id): for tmpl in ( f"https://huggingface.co/{source_id}/raw/main/mergekit_config.yml", f"https://huggingface.co/{source_id}/raw/main/mergekit_config.yaml", ): try: req = urllib.request.Request(tmpl, headers=_hf_auth_headers(tmpl)) with urllib.request.urlopen(req, timeout=15) as resp: text = resp.read().decode("utf-8", "replace") for line in text.splitlines(): line = line.strip() if line.startswith("#"): continue m = re.search(r'base_model\s*:\s*["\']?\s*([A-Za-z0-9._-]+/[A-Za-z0-9._-]+)', line) if m: return m.group(1).strip() except Exception: continue return None def _detect_base(source_id): raw = _hf_api_base_model(source_id) if not raw: return None, False, 0, [] if len(raw) == 1: return raw[0], False, 1, raw tried = _hf_mergekit_base(source_id) if tried: return tried, True, len(raw), raw return raw[0], True, len(raw), raw class TensorRef: __slots__ = ("name", "dtype", "shape") def __init__(self, name, dtype, shape): self.name = name self.dtype = dtype self.shape = tuple(int(x) for x in shape) def __repr__(self): return f"TensorRef({self.name!r}, {self.dtype}, {self.shape})" class SafetensorsSource: kind = "safetensors" def __init__(self, path=None, repo=None, url=None, rev=None): self.path = str(path) if path else None self.repo = repo self.url = url self.rev = rev or "main" self._header = None self._header_size = None self._index = None # weight_map for sharded repos self._shard_headers = {} # shard file -> (header_size, tensors) def _local(self): """(header_size, tensors) for a local file — read once, then cached.""" if self.path and self._header_size is None: with open(self.path, "rb") as f: self._header_size = struct.unpack("= _GGUF_HIGH_TIER: break if not probed: chosen = _pick_gguf(files) url = f"https://huggingface.co/{repo}/resolve/{rev}/{chosen}" try: parsed = _parse_gguf_header(url) head, body = _gguf_unit_census(list(parsed[2].keys())) except SystemExit: raise except Exception: head, body = None, False if head and body: return GgufSource(url=url, rev=rev) _fail(f" ✗ No base model found in '{repo}' ({checked} file(s) checked, none with an output head).") probed.sort(key=lambda t: (t[0], t[1]), reverse=True) best = probed[0][2] if len(best) == 1: _step(f"Using {best[0].rsplit('/', 1)[-1]}") else: _step(f"Using {len(best)} files ({best[0].rsplit('/', 1)[-1]}…)") if len(best) == 1: return GgufSource(url=best[0], rev=rev) return GgufSource(urls=best, rev=rev) def _source_from_repo(repo, rev=None): rev = rev or "main" try: tree = http_get_json(f"https://huggingface.co/api/models/{repo}/tree/{rev}", timeout=20) files = [f.get("path", "") for f in tree if f.get("path")] ggufs = [f for f in files if f.endswith(".gguf")] if ggufs: return _gguf_source_for_repo(repo, rev, ggufs) if any(f.endswith(".safetensors") for f in files): return SafetensorsSource(repo=repo, rev=rev) if files: _fail(f" ✗ '{repo}' has no .safetensors or .gguf files (found: {', '.join(files[:6])}{'…' if len(files) > 6 else ''}).") except SystemExit: raise except Exception: pass return SafetensorsSource(repo=repo, rev=rev) def _pick_gguf(files): """Legacy single-file fallback (pre-tensor-first behavior). Kept for the nothing-resolvable path and external imports; new code goes through _gguf_source_for_repo.""" files = sorted(files) for key in ("fp16", "f16", "q8_0"): for f in files: if key in f.lower(): return f return files[0] def _sanitize_voice_name(name): if not name or not name.strip(): return None name = name.strip() if name.lower().startswith("delta-") or name in (".", ".."): return None bad = set("/\\.: ") for c in name: if c in bad or ord(c) < 32: return None return name _NAME_RULES = "Use letters, numbers, hyphens, underscores (no spaces, /, \\, dots, colons; not starting with delta-)." def _public_source(ref): """Share-safe form of a model reference for voice.json/reports. HF ids pass through; URLs lose userinfo/query/fragment; local paths collapse to basename (directories leak usernames).""" s = (ref or "").strip() if not s or s == "—": return s try: u = urlparse(s) if u.scheme in ("http", "https"): net = u.hostname or "" if u.port: net += f":{u.port}" return f"{u.scheme}://{net}{u.path}" if u.scheme == "file": return Path(u.path).name or s except Exception: pass if (s.startswith(("/", "~", ".")) or "\\" in s or len(s) > 2 and s[1] == ":" or Path(s).expanduser().exists()): return Path(s).name or s return s def _slugify_voice_name(raw, fallback="voice"): """Derive a valid voice name from arbitrary input (filenames, model IDs). Lowercases, maps runs of invalid chars to '-', collapses, trims to 64 chars.""" s = str(raw or "").strip().lower() if not s: return fallback s = re.sub(r"[^a-z0-9_-]+", "-", s) s = re.sub(r"-+", "-", s).strip("-_")[:64] return s or fallback