vlib/cmds_chorus.py

"""Chorus/graft commands: DELLA voice blending (pulse/support/lead).

``voice graft`` is the single-graft primitive (pulse→donor); ``voice chorus``
chains it (pulse→support at 0.8, then →lead at 0.9, or a single 0.9 graft
when support is skipped). Both write named registry voices via the same
``_sing_graft`` core so the CLI can never drift from the tested path.
"""

import hashlib
import json
import sys
from pathlib import Path

from vlib.ui import _fail, _noninteractive, _now_iso, _ok, _say, _step
from vlib.sources import _sanitize_voice_name, _slugify_voice_name, open_source
from vlib.tensors import encode_from_f32, write_safetensors
from vlib.registry import (_assert_not_symlink, _read_meta,
                            _registry_folder, _registry_resolve, log_operation)
from vlib.delta import _choose_voice
from vlib.della import VALID_RESCALES, merge as _della_merge

HOUSE_DENSITY = 0.8
HOUSE_LEAD_DENSITY = 0.9
HOUSE_EPSILON = 0.1
HOUSE_LEAD_EPSILON = 0.05
HOUSE_LAMBDA = 1.0
HOUSE_SEED = 42
HOUSE_RESCALE = "l1"


def _graft_params(args):
    """Coerce graft knobs from argparse; fail loud on bad values."""
    try:
        density = float(getattr(args, "density", HOUSE_DENSITY))
    except (TypeError, ValueError):
        _fail(f"  ✗ --density must be a number, got {getattr(args, 'density', None)!r}.")
    try:
        epsilon = float(getattr(args, "epsilon", HOUSE_EPSILON))
    except (TypeError, ValueError):
        _fail(f"  ✗ --epsilon must be a number, got {getattr(args, 'epsilon', None)!r}.")
    try:
        lambda_ = float(getattr(args, "lambda_", HOUSE_LAMBDA))
    except (TypeError, ValueError):
        _fail(f"  ✗ --lambda must be a number, got {getattr(args, 'lambda_', None)!r}.")
    rescale = str(getattr(args, "rescale", HOUSE_RESCALE) or HOUSE_RESCALE)
    if rescale not in VALID_RESCALES:
        _fail(f"  ✗ --rescale must be {'/'.join(VALID_RESCALES)}, got {rescale!r}.")
    seed = getattr(args, "seed", HOUSE_SEED)
    if seed is not None:
        try:
            seed = int(seed)
        except (TypeError, ValueError):
            _fail(f"  ✗ --seed must be an integer, got {seed!r}.")
    if not 0 < density < 1:
        _fail(f"  ✗ --density must be in (0,1), got {density}.")
    if density + epsilon >= 1 or density - epsilon <= 0:
        _fail(f"  ✗ --density ± --epsilon must stay in (0,1): got density {density:g} epsilon {epsilon:g}.")
    return density, epsilon, lambda_, rescale, seed


def _flag_config(args):
    """Coerce chorus flags to a full config. Shared by all three modes so
    flags mean the same thing everywhere: non-interactive uses them directly,
    --config screens prefill from them, plain applies them silently."""
    def _opt(name, house):
        v = getattr(args, name, None)
        return house if v is None else float(v)
    density = _opt("density", HOUSE_DENSITY)
    lead_density = _opt("lead_density", HOUSE_LEAD_DENSITY)
    epsilon = _opt("epsilon", HOUSE_EPSILON)
    lead_epsilon = _opt("lead_epsilon", HOUSE_LEAD_EPSILON)
    lambda_ = _opt("lambda_", HOUSE_LAMBDA)
    rescale = str(getattr(args, "rescale", HOUSE_RESCALE) or HOUSE_RESCALE)
    if rescale not in VALID_RESCALES:
        _fail(f"  ✗ --rescale must be {'/'.join(VALID_RESCALES)}.")
    seed = getattr(args, "seed", HOUSE_SEED)
    seed = HOUSE_SEED if seed is None else int(seed)
    for _d, _e, _tag in ((density, epsilon, "first graft"),
                         (lead_density, lead_epsilon, "lead graft")):
        if not 0 < _d < 1:
            _fail(f"  ✗ {_tag} density must be in (0,1), got {_d}.")
        if _d + _e >= 1 or _d - _e <= 0:
            _fail(f"  ✗ {_tag} density ± epsilon must stay in (0,1): got {_d:g} ± {_e:g}.")
    return density, lead_density, epsilon, lead_epsilon, lambda_, rescale, seed

