Blame for vlib/delta.py

"""Delta computation (voice minus base) and base resolution."""

from pathlib import Path
import json
import os
import struct
import sys
import time
from vlib.ui import _fail, _noninteractive, _progress_pref, _say, _step, _transform, _warn
from vlib.net import _fmt_eta
from vlib.tensors import _is_delta_safetensors, _is_ggml_quant, decode_to_f32, dtype_bytes, read_safetensors
from vlib.sources import _detect_base, _normalize_hf_repo, _public_source, open_source, pair_role, role_of
from vlib.fetch import _resolve_output_name, _source_arch
from vlib.registry import _assert_not_symlink, _iter_voices, _read_meta


def _choose_voice(args):
    voices = []
    for d in _iter_voices():
        meta = _read_meta(d / "voice.json")
        if meta:
            voices.append((d.name, meta))
    if not voices:
        _fail("  ✗ No voices stored yet. Use 'voice get <model>' to save one.")
    _say("  Available voices:")
    for i, (dirname, meta) in enumerate(voices, 1):
        _say(f"  {i}. {meta.get('name', dirname)} — {meta.get('source_hf_model', meta.get('source', '—'))}")
    try:
        choice = input(f"  Choose voice [1-{len(voices)}]: ").strip()
    except (EOFError, KeyboardInterrupt):
        _say("  ⚠ Cancelled.")
        sys.exit(1)
    if not choice.isdigit() or not (1 <= int(choice) <= len(voices)):
        _fail("  ✗ Invalid choice.")
    return voices[int(choice) - 1][1].get("name", voices[int(choice) - 1][0])


def _resolve_base(args, source_hf):
    base_arg = getattr(args, "base", None)
    if base_arg:
        # allow local file/dir path as base (for manual delta)
        p = Path(base_arg).expanduser()
        if p.exists():
            return str(p.resolve())
        norm = _normalize_hf_repo(base_arg)
        if norm is None:
            if "/" in base_arg or base_arg.startswith(("~", ".")) or base_arg.endswith((".safetensors", ".gguf", ".bin", ".pt")):
                _fail(f'  ✗ "{base_arg}" couldn\'t be found. Pass an existing local path or an HF repo (org/model).')
            _fail(f"  ✗ Invalid HF repo '{base_arg}'. Use org/model or a local path.")
        return norm
    if source_hf not in ("—", "", None) and "/" in source_hf:
        suggested, is_multi, count, _raw = _detect_base(source_hf)
    else:
        suggested, is_multi, count, _raw = None, False, 0, []
    if suggested is None:
        if _noninteractive(args):
            _fail("  ✗ Could not detect the base model. Re-run with --base <org/model or local .safetensors>.")
        try:
            entered = input("  Enter base HF repo (org/model) or local path: ").strip()
        except (EOFError, KeyboardInterrupt):
            _say("  ⚠ Cancelled.")
            sys.exit(1)
        if not entered:
            _fail("  ✗ Invalid base repo.")
        p = Path(entered).expanduser()
        if p.exists():
            return str(p.resolve())
        norm = _normalize_hf_repo(entered)
        if norm is None:
            _fail("  ✗ Invalid base repo.")
        return norm
    if is_multi:
        _say(f"  Base: {suggested} ({count} candidates)")
    else:
        _say(f"  Base detected: {suggested}")
    if _noninteractive(args):
        return suggested
    try:
        entered = input(f"  Press enter to use {suggested}, or type another repo/path: ").strip()
    except (EOFError, KeyboardInterrupt):
        _say("  ⚠ Cancelled.")
        sys.exit(1)
    if not entered:
        return suggested
    p = Path(entered).expanduser()
    if p.exists():
        return str(p.resolve())
    norm = _normalize_hf_repo(entered)
    return norm if norm else suggested


def _delta_gate(args, warn_lines, phrase, fail_hint):
    """Confirm niche delta-on-delta cases. Fails fast when non-interactive."""
    for i, line in enumerate(warn_lines):
        (_warn if i == 0 else _say)(line)
    if _noninteractive(args):
        _fail(f"  ✗ {fail_hint} (re-run interactively and type {phrase} to proceed).")
    try:
        ans = input(f"  Type {phrase} to continue, Enter to cancel: ").strip()
    except (EOFError, KeyboardInterrupt):
        _say("  Cancelled."); sys.exit(1)
    if ans.strip().lower() != phrase.lower():
        _say("  Cancelled."); sys.exit(1)


