diff --git a/.gitignore b/.gitignore index 132e668..1273ff5 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ server_hash.txt proxy/bundle config.toml proxy/config.toml +gui/assets/*_onnx.tar.gz diff --git a/Cargo.lock b/Cargo.lock index b90108e..99ef6c7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -482,7 +482,7 @@ dependencies = [ "bitflags 2.10.0", "cexpr", "clang-sys", - "itertools 0.10.5", + "itertools 0.13.0", "log", "prettyplease", "proc-macro2", @@ -1965,7 +1965,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -2003,7 +2003,7 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "330c60081dcc4c72131f8eb70510f1ac07223e5d4163db481a04a0befcffa412" dependencies = [ - "libloading 0.7.4", + "libloading 0.8.9", ] [[package]] @@ -2218,7 +2218,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4429,7 +4429,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -5586,7 +5586,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.59.0", + "windows-sys 0.60.2", ] [[package]] @@ -6119,7 +6119,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -7257,7 +7257,7 @@ dependencies = [ "getrandom 0.3.4", "once_cell", "rustix", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -8396,7 +8396,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/docker/docker-compose.yaml b/docker/docker-compose.yaml index 9b2fe81..373bbbb 100644 --- a/docker/docker-compose.yaml +++ b/docker/docker-compose.yaml @@ -40,7 +40,7 @@ services: ports: - "4433:4433/tcp" - "4433:4433/udp" - command: ["cargo", "run", "-p", "mumble-web2-proxy"] + command: ["cargo", "run", "-p", "mumble-web2-proxy", "--locked"] network_mode: host mumble-server: diff --git a/gui/build.rs b/gui/build.rs new file mode 100644 index 0000000..ab4c843 --- /dev/null +++ b/gui/build.rs @@ -0,0 +1,38 @@ +use std::path::Path; +use std::process::Command; + +fn main() { + // Define the target directory and file + let assets_dir = "assets"; + let target_file = format!("{}/DeepFilterNet3_ll_onnx.tar.gz", assets_dir); + let target_path = Path::new(&target_file); + + // Check if the file already exists + if target_path.exists() { + println!("cargo:warning=DeepFilterNet model already exists at {}", target_file); + return; + } + + println!("cargo:warning=Downloading DeepFilterNet model to {}...", target_file); + + // Download the file using curl + let url = "https://github.com/Rikorose/DeepFilterNet/raw/refs/heads/main/models/DeepFilterNet3_ll_onnx.tar.gz"; + + let status = Command::new("curl") + .args([ + "-L", // Follow redirects + "-o", &target_file, // Output file + url, + ]) + .status() + .expect("Failed to execute curl command. Make sure curl is installed."); + + if !status.success() { + panic!("Failed to download DeepFilterNet model from {}", url); + } + + println!("cargo:warning=Successfully downloaded DeepFilterNet model to {}", target_file); + + // Rerun this build script if the target file is deleted + println!("cargo:rerun-if-changed={}", target_file); +} diff --git a/gui/src/effects.rs b/gui/src/effects.rs index 338943a..e85c945 100644 --- a/gui/src/effects.rs +++ b/gui/src/effects.rs @@ -1,25 +1,117 @@ use crossbeam::atomic::AtomicCell; +use df::tract::{mut_slice_as_arrayviewmut, slice_as_arrayview}; +use df::tract::{DfParams, DfTract, RuntimeParams}; +use dioxus::prelude::{asset, manganis, Asset}; +use std::cell::RefCell; use std::sync::Arc; +use tracing::{error, info}; + +use crate::imp; + +static DF_MODEL: Asset = asset!("/assets/DeepFilterNet3_ll_onnx.tar.gz"); + +enum DenoisingModelState { + Nothing, + Downloading(Arc>>), + Availible(Box), +} + +fn with_denoising_model( + spawn: &imp::SpawnHandle, + func: impl FnOnce(&mut DfTract) -> O, +) -> Option { + // Using a thread local is super gross, but DfTract is not Send (so it can never leave the current + // thread) while AudioProcessing itself might change threads whenever. + thread_local! { + static STATE: RefCell = const { RefCell::new(DenoisingModelState::Nothing) }; + } + + STATE.with_borrow_mut(|state| match state { + DenoisingModelState::Nothing => { + let cell = Arc::new(AtomicCell::new(None)); + let cell_task = cell.clone(); + *state = DenoisingModelState::Downloading(cell); + spawn.spawn(async move { + let model_bytes = match imp::read_asset_bytes(&DF_MODEL).await { + Ok(b) => b, + Err(e) => { + error!("could not read denoising model from \"{DF_MODEL}\": {e:?}"); + return; + } + }; + let params = match DfParams::from_bytes(&model_bytes) { + Ok(p) => p, + Err(e) => { + error!("could not load denoising model parameters: {e:?}"); + return; + } + }; + cell_task.store(Some(params)); + }); + None + } + DenoisingModelState::Downloading(cell) => { + if let Some(params) = cell.take() { + let mut tract = match DfTract::new(params, &RuntimeParams::default_with_ch(1)) { + Ok(t) => Box::new(t), + Err(e) => { + error!("could not create denoising engine: {e:?}"); + return None; + } + }; + info!("instantiated denoising engine"); + let out = func(&mut tract); + *state = DenoisingModelState::Availible(tract); + Some(out) + } else { + None + } + } + DenoisingModelState::Availible(tract) => Some(func(tract)), + }) +} -#[derive(Default)] pub struct AudioProcessor { - df: Option<::df::DFState>, + denoise: bool, + spawn: imp::SpawnHandle, } impl AudioProcessor { + pub fn new_plain() -> Self { + AudioProcessor { + denoise: false, + spawn: imp::SpawnHandle::current(), + } + } + pub fn new_denoising() -> Self { - let df = ::df::DFState::default(); - AudioProcessor { df: Some(df) } + AudioProcessor { + denoise: true, + spawn: imp::SpawnHandle::current(), + } } } impl AudioProcessor { pub fn process(&mut self, audio: &[f32], output: &mut Vec) { - if let Some(df) = &mut self.df { - let start = output.len(); - output.extend(std::iter::repeat_n(0f32, audio.len())); - df.process_frame(audio, &mut output[start..]); - } else { + let mut finished = false; + if self.denoise { + with_denoising_model(&self.spawn, |df| { + let start = output.len(); + output.extend(std::iter::repeat_n(0f32, audio.len())); + finished = true; + let output = &mut output[start..]; + df.process( + slice_as_arrayview(audio, &[audio.len()]) + .into_shape((1, audio.len())) + .unwrap(), + mut_slice_as_arrayviewmut(output, &[output.len()]) + .into_shape((1, output.len())) + .unwrap(), + ); + }); + } + if !finished { output.extend_from_slice(audio); } } diff --git a/gui/src/imp/desktop.rs b/gui/src/imp/desktop.rs index 6e4a77b..c64464d 100644 --- a/gui/src/imp/desktop.rs +++ b/gui/src/imp/desktop.rs @@ -1,16 +1,15 @@ use crate::app::Command; use crate::effects::{AudioProcessor, AudioProcessorSender}; -use color_eyre::eyre::{eyre, Error}; +use color_eyre::eyre::{eyre, Context, Error}; use cpal::traits::{DeviceTrait, HostTrait}; -use dioxus::hooks::{UnboundedReceiver, UnboundedSender}; +use dioxus::hooks::UnboundedReceiver; use futures::io::{AsyncRead, AsyncWrite}; -use mumble_protocol::control::{ClientControlCodec, ControlPacket}; -use mumble_protocol::Serverbound; +use mumble_protocol::control::ClientControlCodec; use mumble_web2_common::ClientConfig; use std::mem::replace; use std::net::ToSocketAddrs; +use std::sync::Arc; use std::sync::Mutex; -use std::{fmt, io, sync::Arc}; use tokio::net::TcpStream; use tokio_rustls::rustls; use tokio_rustls::rustls::client::danger::{HandshakeSignatureValid, ServerCertVerifier}; @@ -19,8 +18,9 @@ use tokio_rustls::rustls::ClientConfig as RlsClientConfig; use tokio_rustls::rustls::DigitallySignedStruct; use tokio_rustls::TlsConnector; use tokio_util::compat::{TokioAsyncReadCompatExt as _, TokioAsyncWriteCompatExt as _}; -use tracing::{debug, error, info, instrument, warn}; +use tracing::{error, info, instrument, warn}; +pub use tokio::runtime::Handle as SpawnHandle; pub use tokio::task::spawn; pub use tokio::time::sleep; @@ -70,7 +70,7 @@ impl AudioSystem { ) -> Result<(), Error> { let mut encoder = opus::Encoder::new(SAMPLE_RATE, opus::Channels::Mono, opus::Application::Voip)?; - let mut current_processor = AudioProcessor::default(); + let mut current_processor = AudioProcessor::new_plain(); let mut output_buffer = Vec::new(); let processors = self.processors.clone(); let error_callback = move |e: cpal::StreamError| error!("error recording: {e:?}"); @@ -311,3 +311,13 @@ pub fn init_logging() { .with_env_filter(env_filter) .init(); } + +// TODO: once we update to dioxus 0.7, swap this out with the dioxus-asset-resolver crate +pub async fn read_asset_bytes(asset: &dioxus::prelude::Asset) -> color_eyre::Result> { + let cur_exe = std::env::current_exe().unwrap(); + let path = cur_exe + .parent() + .unwrap() + .join(asset.to_string().trim_matches('/')); + Ok(std::fs::read(&path).with_context(|| format!("native path \"{}\"", path.display()))?) +} diff --git a/gui/src/imp/web.rs b/gui/src/imp/web.rs index 9c30f0b..e36d0f6 100644 --- a/gui/src/imp/web.rs +++ b/gui/src/imp/web.rs @@ -8,6 +8,7 @@ use js_sys::Float32Array; use mumble_protocol::control::ClientControlCodec; use mumble_web2_common::ClientConfig; use reqwest::Url; +use std::future::Future; use std::time::Duration; use tracing::level_filters::LevelFilter; use tracing::{debug, error, info, instrument}; @@ -283,7 +284,7 @@ async fn run_encoder_worklet( audio_encoder.configure(&encoder_config); info!("created audio encoder"); - let mut current_processor = AudioProcessor::default(); + let mut current_processor = AudioProcessor::new_plain(); let onmessage: Closure = Closure::new(move |event: MessageEvent| { if let Some(new_processor) = processors.take() { current_processor = new_processor; @@ -444,3 +445,25 @@ pub fn init_logging() { info!("logging initiated"); } + +// TODO: once we update to dioxus 0.7, swap this out with the dioxus-asset-resolver crate +pub async fn read_asset_bytes(asset: &dioxus::prelude::Asset) -> color_eyre::Result> { + let path = asset.to_string(); + let path = path.trim_matches('/'); + Ok(reqwest::get(path).await?.bytes().await?.to_vec()) +} + +pub struct SpawnHandle; + +impl SpawnHandle { + pub fn current() -> Self { + SpawnHandle + } + + pub fn spawn(&self, future: F) + where + F: Future + 'static, + { + spawn(future); + } +} diff --git a/gui/src/lib.rs b/gui/src/lib.rs index 83de958..5b40f62 100644 --- a/gui/src/lib.rs +++ b/gui/src/lib.rs @@ -289,7 +289,7 @@ fn accept_command( if denoise { audio.set_processor(AudioProcessor::new_denoising()); } else { - audio.set_processor(AudioProcessor::default()); + audio.set_processor(AudioProcessor::new_plain()); } } }