def _head_ref(source):
    """Single head tensor ref of a voice source; voices hold exactly one."""
    names = [n for n in source.names() if n != "delta.voice.marker"]
    if not names:
        _fail("  ✗ No tensors found for graft.")
    if len(names) > 1:
        _fail(f"  ✗ Graft needs a single-tensor voice, found {len(names)} tensors "
              f"({', '.join(names[:4])}{'…' if len(names) > 4 else ''}). "
              "Use `voice get <model>` for a head-only voice.")
    return source.ref(names[0])


def _check_overwrite_folder(folder, args):
    if folder.exists() and ((folder / "voice.safetensors").exists()
                             or (folder / "voice.gguf").exists()):
        if _noninteractive(args):
            if not getattr(args, "yes", False):
                _fail(f"  ✗ Voice '{folder.name}' exists. Re-run with --yes to overwrite.")
        else:
            try:
                ans = input(f"  Voice '{folder.name}' exists. Overwrite? [y/N]: ").strip().lower()
            except (EOFError, KeyboardInterrupt):
                _say("  Cancelled.")
                sys.exit(1)
            if ans not in ("y", "yes"):
                _say("  Cancelled.")
                sys.exit(1)


def _resolve_graft_voice(name, role):
    """Registry name (or path) → (voice_id, source). Fails loud when missing."""
    path, meta = _registry_resolve(name)
    if path is None or not path.exists():
        _fail(f"  ✗ {role} voice '{name}' not found. Use 'voice list' or pass a .safetensors path.")
    vid = (meta or {}).get("name") or Path(str(path)).stem
    return vid, open_source(str(path))


def _validate_pair(pulse_ref, donor_ref):
    """Distinct voices + identical head geometry, or fail loud with a table."""
    if tuple(pulse_ref.shape) != tuple(donor_ref.shape):
        _say(f"  {'voice':<12}{'tensor':<28}{'shape'}")
        _say(f"  {'─' * 12}{'─' * 28}{'─' * 20}")
        _say(f"  {'pulse':<12}{pulse_ref.name:<28}{list(pulse_ref.shape)}")
        _say(f"  {'donor':<12}{donor_ref.name:<28}{list(donor_ref.shape)}")
        _fail(f"  ✗ Head shapes differ: {list(pulse_ref.shape)} vs {list(donor_ref.shape)}. "
              "Chorus needs identical head geometry.")
    _say("  Lineage is not verified — voices with different base models still merge, results may vary.")


