Encodec model. (#1771)

* Encodec model.

* Fixes.

* Add the padding functions.

* Get the LSTM bit to work.

* Get the encodec model to generate some tokens (decoder only for now).

* Minor tweak.

* Minor tweak.

candle-examples/examples/encodec/main.rs

@@ -0,0 +1,57 @@
+#[cfg(feature = "mkl")]
+extern crate intel_mkl_src;
+
+#[cfg(feature = "accelerate")]
+extern crate accelerate_src;
+
+use anyhow::Result;
+use candle::{DType, IndexOp};
+use candle_nn::VarBuilder;
+use candle_transformers::models::encodec::{Config, Model};
+use clap::Parser;
+use hf_hub::api::sync::Api;
+
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+struct Args {
+    /// Run on CPU rather than on GPU.
+    #[arg(long)]
+    cpu: bool,
+
+    /// The model weight file, in safetensor format.
+    #[arg(long)]
+    model: Option<String>,
+
+    /// Input file as a safetensors containing the encodec tokens.
+    #[arg(long)]
+    code_file: String,
+
+    /// Output file that will be generated in wav format.
+    #[arg(long)]
+    out: String,
+}
+
+fn main() -> Result<()> {
+    let args = Args::parse();
+    let device = candle_examples::device(args.cpu)?;
+    let model = match args.model {
+        Some(model) => std::path::PathBuf::from(model),
+        None => Api::new()?
+            .model("facebook/encodec_24khz".to_string())
+            .get("model.safetensors")?,
+    };
+    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[model], DType::F32, &device)? };
+    let config = Config::default();
+    let model = Model::new(&config, vb)?;
+
+    let codes = candle::safetensors::load(args.code_file, &device)?;
+    let codes = codes.get("codes").expect("no codes in input file").i(0)?;
+    println!("codes shape: {:?}", codes.shape());
+    let pcm = model.decode(&codes)?;
+    let pcm = pcm.i(0)?.i(0)?.to_vec1::<f32>()?;
+
+    let mut output = std::fs::File::create(&args.out)?;
+    candle_examples::wav::write_pcm_as_wav(&mut output, &pcm, 24_000)?;
+
+    Ok(())
+}

candle-examples/src/lib.rs

@@ -1,6 +1,7 @@
 pub mod coco_classes;
 pub mod imagenet;
 pub mod token_output_stream;
+pub mod wav;
 
 use candle::utils::{cuda_is_available, metal_is_available};
 use candle::{Device, Result, Tensor};

candle-examples/src/wav.rs

@@ -0,0 +1,56 @@
+use std::io::prelude::*;
+
+pub trait Sample {
+    fn to_i16(&self) -> i16;
+}
+
+impl Sample for f32 {
+    fn to_i16(&self) -> i16 {
+        (self.clamp(-1.0, 1.0) * 32767.0) as i16
+    }
+}
+
+impl Sample for f64 {
+    fn to_i16(&self) -> i16 {
+        (self.clamp(-1.0, 1.0) * 32767.0) as i16
+    }
+}
+
+impl Sample for i16 {
+    fn to_i16(&self) -> i16 {
+        *self
+    }
+}
+
+pub fn write_pcm_as_wav<W: Write, S: Sample>(
+    w: &mut W,
+    samples: &[S],
+    sample_rate: u32,
+) -> std::io::Result<()> {
+    let len = 12u32; // header
+    let len = len + 24u32; // fmt
+    let len = len + samples.len() as u32 * 2 + 8; // data
+    let n_channels = 1u16;
+    let bytes_per_second = sample_rate * 2 * n_channels as u32;
+    w.write_all(b"RIFF")?;
+    w.write_all(&(len - 8).to_le_bytes())?; // total length minus 8 bytes
+    w.write_all(b"WAVE")?;
+
+    // Format block
+    w.write_all(b"fmt ")?;
+    w.write_all(&16u32.to_le_bytes())?; // block len minus 8 bytes
+    w.write_all(&1u16.to_le_bytes())?; // PCM
+    w.write_all(&n_channels.to_le_bytes())?; // one channel
+    w.write_all(&sample_rate.to_le_bytes())?;
+    w.write_all(&bytes_per_second.to_le_bytes())?;
+    w.write_all(&2u16.to_le_bytes())?; // 2 bytes of data per sample
+    w.write_all(&16u16.to_le_bytes())?; // bits per sample
+
+    // Data block
+    w.write_all(b"data")?;
+    w.write_all(&(samples.len() as u32 * 2).to_le_bytes())?;
+    for sample in samples.iter() {
+        w.write_all(&sample.to_i16().to_le_bytes())?
+    }
+    Ok(())
+}