diff --git a/Cargo.lock b/Cargo.lock index 83b7476cb..0ec50321d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -74,7 +74,7 @@ dependencies = [ "getrandom 0.2.15", "once_cell", "version_check", - "zerocopy", + "zerocopy 0.7.35", ] [[package]] @@ -2532,6 +2532,18 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "getrandom" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43a49c392881ce6d5c3b8cb70f98717b7c07aabbdff06687b9030dbfbe2725f8" +dependencies = [ + "cfg-if", + "libc", + "wasi 0.13.3+wasi-0.2.2", + "windows-targets 0.52.6", +] + [[package]] name = "getset" version = "0.1.3" @@ -4801,6 +4813,26 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "parler-tts" +version = "0.1.0" +dependencies = [ + "accelerate-src", + "candle-core", + "candle-nn", + "candle-transformers", + "futures-channel", + "hound", + "intel-mkl-src", + "kalosm-common", + "kalosm-model-types", + "rand 0.9.0", + "serde_json", + "thiserror 2.0.9", + "tokenizers", + "tokio", +] + [[package]] name = "password-hash" version = "0.5.0" @@ -5129,7 +5161,7 @@ version = "0.2.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77957b295656769bb8ad2b6a6b09d897d94f05c41b069aede1fcdaa675eaea04" dependencies = [ - "zerocopy", + "zerocopy 0.7.35", ] [[package]] @@ -5402,6 +5434,17 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "rand" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.2", + "zerocopy 0.8.20", +] + [[package]] name = "rand_chacha" version = "0.2.2" @@ -5422,6 +5465,16 @@ dependencies = [ "rand_core 0.6.4", ] +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.2", +] + [[package]] name = "rand_core" version = "0.5.1" @@ -5440,6 +5493,16 @@ dependencies = [ "getrandom 0.2.15", ] +[[package]] +name = "rand_core" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a509b1a2ffbe92afab0e55c8fd99dea1c280e8171bd2d88682bb20bc41cbc2c" +dependencies = [ + "getrandom 0.3.1", + "zerocopy 0.8.20", +] + [[package]] name = "rand_distr" version = "0.4.3" @@ -7341,9 +7404,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.42.0" +version = "1.43.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5cec9b21b0450273377fc97bd4c33a8acffc8c996c987a7c5b319a0083707551" +checksum = "3d61fa4ffa3de412bfea335c6ecff681de2b609ba3c77ef3e00e521813a9ed9e" dependencies = [ "backtrace", "bytes", @@ -7359,9 +7422,9 @@ dependencies = [ [[package]] name = "tokio-macros" -version = "2.4.0" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "693d596312e88961bc67d7f1f97af8a70227d9f90c31bba5806eec004978d752" +checksum = "6e06d43f1345a3bcd39f6a56dbb7dcab2ba47e68e8ac134855e7e2bdbaf8cab8" dependencies = [ "proc-macro2", "quote", @@ -7958,6 +8021,15 @@ version = "0.11.0+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" +[[package]] +name = "wasi" +version = "0.13.3+wasi-0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26816d2e1a4a36a2940b96c5296ce403917633dff8f3440e9b236ed6f6bacad2" +dependencies = [ + "wit-bindgen-rt", +] + [[package]] name = "wasm-bindgen" version = "0.2.99" @@ -8456,6 +8528,15 @@ version = "0.0.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d135d17ab770252ad95e9a872d365cf3090e3be864a34ab46f48555993efc904" +[[package]] +name = "wit-bindgen-rt" +version = "0.33.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3268f3d866458b787f390cf61f4bbb563b922d091359f9608842999eaee3943c" +dependencies = [ + "bitflags 2.6.0", +] + [[package]] name = "write16" version = "1.0.0" @@ -8562,7 +8643,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1b9b4fd18abc82b8136838da5d50bae7bdea537c574d8dc1a34ed098d6c166f0" dependencies = [ "byteorder", - "zerocopy-derive", + "zerocopy-derive 0.7.35", +] + +[[package]] +name = "zerocopy" +version = "0.8.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dde3bb8c68a8f3f1ed4ac9221aad6b10cece3e60a8e2ea54a6a2dec806d0084c" +dependencies = [ + "zerocopy-derive 0.8.20", ] [[package]] @@ -8576,6 +8666,17 @@ dependencies = [ "syn 2.0.91", ] +[[package]] +name = "zerocopy-derive" +version = "0.8.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eea57037071898bf96a6da35fd626f4f27e9cee3ead2a6c703cf09d472b2e700" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.91", +] + [[package]] name = "zerofrom" version = "0.1.5" diff --git a/Cargo.toml b/Cargo.toml index 3812d14a7..e419e8c8e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,7 +25,9 @@ members = [ "interfaces/kalosm-learning", "interfaces/kalosm-learning-macro", "interfaces/kalosm-parse-macro", - "interfaces/kalosm-common", "interfaces/kalosm-model-types", + "interfaces/kalosm-common", + "interfaces/kalosm-model-types", + "models/parler-tts", ] [workspace.dependencies] diff --git a/models/parler-tts/Cargo.toml b/models/parler-tts/Cargo.toml new file mode 100644 index 000000000..69fd49d7c --- /dev/null +++ b/models/parler-tts/Cargo.toml @@ -0,0 +1,59 @@ +[package] +name = "parler-tts" +version = "0.1.0" +edition = "2021" +description = "A simple interface for Parler speech-to-text models in Rust." +license = "MIT/Apache-2.0" +repository = "https://github.com/floneum/floneum" +authors = [ + "Evan Almloff ", + "Miles Murgaw ", +] +keywords = ["ai", "parler", "stt", "speech"] + +[dependencies] +kalosm-common.workspace = true +kalosm-model-types.workspace = true + +candle-core.workspace = true +candle-nn.workspace = true +candle-transformers.workspace = true +tokenizers.workspace = true +thiserror.workspace = true + +serde_json = "1.0.107" +futures-channel = "0.3.31" +rand = "0.9.0" + +hound = { version = "3.5.1", optional = true } +accelerate-src = { version = "0.3.2", optional = true } +intel-mkl-src = { version = "0.8.1", optional = true } + +[dev-dependencies] +tokio = { version = "1.43.0", features = ["full"] } + +[features] +wav = ["dep:hound"] +cuda = ["candle-core/cuda", "candle-nn/cuda", "candle-transformers/cuda"] +cudnn = ["candle-core/cudnn"] +metal = [ + "candle-core/metal", + "candle-nn/metal", + "candle-transformers/metal", + "kalosm-common/metal", +] +accelerate = [ + "dep:accelerate-src", + "candle-core/accelerate", + "candle-nn/accelerate", + "candle-transformers/accelerate", +] +mkl = [ + "dep:intel-mkl-src", + "candle-core/mkl", + "candle-nn/mkl", + "candle-transformers/mkl", +] + +[package.metadata.docs.rs] +features = ["wav"] diff --git a/models/parler-tts/src/builder.rs b/models/parler-tts/src/builder.rs new file mode 100644 index 000000000..6afafd64c --- /dev/null +++ b/models/parler-tts/src/builder.rs @@ -0,0 +1,205 @@ +use kalosm_model_types::{FileSource, ModelLoadingProgress}; +use std::{collections::HashSet, sync::Arc}; + +use crate::{model::ParlerInner, Parler, ParlerDrop, ParlerLoadingError, ParlerMessage}; + +/// The Parler source model to use. +#[derive(Debug, Clone, Copy, PartialEq, Default)] +pub enum ParlerSource { + /// An 880M parameter model. + MiniV1, + /// A 2.3B parameter model. + #[default] + LargeV1, +} + +impl ParlerSource { + pub(crate) fn model_and_revision(&self) -> (&'static str, &'static str) { + match self { + Self::MiniV1 => ("parler-tts/parler-tts-mini-v1", "main"), + Self::LargeV1 => ("parler-tts/parler-tts-large-v1", "main"), + } + } +} + +pub(crate) struct ParlerModelConfig { + model: FileSource, + tokenizer: FileSource, + config: FileSource, +} + +impl ParlerModelConfig { + pub(crate) fn new(model: FileSource, tokenizer: FileSource, config: FileSource) -> Self { + Self { + model, + tokenizer, + config, + } + } +} + +/// A builder for the [`Parler`](crate::Parler) model. +#[derive(Debug, Default)] +pub struct ParlerBuilder { + /// The model to be used. + model: ParlerSource, + + /// The cache location to use for the model (defaults DATA_DIR/kalosm/cache) + cache: kalosm_common::Cache, +} + +impl ParlerBuilder { + /// Create a new [`ParlerBuilder`] with the default configuration. + pub fn new() -> Self { + Self::default() + } + + /// Set the model to be used. + pub fn with_source(mut self, source: ParlerSource) -> Self { + self.model = source; + self + } + + /// Set the cache location to use for the model (defaults DATA_DIR/kalosm/cache) + pub fn with_cache(mut self, cache: kalosm_common::Cache) -> Self { + self.cache = cache; + self + } + + fn get_model_config(&self) -> ParlerModelConfig { + let (model_id, revision) = self.model.model_and_revision(); + let model_file = match self.model { + ParlerSource::MiniV1 => "model.safetensors", + ParlerSource::LargeV1 => "model.safetensors.index.json", + }; + + let model = FileSource::huggingface(model_id, revision, model_file); + let tokenizer = FileSource::huggingface(model_id, revision, "tokenizer.json"); + let config = FileSource::huggingface(model_id, revision, "config.json"); + + ParlerModelConfig::new(model, tokenizer, config) + } + + /// Build the model. + pub async fn build(self) -> Result { + self.build_with_loading_handler(ModelLoadingProgress::multi_bar_loading_indicator()) + .await + } + + /// Build the model with a loading handler. + pub async fn build_with_loading_handler( + self, + mut progress_handler: impl FnMut(ModelLoadingProgress) + Send + Sync + 'static, + ) -> Result { + // Download files + let sources = self.get_model_config(); + let tokenizer_source = sources.tokenizer; + let model_source = sources.model; + let config_source = sources.config; + + // Tokenizer progress + let display_tokenizer_source = format!("Tokenizer ({})", tokenizer_source); + let mut create_progress = + ModelLoadingProgress::downloading_progress(display_tokenizer_source); + let tokenizer_filename = self + .cache + .get(&tokenizer_source, |progress| { + progress_handler(create_progress(progress)) + }) + .await?; + + // Model source progress + let display_model_source = format!("Model ({})", model_source); + let mut create_progress = ModelLoadingProgress::downloading_progress(display_model_source); + let model_file = self + .cache + .get(&model_source, |progress| { + progress_handler(create_progress(progress)) + }) + .await?; + + // Determine if there are multiple model files and download them if required. + let mut model_files = Vec::new(); + match self.model { + ParlerSource::MiniV1 => model_files.push(model_file), + // LargeV1 uses a weight map index file. + ParlerSource::LargeV1 => { + let json: serde_json::Value = + serde_json::from_str(&std::fs::read_to_string(model_file).unwrap()) + .map_err(|e| ParlerLoadingError::ParseModel(e.into()))?; + + let weight_map = match json.get("weight_map") { + Some(serde_json::Value::Object(map)) => map, + _ => { + return Err(ParlerLoadingError::ParseModel( + "expected `weight_map` to be an object".into(), + )) + } + }; + + // De-deplicate files + let mut file_names = HashSet::new(); + for value in weight_map.values() { + if let Some(file) = value.as_str() { + file_names.insert(file); + } + } + + // Download the weight map files form the index. + for file in file_names { + let (model_id, revision) = self.model.model_and_revision(); + let source = FileSource::huggingface(model_id, revision, file); + + let display_model_file = format!("Model ({})", source); + let mut create_progress = + ModelLoadingProgress::downloading_progress(display_model_file); + + let model_file = self + .cache + .get(&source, |progress| { + progress_handler(create_progress(progress)) + }) + .await?; + + model_files.push(model_file); + } + } + }; + + // Config progress + let display_config_source = format!("Config ({})", config_source); + let mut create_progress = ModelLoadingProgress::downloading_progress(display_config_source); + let config = self + .cache + .get(&config_source, |progress| { + progress_handler(create_progress(progress)) + }) + .await?; + + // Start the parler thread + let (tx, rx) = std::sync::mpsc::channel(); + let thread = std::thread::spawn(move || { + let mut model = ParlerInner::new(model_files, tokenizer_filename, config).unwrap(); + while let Ok(message) = rx.recv() { + match message { + ParlerMessage::Kill => return, + ParlerMessage::Generate { + settings, + prompt, + description, + result, + } => { + _ = result.send(model.generate(settings, prompt, description)); + } + } + } + }); + + Ok(Parler { + inner: Arc::new(ParlerDrop { + thread: Some(thread), + sender: tx, + }), + }) + } +} diff --git a/models/parler-tts/src/lib.rs b/models/parler-tts/src/lib.rs new file mode 100644 index 000000000..a5cb47a31 --- /dev/null +++ b/models/parler-tts/src/lib.rs @@ -0,0 +1,159 @@ +//! A simple wrapper around the Parler speech-to-text models. +//! +//! Parler is a speech-to-text AI model that allows you to convert an input prompt into speech audio. +//! Contrary to many other speech-to-text models, Parler uses a description input to determine the speaker. +//! +//! **Audio Output Methods** +//! +//! - Output as a raw PCM vec. +//! - Output as a `wav` file with the `wav` feature. +//! +//! +//! **Common Problems** +//! +//! If the model does not output the full input text in speech, try increasing the temperature and decreasing the top-p setting. +//! This comes with the tradeoff of potentially worse quality output and must be balanced to achieve the desired results. +//! Both settings can be found in [`GenerationSettings`]. +//! +//! ### Example +//! +//! Basic usage of the `parler-tts` crate. There are a variety of configuration options not covered here - +//! namely [`GenerationSettings`]. +//! +//! ```rust,no_run +//! use std::path::PathBuf; +//! use parler_tts::{ParlerBuilder, ParlerSource}; +//! +//! #[tokio::main] +//! async fn main() { +//! let parler = ParlerBuilder::new() +//! .with_source(ParlerSource::MiniV1) +//! .build() +//! .await +//! .unwrap(); +//! +//! let task = parler.generate( +//! "The quick brown fox jumps over the lazy dog.", +//! "Will's voice is monotone yet slightly fast in delivery, with very clear audio.", +//! ); +//! +//! let decoder = task.await.unwrap(); +//! let out = decoder.raw_pcm(); +//! println!("{out:?}"); +//! } +//! ``` +//! +#![warn(missing_docs)] + +#[cfg(feature = "mkl")] +extern crate intel_mkl_src; + +#[cfg(feature = "accelerate")] +extern crate accelerate_src; + +use futures_channel::oneshot; +use kalosm_common::CacheError; +use std::{ + error::Error, + sync::{mpsc, Arc}, +}; + +mod builder; +mod model; +mod task; + +pub use builder::{ParlerBuilder, ParlerSource}; +pub use model::Decoder; +pub use task::{GenerationSeed, GenerationSettings, GenerationTask}; + +/// The Parler speech-to-text model. +#[derive(Clone)] +pub struct Parler { + inner: Arc, +} + +impl Parler { + /// Create a builder for a Parler model. + pub fn builder() -> ParlerBuilder { + ParlerBuilder::default() + } + + /// Create a new default Parler model. + pub async fn new() -> Result { + let model = Self::builder().build().await?; + Ok(model) + } + + /// Generate speech from a prompt and description. + /// + /// This will return a [`GenerationTask`] which is a handle to a generation. + pub fn generate(&self, prompt: impl ToString, description: impl ToString) -> GenerationTask { + GenerationTask { + settings: GenerationSettings::default(), + prompt: prompt.to_string(), + description: description.to_string(), + sender: self.inner.sender.clone(), + receiver: Default::default(), + } + } +} + +struct ParlerDrop { + thread: Option>, + sender: mpsc::Sender, +} + +impl Drop for ParlerDrop { + fn drop(&mut self) { + self.sender.send(ParlerMessage::Kill).unwrap(); + self.thread.take().unwrap().join().unwrap(); + } +} + +enum ParlerMessage { + Kill, + Generate { + settings: GenerationSettings, + prompt: String, + description: String, + result: oneshot::Sender>, + }, +} + +/// An error that can occur when loading a [`Parler`] model. +#[derive(Debug, thiserror::Error)] +pub enum ParlerLoadingError { + /// An error that can occur when trying to load a [`Parler`] model from huggingface or a local file. + #[error("Failed to load model from huggingface or local file: {0}")] + DownloadingError(#[from] CacheError), + + /// An error that can occur when trying to load a [`Parler`] model. + #[error("Failed to parse model json: {0}")] + ParseModel(Box), + + /// An error that can occur when trying to load a [`Parler`] model. + #[error("Failed to load model into device: {0}")] + LoadModel(#[from] candle_core::Error), + /// An error that can occur when trying to load the [`Parler`] tokenizer. + #[error("Failed to load tokenizer: {0}")] + LoadTokenizer(tokenizers::Error), + /// An error that can occur when trying to load the [`Parler`] config. + #[error("Failed to load config: {0}")] + LoadConfig(serde_json::Error), +} + +/// An error that can occur when running a [`Parler`] model. +#[derive(Debug, thiserror::Error)] +pub enum ParlerError { + /// An error that can occur when trying to run a [`Parler`] model. + #[error("Candle error: {0}")] + Candle(#[from] candle_core::Error), + /// An error that can occur when encoding or decoding for a [`Parler`] model. + #[error("Tokenizer error: {0}")] + Tokenizer(tokenizers::Error), + + /// An error that can occur when decoding Parler output into a wav file. + #[cfg(feature = "wav")] + #[error("Wav output error: {0}")] + WavOutput(#[from] hound::Error), +} diff --git a/models/parler-tts/src/model.rs b/models/parler-tts/src/model.rs new file mode 100644 index 000000000..77326f7a9 --- /dev/null +++ b/models/parler-tts/src/model.rs @@ -0,0 +1,141 @@ +use crate::{GenerationSeed, GenerationSettings, ParlerError, ParlerLoadingError}; +use candle_core::{DType, Device, IndexOp, Tensor}; +use candle_nn::VarBuilder; +use candle_transformers::{ + generation::LogitsProcessor, + models::parler_tts::{Config, Model}, +}; +use kalosm_common::accelerated_device_if_available; +use rand::{rngs::ThreadRng, Rng}; +use std::path::PathBuf; +use tokenizers::Tokenizer; + +pub(crate) struct ParlerInner { + device: Device, + tokenizer: Tokenizer, + model: Model, + config: Config, + rng: ThreadRng, +} + +impl ParlerInner { + pub(crate) fn new( + model_files: Vec, + tokenizer_file: PathBuf, + config_file: PathBuf, + ) -> Result { + let device = accelerated_device_if_available()?; + + let tokenizer = tokenizers::Tokenizer::from_file(tokenizer_file) + .map_err(ParlerLoadingError::LoadTokenizer)?; + + let config: Config = serde_json::from_str(&std::fs::read_to_string(config_file).unwrap()) + .map_err(ParlerLoadingError::LoadConfig)?; + + let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_files, DType::F32, &device)? }; + + let model = Model::new(&config, vb)?; + + Ok(Self { + device, + tokenizer, + model, + config, + rng: rand::rng(), + }) + } + + pub(crate) fn generate( + &mut self, + settings: GenerationSettings, + prompt: String, + description: String, + ) -> Result { + // Serialize input into tensor tokens. + let description_tokens = self + .tokenizer + .encode(description, true) + .map_err(ParlerError::Tokenizer)? + .get_ids() + .to_vec(); + let description_tokens = Tensor::new(description_tokens, &self.device)?.unsqueeze(0)?; + + let prompt_tokens = self + .tokenizer + .encode(prompt, true) + .map_err(ParlerError::Tokenizer)? + .get_ids() + .to_vec(); + let prompt_tokens = Tensor::new(prompt_tokens, &self.device)?.unsqueeze(0)?; + + // Generate + let seed = match settings.seed() { + GenerationSeed::Provided(val) => val, + GenerationSeed::Random => self.rng.random(), + }; + + let logit_proc = LogitsProcessor::new(seed, settings.temperature(), settings.top_p()); + let codes = self.model.generate( + &prompt_tokens, + &description_tokens, + logit_proc, + settings.max_steps(), + )?; + + // Decode + let codes = codes.to_dtype(DType::I64)?; + let codes = codes.unsqueeze(0)?; + + let codes = codes.to_device(&self.device)?; + let pcm = self.model.audio_encoder.decode_codes(&codes)?; + + let pcm = pcm.i((0, 0))?; + let pcm = pcm.to_vec1::()?; + + Ok(Decoder::new(pcm, self.config.audio_encoder.sampling_rate)) + } +} + +/// A decoder for Parler output audio data. +pub struct Decoder { + pcm: Vec, + sample_rate: u32, +} + +impl Decoder { + /// Create a new decoder with the specified PCM data and sample rate. + pub fn new(pcm: Vec, sample_rate: u32) -> Self { + Self { pcm, sample_rate } + } + + /// Get a clone of the raw PCM data outputted from the model. + pub fn raw_pcm(&self) -> Vec { + self.pcm.clone() + } + + /// Get the sample rate of the PCM data. + pub fn sample_rate(&self) -> u32 { + self.sample_rate + } + + #[cfg(feature = "wav")] + /// Output the PCM data into a wav file. + pub fn to_wav(&self, path: impl AsRef) -> Result<(), ParlerError> { + use hound::{SampleFormat, WavSpec, WavWriter}; + + let spec = WavSpec { + channels: 1, + sample_rate: self.sample_rate, + bits_per_sample: 16, + sample_format: SampleFormat::Int, + }; + + let mut writer = WavWriter::create(path, spec)?; + for sample in self.pcm.iter() { + writer.write_sample((sample.clamp(-1.0, 1.0) * i16::MAX as f32) as i16)?; + } + writer.finalize()?; + + Ok(()) + } +} diff --git a/models/parler-tts/src/task.rs b/models/parler-tts/src/task.rs new file mode 100644 index 000000000..7192f73ac --- /dev/null +++ b/models/parler-tts/src/task.rs @@ -0,0 +1,154 @@ +use futures_channel::oneshot; +use std::{future::Future, sync::mpsc, task::Poll}; + +use crate::{Decoder, ParlerError, ParlerMessage}; + +/// A generation task. +/// +/// A generation task is your handle to a single generate action for the model. Optionally provide +/// your own settings with the [`GenerationTask::with_settings`] method, and await this task for +/// the results. +pub struct GenerationTask { + pub(crate) settings: GenerationSettings, + pub(crate) prompt: String, + pub(crate) description: String, + pub(crate) sender: mpsc::Sender, + pub(crate) receiver: Option>>, +} + +impl GenerationTask { + /// Provide custom settings for this generation task. + pub fn with_settings(mut self, settings: GenerationSettings) -> Self { + self.settings = settings; + self + } +} + +impl Future for GenerationTask { + type Output = Result; + + fn poll( + self: std::pin::Pin<&mut Self>, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll { + let myself = self.get_mut(); + + // Create the receiver and send start message + if myself.receiver.is_none() { + let (tx, rx) = oneshot::channel(); + myself.receiver = Some(rx); + + _ = myself.sender.send(ParlerMessage::Generate { + settings: myself.settings.clone(), + prompt: myself.prompt.clone(), + description: myself.description.clone(), + result: tx, + }); + } + + let value = myself.receiver.as_mut().unwrap().try_recv().unwrap(); + match value { + Some(val) => Poll::Ready(val), + None => { + cx.waker().clone().wake(); + Poll::Pending + } + } + } +} + +/// The seed to use for a generation. +/// +/// The seed is used for deterministic output. If you want consistent results +/// for the same input, use the same seed. For more varied results, use a random seed. +#[derive(Debug, Clone, Copy)] +pub enum GenerationSeed { + /// Provide your own seed. + Provided(u64), + /// Let us generate a random seed. + Random, +} + +/// Settings for a generation task. +/// +/// These settings can be applied with the [`GenerationTask::with_setings`](GenerationTask::with_settings) method. +#[derive(Debug, Clone)] +pub struct GenerationSettings { + seed: GenerationSeed, + temperature: Option, + top_p: Option, + max_steps: usize, +} + +impl GenerationSettings { + /// Create new generation settings with the default. + pub fn new() -> Self { + Self::default() + } + + /// Get the value of the seed setting. + pub fn seed(&self) -> GenerationSeed { + self.seed + } + + /// Set the value of the seed setting. + /// + /// The seed is used for deterministic output. If you want consistent results + /// for the same input, use the same seed. For more varied results, use a random seed. + pub fn with_seed(mut self, seed: GenerationSeed) -> Self { + self.seed = seed; + self + } + + /// Get the value of the temperature setting. + pub fn temperature(&self) -> Option { + self.temperature + } + + /// Set the value of the temperature setting. + /// + /// Temperature is how creative Parler will tend to be. + /// Value must be at least 0.0 + pub fn with_temperature(mut self, temperature: Option) -> Self { + self.temperature = temperature.map(|val| val.max(0.0)); + self + } + + /// Get the value of the top-p setting. + pub fn top_p(&self) -> Option { + self.top_p + } + + /// Set the value of the top-p setting. + /// + /// Top-p, also known as nucleus sampling, is the cumulative probability threshold for selecting tokens. + /// You can think about it as the threshold required for tokens to be selected by the model. + pub fn with_top_p(mut self, top_p: Option) -> Self { + self.top_p = top_p; + self + } + + /// Get the value of the max steps setting. + pub fn max_steps(&self) -> usize { + self.max_steps + } + + /// Set the value of the max steps setting. + /// + /// More steps will use more RAM or VRAM. + pub fn with_max_steps(mut self, max_steps: usize) -> Self { + self.max_steps = max_steps; + self + } +} + +impl Default for GenerationSettings { + fn default() -> Self { + Self { + seed: GenerationSeed::Random, + temperature: Some(1.0), + top_p: Some(0.1), + max_steps: 512, + } + } +}