def _sing_graft(pulse_id, pulse_src, donor_id, donor_src, out_name,
                density, epsilon, lambda_, rescale, seed, args, chain=None):
    """One DELLA graft pulse→donor into registry voice ``out_name``.

    Shared by ``voice graft`` and ``voice chorus`` so both CLIs run the same
    tensor path. Returns the registry folder.
    """
    import numpy as np
    pulse_ref = _head_ref(pulse_src)
    donor_ref = _head_ref(donor_src)
    if pulse_id.strip().lower() == donor_id.strip().lower():
        _fail("  ✗ Pulse and donor are the same voice — graft needs two distinct voices.")
    from vlib.tensors import _is_ggml_quant as _is_q
    for _role, _ref in (("Pulse", pulse_ref), ("Donor", donor_ref)):
        try:
            _quant = _is_q(_ref.dtype)
        except Exception:
            _quant = False
        if _quant:
            _fail(f"  ✗ {_role} voice is quantized GGUF ({_ref.dtype}) — graft needs float heads. "
                  "Cast it onto a float target first, then graft the result.")
    folder = _registry_folder(out_name)
    _check_overwrite_folder(folder, args)
    _step(f"Grafting {donor_id} onto {pulse_id} (density {density:g})…")
    base_f32 = pulse_src.read_f32(pulse_ref.name)
    donor_f32 = donor_src.read_f32(donor_ref.name)
    merged = _della_merge(np.asarray(base_f32, dtype=np.float32),
                          [np.asarray(donor_f32, dtype=np.float32)],
                          density=density, epsilon=epsilon, lambda_=lambda_,
                          rescale=rescale, seed=seed)
    _assert_not_symlink(folder)
    folder.mkdir(parents=True, exist_ok=True)
    try:
        import os
        os.chmod(folder, 0o700)
    except Exception:
        pass
    _assert_not_symlink(folder)
    out_path = folder / "voice.safetensors"
    tmp = folder / "voice.safetensors.tmp"
    write_safetensors({pulse_ref.name: (encode_from_f32(merged, pulse_ref.dtype),
                                         pulse_ref.dtype,
                                         tuple(pulse_ref.shape))}, str(tmp))
    tmp.replace(out_path)
    try:
        import os
        os.chmod(out_path, 0o600)
    except Exception:
        pass
    cfg = {"density": density, "epsilon": epsilon, "lambda": lambda_,
           "rescale": rescale, "seed": seed}
    meta = {
        "name": out_name.strip().lower(),
        "type": "chorus",
        "method": "della",
        "pulse": pulse_id,
        "donor": donor_id,
        "tensor_name": pulse_ref.name,
        "shape": [int(x) for x in pulse_ref.shape],
        "dtype": pulse_ref.dtype,
        "config": cfg,
        "config_hash": hashlib.sha256(
            json.dumps(cfg, sort_keys=True).encode()).hexdigest()[:12],
    }
    if chain:
        meta["chain"] = chain
    try:
        meta["bytes"] = out_path.stat().st_size
    except Exception:
        pass
    meta["downloaded_at"] = _now_iso()
    jtmp = folder / "voice.json.tmp"
    jtmp.write_text(json.dumps(meta, indent=2) + "\n")
    try:
        import os
        os.chmod(jtmp, 0o600)
    except Exception:
        pass
    jtmp.replace(folder / "voice.json")
    log_operation({"op": "graft", "pulse": pulse_id, "donor": donor_id,
                   "density": density, "epsilon": epsilon, "lambda": lambda_,
                   "rescale": rescale, "seed": seed, "output": str(folder),
                   "status": "ok"})
    return folder


def _prompt_knobs(role, density, epsilon, lambda_, seed, out_name):
    """--config knob screen: Enter accepts each default, typing overrides."""
    if _noninteractive({"non_interactive": True}):
        return density, epsilon, lambda_, seed, out_name
    try:
        raw = input(f"  {role} density [{density:g}]: ").strip()
        if raw:
            density = float(raw)
        raw = input(f"  epsilon [{epsilon:g}]: ").strip()
        if raw:
            epsilon = float(raw)
        raw = input(f"  lambda [{lambda_:g}]: ").strip()
        if raw:
            lambda_ = float(raw)
        raw = input(f"  seed [{seed}]: ").strip()
        if raw:
            seed = int(raw)
        raw = input(f"  output name [{out_name}]: ").strip()
        if raw:
            out_name = raw
    except (EOFError, KeyboardInterrupt):
        _say("  Cancelled.")
        sys.exit(1)
    except ValueError:
        _fail("  ✗ Knob must be a number — re-run `voice chorus --config` to retry.")
    if _sanitize_voice_name(out_name) is None:
        _fail(f"  ✗ Invalid output name '{out_name}'.")
    if not 0 < density < 1:
        _fail(f"  ✗ density must be in (0,1), got {density}.")
    return density, epsilon, lambda_, seed, out_name


def cmd_graft(args):
    pulse_arg = getattr(args, "pulse", None)
    donor_arg = getattr(args, "donor", None)
    if not pulse_arg or not str(pulse_arg).strip():
        pulse_arg = _choose_voice(args)
    if not donor_arg or not str(donor_arg).strip():
        donor_arg = _choose_voice(args)
    pulse_id, pulse_src = _resolve_graft_voice(str(pulse_arg).strip(), "Pulse")
    donor_id, donor_src = _resolve_graft_voice(str(donor_arg).strip(), "Donor")
    density, epsilon, lambda_, rescale, seed = _graft_params(args)
    out = getattr(args, "out", None) or f"chorus-{_slugify_voice_name(pulse_id)}-{_slugify_voice_name(donor_id)}"
    out = str(out).strip()
    if _sanitize_voice_name(out) is None:
        _fail(f"  ✗ Invalid --out name '{out}'.")
    folder = _sing_graft(pulse_id, pulse_src, donor_id, donor_src, out,
                         density, epsilon, lambda_, rescale, seed, args)
    _ok(f"Done → {folder}")
    _say(f"  Use it:   voice cast <your-model.gguf> {folder.name}")
    _say(f"  Hear it:  voice cast <your-model.gguf> {folder.name} --out /tmp/chorus.gguf")


