Skip to main content

rlx_orpheus/decoder/
eager.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4//! SNAC 24 kHz decoder — eager CPU inference from safetensors weights.
5
6use std::path::{Path, PathBuf};
7
8use anyhow::{Context, Result, bail};
9use ndarray::{Array1, Array2, Array3, ArrayView2, s};
10use safetensors::SafeTensors;
11
12use super::config::SnacConfig;
13use super::ops::{
14    conv_transpose1d, conv1d, embed_codes, noise_block, noise_block_zero, repeat_interleave_time,
15    residual_add, snake1d,
16};
17
18/// Output sample rate (Hz).
19pub const SAMPLE_RATE: u32 = 24_000;
20
21/// PCM samples per Orpheus streaming slice (`2048` = one quarter of a 8192-sample SNAC frame).
22pub const SAMPLES_PER_FRAME: usize = 2048;
23
24fn load_f32(st: &SafeTensors<'_>, name: &str) -> Result<Vec<f32>> {
25    let view = st
26        .tensor(name)
27        .with_context(|| format!("Missing weight: {name}"))?;
28    let raw = view.data();
29    use safetensors::tensor::Dtype;
30    Ok(match view.dtype() {
31        Dtype::F32 => {
32            assert!(
33                raw.len() % 4 == 0,
34                "F32 tensor byte length not divisible by 4"
35            );
36            let n = raw.len() / 4;
37            let mut out = Vec::with_capacity(n);
38            #[cfg(target_endian = "little")]
39            unsafe {
40                std::ptr::copy_nonoverlapping(raw.as_ptr(), out.as_mut_ptr() as *mut u8, raw.len());
41                out.set_len(n);
42            }
43            #[cfg(not(target_endian = "little"))]
44            {
45                out.extend(
46                    raw.chunks_exact(4)
47                        .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])),
48                );
49            }
50            out
51        }
52        dt => bail!("Tensor {name}: unsupported dtype {dt:?} (expected F32)"),
53    })
54}
55
56fn as1d(data: Vec<f32>, n: usize) -> Array1<f32> {
57    Array1::from_shape_vec(n, data).expect("1-D shape mismatch")
58}
59
60fn as2d(data: Vec<f32>, rows: usize, cols: usize) -> Array2<f32> {
61    Array2::from_shape_vec((rows, cols), data).expect("2-D shape mismatch")
62}
63
64fn as3d(data: Vec<f32>, d0: usize, d1: usize, d2: usize) -> Array3<f32> {
65    Array3::from_shape_vec((d0, d1, d2), data).expect("3-D shape mismatch")
66}
67
68pub(crate) struct ResidualUnitWeights {
69    pub(crate) snake1_alpha: Array1<f32>,
70    pub(crate) conv1_w: Array3<f32>,
71    pub(crate) conv1_b: Array1<f32>,
72    pub(crate) conv1_pad: usize,
73    pub(crate) conv1_dilation: usize,
74    pub(crate) snake2_alpha: Array1<f32>,
75    pub(crate) conv2_w: Array3<f32>,
76    pub(crate) conv2_b: Array1<f32>,
77    pub(crate) groups: usize,
78}
79
80pub(crate) struct DecoderBlockWeights {
81    pub(crate) snake_alpha: Array1<f32>,
82    pub(crate) upsample_w: Array3<f32>,
83    pub(crate) upsample_b: Array1<f32>,
84    pub(crate) noise_w: Array3<f32>,
85    pub(crate) stride: usize,
86    pub(crate) residual_units: [ResidualUnitWeights; 3],
87}
88
89struct VectorQuantizeWeights {
90    codebook: Array2<f32>,
91    out_proj_w: Array3<f32>,
92    out_proj_b: Array1<f32>,
93    stride: usize,
94}
95
96pub(crate) struct SnacDecoderInner {
97    pub(crate) config: SnacConfig,
98    quantizers: Vec<VectorQuantizeWeights>,
99    pub(crate) init_dw_w: Array3<f32>,
100    pub(crate) init_dw_b: Array1<f32>,
101    pub(crate) init_pw_w: Array3<f32>,
102    pub(crate) init_pw_b: Array1<f32>,
103    pub(crate) blocks: Vec<DecoderBlockWeights>,
104    pub(crate) final_snake_alpha: Array1<f32>,
105    pub(crate) final_conv_w: Array3<f32>,
106    pub(crate) final_conv_b: Array1<f32>,
107}
108
109fn load_alpha(st: &SafeTensors<'_>, key: &str) -> Result<Array1<f32>> {
110    let data = load_f32(st, key)?;
111    let c = data.len();
112    Ok(as1d(data, c))
113}
114
115fn load_residual_unit(
116    st: &SafeTensors<'_>,
117    prefix: &str,
118    dim: usize,
119    groups: usize,
120    dilation: usize,
121) -> Result<ResidualUnitWeights> {
122    let pad = ((7 - 1) * dilation) / 2;
123    Ok(ResidualUnitWeights {
124        snake1_alpha: load_alpha(st, &format!("{prefix}.0.alpha"))?,
125        conv1_w: as3d(
126            load_f32(st, &format!("{prefix}.1.weight"))?,
127            dim,
128            dim / groups,
129            7,
130        ),
131        conv1_b: as1d(load_f32(st, &format!("{prefix}.1.bias"))?, dim),
132        conv1_pad: pad,
133        conv1_dilation: dilation,
134        snake2_alpha: load_alpha(st, &format!("{prefix}.2.alpha"))?,
135        conv2_w: as3d(load_f32(st, &format!("{prefix}.3.weight"))?, dim, dim, 1),
136        conv2_b: as1d(load_f32(st, &format!("{prefix}.3.bias"))?, dim),
137        groups,
138    })
139}
140
141fn residual_unit_forward(x: ArrayView2<f32>, w: &ResidualUnitWeights) -> Array2<f32> {
142    let mut h = snake1d(x, w.snake1_alpha.view());
143    h = conv1d(
144        h.view(),
145        w.conv1_w.view(),
146        Some(w.conv1_b.view()),
147        w.conv1_pad,
148        w.groups,
149        w.conv1_dilation,
150    );
151    h = snake1d(h.view(), w.snake2_alpha.view());
152    h = conv1d(h.view(), w.conv2_w.view(), Some(w.conv2_b.view()), 0, 1, 1);
153    residual_add(x, h.view())
154}
155
156fn load_decoder_block(
157    st: &SafeTensors<'_>,
158    prefix: &str,
159    input_dim: usize,
160    output_dim: usize,
161    stride: usize,
162    depthwise: bool,
163) -> Result<DecoderBlockWeights> {
164    let k = 2 * stride;
165    let ru_groups = if depthwise { output_dim } else { 1 };
166    let block = format!("{prefix}.block");
167    Ok(DecoderBlockWeights {
168        snake_alpha: load_alpha(st, &format!("{block}.0.alpha"))?,
169        upsample_w: as3d(
170            load_f32(st, &format!("{block}.1.weight"))?,
171            input_dim,
172            output_dim,
173            k,
174        ),
175        upsample_b: as1d(load_f32(st, &format!("{block}.1.bias"))?, output_dim),
176        noise_w: as3d(
177            load_f32(st, &format!("{block}.2.linear.weight"))?,
178            output_dim,
179            output_dim,
180            1,
181        ),
182        stride,
183        residual_units: [
184            load_residual_unit(st, &format!("{block}.3.block"), output_dim, ru_groups, 1)?,
185            load_residual_unit(st, &format!("{block}.4.block"), output_dim, ru_groups, 3)?,
186            load_residual_unit(st, &format!("{block}.5.block"), output_dim, ru_groups, 9)?,
187        ],
188    })
189}
190
191fn decoder_block_forward(
192    x: ArrayView2<f32>,
193    w: &DecoderBlockWeights,
194    noise: &mut NoiseState,
195) -> Array2<f32> {
196    let mut h = snake1d(x, w.snake_alpha.view());
197    let padding = w.stride.div_ceil(2);
198    let output_padding = w.stride % 2;
199    h = conv_transpose1d(
200        h.view(),
201        w.upsample_w.view(),
202        Some(w.upsample_b.view()),
203        w.stride,
204        padding,
205        output_padding,
206        1,
207    );
208    if noise.enabled() {
209        let t = h.shape()[1];
210        let n = noise.next_plane(t);
211        h = noise_block(h.view(), w.noise_w.view(), n.view());
212    } else {
213        h = noise_block_zero(h.view());
214    }
215    for ru in &w.residual_units {
216        h = residual_unit_forward(h.view(), ru);
217    }
218    h
219}
220
221fn load_inner(st: &SafeTensors<'_>, config: SnacConfig) -> Result<SnacDecoderInner> {
222    let latent = config.latent_dim();
223    let decoder_dim = config.decoder_dim;
224
225    let quantizers = config
226        .vq_strides
227        .iter()
228        .enumerate()
229        .map(|(i, &stride)| {
230            let prefix = format!("quantizer.quantizers.{i}");
231            Ok(VectorQuantizeWeights {
232                codebook: as2d(
233                    load_f32(st, &format!("{prefix}.codebook.weight"))?,
234                    config.codebook_size,
235                    config.codebook_dim,
236                ),
237                out_proj_w: as3d(
238                    load_f32(st, &format!("{prefix}.out_proj.weight"))?,
239                    latent,
240                    config.codebook_dim,
241                    1,
242                ),
243                out_proj_b: as1d(load_f32(st, &format!("{prefix}.out_proj.bias"))?, latent),
244                stride,
245            })
246        })
247        .collect::<Result<Vec<_>>>()?;
248
249    let init_dw_w = as3d(load_f32(st, "decoder.model.0.weight")?, latent, 1, 7);
250    let init_dw_b = as1d(load_f32(st, "decoder.model.0.bias")?, latent);
251    let init_pw_w = as3d(
252        load_f32(st, "decoder.model.1.weight")?,
253        decoder_dim,
254        latent,
255        1,
256    );
257    let init_pw_b = as1d(load_f32(st, "decoder.model.1.bias")?, decoder_dim);
258
259    let blocks = config
260        .decoder_rates
261        .iter()
262        .enumerate()
263        .map(|(i, &stride)| {
264            let input_dim = decoder_dim / 2_usize.pow(i as u32);
265            let output_dim = decoder_dim / 2_usize.pow(i as u32 + 1);
266            load_decoder_block(
267                st,
268                &format!("decoder.model.{}", i + 2),
269                input_dim,
270                output_dim,
271                stride,
272                config.depthwise,
273            )
274        })
275        .collect::<Result<Vec<_>>>()?;
276
277    let final_dim = decoder_dim / 2_usize.pow(config.decoder_rates.len() as u32);
278    let final_snake_alpha = load_alpha(st, "decoder.model.6.alpha")?;
279    let final_conv_w = as3d(load_f32(st, "decoder.model.7.weight")?, 1, final_dim, 7);
280    let final_conv_b = as1d(load_f32(st, "decoder.model.7.bias")?, 1);
281
282    Ok(SnacDecoderInner {
283        config,
284        quantizers,
285        init_dw_w,
286        init_dw_b,
287        init_pw_w,
288        init_pw_b,
289        blocks,
290        final_snake_alpha,
291        final_conv_w,
292        final_conv_b,
293    })
294}
295
296pub(crate) fn from_codes_inner(
297    inner: &SnacDecoderInner,
298    codes_0: &[i32],
299    codes_1: &[i32],
300    codes_2: &[i32],
301) -> Result<Array2<f32>> {
302    let base_len = codes_0.len();
303    let finest = inner.config.vq_strides[0];
304    let all_codes = [codes_0, codes_1, codes_2];
305    for (i, codes) in all_codes.iter().enumerate() {
306        let stride = inner.config.vq_strides[i];
307        if codes.len() * stride != base_len * finest {
308            bail!(
309                "codes_{i} length {} inconsistent with codes_0={base_len} (need len*stride={})",
310                codes.len(),
311                base_len * finest
312            );
313        }
314    }
315    let mut z_q: Option<Array2<f32>> = None;
316    for (q, codes) in inner.quantizers.iter().zip(all_codes) {
317        validate_codes(codes, inner.config.codebook_size)?;
318        let z_p = embed_codes(codes, q.codebook.view());
319        let mut z_q_i = conv1d(
320            z_p.view(),
321            q.out_proj_w.view(),
322            Some(q.out_proj_b.view()),
323            0,
324            1,
325            1,
326        );
327        if q.stride > 1 {
328            z_q_i = repeat_interleave_time(z_q_i.view(), q.stride);
329        }
330        z_q = Some(match z_q {
331            None => z_q_i,
332            Some(acc) => &acc + &z_q_i,
333        });
334    }
335    Ok(z_q.expect("quantizers non-empty"))
336}
337
338fn validate_codes(codes: &[i32], codebook_size: usize) -> Result<()> {
339    for (i, &c) in codes.iter().enumerate() {
340        if !(0..codebook_size as i32).contains(&c) {
341            bail!("code index {c} at position {i} out of range 0..{codebook_size}");
342        }
343    }
344    Ok(())
345}
346
347/// Per-decode noise planes for SNAC NoiseBlocks (PyTorch `torch.randn(1,1,T)` per block).
348pub(crate) struct NoiseState {
349    block: usize,
350    fixed: Option<Vec<Vec<f32>>>,
351    rng: Option<rand::rngs::StdRng>,
352    enabled: bool,
353}
354
355impl NoiseState {
356    fn disabled() -> Self {
357        Self {
358            block: 0,
359            fixed: None,
360            rng: None,
361            enabled: false,
362        }
363    }
364
365    fn from_reference_dir(dir: &Path) -> Result<Option<Self>> {
366        let mut planes = Vec::new();
367        for i in 0..4 {
368            let path = dir.join(format!("ref_noise_{i}.npy"));
369            if !path.is_file() {
370                return Ok(None);
371            }
372            planes.push(load_npy_f32(&path)?);
373        }
374        Ok(Some(Self {
375            block: 0,
376            fixed: Some(planes),
377            rng: None,
378            enabled: true,
379        }))
380    }
381
382    pub(crate) fn seeded(seed: u64) -> Self {
383        use rand::SeedableRng;
384        Self {
385            block: 0,
386            fixed: None,
387            rng: Some(rand::rngs::StdRng::seed_from_u64(seed)),
388            enabled: true,
389        }
390    }
391
392    fn enabled(&self) -> bool {
393        self.enabled
394    }
395
396    pub(crate) fn next_plane(&mut self, t: usize) -> Array1<f32> {
397        if let Some(fixed) = &self.fixed {
398            let lane = &fixed[self.block.min(fixed.len() - 1)];
399            self.block += 1;
400            debug_assert!(
401                lane.len() >= t,
402                "reference noise lane too short: need {t}, have {}",
403                lane.len()
404            );
405            return Array1::from(lane[..t].to_vec());
406        }
407        let rng = self
408            .rng
409            .as_mut()
410            .expect("seeded noise requires ORPHEUS_SNAC_SEED");
411        use rand_distr::{Distribution, StandardNormal};
412        self.block += 1;
413        let dist = StandardNormal;
414        Array1::from(
415            (0..t)
416                .map(|_| {
417                    let v: f64 = dist.sample(rng);
418                    v as f32
419                })
420                .collect::<Vec<_>>(),
421        )
422    }
423}
424
425pub(crate) fn load_npy_f32(path: &Path) -> Result<Vec<f32>> {
426    let bytes = std::fs::read(path)?;
427    anyhow::ensure!(
428        bytes.starts_with(b"\x93NUMPY"),
429        "not npy: {}",
430        path.display()
431    );
432    let header_len = u16::from_le_bytes(bytes[8..10].try_into()?) as usize;
433    let header = std::str::from_utf8(&bytes[10..10 + header_len])?;
434    let shape_start = header.find('(').context("npy shape")?;
435    let shape_end = header[shape_start..].find(')').context("npy shape")? + shape_start;
436    let shape_str = &header[shape_start + 1..shape_end];
437    let elem_count: usize = if shape_str.trim().is_empty() {
438        0
439    } else {
440        shape_str
441            .split(',')
442            .filter(|s| !s.trim().is_empty())
443            .map(|s| s.trim().parse::<usize>())
444            .collect::<Result<Vec<_>, _>>()?
445            .into_iter()
446            .product()
447    };
448    let data_off = 10 + header_len;
449    let mut out = Vec::with_capacity(elem_count);
450    for chunk in bytes[data_off..].chunks_exact(4).take(elem_count) {
451        out.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]));
452    }
453    Ok(out)
454}
455
456/// Audio samples after the full decoder stack for `t_latent` latent steps.
457#[cfg(feature = "coreml")]
458pub(crate) fn audio_len_for_latent(decoder_rates: &[usize], t_latent: usize) -> usize {
459    let mut t = t_latent;
460    for &stride in decoder_rates {
461        let k = 2 * stride;
462        let padding = stride.div_ceil(2);
463        let output_padding = stride % 2;
464        t = (t - 1) * stride + k - 2 * padding + output_padding;
465    }
466    t
467}
468
469/// Time width after each decoder block upsample (noise plane lengths).
470#[cfg(feature = "coreml")]
471pub(crate) fn decoder_block_times(decoder_rates: &[usize], t_latent: usize) -> [usize; 4] {
472    let mut t = t_latent;
473    let mut out = [0usize; 4];
474    for (slot, &stride) in decoder_rates.iter().enumerate() {
475        let k = 2 * stride;
476        let padding = stride.div_ceil(2);
477        let output_padding = stride % 2;
478        t = (t - 1) * stride + k - 2 * padding + output_padding;
479        out[slot] = t;
480    }
481    out
482}
483
484fn decode_latent(
485    inner: &SnacDecoderInner,
486    z_q: ArrayView2<f32>,
487    noise: &mut NoiseState,
488) -> Array2<f32> {
489    let mut x = if inner.config.depthwise {
490        conv1d(
491            z_q,
492            inner.init_dw_w.view(),
493            Some(inner.init_dw_b.view()),
494            3,
495            inner.config.latent_dim(),
496            1,
497        )
498    } else {
499        conv1d(
500            z_q,
501            inner.init_dw_w.view(),
502            Some(inner.init_dw_b.view()),
503            3,
504            1,
505            1,
506        )
507    };
508    x = conv1d(
509        x.view(),
510        inner.init_pw_w.view(),
511        Some(inner.init_pw_b.view()),
512        0,
513        1,
514        1,
515    );
516    for block in &inner.blocks {
517        x = decoder_block_forward(x.view(), block, noise);
518    }
519    x = snake1d(x.view(), inner.final_snake_alpha.view());
520    x = conv1d(
521        x.view(),
522        inner.final_conv_w.view(),
523        Some(inner.final_conv_b.view()),
524        3,
525        1,
526        1,
527    );
528    x.mapv(|v| v.tanh())
529}
530
531fn config_path_for_weights(path: &Path) -> PathBuf {
532    path.parent()
533        .map(|dir| dir.join("snac_24khz_decoder_config.json"))
534        .unwrap_or_else(|| PathBuf::from("snac_24khz_decoder_config.json"))
535}
536
537/// SNAC neural codec decoder for Orpheus speech tokens.
538pub struct SnacDecoder {
539    inner: SnacDecoderInner,
540    weights_path: PathBuf,
541    noise_seed: Option<u64>,
542    noise_ref_dir: Option<PathBuf>,
543}
544
545impl SnacDecoder {
546    /// Load decoder weights + JSON config from `path` (`.safetensors`).
547    pub fn from_file(path: &Path) -> Result<Self> {
548        if !path.exists() {
549            bail!("SNAC decoder weights not found: {}", path.display());
550        }
551
552        let cfg_path = config_path_for_weights(path);
553        let config = SnacConfig::from_file(&cfg_path).with_context(|| {
554            format!(
555                "SNAC config not found next to weights (expected {})",
556                cfg_path.display()
557            )
558        })?;
559
560        let bytes =
561            std::fs::read(path).with_context(|| format!("read SNAC weights {}", path.display()))?;
562        let st = SafeTensors::deserialize(&bytes)
563            .with_context(|| format!("parse safetensors {}", path.display()))?;
564        let inner = load_inner(&st, config)
565            .with_context(|| format!("load SNAC weights from {}", path.display()))?;
566
567        let noise_seed = std::env::var("ORPHEUS_SNAC_SEED")
568            .ok()
569            .and_then(|s| s.parse().ok());
570        let noise_ref_dir = std::env::var("ORPHEUS_SNAC_REF_DIR")
571            .ok()
572            .filter(|p| !p.is_empty())
573            .map(PathBuf::from)
574            .filter(|p| p.is_dir());
575
576        Ok(Self {
577            inner,
578            weights_path: path.to_path_buf(),
579            noise_seed,
580            noise_ref_dir,
581        })
582    }
583
584    /// Override the PyTorch-style noise seed (`torch.manual_seed` before decode).
585    pub fn with_noise_seed(mut self, seed: u64) -> Self {
586        self.noise_seed = Some(seed);
587        self
588    }
589
590    fn make_noise_state(&self) -> Result<NoiseState> {
591        if !self.inner.config.noise {
592            return Ok(NoiseState::disabled());
593        }
594        if let Some(dir) = &self.noise_ref_dir {
595            if let Some(state) = NoiseState::from_reference_dir(dir)? {
596                return Ok(state);
597            }
598        }
599        if let Some(seed) = self.noise_seed {
600            return Ok(NoiseState::seeded(seed));
601        }
602        // PyTorch SNAC uses per-block `torch.randn`; a fixed seed per decode call
603        // matches upstream closely enough for audible output without ref fixtures.
604        Ok(NoiseState::seeded(42))
605    }
606
607    /// Decode three RVQ codebook streams to mono PCM at [`SAMPLE_RATE`].
608    pub fn decode_codes(
609        &self,
610        codes_0: &[i32],
611        codes_1: &[i32],
612        codes_2: &[i32],
613    ) -> Result<Vec<f32>> {
614        if codes_0.is_empty() {
615            return Ok(Vec::new());
616        }
617        let z_q = from_codes_inner(&self.inner, codes_0, codes_1, codes_2)?;
618        let mut noise = self.make_noise_state()?;
619        let audio = decode_latent(&self.inner, z_q.view(), &mut noise);
620        Ok(audio.slice(s![0, ..]).to_vec())
621    }
622
623    /// Decode interleaved Orpheus frame tokens (7 per frame) via [`crate::tokens::pack_orpheus_codes`].
624    pub fn decode_orpheus_frames(&self, frame_tokens: &[i32]) -> Result<Vec<f32>> {
625        let (c0, c1, c2) = crate::tokens::pack_orpheus_codes(frame_tokens)
626            .ok_or_else(|| anyhow::anyhow!("invalid Orpheus frame token layout"))?;
627        self.decode_codes(&c0, &c1, &c2)
628    }
629
630    #[cfg(feature = "coreml")]
631    pub(crate) fn inner(&self) -> &SnacDecoderInner {
632        &self.inner
633    }
634
635    pub fn config(&self) -> &SnacConfig {
636        &self.inner.config
637    }
638
639    pub fn weights_path(&self) -> &Path {
640        &self.weights_path
641    }
642}
643
644#[cfg(test)]
645mod tests {
646    use super::*;
647    use std::path::PathBuf;
648
649    fn ref_dir() -> Option<PathBuf> {
650        std::env::var("ORPHEUS_SNAC_REF_DIR")
651            .ok()
652            .filter(|p| !p.is_empty())
653            .map(PathBuf::from)
654            .filter(|p| p.is_dir())
655    }
656
657    #[test]
658    fn decode_matches_reference_when_weights_available() {
659        let Some(weights) = super::super::decoder_weights_path_if_available() else {
660            eprintln!("skip decode_matches_reference: set ORPHEUS_SNAC_PATH");
661            return;
662        };
663        let Some(ref_dir) = ref_dir() else {
664            eprintln!(
665                "skip decode_matches_reference: set ORPHEUS_SNAC_REF_DIR=/tmp/rlx-weights/snac"
666            );
667            return;
668        };
669
670        let codes_json: serde_json::Value =
671            serde_json::from_reader(std::fs::File::open(ref_dir.join("ref_codes.json")).unwrap())
672                .unwrap();
673        let c0: Vec<i32> = codes_json["codes_0"][0]
674            .as_array()
675            .unwrap()
676            .iter()
677            .map(|v| v.as_i64().unwrap() as i32)
678            .collect();
679        let c1: Vec<i32> = codes_json["codes_1"][0]
680            .as_array()
681            .unwrap()
682            .iter()
683            .map(|v| v.as_i64().unwrap() as i32)
684            .collect();
685        let c2: Vec<i32> = codes_json["codes_2"][0]
686            .as_array()
687            .unwrap()
688            .iter()
689            .map(|v| v.as_i64().unwrap() as i32)
690            .collect();
691
692        let dec = SnacDecoder::from_file(&weights).expect("SnacDecoder::from_file");
693
694        let z_q = from_codes_inner(&dec.inner, &c0, &c1, &c2).expect("from_codes");
695        let ref_z_q = load_npy_f32(&ref_dir.join("act_z_q.npy")).expect("act_z_q.npy");
696        assert_eq!(z_q.len(), ref_z_q.len());
697        let zq_max = z_q
698            .iter()
699            .zip(ref_z_q.iter())
700            .map(|(a, e)| (a - e).abs())
701            .fold(0.0f32, f32::max);
702        eprintln!("z_q max abs diff vs act_z_q: {zq_max:.6}");
703        assert!(zq_max < 1e-3, "z_q parity failed: max diff {zq_max}");
704
705        let audio = dec.decode_codes(&c0, &c1, &c2).expect("decode_codes");
706        assert_eq!(audio.len(), 8192);
707        let ref_audio = load_npy_f32(&ref_dir.join("ref_decode.npy")).expect("ref_decode.npy");
708        assert_eq!(audio.len(), ref_audio.len());
709        let pcm_max = audio
710            .iter()
711            .zip(ref_audio.iter())
712            .map(|(a, e)| (a - e).abs())
713            .fold(0.0f32, f32::max);
714        eprintln!("pcm max abs diff vs ref_decode: {pcm_max:.6}");
715        assert!(
716            pcm_max < 1e-3,
717            "PCM parity failed: max diff {pcm_max} (need ref_noise_*.npy from export script)"
718        );
719        for (i, &s) in audio.iter().enumerate() {
720            assert!(s.is_finite(), "non-finite sample at {i}: {s}");
721        }
722    }
723
724    #[test]
725    fn decode_orpheus_fixture_center_matches_reference() {
726        let Some(weights) = super::super::decoder_weights_path_if_available() else {
727            eprintln!("skip decode_orpheus_fixture_center: set ORPHEUS_SNAC_PATH");
728            return;
729        };
730        let Some(ref_dir) = ref_dir() else {
731            eprintln!("skip decode_orpheus_fixture_center: set ORPHEUS_SNAC_REF_DIR");
732            return;
733        };
734
735        let fixture = include_str!("../../tests/fixtures/orpheus_hi_codes.txt");
736        let mut lines = fixture
737            .lines()
738            .filter(|l| !l.trim().is_empty() && !l.starts_with('#'));
739        let _count: usize = lines.next().unwrap().trim().parse().unwrap();
740        let codes: Vec<i32> = lines
741            .next()
742            .unwrap()
743            .split_whitespace()
744            .map(|s| s.parse().unwrap())
745            .collect();
746
747        let dec = SnacDecoder::from_file(&weights).expect("SnacDecoder::from_file");
748        let backend = crate::decoder::SnacBackend::Eager(dec);
749        let pcm = crate::decode_orpheus_codes(&backend, &codes).expect("decode_orpheus_codes");
750        assert_eq!(pcm.len(), 4 * super::super::SAMPLES_PER_FRAME);
751
752        let ref_audio = load_npy_f32(&ref_dir.join("ref_decode.npy")).expect("ref_decode.npy");
753        let max_diff = pcm
754            .iter()
755            .zip(ref_audio.iter())
756            .map(|(a, e)| (a - e).abs())
757            .fold(0.0f32, f32::max);
758        eprintln!("fixture full decode max abs diff vs ref_decode: {max_diff:.6}");
759        assert!(
760            max_diff < 1e-3,
761            "fixture PCM full decode parity failed: max diff {max_diff}"
762        );
763    }
764}