tts-voice-creator-clone-and-design-2/core/audio.py

"""Audio helpers: conversion, trim, normalize, loudness, auto-trim."""
from __future__ import annotations

import io
import subprocess
import uuid
import wave
from pathlib import Path

from pydub import AudioSegment

from core.constants import _VOICE_TARGET_DBFS, _VOICE_PEAK_DBFS
from core.registry import TEMP_DIR


def _to_wav_24k(src: Path) -> Path:
    out = TEMP_DIR / f"{src.stem}_24k.wav"
    seg = AudioSegment.from_file(str(src))
    seg = seg.set_frame_rate(24000).set_channels(1).set_sample_width(2)
    seg.export(str(out), format="wav")
    return out


def _to_wav_16k(src: Path) -> Path:
    """16kHz mono WAV — required by Whisper/WhisperX VAD and wav2vec2 alignment."""
    out = TEMP_DIR / f"{src.stem}_16k.wav"
    seg = AudioSegment.from_file(str(src))
    seg = seg.set_frame_rate(16000).set_channels(1).set_sample_width(2)
    seg.export(str(out), format="wav")
    return out


def _trim(src: Path, start_s: float, end_s: float) -> Path:
    seg = AudioSegment.from_file(str(src))
    trimmed = seg[int(start_s * 1000):int(end_s * 1000)]
    trimmed, _ = _normalize_segment(trimmed)
    out = TEMP_DIR / f"{uuid.uuid4().hex}_trimmed.wav"
    trimmed.export(str(out), format="wav")
    return out


def _duration(path: Path) -> float:
    if path.suffix.lower() == ".wav":
        try:
            with wave.open(str(path), "rb") as wf:
                frames = wf.getnframes()
                rate = wf.getframerate()
                if rate:
                    return frames / float(rate)
        except Exception:
            pass

    try:
        seg = AudioSegment.from_file(str(path))
        return len(seg) / 1000.0
    except Exception:
        pass

    probe = subprocess.run(
        [
            "ffprobe", "-v", "error",
            "-show_entries", "format=duration",
            "-of", "default=noprint_wrappers=1:nokey=1",
            str(path),
        ],
        capture_output=True,
        text=True,
        timeout=10,
    )
    if probe.returncode != 0:
        raise RuntimeError(probe.stderr.strip() or "ffprobe failed")
    return float(probe.stdout.strip())


def _normalize_segment(seg: AudioSegment, target_dbfs: float = _VOICE_TARGET_DBFS,
                       peak_dbfs: float = _VOICE_PEAK_DBFS) -> tuple[AudioSegment, dict]:
    before_dbfs = seg.dBFS if seg.dBFS != float("-inf") else None
    before_peak = seg.max_dBFS if seg.max_dBFS != float("-inf") else None
    if before_dbfs is None or before_peak is None:
        return seg, {"before_dbfs": before_dbfs, "after_dbfs": before_dbfs, "gain_db": 0.0, "peak_dbfs": before_peak}

    gain = target_dbfs - before_dbfs
    if before_peak + gain > peak_dbfs:
        gain = peak_dbfs - before_peak
    normalized = seg.apply_gain(gain)
    after_dbfs = normalized.dBFS if normalized.dBFS != float("-inf") else None
    after_peak = normalized.max_dBFS if normalized.max_dBFS != float("-inf") else None
    return normalized, {
        "before_dbfs": round(before_dbfs, 2),
        "after_dbfs": round(after_dbfs, 2) if after_dbfs is not None else None,
        "gain_db": round(gain, 2),
        "peak_dbfs": round(after_peak, 2) if after_peak is not None else None,
    }


def _export_normalized_wav(src: Path, dest: Path, target_dbfs: float = _VOICE_TARGET_DBFS) -> dict:
    seg = AudioSegment.from_file(str(src))
    seg = seg.set_frame_rate(24000).set_channels(1).set_sample_width(2)
    seg, info = _normalize_segment(seg, target_dbfs=target_dbfs)
    seg.export(str(dest), format="wav")
    return info


def _loudness_info(path: Path) -> dict:
    seg = AudioSegment.from_file(str(path))
    dbfs = seg.dBFS if seg.dBFS != float("-inf") else None
    peak = seg.max_dBFS if seg.max_dBFS != float("-inf") else None
    return {
        "dbfs": round(dbfs, 2) if dbfs is not None else None,
        "peak_dbfs": round(peak, 2) if peak is not None else None,
        "target_dbfs": _VOICE_TARGET_DBFS,
    }