def _check_delta_gate(args, voice_path, base_path):
    """Single gate-checker for delta-on-delta niche cases. Computes the
    is_voice/is_base flags internally (GGUF has no marker, so non-safetensors
    sides stay False — matches both the guarded GGUF branch and the
    safetensors branch, where both sides are .safetensors)."""
    is_voice_delta = _is_delta_safetensors(str(voice_path)) if str(voice_path).endswith(".safetensors") else False
    is_base_delta = _is_delta_safetensors(str(base_path)) if str(base_path).endswith(".safetensors") else False
    if not is_voice_delta and is_base_delta:
        _delta_gate(args, ["  The base you gave is a delta-voice (has delta.voice.marker).",
                           "  Deltas are meant to be applied, not used as a base.",
                           "  normal − delta = 2*normal − style — nonsense, not a style."],
                    "UNDO THE UNDO", "delta-on-delta needs explicit confirmation")
    elif is_voice_delta and not is_base_delta:
        _delta_gate(args, ["  The voice you gave is already a delta-voice.",
                           "  Delta-on-normal is usually just style − base. You are doing delta − normal = style − 2*normal."],
                    "DOUBLE DIP", "delta-on-delta needs explicit confirmation")
    elif is_voice_delta and is_base_delta:
        _delta_gate(args, ["  ✨ Congratulations — you have discovered a hidden feature",
                           "     The voice and base are both delta-voices.",
                           "     delta − delta = style difference of style differences."],
                    "BIG BRAIN", "delta-on-delta needs explicit confirmation")


