Skip to main content

av_denoise_vs/
lib.rs

1//! VapourSynth plugin exposing av-denoise as `avd` filters.
2
3mod filter;
4pub mod frames;
5pub mod params;
6
7use anyhow::Error;
8use tracing_subscriber::EnvFilter;
9use vapoursynth::core::CoreRef;
10use vapoursynth::plugins::{Filter, FilterArgument, Metadata};
11use vapoursynth::prelude::{API, Node};
12use vapoursynth::{export_vapoursynth_plugin, make_filter_function};
13
14use crate::filter::Denoise;
15use crate::params::{AlgorithmKind, RawParams};
16
17/// Installs the tracing subscriber that writes the plugin's logs to stderr.
18///
19/// `RUST_LOG` picks what is printed, and without it the plugin logs at `warn` so
20/// an ordinary render stays quiet.
21fn init_logging() {
22    let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn"));
23
24    let _ = tracing_subscriber::fmt()
25        .with_env_filter(filter)
26        .with_writer(std::io::stderr)
27        .try_init();
28}
29
30/// Keeps this plugin's library mapped for the rest of the process.
31///
32/// VapourSynth unloads every plugin library when it frees a core, and
33/// vspipe frees its core right before exiting. The GPU runtime this
34/// plugin builds spawns a device thread per accelerator, plus a
35/// polling thread per stream on the wgpu backends, and those threads
36/// run for the rest of the process. The device thread never blocks. It
37/// spins, yields, then sleeps briefly, over and over. On Windows its
38/// first wake after `FreeLibrary` returns into unmapped code and the
39/// process dies with an access violation, after every frame was
40/// already written. The polling thread parks or waits in the driver,
41/// and dies the same way once anything wakes it.
42///
43/// Pinning the module makes the unload a no-op, so the threads stay
44/// valid until process exit terminates them. On Linux the loader
45/// already refuses to unload a library that registered thread-local
46/// destructors, which is what happens as soon as this plugin's threads
47/// start, so nothing needs doing there. macOS is not covered and has
48/// not been tested.
49///
50/// This runs once, on the first filter creation. That is before any
51/// device thread exists, since only a filter builds a denoiser. The
52/// plugin's init function runs earlier, but the export macro owns its
53/// body and this plugin has no code of its own in it.
54fn pin_plugin_library() {
55    static PIN: std::sync::Once = std::sync::Once::new();
56    PIN.call_once(|| {
57        #[cfg(windows)]
58        pin_plugin_library_windows();
59    });
60}
61
62#[cfg(windows)]
63fn pin_plugin_library_windows() {
64    use std::ffi::c_void;
65
66    const GET_MODULE_HANDLE_EX_FLAG_PIN: u32 = 0x0000_0001;
67    const GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS: u32 = 0x0000_0004;
68
69    #[link(name = "kernel32")]
70    unsafe extern "system" {
71        fn GetModuleHandleExW(flags: u32, module_name: *const u16, module: *mut *mut c_void) -> i32;
72    }
73
74    let address = pin_plugin_library_windows as *const () as *const u16;
75    let mut module: *mut c_void = std::ptr::null_mut();
76    // SAFETY: `address` is a code address inside this library, which is
77    // what `FROM_ADDRESS` asks for, and `module` is a valid out pointer.
78    let ok = unsafe {
79        GetModuleHandleExW(
80            GET_MODULE_HANDLE_EX_FLAG_PIN | GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS,
81            address,
82            &mut module,
83        )
84    };
85    if ok == 0 {
86        tracing::warn!("could not pin the plugin library, the process may crash at exit");
87    }
88}
89
90/// Reads one optional UTF-8 script argument, naming `field` in the error
91/// when the bytes are not valid UTF-8.
92fn opt_string(bytes: Option<&[u8]>, field: &str) -> Result<Option<String>, Error> {
93    bytes
94        .map(|b| String::from_utf8(b.to_vec()).map_err(|_| anyhow::anyhow!("{field} must be valid UTF-8")))
95        .transpose()
96}
97
98/// Reads the optional `accelerators` script argument, a comma-separated
99/// list of accelerator names, into the `Vec<String>` [`RawParams`]
100/// wants.
101///
102/// VapourSynth script arguments have no native string array type that
103/// fits cleanly into `make_filter_function!`'s generated argument
104/// string, so this reuses the plain `data` type and splits it, matching
105/// how `channel_mode` and `device` already take a single string.
106fn opt_accelerators(bytes: Option<&[u8]>) -> Result<Option<Vec<String>>, Error> {
107    let Some(joined) = opt_string(bytes, "accelerators")? else {
108        return Ok(None);
109    };
110
111    let names: Vec<String> = joined
112        .split(',')
113        .map(str::trim)
114        .filter(|s| !s.is_empty())
115        .map(str::to_string)
116        .collect();
117
118    if names.is_empty() {
119        anyhow::bail!("accelerators must name at least one accelerator when set");
120    }
121
122    Ok(Some(names))
123}
124
125/// Reads the optional `motion_compensation` script argument.
126///
127/// VapourSynth script arguments have no native boolean type, so this
128/// takes the plain `int` type every other on/off knob in the wider
129/// VapourSynth ecosystem uses, and reads it the same way: `0` is off,
130/// anything else is on.
131fn opt_bool(value: Option<i64>) -> Option<bool> {
132    value.map(|v| v != 0)
133}
134
135/// Builds a [`RawParams`] from a filter function's raw script arguments.
136#[expect(clippy::too_many_arguments)]
137fn raw_params(
138    strength: Option<f64>,
139    variant: Option<&[u8]>,
140    preset: Option<&[u8]>,
141    prefilter: Option<&[u8]>,
142    channel_mode: Option<&[u8]>,
143    luma_strength: Option<f64>,
144    chroma_strength: Option<f64>,
145    luma_lambda_ht: Option<f64>,
146    chroma_lambda_ht: Option<f64>,
147    luma_mismatch_scale: Option<f64>,
148    chroma_mismatch_scale: Option<f64>,
149    device: Option<&[u8]>,
150    accelerators: Option<&[u8]>,
151    search_radius: Option<i64>,
152    patch_radius: Option<i64>,
153    temporal_radius: Option<i64>,
154    sigma: Option<f64>,
155    sigma_scale: Option<f64>,
156    motion_compensation: Option<i64>,
157    lambda_ht: Option<f64>,
158    lambda_ht_scale: Option<f64>,
159    spatial_radius: Option<i64>,
160    refine: Option<i64>,
161) -> Result<RawParams, Error> {
162    Ok(RawParams {
163        strength,
164        variant: opt_string(variant, "variant")?,
165        preset: opt_string(preset, "preset")?,
166        prefilter: opt_string(prefilter, "prefilter")?,
167        channel_mode: opt_string(channel_mode, "channel_mode")?,
168        luma_strength,
169        chroma_strength,
170        luma_lambda_ht,
171        chroma_lambda_ht,
172        luma_mismatch_scale,
173        chroma_mismatch_scale,
174        device: opt_string(device, "device")?,
175        accelerators: opt_accelerators(accelerators)?,
176        search_radius,
177        patch_radius,
178        temporal_radius,
179        sigma,
180        sigma_scale,
181        motion_compensation: opt_bool(motion_compensation),
182        lambda_ht,
183        lambda_ht_scale,
184        spatial_radius,
185        refine,
186    })
187}
188
189make_filter_function! {
190    NlmeansFunction, "NLMeans"
191
192    #[expect(clippy::too_many_arguments)]
193    fn create_nlmeans<'core>(
194        api: API,
195        core: CoreRef<'core>,
196        clip: Node<'core>,
197        strength: Option<f64>,
198        variant: Option<&[u8]>,
199        preset: Option<&[u8]>,
200        prefilter: Option<&[u8]>,
201        channel_mode: Option<&[u8]>,
202        luma_strength: Option<f64>,
203        chroma_strength: Option<f64>,
204        device: Option<&[u8]>,
205        accelerators: Option<&[u8]>,
206        search_radius: Option<i64>,
207        patch_radius: Option<i64>,
208        temporal_radius: Option<i64>,
209        sigma: Option<f64>,
210        sigma_scale: Option<f64>,
211        motion_compensation: Option<i64>,
212    ) -> Result<Option<Box<dyn Filter<'core> + 'core>>, Error> {
213        let raw = raw_params(
214            strength,
215            variant,
216            preset,
217            prefilter,
218            channel_mode,
219            luma_strength,
220            chroma_strength,
221            None,
222            None,
223            None,
224            None,
225            device,
226            accelerators,
227            search_radius,
228            patch_radius,
229            temporal_radius,
230            sigma,
231            sigma_scale,
232            motion_compensation,
233            None,
234            None,
235            None,
236            None,
237        )?;
238        let filter = Denoise::create(api, core, clip, AlgorithmKind::Nlmeans, &raw)?;
239        Ok(Some(Box::new(filter)))
240    }
241}
242
243make_filter_function! {
244    Nl4dFunction, "NL4D"
245
246    /// Estimates its automatic noise level fresh from each frame's own
247    /// temporal window, rather than smoothing it across the whole
248    /// stream, so a frame denoises to the same pixels no matter what
249    /// order VapourSynth requests frames in. Passing `sigma` pins the
250    /// noise level and skips that estimator entirely.
251    ///
252    /// The first few frames of a clip may differ slightly from the CLI's
253    /// output for the same parameters. The plugin fills a clip's
254    /// leading edge by repeating its first frame across the whole
255    /// temporal window, while the CLI's streaming mode primes a
256    /// narrower repeat before real frames start arriving. The
257    /// difference is bounded, small, and confined to a clip's first
258    /// `2 * temporal_radius` frames.
259    #[expect(clippy::too_many_arguments)]
260    fn create_nl4d<'core>(
261        api: API,
262        core: CoreRef<'core>,
263        clip: Node<'core>,
264        preset: Option<&[u8]>,
265        channel_mode: Option<&[u8]>,
266        luma_strength: Option<f64>,
267        chroma_strength: Option<f64>,
268        luma_lambda_ht: Option<f64>,
269        chroma_lambda_ht: Option<f64>,
270        luma_mismatch_scale: Option<f64>,
271        chroma_mismatch_scale: Option<f64>,
272        device: Option<&[u8]>,
273        accelerators: Option<&[u8]>,
274        temporal_radius: Option<i64>,
275        sigma: Option<f64>,
276        sigma_scale: Option<f64>,
277        lambda_ht: Option<f64>,
278        lambda_ht_scale: Option<f64>,
279        spatial_radius: Option<i64>,
280        refine: Option<i64>,
281    ) -> Result<Option<Box<dyn Filter<'core> + 'core>>, Error> {
282        let raw = raw_params(
283            None,
284            None,
285            preset,
286            None,
287            channel_mode,
288            luma_strength,
289            chroma_strength,
290            luma_lambda_ht,
291            chroma_lambda_ht,
292            luma_mismatch_scale,
293            chroma_mismatch_scale,
294            device,
295            accelerators,
296            None,
297            None,
298            temporal_radius,
299            sigma,
300            sigma_scale,
301            None,
302            lambda_ht,
303            lambda_ht_scale,
304            spatial_radius,
305            refine,
306        )?;
307        let filter = Denoise::create(api, core, clip, AlgorithmKind::Nl4d, &raw)?;
308        Ok(Some(Box::new(filter)))
309    }
310}
311
312export_vapoursynth_plugin! {
313    Metadata {
314        identifier: "com.chillfish8.avdenoise",
315        namespace: "avd",
316        name: "av-denoise",
317        read_only: true,
318    },
319    [NlmeansFunction::new(), Nl4dFunction::new()]
320}