Skip to main content

av_denoise/
denoiser.rs

1use std::collections::VecDeque;
2
3use cubecl::Runtime;
4use cubecl::prelude::ComputeClient;
5
6use crate::accelerate::Accelerator;
7use crate::device::Device;
8use crate::nl4d::{Nl4dDenoiser, Nl4dParams};
9#[cfg(test)]
10use crate::nlmeans::MotionEstimation;
11use crate::nlmeans::{
12    ChannelMode,
13    HqParams,
14    MotionCompensationMode,
15    MotionSearch,
16    NlmDenoiser,
17    NlmParams,
18    Pending,
19    PrefilterMode,
20    hq_default_strength,
21    validate_dimensions,
22};
23use crate::sniff::sniff_best_accelerator;
24
25/// How a [`Denoiser`] should be set up.
26///
27/// Build one with `DenoiserOptions::builder()`. Every field has a
28/// default, so only the parts you care about need naming.
29///
30/// Only the settings every algorithm reads live here. Everything else
31/// belongs to whichever [`Algorithm`] variant actually uses it.
32#[derive(Debug, Clone, bon::Builder)]
33pub struct DenoiserOptions {
34    /// Which channels of the frame to denoise.
35    #[builder(default = ChannelMode::Yuv)]
36    pub channel_mode: ChannelMode,
37    /// Whether to clean each frame on its own or across a temporal
38    /// window.
39    #[builder(default = DenoisingMode::Spacial)]
40    pub mode: DenoisingMode,
41    /// Which algorithm to run, along with the settings only that
42    /// algorithm reads.
43    #[builder(default)]
44    pub algorithm: Algorithm,
45}
46
47/// Which denoising algorithm to run.
48///
49/// Each variant carries its own settings, so a knob one algorithm has no
50/// use for cannot be set on it.
51#[derive(Debug, Copy, Clone, PartialEq)]
52pub enum Algorithm {
53    /// The fast NLMeans path, with fixed weighting and no noise
54    /// measurement.
55    Nlmeans(NlmeansOptions),
56    /// NLMeans with its weighting matched to the measured noise level.
57    ///
58    /// This also uses a different default `strength`, one that adapts to
59    /// the temporal radius and the plane being denoised. See
60    /// [`crate::nlmeans::hq_default_strength`].
61    NlmeansHq(NlmeansHqOptions),
62    /// Groups 8x8 patches across the motion-compensated temporal window
63    /// itself, rather than filtering with NLM first and grouping within
64    /// one frame afterward.
65    ///
66    /// No NLM weighting pass ever runs, so none of the NLM knobs appear
67    /// on [`Nl4dOptions`].
68    Nl4d(Nl4dOptions),
69}
70
71impl Default for Algorithm {
72    fn default() -> Self {
73        Self::Nlmeans(NlmeansOptions::default())
74    }
75}
76
77/// Settings for [`Algorithm::Nlmeans`].
78#[derive(Debug, Copy, Clone, Default, PartialEq)]
79pub struct NlmeansOptions {
80    /// Which reference image the NLM weights are computed against.
81    ///
82    /// `None`, the default, compares patches on the noisy input
83    /// directly. Every other mode costs one extra GPU pass per frame.
84    pub prefilter: PrefilterMode,
85    /// Whether temporal denoising follows motion between frames.
86    ///
87    /// `None`, the default, turns motion compensation off. `Mvtools`
88    /// warps temporal neighbours into line with the centre frame before
89    /// the NLM weighting runs.
90    ///
91    /// Only has an effect when [`DenoiserOptions::mode`] is
92    /// `Temporal { .. }`.
93    pub motion_compensation: MotionCompensationMode,
94    /// Overrides for the NLM search radius, patch radius, strength, and
95    /// self-weight.
96    pub tuning: NlmTuning,
97}
98
99/// Settings for [`Algorithm::NlmeansHq`].
100#[derive(Debug, Copy, Clone, Default, PartialEq)]
101pub struct NlmeansHqOptions {
102    /// Everything the fast path takes, which HQ takes too.
103    pub nlm: NlmeansOptions,
104    /// The noise measurement and confidence weighting HQ adds on top.
105    pub hq: HqParams,
106}
107
108/// Settings for [`Algorithm::Nl4d`].
109///
110/// nl4d runs the HQ front end only for its machinery, the frame ring,
111/// the motion field, and the noise estimate. Nothing weights or averages
112/// patches the NLM way, so the NLM knobs are absent here and the fields
113/// below are the whole surface.
114///
115/// The temporal radius comes from [`DenoiserOptions::mode`], which has to
116/// be `Temporal { .. }`. Motion tracking is always on, because the
117/// grouping kernel reads the motion field and confidence scores it
118/// produces.
119///
120/// `lambda_ht` has a per-plane default. `None` resolves through
121/// [`nl4d_default_lambda_ht`] once the plane being denoised is known.
122/// `lambda_ht_scale` then multiplies whichever value that resolves to.
123///
124/// Every other default comes from [`Nl4dParams::default`].
125#[derive(Debug, Copy, Clone, PartialEq)]
126pub struct Nl4dOptions {
127    /// How motion between frames is tracked.
128    pub motion: MotionSearch,
129    /// A fixed noise standard deviation in `[0, 1]` units, replacing the
130    /// automatic per-frame estimate.
131    ///
132    /// `None`, the default, measures the noise in each pushed frame and
133    /// smooths it over time.
134    pub sigma: Option<f32>,
135    /// A multiplier applied to the measured noise level before anything
136    /// reads it. Defaults to 1.0.
137    ///
138    /// This does nothing when `sigma` pins the noise level, because the
139    /// estimator never runs in that case.
140    pub sigma_scale: f32,
141    /// A multiplier on the per-block mismatch threshold, which sets how
142    /// much extra SAD a block tolerates before its confidence starts to
143    /// fall. Defaults to 1.0.
144    ///
145    /// Higher values tolerate larger mismatches.
146    pub thsad_scale: f32,
147    /// Half-width of the refine window searched around each neighbour
148    /// frame's motion-predicted position, in `1..=4`. Defaults to 2.
149    pub refine: u32,
150    /// Half-width of the spatial candidate window searched in the centre
151    /// frame, in `1..=16`. Defaults to 9.
152    pub spatial_radius: u32,
153    /// Hard-threshold multiplier on the propagated coefficient sigma.
154    /// Higher removes more noise and more fine detail.
155    ///
156    /// `None` resolves through [`nl4d_default_lambda_ht`], which returns
157    /// a different value for luma than for chroma.
158    pub lambda_ht: Option<f32>,
159    /// A multiplier applied to the resolved `lambda_ht`. Defaults to
160    /// 1.0.
161    ///
162    /// It scales an explicit `lambda_ht` and the calibrated per-plane
163    /// default alike, so one value moves both planes together. Has to
164    /// be finite and in `[0.1, 10.0]`.
165    pub lambda_ht_scale: f32,
166    /// The confidence floor below which a whole neighbour block is
167    /// skipped rather than scored, in `[0, 1)`. Defaults to 0.05. Only
168    /// affects how much compute a submit spends, never which candidates
169    /// are admitted once they are scored.
170    pub c_min: f32,
171    /// A multiplier on the mismatch variance a poorly matched temporal
172    /// member carries into the hard threshold. Defaults to 1.0.
173    ///
174    /// The variance grows with the square of this. The mechanism
175    /// saturates well before the top of its accepted range, see
176    /// [`crate::nl4d::Nl4dParams::mismatch_scale`].
177    pub mismatch_scale: f32,
178    /// Whether a temporal member's mismatch variance reaches the
179    /// hard-threshold shrinkage at all. Defaults to `true`. See
180    /// [`crate::nl4d::Nl4dParams::confidence_variance`].
181    pub confidence_variance: bool,
182}
183
184impl Default for Nl4dOptions {
185    fn default() -> Self {
186        let defaults = Nl4dParams::default();
187        let hq = HqParams::default();
188        Self {
189            motion: MotionSearch::default(),
190            sigma: hq.sigma_override,
191            sigma_scale: hq.sigma_scale,
192            thsad_scale: hq.thsad_scale,
193            refine: defaults.refine,
194            spatial_radius: defaults.spatial_radius,
195            // Resolved per plane by `nl4d_default_lambda_ht` at
196            // construction time, once the plane being denoised is
197            // known.
198            lambda_ht: None,
199            lambda_ht_scale: 1.0,
200            c_min: defaults.c_min,
201            mismatch_scale: defaults.mismatch_scale,
202            confidence_variance: defaults.confidence_variance,
203        }
204    }
205}
206
207impl Nl4dOptions {
208    /// The front end's HQ parameters for this configuration.
209    ///
210    /// `temporal_confidence` is always on, because the grouping kernel
211    /// reads the confidence scores it produces. The two strength-related
212    /// switches keep their defaults, since nl4d never runs a weighting
213    /// pass for them to affect.
214    fn to_hq_params(self) -> HqParams {
215        HqParams {
216            sigma_override: self.sigma,
217            sigma_scale: self.sigma_scale,
218            thsad_scale: self.thsad_scale,
219            temporal_confidence: true,
220            ..HqParams::default()
221        }
222    }
223}
224
225/// The default `lambda_ht` for nl4d's hard-threshold stage, per plane.
226///
227/// `lambda_ht` is how many standard deviations of estimated noise a
228/// transform coefficient has to clear to survive. Raising it removes more
229/// noise and more fine detail with it, so the value is a trade rather
230/// than an optimum.
231///
232/// Luma gets 5.3, picked by eye from rendered comparisons on real grain
233/// and deliberately biased toward keeping detail. Higher values remove
234/// visibly more noise, but not enough to be worth what they cost in
235/// texture.
236///
237/// `ChannelMode::Yuv` reads the luma value, on the same "a fused pass is
238/// dominated by luma" assumption [`hq_default_strength`]
239/// makes for its own Yuv case.
240///
241/// Chroma gets 4.2, picked the same way from the chroma residuals with
242/// luma pinned at 5.3.
243pub fn nl4d_default_lambda_ht(channels: ChannelMode) -> f32 {
244    match channels {
245        ChannelMode::Luma | ChannelMode::Yuv => 5.3,
246        ChannelMode::Chroma => 4.2,
247    }
248}
249
250/// Resolves `Nl4dOptions.lambda_ht` for one plane, falling back to
251/// [`nl4d_default_lambda_ht`] when the caller left it unset, then
252/// applies `lambda_ht_scale`.
253///
254/// The scale multiplies an explicit value and the calibrated default
255/// alike, so it moves both planes together whether or not one of them
256/// is pinned.
257///
258/// The range check lives here rather than in [`Nl4dParams`],
259/// which only ever sees the product. A scale of 0 would surface there as
260/// a complaint about `lambda_ht`, naming a knob the caller never set.
261fn resolve_lambda_ht(opts: &Nl4dOptions, channels: ChannelMode) -> Result<f32, String> {
262    if !(opts.lambda_ht_scale.is_finite() && (0.1..=10.0).contains(&opts.lambda_ht_scale)) {
263        return Err(format!(
264            "lambda_ht_scale must be finite and in [0.1, 10.0], got {}",
265            opts.lambda_ht_scale
266        ));
267    }
268
269    let lambda_ht = opts.lambda_ht.unwrap_or_else(|| nl4d_default_lambda_ht(channels));
270
271    Ok(lambda_ht * opts.lambda_ht_scale)
272}
273
274/// Whether a frame is cleaned on its own or alongside its neighbours.
275#[derive(Debug, Copy, Clone, Eq, PartialEq)]
276pub enum DenoisingMode {
277    /// Cleans each frame using only its own pixels.
278    Spacial,
279    /// Cleans each frame using a window of `2 * radius + 1` frames.
280    Temporal { radius: u32 },
281}
282
283/// NLM tuning knobs.
284///
285/// Every field is optional. Whatever is left unset falls back to the
286/// library default.
287#[derive(Debug, Copy, Clone, Default, PartialEq)]
288pub struct NlmTuning {
289    pub search_radius: Option<u32>,
290    pub patch_radius: Option<u32>,
291    pub strength: Option<f32>,
292    pub self_weight: Option<f32>,
293}
294
295impl DenoiserOptions {
296    /// Turns this option set into the low-level [`NlmParams`] a backend
297    /// denoiser is built from.
298    ///
299    /// For nl4d this describes the front end only, since nl4d's own
300    /// grouping stage is configured from [`Nl4dOptions`] separately in
301    /// [`build_engine`].
302    ///
303    /// Whichever default `strength` applies is folded in here. For the
304    /// HQ algorithm that comes from
305    /// [`crate::nlmeans::hq_default_strength`].
306    ///
307    /// This is public so callers building per-plane options, and tests,
308    /// can read the resolved values without building a real `Denoiser`.
309    #[doc(hidden)]
310    pub fn to_nlm_params(&self) -> NlmParams {
311        let temporal_radius = match self.mode {
312            DenoisingMode::Spacial => 0,
313            DenoisingMode::Temporal { radius } => radius,
314        };
315
316        match self.algorithm {
317            Algorithm::Nlmeans(opts) => self.nlm_params_for(opts, None, temporal_radius),
318            Algorithm::NlmeansHq(opts) => self.nlm_params_for(opts.nlm, Some(opts.hq), temporal_radius),
319            // nl4d never runs a weighting pass, so `strength`,
320            // `search_radius`, `patch_radius`, and `self_weight` stay at
321            // their library defaults and no prefilter is built.
322            Algorithm::Nl4d(opts) => NlmParams {
323                channels: self.channel_mode,
324                motion_compensation: opts.motion.into(),
325                temporal_radius,
326                hq: Some(opts.to_hq_params()),
327                ..NlmParams::default()
328            },
329        }
330    }
331
332    /// [`Self::to_nlm_params`] for whichever of the two NLM algorithms
333    /// is running, with `hq` set only for the quality one.
334    fn nlm_params_for(&self, opts: NlmeansOptions, hq: Option<HqParams>, temporal_radius: u32) -> NlmParams {
335        // An explicit `strength` always wins, whether it came straight
336        // from `NlmTuning` or from a per-plane override the caller
337        // already folded in.
338        //
339        // Otherwise the default depends on `auto_strength`. With it on,
340        // HQ reads `strength` as a multiplier on the measured noise
341        // level, so it needs its own calibrated default rather than the
342        // fast path's absolute FFmpeg-style one. That calibrated default
343        // also varies with the temporal radius and with the plane
344        // `channel_mode` names, because each per-plane `Denoiser`
345        // carries its own channel mode.
346        //
347        // With auto-strength off, HQ reads `strength` as an absolute
348        // value just like the fast path, so it falls back to the same
349        // absolute default.
350        let strength = opts.tuning.strength.unwrap_or(match hq {
351            Some(hq) if hq.auto_strength => hq_default_strength(self.channel_mode, temporal_radius),
352            _ => NlmParams::default().strength,
353        });
354
355        let defaults = NlmParams::default();
356        NlmParams {
357            channels: self.channel_mode,
358            prefilter: opts.prefilter,
359            motion_compensation: opts.motion_compensation,
360            temporal_radius,
361            hq,
362            strength,
363            search_radius: opts.tuning.search_radius.unwrap_or(defaults.search_radius),
364            patch_radius: opts.tuning.patch_radius.unwrap_or(defaults.patch_radius),
365            self_weight: opts.tuning.self_weight.unwrap_or(defaults.self_weight),
366        }
367    }
368}
369
370/// Errors reported by the high-level [`Denoiser`].
371#[derive(Debug, thiserror::Error)]
372pub enum DenoiserError {
373    /// An earlier denoised frame has not been collected yet, so pushing
374    /// again would overwrite it in the double-buffered output slot.
375    ///
376    /// Call [`Denoiser::recv_frame`] or [`Denoiser::try_recv_frame`],
377    /// then retry the same `push_frame` call.
378    #[error("denoiser queue is full, collect the pending frame before pushing more")]
379    QueueFull,
380    /// None of the accelerators in the priority list could be started.
381    #[error("no accelerator from the priority list is available")]
382    NoAcceleratorAvailable,
383    /// Anything else, wrapping the internal `anyhow` errors raised by
384    /// kernel dispatch and readback.
385    #[error(transparent)]
386    Other(#[from] anyhow::Error),
387}
388
389/// Either denoiser a `Backend` runtime arm can hold.
390///
391/// This keeps `Backend`'s own match arms at one line each. Without it,
392/// adding a second denoiser type would multiply the runtime arms instead
393/// of fanning out once here.
394enum Engine<R: Runtime> {
395    Nlm(Box<NlmDenoiser<R>>),
396    Nl4d(Box<Nl4dDenoiser<R>>),
397}
398
399impl<R: Runtime> Engine<R> {
400    fn push_frame(&mut self, frame: &[f32]) {
401        match self {
402            Self::Nlm(d) => d.push_frame(frame),
403            Self::Nl4d(d) => d.push_frame(frame),
404        }
405    }
406
407    fn denoise_submit(&mut self) -> Result<Option<Pending<R>>, anyhow::Error> {
408        match self {
409            Self::Nlm(d) => d.denoise_submit(),
410            // `Nl4dDenoiser::denoise_submit` already returns
411            // `DenoiserError` rather than `anyhow::Error`, so this leans
412            // on `DenoiserError`'s own `anyhow::Error` conversion instead
413            // of re-wrapping it.
414            Self::Nl4d(d) => d.denoise_submit().map_err(anyhow::Error::from),
415        }
416    }
417
418    fn flush(&mut self, sink: impl FnMut(&[f32])) -> Result<(), anyhow::Error> {
419        match self {
420            Self::Nlm(d) => d.flush(sink),
421            Self::Nl4d(d) => d.flush(sink).map_err(anyhow::Error::from),
422        }
423    }
424}
425
426/// Builds whichever [`Engine`] `algorithm` calls for.
427///
428/// `Algorithm::Nl4d` carries its own grouping tuning, which is not part
429/// of `NlmParams`, so it is read from `algorithm` directly rather than
430/// from `params`. This is also where an unset `lambda_ht` picks up its
431/// calibrated per-plane default (`resolve_lambda_ht`), the same way
432/// `to_nlm_params` resolves HQ's calibrated `strength`, since this is
433/// the first point construction has both `opts` and `params.channels`
434/// together.
435fn build_engine<R: Runtime>(
436    client: &ComputeClient<R>,
437    algorithm: &Algorithm,
438    params: NlmParams,
439    width: u32,
440    height: u32,
441) -> Result<Engine<R>, DenoiserError> {
442    match algorithm {
443        Algorithm::Nl4d(opts) => {
444            // nl4d groups patches across neighbouring frames, so there
445            // is nothing for it to do without a temporal window.
446            if params.temporal_radius == 0 {
447                return Err(DenoiserError::Other(anyhow::anyhow!(
448                    "nl4d needs a temporal window, set DenoiserOptions::mode to \
449                     DenoisingMode::Temporal"
450                )));
451            }
452
453            let lambda_ht = resolve_lambda_ht(opts, params.channels)
454                .map_err(|e| DenoiserError::Other(anyhow::anyhow!(e)))?;
455            let nl4d_params = Nl4dParams {
456                temporal_radius: params.temporal_radius,
457                nlm: params,
458                refine: opts.refine,
459                spatial_radius: opts.spatial_radius,
460                lambda_ht,
461                c_min: opts.c_min,
462                mismatch_scale: opts.mismatch_scale,
463                confidence_variance: opts.confidence_variance,
464            };
465            let denoiser = Nl4dDenoiser::new(client, nl4d_params, width, height)
466                .map_err(|e| DenoiserError::Other(anyhow::anyhow!(e)))?;
467            Ok(Engine::Nl4d(Box::new(denoiser)))
468        },
469        Algorithm::Nlmeans(_) | Algorithm::NlmeansHq(_) => Ok(Engine::Nlm(Box::new(NlmDenoiser::new(
470            client, params, width, height,
471        )))),
472    }
473}
474
475enum Backend {
476    #[cfg(feature = "cuda")]
477    Cuda(Engine<cubecl::cuda::CudaRuntime>),
478    #[cfg(feature = "rocm")]
479    Rocm(Engine<cubecl::hip::HipRuntime>),
480    #[cfg(any(feature = "vulkan", feature = "metal"))]
481    Wgpu(Engine<cubecl::wgpu::WgpuRuntime>),
482}
483
484enum BackendPending {
485    #[cfg(feature = "cuda")]
486    Cuda(Pending<cubecl::cuda::CudaRuntime>),
487    #[cfg(feature = "rocm")]
488    Rocm(Pending<cubecl::hip::HipRuntime>),
489    #[cfg(any(feature = "vulkan", feature = "metal"))]
490    Wgpu(Pending<cubecl::wgpu::WgpuRuntime>),
491}
492
493impl BackendPending {
494    fn wait(self) -> Result<Vec<f32>, anyhow::Error> {
495        match self {
496            #[cfg(feature = "cuda")]
497            Self::Cuda(p) => p.wait(),
498            #[cfg(feature = "rocm")]
499            Self::Rocm(p) => p.wait(),
500            #[cfg(any(feature = "vulkan", feature = "metal"))]
501            Self::Wgpu(p) => p.wait(),
502        }
503    }
504}
505
506/// How many readbacks the high-level [`Denoiser`] keeps in flight at
507/// once.
508///
509/// This has to match the backend's output-handle count, which is two.
510/// Going past it would reuse the oldest pending frame's output handle
511/// and quietly corrupt the results.
512pub const MAX_PENDING: usize = 2;
513
514/// A stateful denoiser that cleans a stream of frames.
515///
516/// Push frames in order with [`push_frame`](Self::push_frame) and
517/// collect the cleaned ones with [`recv_frame`](Self::recv_frame) or
518/// [`try_recv_frame`](Self::try_recv_frame).
519///
520/// At the end of the stream call [`flush`](Self::flush) to drain
521/// whatever temporal context is left.
522///
523/// Frames are `f32` values in `[0, 1]`, laid out as
524/// `width * height * channels`.
525///
526/// ```no_run
527/// use av_denoise::accelerate::Accelerator;
528/// use av_denoise::{ChannelMode, Denoiser, DenoiserOptions, DenoisingMode, Device};
529///
530/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
531/// let options = DenoiserOptions::builder()
532///     .channel_mode(ChannelMode::Luma)
533///     .mode(DenoisingMode::Temporal { radius: 2 })
534///     .build();
535///
536/// let mut denoiser = Denoiser::create(
537///     &[Accelerator::Vulkan],
538///     &Device::Default,
539///     1920,
540///     1080,
541///     options,
542/// )?;
543///
544/// let frames: Vec<Vec<f32>> = read_my_frames();
545/// let mut cleaned: Vec<Vec<f32>> = Vec::new();
546///
547/// for frame in &frames {
548///     denoiser.push_frame(frame)?;
549///
550///     // Temporal denoising runs a few frames behind the input, so
551///     // there is not always one ready to collect.
552///     if let Some(out) = denoiser.recv_frame()? {
553///         cleaned.push(out);
554///     }
555/// }
556///
557/// // Drain the frames still inside the temporal window.
558/// denoiser.flush(|out| cleaned.push(out))?;
559/// # Ok(())
560/// # }
561/// # fn read_my_frames() -> Vec<Vec<f32>> { Vec::new() }
562/// ```
563pub struct Denoiser {
564    backend: Backend,
565    pending: VecDeque<BackendPending>,
566    accelerator: Accelerator,
567    width: u32,
568    height: u32,
569    channels: u32,
570    temporal_radius: u32,
571    frames_pushed: u32,
572}
573
574impl Denoiser {
575    /// Tries each accelerator in `accelerators` in order and builds a
576    /// denoiser on the first one that works.
577    ///
578    /// `device` picks a non-default device on the chosen runtime.
579    ///
580    /// # Thread stack size
581    ///
582    /// cubecl spawns its own per-device worker thread, named
583    /// `DS{U,D}-…`, and runs GPU kernel codegen on it. That thread gets
584    /// Rust's default stack, which is `RUST_MIN_STACK` or 2 MiB when
585    /// that is unset.
586    ///
587    /// The windowed NLM kernels unroll their body
588    /// `(2 * search_radius + 1)^2` times, so a `search_radius` of about
589    /// 5 or more can overflow the 2 MiB default and abort the process.
590    ///
591    /// Callers using a `search_radius` above 4 should set
592    /// `RUST_MIN_STACK` to at least 16 MiB before any cubecl thread
593    /// spawns, usually right at the top of `main`.
594    ///
595    /// ```no_run
596    /// if std::env::var_os("RUST_MIN_STACK").is_none() {
597    ///     // SAFETY: single-threaded at startup.
598    ///     unsafe { std::env::set_var("RUST_MIN_STACK", "16777216") };
599    /// }
600    /// ```
601    pub fn create(
602        accelerators: &[Accelerator],
603        device: &Device,
604        width: u32,
605        height: u32,
606        options: DenoiserOptions,
607    ) -> Result<Self, DenoiserError> {
608        let accelerator =
609            sniff_best_accelerator(accelerators, device).ok_or(DenoiserError::NoAcceleratorAvailable)?;
610
611        let params = options.to_nlm_params();
612        params.validate()?;
613        validate_dimensions(width, height)?;
614
615        let channels = params.channels.count();
616        let temporal_radius = params.temporal_radius;
617        let backend = build_backend(accelerator, device, &options.algorithm, params, width, height)?;
618
619        Ok(Self {
620            backend,
621            pending: VecDeque::with_capacity(MAX_PENDING),
622            accelerator,
623            width,
624            height,
625            channels,
626            temporal_radius,
627            frames_pushed: 0,
628        })
629    }
630
631    /// The accelerator [`sniff_best_accelerator`] picked.
632    pub fn selected_accelerator(&self) -> Accelerator {
633        self.accelerator
634    }
635
636    /// The width passed at construction.
637    pub fn width(&self) -> u32 {
638        self.width
639    }
640
641    /// The height passed at construction.
642    pub fn height(&self) -> u32 {
643        self.height
644    }
645
646    /// Uploads one frame into the temporal window.
647    ///
648    /// `frame` holds `width * height * channels` `f32` values in
649    /// `[0, 1]`.
650    ///
651    /// Once the window is full and the pipeline has room, this also
652    /// starts the kernels for the next denoised frame.
653    ///
654    /// Up to `MAX_PENDING` outputs can be in flight at once, so the GPU
655    /// runs one frame's kernels while the previous frame's readback is
656    /// still travelling. At that ceiling this returns
657    /// [`DenoiserError::QueueFull`], and the caller has to drain a frame
658    /// with [`Self::recv_frame`] before pushing more.
659    pub fn push_frame(&mut self, frame: &[f32]) -> Result<(), DenoiserError> {
660        // After `temporal_radius` real pushes the leading-edge mirror
661        // has primed the window, so the next push produces a pending
662        // frame. From then on every push takes a pending slot.
663        let window_full = self.frames_pushed > self.temporal_radius;
664        if window_full && self.pending.len() >= MAX_PENDING {
665            return Err(DenoiserError::QueueFull);
666        }
667
668        match &mut self.backend {
669            #[cfg(feature = "cuda")]
670            Backend::Cuda(d) => {
671                d.push_frame(frame);
672                if let Some(p) = d.denoise_submit()? {
673                    self.pending.push_back(BackendPending::Cuda(p));
674                }
675            },
676            #[cfg(feature = "rocm")]
677            Backend::Rocm(d) => {
678                d.push_frame(frame);
679                if let Some(p) = d.denoise_submit()? {
680                    self.pending.push_back(BackendPending::Rocm(p));
681                }
682            },
683            #[cfg(any(feature = "vulkan", feature = "metal"))]
684            Backend::Wgpu(d) => {
685                d.push_frame(frame);
686                if let Some(p) = d.denoise_submit()? {
687                    self.pending.push_back(BackendPending::Wgpu(p));
688                }
689            },
690        }
691
692        self.frames_pushed = self.frames_pushed.saturating_add(1);
693        Ok(())
694    }
695
696    /// Blocks until the in-flight denoise finishes and returns the
697    /// cleaned frame.
698    ///
699    /// Returns `Ok(None)` when nothing is in flight, which happens while
700    /// the temporal window is still filling up.
701    pub fn recv_frame(&mut self) -> Result<Option<Vec<f32>>, DenoiserError> {
702        let Some(pending) = self.pending.pop_front() else {
703            return Ok(None);
704        };
705        Ok(Some(pending.wait()?))
706    }
707
708    /// Collects the in-flight denoise if one is ready.
709    ///
710    /// This can still block for a moment while the runtime confirms the
711    /// readback has landed. When the kernels have already finished the
712    /// wait is effectively nothing.
713    pub fn try_recv_frame(&mut self) -> Result<Option<Vec<f32>>, DenoiserError> {
714        self.recv_frame()
715    }
716
717    /// Drains the in-flight frames and the trailing temporal tail,
718    /// handing each frame it produces to `sink`.
719    ///
720    /// The tail is padded by repeating the last pushed frame.
721    ///
722    /// On success the denoiser is ready for a fresh, unrelated stream of
723    /// the same size and parameters. Pushing again after a flush starts
724    /// a new temporal window from scratch, and flushing more than once
725    /// is fine.
726    ///
727    /// If `flush` returns `Err` the denoiser is in an undefined state
728    /// and should be dropped rather than reused.
729    pub fn flush(&mut self, mut sink: impl FnMut(Vec<f32>)) -> Result<(), DenoiserError> {
730        // Drain the whole pending pipeline, up to MAX_PENDING frames,
731        // before submitting the trailing-tail mirrors.
732        while let Some(frame) = self.recv_frame()? {
733            sink(frame);
734        }
735
736        let pixels = (self.width * self.height) as usize;
737        let channels = self.channels as usize;
738        let scratch_cap = pixels * channels;
739
740        match &mut self.backend {
741            #[cfg(feature = "cuda")]
742            Backend::Cuda(d) => d.flush(|slice| {
743                let mut v = Vec::with_capacity(scratch_cap);
744                v.extend_from_slice(slice);
745                sink(v);
746            })?,
747            #[cfg(feature = "rocm")]
748            Backend::Rocm(d) => d.flush(|slice| {
749                let mut v = Vec::with_capacity(scratch_cap);
750                v.extend_from_slice(slice);
751                sink(v);
752            })?,
753            #[cfg(any(feature = "vulkan", feature = "metal"))]
754            Backend::Wgpu(d) => d.flush(|slice| {
755                let mut v = Vec::with_capacity(scratch_cap);
756                v.extend_from_slice(slice);
757                sink(v);
758            })?,
759        }
760
761        // The backend has already reset its own stream indices. Reset
762        // the outer push counter too, so the next push re-arms the
763        // window-priming check at the top of `push_frame`.
764        self.frames_pushed = 0;
765
766        Ok(())
767    }
768}
769
770fn build_backend(
771    accel: Accelerator,
772    device: &Device,
773    algorithm: &Algorithm,
774    params: NlmParams,
775    width: u32,
776    height: u32,
777) -> Result<Backend, DenoiserError> {
778    match accel {
779        #[cfg(feature = "cuda")]
780        Accelerator::Cuda => {
781            let dev = device.to_cuda()?;
782            let client = <cubecl::cuda::CudaRuntime as Runtime>::client(&dev);
783            Ok(Backend::Cuda(build_engine(
784                &client, algorithm, params, width, height,
785            )?))
786        },
787        #[cfg(feature = "rocm")]
788        Accelerator::Rocm => {
789            let dev = device.to_amd()?;
790            let client = <cubecl::hip::HipRuntime as Runtime>::client(&dev);
791            Ok(Backend::Rocm(build_engine(
792                &client, algorithm, params, width, height,
793            )?))
794        },
795        #[cfg(feature = "vulkan")]
796        Accelerator::Vulkan => {
797            let dev = device.to_wgpu()?;
798            let client = <cubecl::wgpu::WgpuRuntime as Runtime>::client(&dev);
799            Ok(Backend::Wgpu(build_engine(
800                &client, algorithm, params, width, height,
801            )?))
802        },
803        #[cfg(feature = "metal")]
804        Accelerator::Metal => {
805            let dev = device.to_wgpu()?;
806            let client = <cubecl::wgpu::WgpuRuntime as Runtime>::client(&dev);
807            Ok(Backend::Wgpu(build_engine(
808                &client, algorithm, params, width, height,
809            )?))
810        },
811        // Keeps the match exhaustive on docs.rs, where `cfg(docsrs)`
812        // widens the `Accelerator` enum to include variants whose
813        // backend feature is not enabled. Never reached at runtime.
814        #[cfg(docsrs)]
815        #[allow(unreachable_patterns)]
816        _ => unreachable!(),
817    }
818}
819
820#[cfg(test)]
821mod options_tests {
822    use super::*;
823
824    /// `Algorithm::NlmeansHq` with `hq` overridden and everything else
825    /// left at its default.
826    fn hq(hq: HqParams) -> Algorithm {
827        Algorithm::NlmeansHq(NlmeansHqOptions {
828            hq,
829            ..NlmeansHqOptions::default()
830        })
831    }
832
833    /// `Algorithm::Nlmeans` with `tuning` overridden.
834    fn fast_tuned(tuning: NlmTuning) -> Algorithm {
835        Algorithm::Nlmeans(NlmeansOptions {
836            tuning,
837            ..NlmeansOptions::default()
838        })
839    }
840
841    #[test]
842    fn nl4d_default_lambda_ht_differs_between_luma_and_chroma() {
843        let luma = nl4d_default_lambda_ht(ChannelMode::Luma);
844        let chroma = nl4d_default_lambda_ht(ChannelMode::Chroma);
845
846        assert!((luma - 5.3).abs() < f32::EPSILON);
847        assert!((chroma - 4.2).abs() < f32::EPSILON);
848        assert!(
849            (chroma - luma).abs() > f32::EPSILON,
850            "the two planes should not resolve to the same default"
851        );
852    }
853
854    #[test]
855    fn nl4d_default_lambda_ht_yuv_reads_the_luma_value() {
856        let yuv = nl4d_default_lambda_ht(ChannelMode::Yuv);
857        let luma = nl4d_default_lambda_ht(ChannelMode::Luma);
858
859        assert!((yuv - luma).abs() < f32::EPSILON);
860    }
861
862    #[test]
863    fn resolve_lambda_ht_unset_uses_the_per_plane_default() {
864        let opts = Nl4dOptions::default();
865
866        let luma = resolve_lambda_ht(&opts, ChannelMode::Luma).expect("the default scale is in range");
867        let chroma = resolve_lambda_ht(&opts, ChannelMode::Chroma).expect("the default scale is in range");
868
869        assert!((luma - 5.3).abs() < f32::EPSILON, "got {luma}");
870        assert!((chroma - 4.2).abs() < f32::EPSILON, "got {chroma}");
871    }
872
873    #[test]
874    fn resolve_lambda_ht_explicit_value_overrides_every_plane() {
875        let opts = Nl4dOptions {
876            lambda_ht: Some(4.4),
877            ..Nl4dOptions::default()
878        };
879
880        for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
881            let got = resolve_lambda_ht(&opts, channels).expect("the default scale is in range");
882            assert!(
883                (got - 4.4).abs() < f32::EPSILON,
884                "channels {channels:?} got {got}"
885            );
886        }
887    }
888
889    #[test]
890    fn resolve_lambda_ht_default_scale_leaves_the_value_alone() {
891        let opts = Nl4dOptions::default();
892
893        for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
894            let got = resolve_lambda_ht(&opts, channels).expect("the default scale is in range");
895            let want = nl4d_default_lambda_ht(channels);
896            assert!(
897                (got - want).abs() < f32::EPSILON,
898                "channels {channels:?} got {got}"
899            );
900        }
901    }
902
903    #[test]
904    fn resolve_lambda_ht_scale_multiplies_the_per_plane_default() {
905        let opts = Nl4dOptions {
906            lambda_ht_scale: 1.1,
907            ..Nl4dOptions::default()
908        };
909
910        for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
911            let got = resolve_lambda_ht(&opts, channels).expect("1.1 is in range");
912            let want = nl4d_default_lambda_ht(channels) * 1.1;
913            assert!(
914                (got - want).abs() < 1e-5,
915                "channels {channels:?} got {got}, want {want}"
916            );
917        }
918    }
919
920    /// The scale is not limited to the defaults. Pinning one plane and
921    /// scaling both is the combination this exists for.
922    #[test]
923    fn resolve_lambda_ht_scale_multiplies_an_explicit_value() {
924        let opts = Nl4dOptions {
925            lambda_ht: Some(5.0),
926            lambda_ht_scale: 0.9,
927            ..Nl4dOptions::default()
928        };
929
930        let got = resolve_lambda_ht(&opts, ChannelMode::Luma).expect("0.9 is in range");
931        assert!((got - 4.5).abs() < 1e-5, "got {got}");
932    }
933
934    #[test]
935    fn resolve_lambda_ht_rejects_an_out_of_range_scale() {
936        for bad in [0.0, -1.0, 0.05, 10.5, f32::NAN, f32::INFINITY] {
937            let opts = Nl4dOptions {
938                lambda_ht_scale: bad,
939                ..Nl4dOptions::default()
940            };
941            let err = resolve_lambda_ht(&opts, ChannelMode::Luma).unwrap_err();
942            assert!(
943                err.contains("lambda_ht_scale"),
944                "lambda_ht_scale={bad} should be rejected, got {err}"
945            );
946        }
947    }
948
949    #[test]
950    fn the_default_algorithm_is_the_fast_nlmeans_path() {
951        let opts = DenoiserOptions::builder().build();
952        assert_eq!(opts.algorithm, Algorithm::Nlmeans(NlmeansOptions::default()));
953    }
954
955    #[test]
956    fn spatial_mode_maps_to_zero_temporal_radius() {
957        let opts = DenoiserOptions::builder()
958            .channel_mode(ChannelMode::Yuv)
959            .mode(DenoisingMode::Spacial)
960            .build();
961        let params = opts.to_nlm_params();
962
963        assert_eq!(params.temporal_radius, 0);
964        assert_eq!(params.channels, ChannelMode::Yuv);
965    }
966
967    #[test]
968    fn temporal_mode_propagates_radius() {
969        let opts = DenoiserOptions::builder()
970            .mode(DenoisingMode::Temporal { radius: 3 })
971            .build();
972        let params = opts.to_nlm_params();
973
974        assert_eq!(params.temporal_radius, 3);
975    }
976
977    #[test]
978    fn prefilter_passthrough() {
979        let opts = DenoiserOptions::builder()
980            .algorithm(Algorithm::Nlmeans(NlmeansOptions {
981                prefilter: PrefilterMode::Bilateral {
982                    sigma_s: 3.0,
983                    sigma_r: 0.02,
984                },
985                ..NlmeansOptions::default()
986            }))
987            .build();
988        let params = opts.to_nlm_params();
989
990        assert!(matches!(params.prefilter, PrefilterMode::Bilateral { .. }));
991    }
992
993    #[test]
994    fn hq_unset_prefilter_defaults_to_none() {
995        let opts = DenoiserOptions::builder()
996            .algorithm(hq(HqParams::default()))
997            .build();
998        let params = opts.to_nlm_params();
999
1000        assert!(matches!(params.prefilter, PrefilterMode::None));
1001    }
1002
1003    #[test]
1004    fn fast_unset_prefilter_defaults_to_none() {
1005        let opts = DenoiserOptions::builder()
1006            .algorithm(Algorithm::Nlmeans(NlmeansOptions::default()))
1007            .build();
1008        let params = opts.to_nlm_params();
1009
1010        assert!(matches!(params.prefilter, PrefilterMode::None));
1011    }
1012
1013    #[test]
1014    fn hq_unset_strength_defaults_to_hq_default_strength() {
1015        // Default channel_mode is Yuv, default mode is Spacial (radius 0).
1016        let opts = DenoiserOptions::builder()
1017            .algorithm(hq(HqParams::default()))
1018            .build();
1019        let params = opts.to_nlm_params();
1020
1021        let expected = hq_default_strength(ChannelMode::Yuv, 0);
1022        assert!((params.strength - expected).abs() < f32::EPSILON);
1023    }
1024
1025    #[test]
1026    fn hq_no_auto_strength_falls_back_to_the_legacy_absolute_default() {
1027        // `effective_strength_with` only reads `strength` as a
1028        // multiplier on the measured sigma when `auto_strength` is true.
1029        // With it false, `strength` is an FFmpeg-style absolute value,
1030        // so the fallback has to be the fast path's absolute default
1031        // rather than a calibrated multiplier from
1032        // `hq_default_strength`.
1033        let opts = DenoiserOptions::builder()
1034            .algorithm(hq(HqParams {
1035                auto_strength: false,
1036                ..HqParams::default()
1037            }))
1038            .build();
1039        let params = opts.to_nlm_params();
1040
1041        let expected = NlmParams::default().strength;
1042        assert!(
1043            (params.strength - expected).abs() < f32::EPSILON,
1044            "expected the legacy absolute default {expected}, got {}, which looks like the \
1045             auto-strength multiplier table leaking through",
1046            params.strength
1047        );
1048    }
1049
1050    #[test]
1051    fn hq_luma_r4_uses_measured_table_value() {
1052        let opts = DenoiserOptions::builder()
1053            .channel_mode(ChannelMode::Luma)
1054            .mode(DenoisingMode::Temporal { radius: 4 })
1055            .algorithm(hq(HqParams::default()))
1056            .build();
1057        let params = opts.to_nlm_params();
1058
1059        assert!((params.strength - 0.35).abs() < f32::EPSILON);
1060    }
1061
1062    #[test]
1063    fn hq_chroma_r4_uses_measured_table_value() {
1064        let opts = DenoiserOptions::builder()
1065            .channel_mode(ChannelMode::Chroma)
1066            .mode(DenoisingMode::Temporal { radius: 4 })
1067            .algorithm(hq(HqParams::default()))
1068            .build();
1069        let params = opts.to_nlm_params();
1070
1071        assert!((params.strength - 0.70).abs() < f32::EPSILON);
1072    }
1073
1074    #[test]
1075    fn hq_yuv_r8_uses_measured_table_value() {
1076        let opts = DenoiserOptions::builder()
1077            .channel_mode(ChannelMode::Yuv)
1078            .mode(DenoisingMode::Temporal { radius: 8 })
1079            .algorithm(hq(HqParams::default()))
1080            .build();
1081        let params = opts.to_nlm_params();
1082
1083        assert!((params.strength - 0.30).abs() < f32::EPSILON);
1084    }
1085
1086    #[test]
1087    fn hq_spacial_mode_uses_radius_zero_table_values() {
1088        for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
1089            let opts = DenoiserOptions::builder()
1090                .channel_mode(channels)
1091                .mode(DenoisingMode::Spacial)
1092                .algorithm(hq(HqParams::default()))
1093                .build();
1094            let params = opts.to_nlm_params();
1095
1096            let expected = hq_default_strength(channels, 0);
1097            assert!(
1098                (params.strength - expected).abs() < f32::EPSILON,
1099                "for channels {channels:?} expected {expected}, got {}",
1100                params.strength
1101            );
1102        }
1103    }
1104
1105    #[test]
1106    fn hq_explicit_strength_wins_over_the_table_for_every_plane() {
1107        for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
1108            let opts = DenoiserOptions::builder()
1109                .channel_mode(channels)
1110                .mode(DenoisingMode::Temporal { radius: 4 })
1111                .algorithm(Algorithm::NlmeansHq(NlmeansHqOptions {
1112                    nlm: NlmeansOptions {
1113                        tuning: NlmTuning {
1114                            strength: Some(0.99),
1115                            ..NlmTuning::default()
1116                        },
1117                        ..NlmeansOptions::default()
1118                    },
1119                    hq: HqParams::default(),
1120                }))
1121                .build();
1122            let params = opts.to_nlm_params();
1123
1124            assert!(
1125                (params.strength - 0.99).abs() < f32::EPSILON,
1126                "for channels {channels:?} the explicit strength was overridden by the table"
1127            );
1128        }
1129    }
1130
1131    #[test]
1132    fn fast_unset_strength_defaults_to_legacy_default() {
1133        let opts = DenoiserOptions::builder()
1134            .algorithm(Algorithm::Nlmeans(NlmeansOptions::default()))
1135            .build();
1136        let params = opts.to_nlm_params();
1137
1138        assert!((params.strength - 1.2).abs() < f32::EPSILON);
1139    }
1140
1141    #[test]
1142    fn nl4d_options_default_matches_nl4d_params_default() {
1143        let opts = Nl4dOptions::default();
1144        let params = crate::nl4d::Nl4dParams::default();
1145
1146        assert_eq!(opts.refine, params.refine);
1147        assert_eq!(opts.spatial_radius, params.spatial_radius);
1148        assert!((opts.c_min - params.c_min).abs() < f32::EPSILON);
1149        assert_eq!(opts.confidence_variance, params.confidence_variance);
1150        // The two `lambda_ht` fields hold different things, so they are
1151        // not compared. `opts.lambda_ht` stays `None` and is deferred to
1152        // `nl4d_default_lambda_ht` once the plane is known (see
1153        // `resolve_lambda_ht_unset_uses_the_per_plane_default` above),
1154        // while `params.lambda_ht` is a concrete default mirroring the
1155        // Luma/Yuv value.
1156        assert_eq!(opts.lambda_ht, None);
1157        assert!((params.lambda_ht - nl4d_default_lambda_ht(ChannelMode::Yuv)).abs() < f32::EPSILON);
1158    }
1159
1160    /// nl4d takes its own noise and confidence knobs rather than a whole
1161    /// [`HqParams`], so the three it does take have to reach the front
1162    /// end and the rest have to arrive at their defaults.
1163    #[test]
1164    fn nl4d_builds_the_front_ends_hq_params_from_its_own_fields() {
1165        let opts = DenoiserOptions::builder()
1166            .mode(DenoisingMode::Temporal { radius: 2 })
1167            .algorithm(Algorithm::Nl4d(Nl4dOptions {
1168                sigma: Some(0.02),
1169                sigma_scale: 1.3,
1170                thsad_scale: 0.8,
1171                ..Nl4dOptions::default()
1172            }))
1173            .build();
1174        let params = opts.to_nlm_params();
1175
1176        let hq = params.hq.expect("nl4d always runs the hq front end");
1177        assert_eq!(hq.sigma_override, Some(0.02));
1178        assert!((hq.sigma_scale - 1.3).abs() < f32::EPSILON);
1179        assert!((hq.thsad_scale - 0.8).abs() < f32::EPSILON);
1180        assert!(
1181            hq.temporal_confidence,
1182            "the grouping kernel reads the confidence scores, so this cannot be off"
1183        );
1184    }
1185
1186    /// The temporal radius has one source now, `mode`, so nl4d cannot
1187    /// disagree with the front end's ring about how wide the window is.
1188    #[test]
1189    fn nl4d_reads_its_temporal_radius_from_the_denoising_mode() {
1190        for radius in [1u32, 4, 8] {
1191            let opts = DenoiserOptions::builder()
1192                .mode(DenoisingMode::Temporal { radius })
1193                .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1194                .build();
1195
1196            assert_eq!(opts.to_nlm_params().temporal_radius, radius);
1197        }
1198    }
1199
1200    /// nl4d never runs an NLM weighting pass, so a prefilter would cost
1201    /// a GPU pass per frame producing a reference image nothing reads.
1202    #[test]
1203    fn nl4d_never_builds_a_prefilter() {
1204        let opts = DenoiserOptions::builder()
1205            .mode(DenoisingMode::Temporal { radius: 2 })
1206            .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1207            .build();
1208
1209        assert!(matches!(opts.to_nlm_params().prefilter, PrefilterMode::None));
1210    }
1211
1212    /// Nothing in the nl4d path reads `strength`, so it stays at the
1213    /// library default rather than picking up HQ's calibrated table.
1214    #[test]
1215    fn nl4d_leaves_the_nlm_weighting_knobs_at_their_defaults() {
1216        let defaults = NlmParams::default();
1217        let opts = DenoiserOptions::builder()
1218            .channel_mode(ChannelMode::Luma)
1219            .mode(DenoisingMode::Temporal { radius: 4 })
1220            .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1221            .build();
1222        let params = opts.to_nlm_params();
1223
1224        assert!((params.strength - defaults.strength).abs() < f32::EPSILON);
1225        assert_eq!(params.search_radius, defaults.search_radius);
1226        assert_eq!(params.patch_radius, defaults.patch_radius);
1227        assert!((params.self_weight - defaults.self_weight).abs() < f32::EPSILON);
1228    }
1229
1230    /// nl4d always tracks motion, so its `MotionSearch` reaches the
1231    /// front end as an active `Mvtools` mode.
1232    #[test]
1233    fn nl4d_motion_search_becomes_an_active_mvtools_mode() {
1234        let opts = DenoiserOptions::builder()
1235            .mode(DenoisingMode::Temporal { radius: 2 })
1236            .algorithm(Algorithm::Nl4d(Nl4dOptions {
1237                motion: MotionSearch {
1238                    blksize: 32,
1239                    overlap: 16,
1240                    search_radius: 6,
1241                    pyramid_levels: 1,
1242                    estimation: MotionEstimation::Direct,
1243                },
1244                ..Nl4dOptions::default()
1245            }))
1246            .build();
1247        let params = opts.to_nlm_params();
1248
1249        assert!(matches!(
1250            params.motion_compensation,
1251            MotionCompensationMode::Mvtools {
1252                blksize: 32,
1253                overlap: 16,
1254                search_radius: 6,
1255                pyramid_levels: 1,
1256                estimation: MotionEstimation::Direct,
1257            }
1258        ));
1259    }
1260
1261    #[test]
1262    fn nl4d_motion_search_defaults_match_the_front_ends_own_defaults() {
1263        let opts = DenoiserOptions::builder()
1264            .mode(DenoisingMode::Temporal { radius: 2 })
1265            .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1266            .build();
1267        let params = opts.to_nlm_params();
1268
1269        assert_eq!(
1270            params.motion_compensation,
1271            crate::nl4d::Nl4dParams::default().nlm.motion_compensation
1272        );
1273    }
1274
1275    #[test]
1276    fn motion_compensation_passthrough() {
1277        let opts = DenoiserOptions::builder()
1278            .mode(DenoisingMode::Temporal { radius: 1 })
1279            .algorithm(Algorithm::Nlmeans(NlmeansOptions {
1280                motion_compensation: MotionCompensationMode::Mvtools {
1281                    blksize: 16,
1282                    overlap: 8,
1283                    search_radius: 4,
1284                    pyramid_levels: 2,
1285                    estimation: MotionEstimation::Direct,
1286                },
1287                ..NlmeansOptions::default()
1288            }))
1289            .build();
1290        let params = opts.to_nlm_params();
1291
1292        assert!(matches!(
1293            params.motion_compensation,
1294            MotionCompensationMode::Mvtools {
1295                blksize: 16,
1296                overlap: 8,
1297                search_radius: 4,
1298                pyramid_levels: 2,
1299                ..
1300            }
1301        ));
1302    }
1303
1304    #[test]
1305    fn motion_compensation_defaults_to_none() {
1306        let opts = DenoiserOptions::builder().build();
1307        let params = opts.to_nlm_params();
1308        assert!(matches!(params.motion_compensation, MotionCompensationMode::None));
1309    }
1310
1311    #[test]
1312    fn nlm_tuning_overrides_individual_fields() {
1313        let defaults = NlmParams::default();
1314        let opts = DenoiserOptions::builder()
1315            .algorithm(fast_tuned(NlmTuning {
1316                search_radius: Some(7),
1317                patch_radius: None,
1318                strength: Some(2.5),
1319                self_weight: None,
1320            }))
1321            .build();
1322        let params = opts.to_nlm_params();
1323
1324        assert_eq!(params.search_radius, 7);
1325        assert_eq!(params.patch_radius, defaults.patch_radius);
1326        assert!((params.strength - 2.5).abs() < f32::EPSILON);
1327        assert!((params.self_weight - defaults.self_weight).abs() < f32::EPSILON);
1328    }
1329}
1330
1331#[cfg(all(test, feature = "vulkan"))]
1332mod tests {
1333    use super::*;
1334
1335    fn opts(mode: DenoisingMode) -> DenoiserOptions {
1336        DenoiserOptions::builder()
1337            .channel_mode(ChannelMode::Luma)
1338            .mode(mode)
1339            .build()
1340    }
1341
1342    fn frame(w: u32, h: u32) -> Vec<f32> {
1343        vec![0.5f32; (w * h) as usize]
1344    }
1345
1346    #[test]
1347    fn spatial_denoise_roundtrip() {
1348        let mut d = Denoiser::create(
1349            &[Accelerator::Vulkan],
1350            &Device::Default,
1351            16,
1352            16,
1353            opts(DenoisingMode::Spacial),
1354        )
1355        .expect("denoiser construction failed");
1356        assert_eq!(d.selected_accelerator(), Accelerator::Vulkan);
1357
1358        d.push_frame(&frame(16, 16)).expect("push failed");
1359        let out = d.recv_frame().expect("recv failed").expect("no frame");
1360        assert_eq!(out.len(), 16 * 16);
1361    }
1362
1363    #[test]
1364    fn nl4d_algorithm_round_trips_through_the_facade() {
1365        let opts = DenoiserOptions::builder()
1366            .channel_mode(ChannelMode::Luma)
1367            .mode(DenoisingMode::Temporal { radius: 2 })
1368            .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1369            .build();
1370        let mut d = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, opts)
1371            .expect("nl4d denoiser construction failed");
1372        assert_eq!(d.selected_accelerator(), Accelerator::Vulkan);
1373
1374        // temporal_radius is 2, so a single push does not fill the
1375        // window yet, the same convention every temporal algorithm here
1376        // follows.
1377        d.push_frame(&frame(16, 16)).expect("push failed");
1378        assert!(d.recv_frame().expect("recv failed").is_none());
1379
1380        let mut out = Vec::new();
1381        d.flush(|f| out.push(f)).expect("flush failed");
1382        assert_eq!(out.len(), 1, "expected exactly one output for one pushed frame");
1383        assert_eq!(out[0].len(), 16 * 16);
1384    }
1385
1386    /// nl4d groups patches across neighbouring frames, so a spatial
1387    /// mode leaves it nothing to do. The temporal radius has one source
1388    /// now, `mode`, so this is the only way to ask for that.
1389    #[test]
1390    fn nl4d_rejects_a_spatial_denoising_mode() {
1391        let opts = DenoiserOptions::builder()
1392            .channel_mode(ChannelMode::Luma)
1393            .mode(DenoisingMode::Spacial)
1394            .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1395            .build();
1396        let result = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, opts);
1397
1398        match result {
1399            Err(DenoiserError::Other(e)) => assert!(
1400                e.to_string().contains("temporal window"),
1401                "unexpected error message: {e}"
1402            ),
1403            Err(other) => panic!("expected DenoiserError::Other, got {other:?}"),
1404            Ok(_) => panic!("expected a rejection, got Ok"),
1405        }
1406    }
1407
1408    #[test]
1409    fn invalid_params_surface_as_error() {
1410        let bad = DenoiserOptions::builder()
1411            .algorithm(Algorithm::Nlmeans(NlmeansOptions {
1412                tuning: NlmTuning {
1413                    strength: Some(0.0),
1414                    ..NlmTuning::default()
1415                },
1416                ..NlmeansOptions::default()
1417            }))
1418            .build();
1419        let result = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, bad);
1420
1421        match result {
1422            Err(DenoiserError::Other(_)) => {},
1423            Err(other) => panic!("expected DenoiserError::Other, got {other:?}"),
1424            Ok(_) => panic!("expected validation error, got Ok"),
1425        }
1426    }
1427
1428    #[test]
1429    fn tiny_frame_dimensions_surface_as_error() {
1430        let result = Denoiser::create(
1431            &[Accelerator::Vulkan],
1432            &Device::Default,
1433            2,
1434            2,
1435            opts(DenoisingMode::Spacial),
1436        );
1437
1438        match result {
1439            Err(DenoiserError::Other(e)) => {
1440                assert!(
1441                    e.to_string().contains("supported minimum"),
1442                    "unexpected error message: {e}"
1443                );
1444            },
1445            Err(other) => panic!("expected DenoiserError::Other, got {other:?}"),
1446            Ok(_) => panic!("expected dimension validation error, got Ok"),
1447        }
1448    }
1449
1450    #[test]
1451    fn push_after_pending_returns_queue_full() {
1452        let mut d = Denoiser::create(
1453            &[Accelerator::Vulkan],
1454            &Device::Default,
1455            16,
1456            16,
1457            opts(DenoisingMode::Spacial),
1458        )
1459        .unwrap();
1460
1461        // The pipeline is two deep because the output handles are
1462        // double-buffered, so the first two pushes both submit. The
1463        // third would overwrite the oldest pending frame's output slot,
1464        // so it is rejected with QueueFull.
1465        d.push_frame(&frame(16, 16)).unwrap();
1466        d.push_frame(&frame(16, 16)).unwrap();
1467        let err = d.push_frame(&frame(16, 16)).expect_err("expected QueueFull");
1468        assert!(matches!(err, DenoiserError::QueueFull));
1469
1470        let out = d.recv_frame().unwrap().unwrap();
1471        assert_eq!(out.len(), 16 * 16);
1472
1473        // After draining one slot the next push must succeed.
1474        d.push_frame(&frame(16, 16)).expect("push after drain failed");
1475    }
1476
1477    fn frame_filled(w: u32, h: u32, value: f32) -> Vec<f32> {
1478        vec![value; (w * h) as usize]
1479    }
1480
1481    /// Pushes `n` frames of the given value, receiving along the way to
1482    /// keep the in-flight pipeline below `MAX_PENDING`.
1483    fn push_n_with_drain(d: &mut Denoiser, n: usize, value: f32, out: &mut Vec<Vec<f32>>) {
1484        for _ in 0..n {
1485            loop {
1486                match d.push_frame(&frame_filled(16, 16, value)) {
1487                    Ok(()) => break,
1488                    Err(DenoiserError::QueueFull) => {
1489                        let f = d
1490                            .recv_frame()
1491                            .expect("recv ok")
1492                            .expect("queue full but recv yielded none");
1493                        out.push(f);
1494                    },
1495                    Err(e) => panic!("unexpected push error: {e:?}"),
1496                }
1497            }
1498        }
1499    }
1500
1501    #[test]
1502    fn flush_leaves_denoiser_reusable_spatial() {
1503        let mut d = Denoiser::create(
1504            &[Accelerator::Vulkan],
1505            &Device::Default,
1506            16,
1507            16,
1508            opts(DenoisingMode::Spacial),
1509        )
1510        .unwrap();
1511
1512        let mut batch_a = Vec::new();
1513        push_n_with_drain(&mut d, 5, 0.25, &mut batch_a);
1514        d.flush(|f| batch_a.push(f)).expect("first flush failed");
1515        assert_eq!(batch_a.len(), 5);
1516
1517        // After flush the pipeline must be empty.
1518        assert!(d.recv_frame().unwrap().is_none());
1519
1520        let mut batch_b = Vec::new();
1521        push_n_with_drain(&mut d, 5, 0.75, &mut batch_b);
1522        d.flush(|f| batch_b.push(f)).expect("second flush failed");
1523        assert_eq!(batch_b.len(), 5);
1524
1525        for v in batch_b.iter().flatten() {
1526            assert!((v - 0.75).abs() < 0.1, "batch_b carried state from batch_a: {v}");
1527        }
1528        for v in batch_a.iter().flatten() {
1529            assert!((v - 0.25).abs() < 0.1, "batch_a value unexpectedly drifted: {v}");
1530        }
1531    }
1532
1533    #[test]
1534    fn flush_leaves_denoiser_reusable_temporal() {
1535        let mut d = Denoiser::create(
1536            &[Accelerator::Vulkan],
1537            &Device::Default,
1538            16,
1539            16,
1540            opts(DenoisingMode::Temporal { radius: 1 }),
1541        )
1542        .unwrap();
1543
1544        let mut batch_a = Vec::new();
1545        push_n_with_drain(&mut d, 5, 0.25, &mut batch_a);
1546        d.flush(|f| batch_a.push(f)).expect("first flush failed");
1547        assert_eq!(batch_a.len(), 5, "expected 5 frames from first batch");
1548
1549        // The temporal window must be empty after a flush, so the first
1550        // push of the new stream should not produce a pending frame.
1551        // With r=1 the window needs 3 frames before `denoise_submit`
1552        // fires.
1553        assert!(d.recv_frame().unwrap().is_none());
1554        d.push_frame(&frame_filled(16, 16, 0.75)).unwrap();
1555        assert!(
1556            d.recv_frame().unwrap().is_none(),
1557            "first push of new temporal stream should not produce output yet"
1558        );
1559
1560        // Push 4 more frames (5 total in batch B) with drain.
1561        let mut batch_b = Vec::new();
1562        push_n_with_drain(&mut d, 4, 0.75, &mut batch_b);
1563        d.flush(|f| batch_b.push(f)).expect("second flush failed");
1564        assert_eq!(batch_b.len(), 5, "expected 5 frames from second batch");
1565
1566        for v in batch_b.iter().flatten() {
1567            assert!((v - 0.75).abs() < 0.1, "batch_b carried state from batch_a: {v}");
1568        }
1569    }
1570
1571    #[test]
1572    fn flush_emits_exactly_n_outputs_for_small_n() {
1573        // With temporal radius R=2 the window is 5 frames. Pushing fewer
1574        // than R+1 frames means the window never fills while pushing, so
1575        // flush must still emit one output per pushed frame rather than
1576        // R+1 of them.
1577        for n in 1..=5usize {
1578            let mut d = Denoiser::create(
1579                &[Accelerator::Vulkan],
1580                &Device::Default,
1581                16,
1582                16,
1583                opts(DenoisingMode::Temporal { radius: 2 }),
1584            )
1585            .unwrap();
1586
1587            let mut out = Vec::new();
1588            push_n_with_drain(&mut d, n, 0.5, &mut out);
1589            d.flush(|f| out.push(f)).expect("flush failed");
1590            assert_eq!(
1591                out.len(),
1592                n,
1593                "expected {n} outputs for {n} pushes, got {}",
1594                out.len()
1595            );
1596        }
1597    }
1598}