Initial commit

ackit/encoder.py

@@ -0,0 +1,155 @@
+"""Encode audio into game-compatible WEMs (portable: libvorbis via libsndfile, no Wwise, no ffmpeg).
+
+Pipeline: float32 PCM -> libvorbis (Ogg) -> `builder.ogg_to_wem` (Wwise 2011.1 layout, game codebook
+IDs, FNV-1 hash, decode-alloc size). The quality is chosen so that every codebook libvorbis uses exists
+in the game's built-in library (docs/formats/wem.md); setups the game itself ships are preferred.
+"""
+import io
+import json
+import os
+
+import numpy as np
+import soundfile as sf
+
+from . import audio
+from .builder import ogg_to_wem
+from .convert import library
+from .known_setups import lookup_setup
+from .ogg import read_packets
+from .paths import user_dir
+from .wwvorbis import V2W, transcode_setup
+
+QUALITIES = [q / 2 for q in range(20, -1, -1)]  # 10.0 .. 0.0 (libvorbis q; libsndfile cannot do q<0)
+AUTO_MAX_Q_KNOWN = 6.0   # highest q considered when looking for a setup the game ships
+AUTO_MAX_Q = 4.0         # otherwise, highest compatible q (game files sit around q3-5)
+FALLBACK_RATES = (32000, 28000, 24000)          # e.g. 6 channels at 44.1/48 kHz
+FIT_RATES = (24000, 22050, 16000, 12000, 11025, 8000)
+
+
+class EncodeError(ValueError):
+    pass
+
+
+def vorbis_ogg(x, rate, quality):
+    """Encode float32 PCM to Ogg Vorbis bytes at libvorbis quality `quality` (0..10)."""
+    b = io.BytesIO()
+    level = min(1.0, max(0.0, 1.0 - quality / 10.0))
+    x = np.ascontiguousarray(audio._as2d(x))
+    # write in blocks: libsndfile's Vorbis writer overflows the C stack on very large single writes
+    with sf.SoundFile(b, "w", rate, x.shape[1], format="OGG", subtype="VORBIS", compression_level=level) as f:
+        for i in range(0, len(x), 4096):
+            f.write(x[i:i + 4096])
+    return b.getvalue()
+
+
+# --- setup probing (which q values are usable at a rate/channel count) ------------------------
+_cache = None
+
+
+def _cache_path():
+    return os.path.join(user_dir("cache"), "vorbis_probe.json")
+
+
+def _probe_cache():
+    global _cache
+    if _cache is None:
+        try:
+            c = json.load(open(_cache_path()))
+            _cache = c if c.get("_lib") == sf.__libsndfile_version__ else None
+        except (OSError, ValueError):
+            _cache = None
+        _cache = _cache or {"_lib": sf.__libsndfile_version__}
+    return _cache
+
+
+def probe_setup(rate, channels, quality):
+    """(stripped Wwise setup or None, error) for libvorbis at these settings (cached)."""
+    key = "%d,%d,%.1f" % (rate, channels, quality)
+    c = _probe_cache()
+    if key not in c:
+        noise = (np.random.default_rng(7).standard_normal((rate, channels)) * 0.1).astype(np.float32)
+        pk = read_packets(vorbis_ogg(noise, rate, quality))
+        try:
+            s, _ = transcode_setup(pk[2][0], V2W, channels, library())
+            c[key] = [s.hex(), None]
+        except ValueError as e:
+            c[key] = [None, str(e)]
+        try:
+            json.dump(c, open(_cache_path(), "w"))
+        except OSError:
+            pass
+    s, err = c[key]
+    return (bytes.fromhex(s) if s else None), err
+
+
+def choose_quality(rate, channels, quality="auto"):
+    """-> (quality, known) for a rate/channel count, or raise EncodeError."""
+    if quality != "auto":
+        q = float(quality)
+        s, err = probe_setup(rate, channels, q)
+        if s is None:
+            raise EncodeError("quality %s cannot be used at %d Hz / %d ch: %s" % (quality, rate, channels, err))
+        return q, bool(lookup_setup(s))
+    compatible = None
+    for q in [x for x in QUALITIES if x <= AUTO_MAX_Q_KNOWN]:
+        s, _ = probe_setup(rate, channels, q)
+        if s is None:
+            continue
+        if lookup_setup(s):
+            return q, True
+        if compatible is None and q <= AUTO_MAX_Q:
+            compatible = q
+    if compatible is None:
+        raise EncodeError("no libvorbis quality yields game-compatible codebooks at %d Hz / %d ch" % (rate, channels))
+    return compatible, False
+
+
+def pick_format(rate, channels, quality="auto"):
+    """Find a usable (rate, quality). Falls back to lower rates when needed (Wwise resamples every
+    source, so playback speed is unaffected)."""
+    try:
+        return (rate,) + choose_quality(rate, channels, quality)
+    except EncodeError as first:
+        for r in FALLBACK_RATES:
+            if r < rate:
+                try:
+                    return (r,) + choose_quality(r, channels, quality)
+                except EncodeError:
+                    pass
+        raise first
+
+
+def encode(x, src_rate, rate, channels, quality="auto"):
+    """PCM -> WEM bytes in the requested format. Returns (wem, info)."""
+    r, q, known = pick_format(rate, channels, quality)
+    pcm = audio.conform(x, src_rate, r, channels)
+    if len(pcm) < 64:
+        raise EncodeError("the audio is empty")
+    wem = ogg_to_wem(vorbis_ogg(pcm, r, q))
+    return wem, dict(rate=r, quality=q, known_setup=known, size=len(wem), channels=channels,
+                     seconds=len(pcm) / float(r), requested_rate=rate)
+
+
+def encode_fit(x, src_rate, rate, channels, limit=None, quality="auto", progress=None):
+    """Like `encode`, but if the result is larger than `limit` bytes lower the quality step by step,
+    then the sample rate, until it fits. Raises EncodeError if it cannot fit."""
+    wem, info = encode(x, src_rate, rate, channels, quality)
+    if limit is None or len(wem) <= limit:
+        return wem, info
+    first = len(wem)
+    for r in [info["rate"]] + [f for f in FIT_RATES if f < info["rate"]]:
+        pcm = audio.conform(x, src_rate, r, channels)
+        for q in QUALITIES:
+            if r == info["rate"] and q >= info["quality"]:
+                continue
+            if probe_setup(r, channels, q)[0] is None:
+                continue
+            if progress:
+                progress("shrinking: %d Hz, quality %.1f" % (r, q))
+            w = ogg_to_wem(vorbis_ogg(pcm, r, q))
+            if len(w) <= limit:
+                info.update(rate=r, quality=q, size=len(w), shrunk_from=first,
+                            known_setup=bool(lookup_setup(probe_setup(r, channels, q)[0])))
+                return w, info
+    raise EncodeError("the audio is too long for this sound even at the lowest quality "
+                      "(%d bytes, maximum %d). Shorten it." % (first, limit))