def _pick_from(voices, exclude=()):
    """Numbered pick from a fixed registry list (excludes by lowercase id)."""
    shown = [(vid, src) for vid, src in voices
             if vid.strip().lower() not in {e.strip().lower() for e in exclude}]
    if not shown:
        _fail("  ✗ No more voices to choose from.")
    _say("  Available voices:")
    for i, (vid, _src) in enumerate(shown, 1):
        _say(f"  {i}. {vid}")
    try:
        choice = input(f"  Choose [1-{len(shown)}]: ").strip()
    except (EOFError, KeyboardInterrupt):
        _say("  Cancelled.")
        sys.exit(1)
    if not choice.isdigit() or not (1 <= int(choice) <= len(shown)):
        _fail("  ✗ Invalid choice.")
    return shown[int(choice) - 1]


def _registry_voices():
    """All registry voices as (id, source) pairs; fails loud when empty."""
    from vlib.registry import _iter_voices
    voices = []
    for d in _iter_voices():
        meta = _read_meta(d / "voice.json")
        if not meta:
            continue
        vid = meta.get("name", d.name)
        try:
            src = open_source(str(d / ("voice.gguf" if (d / "voice.gguf").exists()
                                        else "voice.safetensors")))
        except Exception as e:
            _fail(f"  ✗ Voice '{vid}' is unreadable ({e}).")
        voices.append((vid, src))
    if not voices:
        _fail("  ✗ No voices stored yet. Use 'voice get <model>' to save one.")
    return voices


