load an actual denoising model

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<AtomicCell<Option<DfParams>>>),
+    Availible(Box<DfTract>),
+}
+
+fn with_denoising_model<O>(
+    spawn: &imp::SpawnHandle,
+    func: impl FnOnce(&mut DfTract) -> O,
+) -> Option<O> {
+    // 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<DenoisingModelState> = 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<f32>) {
-        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);
         }
     }