def _compute_delta(voice_path, base_path, delta_folder, delta_name, args=None):
    # Fast path: both safetensors (no dequant, chunked, efficient for 1.4GB).
    # Slow path: either side GGUF quant -> dequant once to F32 (unavoidable), direct numpy (efficient for <50MB tests).
    _v_gguf = _b_gguf = False
    # Chunk seam resolves through the voice facade (not vlib.tensors): the suite
    # patches voice._DELTA_CHUNK_ELEMS / voice._gguf_raw_slice_to_f32, so the
    # slow path must honor the facade namespace. Local import: voice imports us.
    import voice as _facade
    try:
        v_hdr, v_tensors = read_safetensors(str(voice_path))
    except Exception:
        _v_gguf = True
        v_hdr = v_tensors = None
    try:
        b_hdr, b_tensors = read_safetensors(str(base_path))
    except Exception:
        _b_gguf = True
        b_hdr = b_tensors = None
    if _v_gguf or _b_gguf:
        import numpy as np
        # Resolve tensor names via open_source (handles GGUF + safetensors uniformly).
        v_src = open_source(str(voice_path))
        b_src = open_source(str(base_path))
        v_names = [n for n in v_src.names() if n != "delta.voice.marker"]
        b_names = [n for n in b_src.names() if n != "delta.voice.marker"]
        if not v_names or not b_names:
            _fail("  ✗ No tensors found for delta.")
        # Match by exact name: every voice tensor needs a same-named base tensor.
        # Output head first for deterministic order, then voice order.
        try:
            arch, cfg = _source_arch(v_src)
            _cand = _resolve_output_name(v_src, cfg, v_names)
        except Exception:
            _cand = None
        ordered = ([_cand] if _cand in v_names else []) + [n for n in v_names if n != _cand]
        missing = [n for n in ordered if n not in b_names]
        fallback = {}
        for m in missing:
            if role_of(m) not in ("head", "embed"):
                continue
            b = pair_role(m, tuple(v_src.ref(m).shape), b_names,
                          lambda n: tuple(b_src.ref(n).shape))
            if b is not None:
                fallback[m] = b
        missing = [n for n in missing if n not in fallback]
        if missing:
            shown = ", ".join(missing[:5]) + ("…" if len(missing) > 5 else "")
            _fail(f"  ✗ {len(missing)} voice tensor(s) have no match in base: {shown}. Delta needs same names on both sides.")
        pairs = [(n, fallback.get(n, n)) for n in ordered]
        _check_delta_gate(args, voice_path, base_path)
        _assert_not_symlink(delta_folder)
        delta_folder.mkdir(parents=True, exist_ok=True)
        try:
            os.chmod(delta_folder, 0o700)
        except Exception:
            pass
        out_path = delta_folder / "voice.safetensors"
        tmp_out = delta_folder / "voice.safetensors.tmp"
        # Header first (offsets need shapes only), then stream one pair at a
        # time so packs never hold every tensor in RAM at once.
        header, off = {}, 0
        for vname, bname in pairs:
            vref, bref = v_src.ref(vname), b_src.ref(bname)
            if tuple(vref.shape) != tuple(bref.shape):
                _fail(f"  ✗ Voice tensor '{vname}' shape {tuple(vref.shape)} doesn't fit base {tuple(bref.shape)}.")
            nbytes = 2
            for d in vref.shape:
                nbytes *= int(d)
            header[vname] = {"dtype": "F16", "shape": [int(x) for x in vref.shape], "data_offsets": [off, off + nbytes]}
            off += nbytes
        header["delta.voice.marker"] = {"dtype": "U8", "shape": [1], "data_offsets": [off, off + 1]}
        header["__metadata__"] = {"type": "delta", "base": _public_source(str(base_path)), "voice": _public_source(str(voice_path))}
        hj = json.dumps(header).encode("utf-8")
        with open(tmp_out, "wb") as out:
            out.write(struct.pack("<Q", len(hj)))
            out.write(hj)
            for vname, bname in pairs:
                vref, bref = v_src.ref(vname), b_src.ref(bname)
                _warn(f"  Dequanting {vname} ({vref.dtype}) and base ({bref.dtype}) once to F32 for delta math (unavoidable for GGUF quant).")
                elems = 1
                for d in vref.shape:
                    elems *= int(d)
                # Within a pair, decode/dequant in _DELTA_CHUNK_ELEMS slices so peak
                # stays O(chunk): plain sides stream raw slices; quant sides stream
                # block-aligned raw slices (GGML_QUANT_SIZES units decode independently).
                from gguf.constants import GGMLQuantizationType as _Q
                from gguf.constants import GGML_QUANT_SIZES as _QS
                v_q = _is_ggml_quant(vref.dtype)
                b_q = _is_ggml_quant(bref.dtype)
                v_qt = _Q[vref.dtype] if v_q else None
                b_qt = _Q[bref.dtype] if b_q else None
                v_bpe = 1 if v_q else dtype_bytes(vref.dtype)
                b_bpe = 1 if b_q else dtype_bytes(bref.dtype)

                def _quant_slice(src, name, qt, at, n):
                    bs, ts = _QS[qt]
                    b0 = at // bs
                    b1 = (at + n + bs - 1) // bs
                    seg = src.read_raw_slice(name, b0 * ts, (b1 - b0) * ts)
                    return _facade._gguf_raw_slice_to_f32(seg, qt, at - b0 * bs, n)

                done = 0
                while done < elems:
                    cur = min(_facade._DELTA_CHUNK_ELEMS, elems - done)
                    if v_q:
                        v_f32 = _quant_slice(v_src, vname, v_qt, done, cur)
                    else:
                        v_f32 = decode_to_f32(v_src.read_raw_slice(vname, done * v_bpe, cur * v_bpe),
                                              vref.dtype, (cur,))
                    if b_q:
                        b_f32 = _quant_slice(b_src, bname, b_qt, done, cur)
                    else:
                        b_f32 = decode_to_f32(b_src.read_raw_slice(bname, done * b_bpe, cur * b_bpe),
                                              bref.dtype, (cur,))
                    out.write((v_f32.astype(np.float32) - b_f32.astype(np.float32)).astype("float16").tobytes())
                    done += cur
            out.write(b"\x01")
            out.flush()
            try:
                os.fsync(out.fileno())
            except Exception:
                pass
        tmp_out.replace(out_path)
        return out_path
    # Original chunked safetensors-safetensors path below (no dequant beyond F32 math per chunk).
    # Match by exact name: every voice tensor needs a same-named base tensor.
    v_keys = [k for k in v_tensors.keys() if k != "delta.voice.marker"]
    if not v_keys:
        _fail("  ✗ No tensors found for delta.")
    missing = [k for k in v_keys if k not in b_tensors]
    fallback = {}
    for m in missing:
        if role_of(m) not in ("head", "embed"):
            continue
        b = pair_role(m, tuple(v_tensors[m]["shape"]), list(b_tensors.keys()),
                      lambda n: tuple(b_tensors[n]["shape"]))
        if b is not None:
            fallback[m] = b
    missing = [k for k in missing if k not in fallback]
    if missing:
        shown = ", ".join(missing[:5]) + ("…" if len(missing) > 5 else "")
        _fail(f"  ✗ {len(missing)} voice tensor(s) have no match in base: {shown}. Delta needs same names on both sides.")
    pairs = [(k, fallback.get(k, k)) for k in v_keys]
    _check_delta_gate(args, voice_path, base_path)
    with open(voice_path, "rb") as f:
        v_hs = struct.unpack("<Q", f.read(8))[0]
    with open(base_path, "rb") as f:
        b_hs = struct.unpack("<Q", f.read(8))[0]
    specs = []
    header, off = {}, 0
    for vname, bname in pairs:
        vinfo, binfo = v_tensors[vname], b_tensors[bname]
        vshape, bshape = tuple(vinfo["shape"]), tuple(binfo["shape"])
        if vshape != bshape:
            _fail(f"  ✗ Voice tensor '{vname}' shape {vshape} doesn't fit base {bshape}.")
        elems = 1
        for d in vshape:
            elems *= int(d)
        nbytes = elems * 2
        header[vname] = {"dtype": "F16", "shape": list(vshape), "data_offsets": [off, off + nbytes]}
        off += nbytes
        specs.append({"name": vname, "vinfo": vinfo, "binfo": binfo,
                      "v_start": 8 + v_hs + vinfo["data_offsets"][0],
                      "b_start": 8 + b_hs + binfo["data_offsets"][0],
                      "v_bpe": dtype_bytes(vinfo["dtype"]), "b_bpe": dtype_bytes(binfo["dtype"]),
                      "elems": elems})
    total_elems = off // 2
    marker_name = "delta.voice.marker"
    header[marker_name] = {"dtype": "U8", "shape": [1], "data_offsets": [off, off + 1]}
    header["__metadata__"] = {"type": "delta", "base": _public_source(str(base_path)), "voice": _public_source(str(voice_path))}
    hj = json.dumps(header).encode("utf-8")
    _assert_not_symlink(delta_folder)
    delta_folder.mkdir(parents=True, exist_ok=True)
    try:
        os.chmod(delta_folder, 0o700)
    except Exception:
        pass
    out_path = delta_folder / "voice.safetensors"
    tmp_out = delta_folder / "voice.safetensors.tmp"
    chunk = 1 << 20
    total_bytes = total_elems * 2
    t0 = time.time()
    last = [0.0]
    last_pct = [-1.0]
    tiny = total_bytes < 100000
    def _delta_progress(force=False):
        now = time.time()
        if not force and now - last[0] < 0.25 and processed < total_elems:
            return
        last[0] = now
        pct = round(processed / total_elems * 100) if total_elems else 100
        if pct == last_pct[0] and not force:
            return
        last_pct[0] = pct
        if tiny:
            msg = f"  → crafting delta {pct:3.0f}% · {processed}/{total_elems} elems"
        else:
            elapsed = now - t0 or 1e-6
            mbps = (processed * 2) / elapsed / 1e6
            remain = total_bytes - processed * 2
            eta = _fmt_eta(remain, mbps)
            msg = f"  → crafting delta {pct:3.0f}% · {processed*2/1e6:.0f}/{total_bytes/1e6:.0f} MB · {mbps:.1f} MB/s · {eta}  "
        sys.stdout.write(_transform(msg) + ("\n" if _progress_pref() else "\r"))
        sys.stdout.flush()
    with open(voice_path, "rb") as fv, open(base_path, "rb") as fb, open(tmp_out, "wb") as out:
        out.write(struct.pack("<Q", len(hj)))
        out.write(hj)
        processed = 0
        _delta_progress(force=True)
        for spec in specs:
            fv.seek(spec["v_start"])
            fb.seek(spec["b_start"])
            vinfo, binfo = spec["vinfo"], spec["binfo"]
            v_bpe, b_bpe = spec["v_bpe"], spec["b_bpe"]
            done = 0
            while done < spec["elems"]:
                cur = min(chunk, spec["elems"] - done)
                v_raw = fv.read(cur * v_bpe)
                b_raw = fb.read(cur * b_bpe)
                if len(v_raw) != cur * v_bpe or len(b_raw) != cur * b_bpe:
                    sys.stdout.write("\n")
                    _fail(f"  ✗ Short read during delta compute ({spec['name']}).")
                v_f32 = decode_to_f32(v_raw, vinfo["dtype"], (cur,))
                b_f32 = decode_to_f32(b_raw, binfo["dtype"], (cur,))
                delta_f16 = (v_f32 - b_f32).astype("float16")
                out.write(delta_f16.tobytes())
                done += cur
                processed += cur
                _delta_progress()
        _delta_progress()
        sys.stdout.write("\n")
        # marker byte for auto-detect (delta.voice.marker)
        out.write(b"\x01")
        out.flush()
        try:
            os.fsync(out.fileno())
        except Exception:
            pass
    tmp_out.replace(out_path)
    return out_path