def _auto_trim_bounds(path: Path) -> dict:
    seg = AudioSegment.from_file(str(path)).set_channels(1)
    dur_ms = len(seg)
    if dur_ms <= 20_000:
        return {
            "start": 0.0,
            "end": dur_ms / 1000.0,
            "duration": dur_ms / 1000.0,
            "reason": "Audio is already short enough.",
        }

    chunk_ms = 250
    chunks = []
    overall_db = seg.dBFS if seg.dBFS != float("-inf") else -60.0
    speech_floor = max(overall_db - 18.0, -45.0)

    for pos in range(0, dur_ms, chunk_ms):
        ch = seg[pos:pos + chunk_ms]
        db = ch.dBFS if ch.dBFS != float("-inf") else -80.0
        max_db = ch.max_dBFS if ch.max_dBFS != float("-inf") else -80.0
        chunks.append({"db": db, "speech": db >= speech_floor, "clipped": max_db > -1.0})

    def score_window(start_ms: int, length_ms: int) -> tuple[float, dict]:
        first = max(0, start_ms // chunk_ms)
        last = min(len(chunks), (start_ms + length_ms + chunk_ms - 1) // chunk_ms)
        win = chunks[first:last]
        if not win:
            return -9999.0, {}
        speech_ratio = sum(1 for c in win if c["speech"]) / len(win)
        silence_ratio = 1.0 - speech_ratio
        clip_ratio = sum(1 for c in win if c["clipped"]) / len(win)
        speech_dbs = [c["db"] for c in win if c["speech"]]
        avg_db = sum(speech_dbs) / len(speech_dbs) if speech_dbs else -80.0
        variance = sum((x - avg_db) ** 2 for x in speech_dbs) / len(speech_dbs) if speech_dbs else 100.0
        loudness_penalty = abs(avg_db - (-20.0)) * 1.7
        steadiness_penalty = min(18.0, variance ** 0.5 * 1.4)
        duration_s = length_ms / 1000.0
        duration_penalty = abs(duration_s - 12.0) * 0.9
        score = (
            speech_ratio * 100.0
            - silence_ratio * 55.0
            - clip_ratio * 85.0
            - loudness_penalty
            - steadiness_penalty
            - duration_penalty
        )
        return score, {
            "speech_ratio": speech_ratio,
            "silence_ratio": silence_ratio,
            "clip_ratio": clip_ratio,
            "avg_db": avg_db,
        }

    best = None
    window_lengths = [8_000, 10_000, 12_000, 15_000, 18_000]
    for length_ms in window_lengths:
        if length_ms > dur_ms:
            continue
        for start_ms in range(0, dur_ms - length_ms + 1, 500):
            score, metrics = score_window(start_ms, length_ms)
            if best is None or score > best["score"]:
                best = {"start_ms": start_ms, "length_ms": length_ms, "score": score, "metrics": metrics}

    if best is None:
        end_ms = min(dur_ms, 12_000)
        return {"start": 0.0, "end": end_ms / 1000.0, "duration": end_ms / 1000.0,
                "reason": "Using the beginning because no stable speech window was found."}

    start_s = best["start_ms"] / 1000.0
    end_s = (best["start_ms"] + best["length_ms"]) / 1000.0
    m = best["metrics"]
    return {
        "start": round(start_s, 2),
        "end": round(end_s, 2),
        "duration": round(end_s - start_s, 2),
        "score": round(best["score"], 2),
        "reason": (
            f"Selected {end_s - start_s:.1f}s with "
            f"{m.get('speech_ratio', 0) * 100:.0f}% speech, "
            f"{m.get('silence_ratio', 0) * 100:.0f}% silence, "
            f"avg {m.get('avg_db', -80):.1f} dBFS."
        ),
    }


def _audio_segment_from_wav(audio: bytes) -> AudioSegment:
    segment = AudioSegment.from_file(io.BytesIO(audio), format="wav")
    return segment.set_channels(1).set_sample_width(2).set_frame_rate(24000)


def _wav_bytes_from_segment(segment: AudioSegment) -> bytes:
    out = io.BytesIO()
    segment.export(out, format="wav")
    return out.getvalue()