omnius 1.0.0

package/voices/personaplex/clone-voice.py

@@ -0,0 +1,384 @@
+#!/usr/bin/env python3
+"""
+clone-voice.py — High-fidelity PersonaPlex voice cloning.
+
+Applies LuxTTS-inspired preprocessing before embedding extraction:
+  1. Resample to 24kHz mono
+  2. Noise reduction (spectral gating)
+  3. Silence trimming (energy-based VAD)
+  4. LUFS normalization to -20 dBFS (tuned for PersonaPlex)
+  5. Duration optimization (trim to 4-8s sweet spot)
+  6. Multi-segment embedding averaging for long clips
+
+Usage:
+  python clone-voice.py --input voice.wav --name MyVoice [--device cuda]
+  python clone-voice.py --input voice.wav --name MyVoice --segments 3  # multi-segment averaging
+"""
+
+import argparse
+import os
+import sys
+import logging
+import numpy as np
+
+logging.basicConfig(level=logging.INFO, format="%(message)s")
+log = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Audio preprocessing (LuxTTS-inspired)
+# ---------------------------------------------------------------------------
+
+def preprocess_audio(input_path: str, target_sr: int = 24000,
+                     target_lufs: float = -20.0,
+                     min_duration: float = 2.0,
+                     max_duration: float = 8.0,
+                     denoise: bool = True) -> "torch.Tensor":
+    """
+    Full preprocessing pipeline:
+      1. Load + resample to 24kHz mono
+      2. Noise reduction via spectral gating
+      3. Silence trimming (leading/trailing)
+      4. LUFS normalization
+      5. Duration clipping to sweet spot
+    Returns: [1, T] tensor at target_sr
+    """
+    import torch
+    import torchaudio
+
+    log.info(f"  Loading: {input_path}")
+    wav, sr = torchaudio.load(input_path)
+
+    # Stereo → mono
+    if wav.shape[0] > 1:
+        wav = wav.mean(dim=0, keepdim=True)
+        log.info(f"  Converted stereo → mono")
+
+    # Resample
+    if sr != target_sr:
+        wav = torchaudio.transforms.Resample(sr, target_sr)(wav)
+        sr = target_sr
+        log.info(f"  Resampled to {target_sr}Hz")
+
+    audio_np = wav.squeeze().numpy()
+    orig_duration = len(audio_np) / sr
+
+    # Step 1: Noise reduction
+    if denoise:
+        try:
+            import noisereduce as nr
+            log.info(f"  Denoising (spectral gating)...")
+            audio_np = nr.reduce_noise(
+                y=audio_np,
+                sr=sr,
+                prop_decrease=0.7,  # moderate — preserve voice character
+                n_fft=2048,
+                hop_length=512,
+            )
+            log.info(f"  Denoised")
+        except ImportError:
+            log.info(f"  noisereduce not available, skipping denoise")
+
+    # Step 2: Silence trimming (energy-based)
+    frame_length = int(0.025 * sr)  # 25ms frames
+    hop = int(0.010 * sr)           # 10ms hop
+    energy = []
+    for i in range(0, len(audio_np) - frame_length, hop):
+        frame = audio_np[i:i + frame_length]
+        energy.append(np.sqrt(np.mean(frame ** 2)))
+    energy = np.array(energy)
+
+    if len(energy) > 0:
+        # Threshold: 5% of peak energy (aggressive silence removal)
+        threshold = max(np.percentile(energy, 10), np.max(energy) * 0.05)
+        voiced = np.where(energy > threshold)[0]
+        if len(voiced) > 0:
+            start_sample = max(0, voiced[0] * hop - int(0.1 * sr))   # 100ms margin
+            end_sample = min(len(audio_np), (voiced[-1] + 1) * hop + int(0.1 * sr))
+            audio_np = audio_np[start_sample:end_sample]
+            trimmed_dur = len(audio_np) / sr
+            if trimmed_dur < orig_duration - 0.2:
+                log.info(f"  Trimmed silence: {orig_duration:.1f}s → {trimmed_dur:.1f}s")
+
+    # Step 3: LUFS normalization
+    try:
+        import pyloudnorm as pyln
+        meter = pyln.Meter(sr)
+        current_lufs = meter.integrated_loudness(audio_np)
+        if not np.isinf(current_lufs) and not np.isnan(current_lufs):
+            audio_np = pyln.normalize.loudness(audio_np, current_lufs, target_lufs)
+            log.info(f"  Normalized: {current_lufs:.1f} → {target_lufs:.1f} LUFS")
+    except Exception as e:
+        # Fallback: simple RMS normalization
+        rms = np.sqrt(np.mean(audio_np ** 2))
+        if rms > 0:
+            target_rms = 10 ** (target_lufs / 20)  # approximate
+            audio_np = audio_np * (target_rms / rms)
+            log.info(f"  RMS-normalized (fallback)")
+
+    # Step 4: Duration clipping
+    duration = len(audio_np) / sr
+    if duration > max_duration:
+        # Take the most energetic segment, not just the start
+        segment_samples = int(max_duration * sr)
+        # Sliding window energy to find the best segment
+        window_energy = []
+        step = int(0.5 * sr)  # 500ms steps
+        for i in range(0, len(audio_np) - segment_samples, step):
+            seg = audio_np[i:i + segment_samples]
+            window_energy.append((i, np.sqrt(np.mean(seg ** 2))))
+        if window_energy:
+            best_start = max(window_energy, key=lambda x: x[1])[0]
+            audio_np = audio_np[best_start:best_start + segment_samples]
+            log.info(f"  Selected best {max_duration:.0f}s segment (from {duration:.1f}s)")
+    elif duration < min_duration:
+        log.warning(f"  Audio too short ({duration:.1f}s < {min_duration:.0f}s minimum)")
+
+    final_duration = len(audio_np) / sr
+    log.info(f"  Final: {final_duration:.1f}s, {sr}Hz mono")
+
+    wav_out = torch.from_numpy(audio_np).float().unsqueeze(0)  # [1, T]
+    return wav_out
+
+
+# ---------------------------------------------------------------------------
+# Multi-segment embedding averaging
+# ---------------------------------------------------------------------------
+
+def clone_voice_multiseg(input_wav: str, output_name: str, device: str = "cuda",
+                         hf_repo: str = "nvidia/personaplex-7b-v1",
+                         cpu_offload: bool = False,
+                         n_segments: int = 1,
+                         target_lufs: float = -20.0,
+                         max_seg_duration: float = 8.0):
+    """
+    Clone voice with optional multi-segment averaging.
+
+    For n_segments > 1 on long clips:
+      - Split preprocessed audio into overlapping segments
+      - Extract embeddings from each segment independently
+      - Average embeddings in latent space
+      - Use final segment's KV-cache (most voice information)
+    """
+    import torch
+    import torchaudio
+    from huggingface_hub import hf_hub_download
+    from moshi.models import loaders
+    from moshi.models.lm import LMGen
+
+    voices_dir = os.path.join(os.path.dirname(__file__), "custom_voices")
+    os.makedirs(voices_dir, exist_ok=True)
+    output_pt = os.path.join(voices_dir, f"{output_name}.pt")
+
+    if os.path.exists(output_pt):
+        log.info(f"Voice '{output_name}' already exists at {output_pt}")
+        log.info("Delete it first if you want to re-clone.")
+        return output_pt
+
+    # ── Preprocessing ────────────────────────────────────────────────────
+    log.info("Step 1: Preprocessing audio...")
+    # For multi-segment, preprocess with longer max duration
+    effective_max = max_seg_duration if n_segments == 1 else max_seg_duration * n_segments
+    preprocessed = preprocess_audio(
+        input_wav,
+        target_lufs=target_lufs,
+        max_duration=effective_max,
+        denoise=True,
+    )
+
+    # Save preprocessed audio for inspection
+    prep_path = os.path.join(voices_dir, f"{output_name}_preprocessed.wav")
+    torchaudio.save(prep_path, preprocessed, 24000)
+    log.info(f"  Saved preprocessed audio: {prep_path}")
+
+    # ── Split into segments if requested ─────────────────────────────────
+    sr = 24000
+    total_samples = preprocessed.shape[1]
+    total_duration = total_samples / sr
+
+    if n_segments > 1 and total_duration > max_seg_duration:
+        seg_samples = int(max_seg_duration * sr)
+        overlap = int(1.0 * sr)  # 1s overlap
+        stride = max(seg_samples - overlap, int(2.0 * sr))
+        segments = []
+        for i in range(0, total_samples - seg_samples + 1, stride):
+            seg = preprocessed[:, i:i + seg_samples]
+            segments.append(seg)
+            if len(segments) >= n_segments:
+                break
+        # If we didn't get enough, just use what we have
+        if len(segments) == 0:
+            segments = [preprocessed]
+        log.info(f"  Split into {len(segments)} segments ({max_seg_duration:.0f}s each, 1s overlap)")
+    else:
+        segments = [preprocessed]
+
+    # ── Load models ──────────────────────────────────────────────────────
+    log.info("\nStep 2: Loading models...")
+    mimi_weight = hf_hub_download(hf_repo, loaders.MIMI_NAME)
+    mimi = loaders.get_mimi(mimi_weight, device)
+    mimi.streaming_forever(1)
+
+    moshi_weight = hf_hub_download(hf_repo, loaders.MOSHI_NAME)
+    lm = loaders.get_moshi_lm(moshi_weight, device=device, cpu_offload=cpu_offload)
+    lm.eval()
+
+    frame_size = int(mimi.sample_rate / mimi.frame_rate)
+    other_mimi = loaders.get_mimi(mimi_weight, device)
+    other_mimi.streaming_forever(1)
+
+    log.info("  Warming up...")
+    from moshi.offline import warmup
+    lm_gen = LMGen(
+        lm,
+        audio_silence_frame_cnt=int(0.5 * mimi.frame_rate),
+        sample_rate=mimi.sample_rate,
+        device=device,
+        frame_rate=mimi.frame_rate,
+        save_voice_prompt_embeddings=True,
+        use_sampling=False,
+        temp=0.8,
+        temp_text=0.7,
+        top_k=250,
+        top_k_text=25,
+    )
+    lm_gen.streaming_forever(1)
+    warmup(mimi, other_mimi, lm_gen, device, frame_size)
+
+    # ── Extract embeddings per segment ───────────────────────────────────
+    all_embeddings = []
+    final_cache = None
+
+    for seg_idx, seg_audio in enumerate(segments):
+        seg_dur = seg_audio.shape[1] / sr
+        log.info(f"\nStep 3.{seg_idx + 1}: Extracting embeddings from segment {seg_idx + 1}/{len(segments)} ({seg_dur:.1f}s)...")
+
+        # Write temp WAV for this segment
+        tmp_seg = os.path.join(voices_dir, f"_tmp_seg_{seg_idx}.wav")
+        torchaudio.save(tmp_seg, seg_audio, sr)
+
+        # Reset state for each segment
+        mimi.reset_streaming()
+        other_mimi.reset_streaming()
+        lm_gen.reset_streaming()
+
+        # Load and process
+        lm_gen.load_voice_prompt(tmp_seg)
+
+        # Trick save path so .pt goes where we want
+        lm_gen.voice_prompt = os.path.join(voices_dir, f"_tmp_seg_{seg_idx}.wav")
+        lm_gen._step_voice_prompt(mimi)
+
+        # Collect embeddings
+        auto_saved = os.path.join(voices_dir, f"_tmp_seg_{seg_idx}.pt")
+        if os.path.exists(auto_saved):
+            state = torch.load(auto_saved, map_location="cpu", weights_only=False)
+            all_embeddings.append(state["embeddings"])
+            final_cache = state["cache"]
+            os.remove(auto_saved)
+            log.info(f"    Extracted {state['embeddings'].shape[0]} frames")
+        else:
+            log.warning(f"    Segment {seg_idx + 1} failed to produce embeddings")
+
+        # Clean up temp
+        os.remove(tmp_seg)
+
+    if not all_embeddings:
+        log.error("No embeddings extracted!")
+        return None
+
+    # ── Average embeddings across segments ───────────────────────────────
+    if len(all_embeddings) > 1:
+        log.info(f"\nStep 4: Averaging {len(all_embeddings)} segment embeddings...")
+        # Pad/truncate to same length, then average
+        min_frames = min(e.shape[0] for e in all_embeddings)
+        truncated = [e[:min_frames] for e in all_embeddings]
+        stacked = torch.stack(truncated, dim=0)  # [N_segs, min_frames, 1, 1, 4096]
+        averaged = stacked.mean(dim=0)            # [min_frames, 1, 1, 4096]
+        log.info(f"  Averaged: {averaged.shape[0]} frames from {len(all_embeddings)} segments")
+    else:
+        averaged = all_embeddings[0]
+
+    # ── Save final .pt ───────────────────────────────────────────────────
+    torch.save({
+        "embeddings": averaged.detach().cpu(),
+        "cache": final_cache.detach().cpu() if final_cache is not None else torch.zeros(1, 17, 4, dtype=torch.int64),
+    }, output_pt)
+
+    # Verify
+    state = torch.load(output_pt, map_location="cpu", weights_only=False)
+    emb_shape = state["embeddings"].shape
+    log.info(f"\nVoice cloned successfully!")
+    log.info(f"  Output: {output_pt}")
+    log.info(f"  Embeddings: {emb_shape} ({emb_shape[0]} frames, ~{emb_shape[0] / 12.5:.1f}s)")
+    log.info(f"  Segments averaged: {len(all_embeddings)}")
+    log.info(f"  Preprocessing: denoise + silence trim + {target_lufs:.0f} LUFS normalize")
+
+    return output_pt
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="High-fidelity PersonaPlex voice cloning"
+    )
+    parser.add_argument(
+        "--input", "-i", required=True,
+        help="Input audio file (WAV, MP3, FLAC — any format torchaudio supports)"
+    )
+    parser.add_argument(
+        "--name", "-n", required=True,
+        help="Name for the cloned voice (e.g. 'MyVoice')"
+    )
+    parser.add_argument(
+        "--device", "-d", default="cuda",
+        help="Device to run on (default: cuda)"
+    )
+    parser.add_argument(
+        "--cpu-offload", action="store_true",
+        help="Offload to CPU if GPU memory is insufficient"
+    )
+    parser.add_argument(
+        "--hf-repo", default="nvidia/personaplex-7b-v1",
+        help="HuggingFace model repo"
+    )
+    parser.add_argument(
+        "--segments", "-s", type=int, default=1,
+        help="Number of segments to average (1 = single pass, 3+ = multi-segment averaging for longer clips)"
+    )
+    parser.add_argument(
+        "--lufs", type=float, default=-20.0,
+        help="Target LUFS normalization (default: -20, PersonaPlex default is -24)"
+    )
+    parser.add_argument(
+        "--max-duration", type=float, default=8.0,
+        help="Max duration per segment in seconds (default: 8, built-in voices use ~4)"
+    )
+    parser.add_argument(
+        "--no-denoise", action="store_true",
+        help="Skip noise reduction (if reference is already clean)"
+    )
+
+    args = parser.parse_args()
+
+    if not os.path.exists(args.input):
+        print(f"Error: Input file not found: {args.input}")
+        sys.exit(1)
+
+    import torch
+    with torch.no_grad():
+        result = clone_voice_multiseg(
+            input_wav=args.input,
+            output_name=args.name,
+            device=args.device,
+            hf_repo=args.hf_repo,
+            cpu_offload=args.cpu_offload,
+            n_segments=args.segments,
+            target_lufs=args.lufs,
+            max_seg_duration=args.max_duration,
+        )
+
+    sys.exit(0 if result else 1)
+
+
+if __name__ == "__main__":
+    main()