def cmd_chorus(args):
    config = bool(getattr(args, "config", False))
    if _noninteractive(args):
        # Flags override defaults silently; positionals answer the prompts.
        # --config has no staged screens here (no TTY): flags ARE the config.
        pulse_arg = getattr(args, "pulse", None)
        donor_args = getattr(args, "donors", None) or []
        if not pulse_arg or not donor_args:
            _fail("  ✗ Non-interactive chorus needs `voice chorus <pulse> <support> <lead>` "
                  "(pass one donor for a duet).")
        if len(donor_args) > 2:
            _fail("  ✗ Chorus takes at most two donors (support + lead).")
        pulse_id, pulse_src = _resolve_graft_voice(str(pulse_arg).strip(), "Pulse")
        ids_srcs = [_resolve_graft_voice(str(d).strip(), "Donor") for d in donor_args]
        density, lead_density, epsilon, lead_epsilon, lambda_, rescale, seed = _flag_config(args)
        out = getattr(args, "out", None)
        if out is not None:
            out = str(out).strip()
            if _sanitize_voice_name(out) is None:
                _fail(f"  ✗ Invalid --out name '{out}'.")
        if len(ids_srcs) == 1:
            (lead_id, lead_src) = ids_srcs[0]
            name = str(out or f"chorus-{_slugify_voice_name(pulse_id)}-{_slugify_voice_name(lead_id)}").strip()
            folder = _sing_graft(pulse_id, pulse_src, lead_id, lead_src, name,
                                 lead_density, lead_epsilon, lambda_, rescale, seed,
                                 args)
        else:
            (support_id, support_src), (lead_id, lead_src) = ids_srcs
            duet = str(f"chorus-{_slugify_voice_name(pulse_id)}-{_slugify_voice_name(support_id)}").strip()
            folder = _sing_graft(pulse_id, pulse_src, support_id, support_src, duet,
                                 density, epsilon, lambda_, rescale, seed, args)
            _say(f"  Duet → {folder.name} (audition it: voice cast <model.gguf> {folder.name})")
            duet_src = open_source(str(folder / "voice.safetensors"))
            name = str(out or f"chorus-{_slugify_voice_name(pulse_id)}-{_slugify_voice_name(lead_id)}").strip()
            folder = _sing_graft(folder.name, duet_src, lead_id, lead_src, name,
                                 lead_density, lead_epsilon, lambda_, rescale, seed + 1,
                                 args, chain=[duet, name])
        _ok(f"Done → {folder}")
        _say(f"  Hear it:  voice cast <your-model.gguf> {folder.name} --out /tmp/chorus.gguf")
        return

    voices = _registry_voices()
    _say("  Choose pulse:")
    pulse_id, pulse_src = _pick_from(voices)
    _say("  Choose support (Enter to skip):")
    shown = [(vid, src) for vid, src in voices
             if vid.strip().lower() != pulse_id.strip().lower()]
    for i, (vid, _src) in enumerate(shown, 1):
        _say(f"  {i}. {vid}")
    try:
        raw = input(f"  Choose support [1-{len(shown)}] (Enter to skip): ").strip()
    except (EOFError, KeyboardInterrupt):
        _say("  Cancelled.")
        sys.exit(1)
    support = None
    if raw:
        if not raw.isdigit() or not (1 <= int(raw) <= len(shown)):
            _fail("  ✗ Invalid choice.")
        support = shown[int(raw) - 1]
    _say("  Choose lead:")
    taken = {pulse_id.strip().lower()} | ({support[0].strip().lower()} if support else set())
    lead_id, lead_src = _pick_from(voices, exclude=taken)

    _fd, _ld, _fe, _le, _fl, _fr, _fs = _flag_config(args)
    if support is None:
        # Duet: single graft at lead density. Plain applies flags silently;
        # --config screens prefill from flags, Enter keeps them.
        if config:
            out = f"chorus-{_slugify_voice_name(pulse_id)}-{_slugify_voice_name(lead_id)}"
            density, epsilon, lambda_, seed, out = _prompt_knobs(
                "Lead", _ld, _le, _fl, _fs, out)
            rescale = _fr
        else:
            density, epsilon, lambda_, seed, rescale = _ld, _le, _fl, _fs, _fr
            out = f"chorus-{_slugify_voice_name(pulse_id)}-{_slugify_voice_name(lead_id)}"
        _say(f"  {pulse_id} pulse + {lead_id} lead → {out}")
        folder = _sing_graft(pulse_id, pulse_src, lead_id, lead_src, out,
                             density, epsilon, lambda_, rescale, seed, args)
        _ok(f"Done → {folder}")
        _say(f"  Hear it:  voice cast <your-model.gguf> {folder.name} --out /tmp/chorus.gguf")
        return

    support_id, support_src = support
    _say(f"  {pulse_id} pulse + {support_id} support + {lead_id} lead → chorus-{_slugify_voice_name(pulse_id)}-{_slugify_voice_name(lead_id)}")
    if config:
        duet = f"chorus-{_slugify_voice_name(pulse_id)}-{_slugify_voice_name(support_id)}"
        density, epsilon, lambda_, seed, duet = _prompt_knobs(
            "Support", _fd, _fe, _fl, _fs, duet)
        folder = _sing_graft(pulse_id, pulse_src, support_id, support_src, duet,
                             density, epsilon, lambda_, _fr, seed, args)
        _say(f"  Duet → {folder.name} — cast it and listen before the lead lands: "
             f"voice cast <model.gguf> {folder.name}")
        final = f"chorus-{_slugify_voice_name(pulse_id)}-{_slugify_voice_name(lead_id)}"
        density2, epsilon2, lambda2, seed2, final = _prompt_knobs(
            "Lead", _ld, _le, lambda_, seed + 1, final)
        duet_src = open_source(str(folder / "voice.safetensors"))
        folder = _sing_graft(folder.name, duet_src, lead_id, lead_src, final,
                             density2, epsilon2, lambda2, _fr, seed2,
                             args, chain=[duet, final])
    else:
        duet = f"chorus-{_slugify_voice_name(pulse_id)}-{_slugify_voice_name(support_id)}"
        folder = _sing_graft(pulse_id, pulse_src, support_id, support_src, duet,
                             _fd, _fe, _fl, _fr, _fs, args)
        duet_src = open_source(str(folder / "voice.safetensors"))
        final = f"chorus-{_slugify_voice_name(pulse_id)}-{_slugify_voice_name(lead_id)}"
        folder = _sing_graft(folder.name, duet_src, lead_id, lead_src, final,
                             _ld, _le, _fl, _fr, _fs + 1, args,
                             chain=[duet, final])
    _ok(f"Done → {folder}")
    _say(f"  Hear it:  voice cast <your-model.gguf> {folder.name} --out /tmp/chorus.gguf")