rlx_tsac/native/
mod.rs

1//! Native TSAC inference (vendored [tsac-ng](https://github.com/Hope2333/tsac-ng)).
2//!
3//! Maps RLX [`Device`] tags to tsac-ng backends (CPU / CUDA / HIP / Vulkan).
4
5pub(crate) mod ffi;
6
7use crate::TsacOptions;
8use crate::device::resolve_codec_device;
9use crate::download::{dac_model_path, weights_available};
10use anyhow::{Context, Result, bail};
11use rlx_runtime::Device;
12use std::ffi::CString;
13use std::path::{Path, PathBuf};
14use std::ptr;
15
16pub use ffi::TsacBackend;
17
18pub struct NativeCodec {
19    install_dir: PathBuf,
20    ctx: *mut ffi::TSACContext,
21    device: Device,
22    /// RLX-graph decoder (cpu/metal/mlx/wgpu/…). Present when a GPU backend was
23    /// requested and is available; decode then runs natively on that backend,
24    /// bit-exact with the vendored C decoder. Encode stays on the C path.
25    rlx_decoder: Option<crate::rlx_decode::RlxDecoder>,
26}
27
28impl NativeCodec {
29    pub fn open(install_dir: impl AsRef<Path>, options: &TsacOptions) -> Result<Self> {
30        let install_dir = install_dir.as_ref().to_path_buf();
31        if !weights_available(&install_dir) {
32            bail!(
33                "missing TSAC weights under {} — run `just fetch-tsac`",
34                install_dir.display()
35            );
36        }
37        let device = resolve_codec_device(options.device);
38        let backend = rlx_device_to_tsac_backend(device);
39        let model_dir = CString::new(install_dir.to_string_lossy().into_owned())
40            .context("install dir utf-8")?;
41        let ctx = unsafe { ffi::tsac_init(backend, 0, model_dir.as_ptr()) };
42        if ctx.is_null() {
43            bail!("tsac_init failed for backend {backend:?} on {device:?}");
44        }
45        // Build the RLX-graph decoder for GPU backends (weights are copied out of
46        // ctx, so it stays valid for the lifetime of this codec).
47        let rlx_decoder = match rlx_graph_device(options.device) {
48            Some(dev) => crate::rlx_decode::RlxDecoder::from_ctx(ctx, dev).ok(),
49            None => None,
50        };
51        Ok(Self {
52            install_dir,
53            ctx,
54            device,
55            rlx_decoder,
56        })
57    }
58
59    pub fn device(&self) -> Device {
60        self.device
61    }
62
63    /// Raw context pointer — used by the RLX graph backend to pull bit-exact
64    /// dequantized weights via the C `dequant_weights` bridge.
65    pub(crate) fn ctx_raw(&self) -> *mut ffi::TSACContext {
66        self.ctx
67    }
68
69    pub fn encode_file(&self, in_wav: &Path, out_tsac: &Path, options: &TsacOptions) -> Result<()> {
70        let in_c = path_to_c(in_wav)?;
71        let out_c = path_to_c(out_tsac)?;
72        let n_codebooks = options.quality.unwrap_or(12) as i32;
73        let rc = unsafe {
74            ffi::tsac_compress_file(self.ctx, in_c.as_ptr(), out_c.as_ptr(), n_codebooks)
75        };
76        if rc != ffi::TSAC_OK {
77            bail!("tsac_compress_file failed ({rc})");
78        }
79        Ok(())
80    }
81
82    pub fn decode_file(&self, in_tsac: &Path, out_wav: &Path) -> Result<()> {
83        // GPU backends decode through the RLX graph (bit-exact with C).
84        if let Some(rlx) = &self.rlx_decoder {
85            return rlx.decode_file(in_tsac, out_wav);
86        }
87        let in_c = path_to_c(in_tsac)?;
88        let out_c = path_to_c(out_wav)?;
89        let rc = unsafe { ffi::tsac_decompress_file(self.ctx, in_c.as_ptr(), out_c.as_ptr()) };
90        if rc != ffi::TSAC_OK {
91            bail!("tsac_decompress_file failed ({rc})");
92        }
93        Ok(())
94    }
95
96    pub fn model_hint(&self) -> PathBuf {
97        dac_model_path(&self.install_dir, false)
98    }
99}
100
101impl Drop for NativeCodec {
102    fn drop(&mut self) {
103        if !self.ctx.is_null() {
104            unsafe { ffi::tsac_free(self.ctx) };
105            self.ctx = ptr::null_mut();
106        }
107    }
108}
109
110unsafe impl Send for NativeCodec {}
111
112/// The rlx-runtime device to run the graph decoder on for a requested device, or
113/// `None` to keep the C decode path (CPU, or backend unavailable at runtime).
114fn rlx_graph_device(requested: Device) -> Option<Device> {
115    match requested {
116        Device::Metal
117        | Device::Mlx
118        | Device::Gpu
119        | Device::Cuda
120        | Device::Rocm
121        | Device::Vulkan
122            if rlx_runtime::is_available(requested) =>
123        {
124            Some(requested)
125        }
126        _ => None,
127    }
128}
129
130fn path_to_c(path: &Path) -> Result<CString> {
131    CString::new(path.to_string_lossy().into_owned()).context("path utf-8")
132}
133
134fn rlx_device_to_tsac_backend(device: Device) -> TsacBackend {
135    match device {
136        Device::Cuda => TsacBackend::Cuda,
137        Device::Rocm => TsacBackend::Hip,
138        Device::Vulkan | Device::Gpu => TsacBackend::Vulkan,
139        // Metal / MLX: MoltenVK when vulkan feature is on; otherwise CPU SIMD.
140        Device::Metal | Device::Mlx => {
141            if cfg!(feature = "vulkan") {
142                TsacBackend::Vulkan
143            } else {
144                TsacBackend::Cpu
145            }
146        }
147        _ => TsacBackend::Cpu,
148    }
149}
rlx_tsac/native/mod.rs

rlx_tsac/native/
mod.rs