package/voices/personaplex/dequant-loader.py

@@ -0,0 +1,174 @@
+#!/usr/bin/env python3
+"""
+dequant-loader.py — Pre-dequantize quantized PersonaPlex weights to bf16 cache.
+
+For NF4 (INT4) or TurboQuant 2-bit weights, dequantizes to a temporary
+bf16 safetensors file that moshi.server can load natively.
+
+Usage:
+  python dequant-loader.py --input model-nf4.safetensors --output /tmp/model-bf16.safetensors
+  python dequant-loader.py --input model-turbo2bit.safetensors --output /tmp/model-bf16.safetensors
+
+The output file can then be passed to moshi.server via --moshi-weight.
+"""
+
+import os, sys, math, time
+import torch
+from safetensors.torch import load_file, save_file
+
+NF2_CENTROIDS = torch.tensor([-1.5104, -0.4528, 0.4528, 1.5104])
+
+
+def fast_wht(x):
+    """Vectorized Walsh-Hadamard Transform."""
+    n = x.shape[-1]
+    h = 1
+    while h < n:
+        x_view = x.view(*x.shape[:-1], -1, 2, h)
+        a = x_view[..., 0, :].clone()
+        b = x_view[..., 1, :].clone()
+        x_view[..., 0, :] = a + b
+        x_view[..., 1, :] = a - b
+        x = x_view.reshape(*x.shape)
+        h *= 2
+    return x / math.sqrt(n)
+
+
+def detect_format(state):
+    """Detect if weights are NF4 (INT4), TurboQuant 2-bit, or plain."""
+    has_scales = any(k.endswith(".__scales__") for k in state)
+    has_packed = any(k.endswith(".packed") for k in state)
+    if has_packed:
+        return "turbo2bit"
+    if has_scales:
+        return "nf4"
+    return "plain"
+
+
+def dequant_nf4(state):
+    """Dequantize INT4 NF4 weights."""
+    result = {}
+    processed = set()
+
+    for name in list(state.keys()):
+        if name.endswith(".__scales__") or name.endswith(".__shape__") or name.endswith(".__numel__"):
+            continue
+        if name in processed:
+            continue
+
+        scales_key = f"{name}.__scales__"
+        if scales_key in state:
+            packed = state[name]
+            scales = state[scales_key].float()
+            shape = state[f"{name}.__shape__"].tolist()
+            numel = state[f"{name}.__numel__"].item()
+            group_size = 64
+
+            lo = (packed & 0x0F).to(torch.int8) - 8
+            hi = ((packed >> 4) & 0x0F).to(torch.int8) - 8
+            unpacked = torch.zeros(packed.numel() * 2, dtype=torch.float32)
+            unpacked[0::2] = lo.float()
+            unpacked[1::2] = hi.float()
+
+            n_groups = scales.numel()
+            groups = unpacked[:n_groups * group_size].reshape(n_groups, group_size)
+            deq = (groups * scales.unsqueeze(1)).reshape(-1)[:numel]
+
+            orig_shape = [s for s in shape if s > 0]
+            result[name] = deq.reshape(orig_shape).to(torch.bfloat16)
+            processed.add(name)
+        else:
+            result[name] = state[name].to(torch.bfloat16)
+            processed.add(name)
+
+    return result
+
+
+def dequant_turbo2bit(state):
+    """Dequantize TurboQuant 2-bit (NF2 + WHT) weights."""
+    result = {}
+    processed = set()
+
+    for name in list(state.keys()):
+        if any(name.endswith(f".{s}") for s in ["packed", "scales", "shape", "numel", "gs", "np2"]):
+            continue
+        if name in processed:
+            continue
+
+        packed_key = f"{name}.packed"
+        if packed_key in state:
+            gs = state[f"{name}.gs"].item()
+            gs_pow2 = state[f"{name}.np2"].item()
+            numel = state[f"{name}.numel"].item()
+            shape = [s for s in state[f"{name}.shape"].tolist() if s > 0]
+            scales = state[f"{name}.scales"].float()
+            packed = state[packed_key]
+            n_groups = scales.numel()
+
+            # Unpack 2-bit
+            p = packed.reshape(n_groups, gs // 4)
+            codes = torch.zeros(n_groups, gs, dtype=torch.long)
+            for i in range(4):
+                codes[:, i::4] = (p >> (2 * i)) & 0x03
+
+            dequant = NF2_CENTROIDS[codes]
+
+            # Inverse WHT
+            if gs_pow2 > gs:
+                dequant = torch.cat([dequant, torch.zeros(n_groups, gs_pow2 - gs)], dim=1)
+            dequant = fast_wht(dequant)
+            dequant = dequant[:, :gs]
+
+            dequant = dequant * scales.unsqueeze(1)
+            result[name] = dequant.reshape(-1)[:numel].reshape(shape).to(torch.bfloat16)
+            processed.add(name)
+        else:
+            result[name] = state[name].to(torch.bfloat16)
+            processed.add(name)
+
+    return result
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description="Dequantize PersonaPlex weights to bf16")
+    parser.add_argument("--input", "-i", required=True, help="Quantized safetensors file")
+    parser.add_argument("--output", "-o", required=True, help="Output bf16 safetensors file")
+    parser.add_argument("--device", "-d", default="cpu", help="Device for dequantization")
+    args = parser.parse_args()
+
+    if not os.path.exists(args.input):
+        print(f"Error: {args.input} not found")
+        sys.exit(1)
+
+    # Skip if output already exists and is newer than input
+    if os.path.exists(args.output) and os.path.getmtime(args.output) > os.path.getmtime(args.input):
+        print(f"Cached: {args.output} is up to date")
+        sys.exit(0)
+
+    print(f"Loading {args.input}...")
+    t0 = time.time()
+    state = load_file(args.input, device=args.device)
+
+    fmt = detect_format(state)
+    print(f"Format: {fmt}")
+
+    if fmt == "nf4":
+        result = dequant_nf4(state)
+    elif fmt == "turbo2bit":
+        result = dequant_turbo2bit(state)
+    else:
+        print("Already plain bf16/fp16 — copying")
+        result = {k: v.to(torch.bfloat16) for k, v in state.items()}
+
+    t1 = time.time()
+    print(f"Dequantized {len(result)} tensors in {t1-t0:.1f}s")
+
+    print(f"Saving to {args.output}...")
+    save_file(result, args.output)
+    size_gb = os.path.getsize(args.output) / 1024**3
+    print(f"Done: {size_gb:.2f} GB")
+
+
+if __name__ == "__main__":
+    main()