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 Depth,
14 HqParams,
15 MotionCompensationMode,
16 MotionSearch,
17 NlmDenoiser,
18 NlmParams,
19 Pending,
20 PrefilterMode,
21 TryWait,
22 hq_default_strength,
23 validate_dimensions,
24};
25use crate::sniff::sniff_best_accelerator;
26
27#[derive(Debug, Clone, bon::Builder)]
35pub struct DenoiserOptions {
36 #[builder(default = ChannelMode::Yuv)]
38 pub channel_mode: ChannelMode,
39 #[builder(default = DenoisingMode::Spacial)]
42 pub mode: DenoisingMode,
43 #[builder(default)]
46 pub algorithm: Algorithm,
47 #[builder(default = OutputFormat::F32)]
49 pub output_format: OutputFormat,
50}
51
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub enum OutputFormat {
60 F32,
62 Wire { depth: Depth },
64}
65
66#[derive(Debug, Clone, PartialEq)]
68pub enum FrameOutput {
69 F32(Vec<f32>),
71 Wire(Vec<u8>),
74}
75
76impl FrameOutput {
77 pub fn into_f32(self) -> Option<Vec<f32>> {
79 match self {
80 Self::F32(v) => Some(v),
81 Self::Wire(_) => None,
82 }
83 }
84
85 pub fn into_wire(self) -> Option<Vec<u8>> {
87 match self {
88 Self::Wire(v) => Some(v),
89 Self::F32(_) => None,
90 }
91 }
92
93 pub fn as_f32(&self) -> Option<&[f32]> {
96 match self {
97 Self::F32(v) => Some(v),
98 Self::Wire(_) => None,
99 }
100 }
101
102 pub fn as_wire(&self) -> Option<&[u8]> {
105 match self {
106 Self::Wire(v) => Some(v),
107 Self::F32(_) => None,
108 }
109 }
110}
111
112#[derive(Debug, Copy, Clone, PartialEq)]
117pub enum Algorithm {
118 Nlmeans(NlmeansOptions),
121 NlmeansHq(NlmeansHqOptions),
127 Nl4d(Nl4dOptions),
134}
135
136impl Default for Algorithm {
137 fn default() -> Self {
138 Self::Nlmeans(NlmeansOptions::default())
139 }
140}
141
142#[derive(Debug, Copy, Clone, Default, PartialEq)]
144pub struct NlmeansOptions {
145 pub prefilter: PrefilterMode,
150 pub motion_compensation: MotionCompensationMode,
159 pub tuning: NlmTuning,
162}
163
164#[derive(Debug, Copy, Clone, Default, PartialEq)]
166pub struct NlmeansHqOptions {
167 pub nlm: NlmeansOptions,
169 pub hq: HqParams,
171}
172
173#[derive(Debug, Copy, Clone, PartialEq)]
191pub struct Nl4dOptions {
192 pub motion: MotionSearch,
194 pub sigma: Option<f32>,
200 pub sigma_scale: f32,
206 pub thsad_scale: f32,
212 pub refine: u32,
215 pub spatial_radius: u32,
218 pub lambda_ht: Option<f32>,
224 pub lambda_ht_scale: f32,
231 pub c_min: f32,
236 pub mismatch_scale: f32,
243 pub confidence_variance: bool,
247 pub windowed_noise_estimation: bool,
257}
258
259impl Default for Nl4dOptions {
260 fn default() -> Self {
261 let defaults = Nl4dParams::default();
262 let hq = HqParams::default();
263 Self {
264 motion: MotionSearch::default(),
265 sigma: hq.sigma_override,
266 sigma_scale: hq.sigma_scale,
267 thsad_scale: hq.thsad_scale,
268 refine: defaults.refine,
269 spatial_radius: defaults.spatial_radius,
270 lambda_ht: None,
274 lambda_ht_scale: 1.0,
275 c_min: defaults.c_min,
276 mismatch_scale: defaults.mismatch_scale,
277 confidence_variance: defaults.confidence_variance,
278 windowed_noise_estimation: false,
279 }
280 }
281}
282
283impl Nl4dOptions {
284 fn to_hq_params(self) -> HqParams {
291 HqParams {
292 sigma_override: self.sigma,
293 sigma_scale: self.sigma_scale,
294 thsad_scale: self.thsad_scale,
295 temporal_confidence: true,
296 windowed_noise_estimation: self.windowed_noise_estimation,
297 ..HqParams::default()
298 }
299 }
300}
301
302pub fn nl4d_default_lambda_ht(channels: ChannelMode) -> f32 {
321 match channels {
322 ChannelMode::Luma | ChannelMode::Yuv => 5.3,
323 ChannelMode::Chroma => 4.2,
324 }
325}
326
327fn resolve_lambda_ht(opts: &Nl4dOptions, channels: ChannelMode) -> Result<f32, String> {
339 if !(opts.lambda_ht_scale.is_finite() && (0.1..=10.0).contains(&opts.lambda_ht_scale)) {
340 return Err(format!(
341 "lambda_ht_scale must be finite and in [0.1, 10.0], got {}",
342 opts.lambda_ht_scale
343 ));
344 }
345
346 let lambda_ht = opts.lambda_ht.unwrap_or_else(|| nl4d_default_lambda_ht(channels));
347
348 Ok(lambda_ht * opts.lambda_ht_scale)
349}
350
351#[derive(Debug, Copy, Clone, Default, PartialEq, Eq, strum_macros::EnumString)]
362#[strum(ascii_case_insensitive)]
363pub enum Preset {
364 Veryfast,
366 Fast,
368 #[default]
370 Base,
371 Slow,
373 Veryslow,
375}
376
377#[derive(Debug, Copy, Clone, PartialEq, Eq, strum_macros::EnumString)]
379#[strum(ascii_case_insensitive)]
380pub enum NlmeansVariant {
381 Fast,
383 Hq,
386}
387
388pub fn nlmeans_variant_for(preset: Preset) -> NlmeansVariant {
390 match preset {
391 Preset::Veryfast => NlmeansVariant::Fast,
392 Preset::Fast | Preset::Base | Preset::Slow | Preset::Veryslow => NlmeansVariant::Hq,
393 }
394}
395
396pub fn nlmeans_temporal_radius_for(preset: Preset) -> u32 {
399 match preset {
400 Preset::Veryfast => 0,
401 Preset::Fast => 1,
402 Preset::Base => 2,
403 Preset::Slow => 4,
404 Preset::Veryslow => 8,
405 }
406}
407
408pub fn nlmeans_search_radius_for(preset: Preset) -> u32 {
411 match preset {
412 Preset::Veryfast | Preset::Fast | Preset::Base => 2,
413 Preset::Slow | Preset::Veryslow => 4,
414 }
415}
416
417pub fn nl4d_temporal_radius_for(preset: Preset) -> u32 {
423 match preset {
424 Preset::Veryfast | Preset::Fast => 1,
425 Preset::Base => 2,
426 Preset::Slow => 4,
427 Preset::Veryslow => 8,
428 }
429}
430
431pub fn nl4d_spatial_radius_for(preset: Preset) -> u32 {
442 match preset {
443 Preset::Veryfast => 6,
444 Preset::Fast | Preset::Base | Preset::Slow | Preset::Veryslow => {
445 Nl4dOptions::default().spatial_radius
446 },
447 }
448}
449
450#[derive(Debug, Copy, Clone, Eq, PartialEq)]
452pub enum DenoisingMode {
453 Spacial,
455 Temporal { radius: u32 },
457}
458
459#[derive(Debug, Copy, Clone, Default, PartialEq)]
464pub struct NlmTuning {
465 pub search_radius: Option<u32>,
466 pub patch_radius: Option<u32>,
467 pub strength: Option<f32>,
468 pub self_weight: Option<f32>,
469}
470
471impl DenoiserOptions {
472 #[doc(hidden)]
486 pub fn to_nlm_params(&self) -> NlmParams {
487 let temporal_radius = match self.mode {
488 DenoisingMode::Spacial => 0,
489 DenoisingMode::Temporal { radius } => radius,
490 };
491
492 match self.algorithm {
493 Algorithm::Nlmeans(opts) => self.nlm_params_for(opts, None, temporal_radius),
494 Algorithm::NlmeansHq(opts) => self.nlm_params_for(opts.nlm, Some(opts.hq), temporal_radius),
495 Algorithm::Nl4d(opts) => NlmParams {
499 channels: self.channel_mode,
500 motion_compensation: opts.motion.into(),
501 temporal_radius,
502 hq: Some(opts.to_hq_params()),
503 ..NlmParams::default()
504 },
505 }
506 }
507
508 fn nlm_params_for(&self, opts: NlmeansOptions, hq: Option<HqParams>, temporal_radius: u32) -> NlmParams {
511 let strength = opts.tuning.strength.unwrap_or(match hq {
527 Some(hq) if hq.auto_strength => hq_default_strength(self.channel_mode, temporal_radius),
528 _ => NlmParams::default().strength,
529 });
530
531 let defaults = NlmParams::default();
532 NlmParams {
533 channels: self.channel_mode,
534 prefilter: opts.prefilter,
535 motion_compensation: opts.motion_compensation,
536 temporal_radius,
537 hq,
538 strength,
539 search_radius: opts.tuning.search_radius.unwrap_or(defaults.search_radius),
540 patch_radius: opts.tuning.patch_radius.unwrap_or(defaults.patch_radius),
541 self_weight: opts.tuning.self_weight.unwrap_or(defaults.self_weight),
542 }
543 }
544}
545
546#[derive(Debug, thiserror::Error)]
548pub enum DenoiserError {
549 #[error("denoiser queue is full, collect the pending frame before pushing more")]
555 QueueFull,
556 #[error("denoiser failed earlier, reset the stream before using it again")]
561 Poisoned,
562 #[error("no accelerator from the priority list is available")]
564 NoAcceleratorAvailable,
565 #[error(transparent)]
568 Other(#[from] anyhow::Error),
569}
570
571enum Engine<R: Runtime> {
577 Nlm(Box<NlmDenoiser<R>>),
578 Nl4d(Box<Nl4dDenoiser<R>>),
579}
580
581impl<R: Runtime> Engine<R> {
582 fn is_nl4d(&self) -> bool {
583 matches!(self, Self::Nl4d(_))
584 }
585
586 fn push_frame(&mut self, frame: &[f32]) {
587 match self {
588 Self::Nlm(d) => d.push_frame(frame),
589 Self::Nl4d(d) => d.push_frame(frame),
590 }
591 }
592
593 fn push_frame_wire(&mut self, planes: &[&[u8]], depth: Depth) {
594 match self {
595 Self::Nlm(d) => d.push_frame_wire(planes, depth),
596 Self::Nl4d(d) => d.push_frame_wire(planes, depth),
597 }
598 }
599
600 fn denoise_submit(&mut self) -> Result<Option<Pending<R>>, anyhow::Error> {
601 match self {
602 Self::Nlm(d) => d.denoise_submit(),
603 Self::Nl4d(d) => d.denoise_submit().map_err(anyhow::Error::from),
608 }
609 }
610
611 #[cfg(test)]
612 fn wire_outputs(&self) -> Option<&[cubecl::server::Handle; 2]> {
613 match self {
614 Self::Nlm(d) => d.wire_outputs_for_test(),
615 Self::Nl4d(d) => d.wire_outputs_for_test(),
616 }
617 }
618
619 fn flush(&mut self, sink: impl FnMut(&FrameOutput)) -> Result<(), anyhow::Error> {
620 match self {
621 Self::Nlm(d) => d.flush(sink),
622 Self::Nl4d(d) => d.flush(sink).map_err(anyhow::Error::from),
623 }
624 }
625
626 fn reset_stream(&mut self) {
627 match self {
628 Self::Nlm(d) => d.reset_stream_state(),
629 Self::Nl4d(d) => d.reset_stream(),
630 }
631 }
632}
633
634fn build_engine<R: Runtime>(
644 client: &ComputeClient<R>,
645 algorithm: &Algorithm,
646 params: NlmParams,
647 width: u32,
648 height: u32,
649 output_format: OutputFormat,
650) -> Result<Engine<R>, DenoiserError> {
651 match algorithm {
652 Algorithm::Nl4d(opts) => {
653 if params.temporal_radius == 0 {
656 return Err(DenoiserError::Other(anyhow::anyhow!(
657 "nl4d needs a temporal window, set DenoiserOptions::mode to \
658 DenoisingMode::Temporal"
659 )));
660 }
661
662 let lambda_ht = resolve_lambda_ht(opts, params.channels)
663 .map_err(|e| DenoiserError::Other(anyhow::anyhow!(e)))?;
664 let nl4d_params = Nl4dParams {
665 temporal_radius: params.temporal_radius,
666 nlm: params,
667 refine: opts.refine,
668 spatial_radius: opts.spatial_radius,
669 lambda_ht,
670 c_min: opts.c_min,
671 mismatch_scale: opts.mismatch_scale,
672 confidence_variance: opts.confidence_variance,
673 };
674 let denoiser =
675 Nl4dDenoiser::with_output_format(client, nl4d_params, width, height, output_format)
676 .map_err(|e| DenoiserError::Other(anyhow::anyhow!(e)))?;
677 Ok(Engine::Nl4d(Box::new(denoiser)))
678 },
679 Algorithm::Nlmeans(_) | Algorithm::NlmeansHq(_) => Ok(Engine::Nlm(Box::new(
680 NlmDenoiser::with_output_format(client, params, width, height, output_format),
681 ))),
682 }
683}
684
685enum Backend {
686 #[cfg(feature = "cuda")]
687 Cuda(Engine<cubecl::cuda::CudaRuntime>),
688 #[cfg(feature = "rocm")]
689 Rocm(Engine<cubecl::hip::HipRuntime>),
690 #[cfg(any(feature = "vulkan", feature = "metal"))]
691 Wgpu(Engine<cubecl::wgpu::WgpuRuntime>),
692}
693
694impl Backend {
695 fn is_nl4d(&self) -> bool {
696 match self {
697 #[cfg(feature = "cuda")]
698 Self::Cuda(e) => e.is_nl4d(),
699 #[cfg(feature = "rocm")]
700 Self::Rocm(e) => e.is_nl4d(),
701 #[cfg(any(feature = "vulkan", feature = "metal"))]
702 Self::Wgpu(e) => e.is_nl4d(),
703 }
704 }
705
706 #[cfg(test)]
707 fn wire_outputs(&self) -> Option<&[cubecl::server::Handle; 2]> {
708 match self {
709 #[cfg(feature = "cuda")]
710 Self::Cuda(e) => e.wire_outputs(),
711 #[cfg(feature = "rocm")]
712 Self::Rocm(e) => e.wire_outputs(),
713 #[cfg(any(feature = "vulkan", feature = "metal"))]
714 Self::Wgpu(e) => e.wire_outputs(),
715 }
716 }
717}
718
719enum BackendPending {
720 #[cfg(feature = "cuda")]
721 Cuda(Pending<cubecl::cuda::CudaRuntime>),
722 #[cfg(feature = "rocm")]
723 Rocm(Pending<cubecl::hip::HipRuntime>),
724 #[cfg(any(feature = "vulkan", feature = "metal"))]
725 Wgpu(Pending<cubecl::wgpu::WgpuRuntime>),
726}
727
728impl BackendPending {
729 fn wait(self) -> Result<FrameOutput, anyhow::Error> {
730 match self {
731 #[cfg(feature = "cuda")]
732 Self::Cuda(p) => p.wait(),
733 #[cfg(feature = "rocm")]
734 Self::Rocm(p) => p.wait(),
735 #[cfg(any(feature = "vulkan", feature = "metal"))]
736 Self::Wgpu(p) => p.wait(),
737 }
738 }
739
740 fn try_wait(self) -> Result<Result<FrameOutput, Self>, anyhow::Error> {
743 match self {
744 #[cfg(feature = "cuda")]
745 Self::Cuda(p) => match p.try_wait()? {
746 TryWait::Ready(frame) => Ok(Ok(frame)),
747 TryWait::NotReady(p) => Ok(Err(Self::Cuda(p))),
748 },
749 #[cfg(feature = "rocm")]
750 Self::Rocm(p) => match p.try_wait()? {
751 TryWait::Ready(frame) => Ok(Ok(frame)),
752 TryWait::NotReady(p) => Ok(Err(Self::Rocm(p))),
753 },
754 #[cfg(any(feature = "vulkan", feature = "metal"))]
755 Self::Wgpu(p) => match p.try_wait()? {
756 TryWait::Ready(frame) => Ok(Ok(frame)),
757 TryWait::NotReady(p) => Ok(Err(Self::Wgpu(p))),
758 },
759 }
760 }
761}
762
763pub const MAX_PENDING: usize = 2;
770
771#[derive(Debug, Clone, Copy, PartialEq, Eq)]
782pub struct WindowSpan {
783 pub behind: usize,
785 pub ahead: usize,
787}
788
789impl WindowSpan {
790 pub fn frame_count(&self) -> usize {
793 self.behind + 1 + self.ahead
794 }
795}
796
797pub struct Denoiser {
848 backend: Backend,
849 pending: VecDeque<BackendPending>,
850 accelerator: Accelerator,
851 width: u32,
852 height: u32,
853 temporal_radius: u32,
854 output_format: OutputFormat,
855 frames_pushed: u32,
856 poisoned: bool,
862}
863
864impl Denoiser {
865 pub fn create(
885 accelerators: &[Accelerator],
886 device: &Device,
887 width: u32,
888 height: u32,
889 options: DenoiserOptions,
890 ) -> Result<Self, DenoiserError> {
891 let accelerator =
892 sniff_best_accelerator(accelerators, device).ok_or(DenoiserError::NoAcceleratorAvailable)?;
893
894 let params = options.to_nlm_params();
895 params.validate()?;
896 validate_dimensions(width, height)?;
897
898 let temporal_radius = params.temporal_radius;
899 let backend = build_backend(
900 accelerator,
901 device,
902 &options.algorithm,
903 params,
904 width,
905 height,
906 options.output_format,
907 )?;
908
909 Ok(Self {
910 backend,
911 pending: VecDeque::with_capacity(MAX_PENDING),
912 accelerator,
913 width,
914 height,
915 temporal_radius,
916 output_format: options.output_format,
917 frames_pushed: 0,
918 poisoned: false,
919 })
920 }
921
922 pub fn selected_accelerator(&self) -> Accelerator {
924 self.accelerator
925 }
926
927 pub fn width(&self) -> u32 {
929 self.width
930 }
931
932 pub fn height(&self) -> u32 {
934 self.height
935 }
936
937 pub fn temporal_radius(&self) -> u32 {
939 self.temporal_radius
940 }
941
942 pub fn output_format(&self) -> OutputFormat {
944 self.output_format
945 }
946
947 pub fn window_span(&self) -> WindowSpan {
965 let radius = self.temporal_radius as usize;
966 let span = if self.backend.is_nl4d() {
967 2 * radius
968 } else {
969 radius
970 };
971 WindowSpan {
972 behind: span,
973 ahead: span,
974 }
975 }
976
977 pub fn push_frame(&mut self, frame: &[f32]) -> Result<(), DenoiserError> {
996 if self.poisoned {
997 return Err(DenoiserError::Poisoned);
998 }
999 self.push_frame_inner(frame).inspect_err(|err| {
1000 if !matches!(err, DenoiserError::QueueFull) {
1001 self.poisoned = true;
1002 }
1003 })
1004 }
1005
1006 fn push_frame_inner(&mut self, frame: &[f32]) -> Result<(), DenoiserError> {
1007 let window_full = self.frames_pushed > self.temporal_radius;
1011 if window_full && self.pending.len() >= MAX_PENDING {
1012 return Err(DenoiserError::QueueFull);
1013 }
1014
1015 match &mut self.backend {
1016 #[cfg(feature = "cuda")]
1017 Backend::Cuda(d) => {
1018 d.push_frame(frame);
1019 if let Some(p) = d.denoise_submit()? {
1020 self.pending.push_back(BackendPending::Cuda(p));
1021 }
1022 },
1023 #[cfg(feature = "rocm")]
1024 Backend::Rocm(d) => {
1025 d.push_frame(frame);
1026 if let Some(p) = d.denoise_submit()? {
1027 self.pending.push_back(BackendPending::Rocm(p));
1028 }
1029 },
1030 #[cfg(any(feature = "vulkan", feature = "metal"))]
1031 Backend::Wgpu(d) => {
1032 d.push_frame(frame);
1033 if let Some(p) = d.denoise_submit()? {
1034 self.pending.push_back(BackendPending::Wgpu(p));
1035 }
1036 },
1037 }
1038
1039 self.frames_pushed = self.frames_pushed.saturating_add(1);
1040 Ok(())
1041 }
1042
1043 pub fn push_frame_wire(&mut self, planes: &[&[u8]], depth: Depth) -> Result<(), DenoiserError> {
1052 if self.poisoned {
1053 return Err(DenoiserError::Poisoned);
1054 }
1055 self.push_frame_wire_inner(planes, depth).inspect_err(|err| {
1056 if !matches!(err, DenoiserError::QueueFull) {
1057 self.poisoned = true;
1058 }
1059 })
1060 }
1061
1062 fn push_frame_wire_inner(&mut self, planes: &[&[u8]], depth: Depth) -> Result<(), DenoiserError> {
1063 let window_full = self.frames_pushed > self.temporal_radius;
1065 if window_full && self.pending.len() >= MAX_PENDING {
1066 return Err(DenoiserError::QueueFull);
1067 }
1068
1069 match &mut self.backend {
1070 #[cfg(feature = "cuda")]
1071 Backend::Cuda(d) => {
1072 d.push_frame_wire(planes, depth);
1073 if let Some(p) = d.denoise_submit()? {
1074 self.pending.push_back(BackendPending::Cuda(p));
1075 }
1076 },
1077 #[cfg(feature = "rocm")]
1078 Backend::Rocm(d) => {
1079 d.push_frame_wire(planes, depth);
1080 if let Some(p) = d.denoise_submit()? {
1081 self.pending.push_back(BackendPending::Rocm(p));
1082 }
1083 },
1084 #[cfg(any(feature = "vulkan", feature = "metal"))]
1085 Backend::Wgpu(d) => {
1086 d.push_frame_wire(planes, depth);
1087 if let Some(p) = d.denoise_submit()? {
1088 self.pending.push_back(BackendPending::Wgpu(p));
1089 }
1090 },
1091 }
1092
1093 self.frames_pushed = self.frames_pushed.saturating_add(1);
1094 Ok(())
1095 }
1096
1097 pub fn push_frame_wire_priming(&mut self, planes: &[&[u8]], depth: Depth) -> Result<(), DenoiserError> {
1102 if self.poisoned {
1103 return Err(DenoiserError::Poisoned);
1104 }
1105 match &mut self.backend {
1106 #[cfg(feature = "cuda")]
1107 Backend::Cuda(d) => d.push_frame_wire(planes, depth),
1108 #[cfg(feature = "rocm")]
1109 Backend::Rocm(d) => d.push_frame_wire(planes, depth),
1110 #[cfg(any(feature = "vulkan", feature = "metal"))]
1111 Backend::Wgpu(d) => d.push_frame_wire(planes, depth),
1112 }
1113
1114 self.frames_pushed = self.frames_pushed.saturating_add(1);
1115 Ok(())
1116 }
1117
1118 pub fn push_frame_priming(&mut self, frame: &[f32]) -> Result<(), DenoiserError> {
1130 if self.poisoned {
1131 return Err(DenoiserError::Poisoned);
1132 }
1133 match &mut self.backend {
1134 #[cfg(feature = "cuda")]
1135 Backend::Cuda(d) => d.push_frame(frame),
1136 #[cfg(feature = "rocm")]
1137 Backend::Rocm(d) => d.push_frame(frame),
1138 #[cfg(any(feature = "vulkan", feature = "metal"))]
1139 Backend::Wgpu(d) => d.push_frame(frame),
1140 }
1141
1142 self.frames_pushed = self.frames_pushed.saturating_add(1);
1143 Ok(())
1144 }
1145
1146 pub fn reset_stream(&mut self) {
1153 self.pending.clear();
1154 self.frames_pushed = 0;
1155 self.poisoned = false;
1156
1157 match &mut self.backend {
1158 #[cfg(feature = "cuda")]
1159 Backend::Cuda(d) => d.reset_stream(),
1160 #[cfg(feature = "rocm")]
1161 Backend::Rocm(d) => d.reset_stream(),
1162 #[cfg(any(feature = "vulkan", feature = "metal"))]
1163 Backend::Wgpu(d) => d.reset_stream(),
1164 }
1165 }
1166
1167 pub fn recv_frame(&mut self) -> Result<Option<FrameOutput>, DenoiserError> {
1176 if self.poisoned {
1177 return Err(DenoiserError::Poisoned);
1178 }
1179 self.recv_frame_inner().inspect_err(|_| self.poisoned = true)
1180 }
1181
1182 fn recv_frame_inner(&mut self) -> Result<Option<FrameOutput>, DenoiserError> {
1183 let Some(pending) = self.pending.pop_front() else {
1184 return Ok(None);
1185 };
1186 Ok(Some(pending.wait()?))
1187 }
1188
1189 pub fn try_recv_frame(&mut self) -> Result<Option<FrameOutput>, DenoiserError> {
1204 if self.poisoned {
1205 return Err(DenoiserError::Poisoned);
1206 }
1207 self.try_recv_frame_inner().inspect_err(|_| self.poisoned = true)
1208 }
1209
1210 fn try_recv_frame_inner(&mut self) -> Result<Option<FrameOutput>, DenoiserError> {
1211 let Some(pending) = self.pending.pop_front() else {
1212 return Ok(None);
1213 };
1214
1215 match pending.try_wait()? {
1216 Ok(frame) => Ok(Some(frame)),
1217 Err(pending) => {
1218 self.pending.push_front(pending);
1219 Ok(None)
1220 },
1221 }
1222 }
1223
1224 pub fn flush(&mut self, sink: impl FnMut(FrameOutput)) -> Result<(), DenoiserError> {
1237 if self.poisoned {
1238 return Err(DenoiserError::Poisoned);
1239 }
1240 self.flush_inner(sink).inspect_err(|_| self.poisoned = true)
1241 }
1242
1243 fn flush_inner(&mut self, mut sink: impl FnMut(FrameOutput)) -> Result<(), DenoiserError> {
1244 while let Some(frame) = self.recv_frame_inner()? {
1249 sink(frame);
1250 }
1251
1252 match &mut self.backend {
1256 #[cfg(feature = "cuda")]
1257 Backend::Cuda(d) => d.flush(|frame| sink(frame.clone()))?,
1258 #[cfg(feature = "rocm")]
1259 Backend::Rocm(d) => d.flush(|frame| sink(frame.clone()))?,
1260 #[cfg(any(feature = "vulkan", feature = "metal"))]
1261 Backend::Wgpu(d) => d.flush(|frame| sink(frame.clone()))?,
1262 }
1263
1264 self.frames_pushed = 0;
1268
1269 Ok(())
1270 }
1271
1272 #[cfg(test)]
1276 pub(crate) fn poison_for_test(&mut self) {
1277 self.poisoned = true;
1278 }
1279
1280 #[cfg(test)]
1283 pub(crate) fn wire_outputs_for_test(&self) -> Option<&[cubecl::server::Handle; 2]> {
1284 self.backend.wire_outputs()
1285 }
1286}
1287
1288fn build_backend(
1289 accel: Accelerator,
1290 device: &Device,
1291 algorithm: &Algorithm,
1292 params: NlmParams,
1293 width: u32,
1294 height: u32,
1295 output_format: OutputFormat,
1296) -> Result<Backend, DenoiserError> {
1297 match accel {
1298 #[cfg(feature = "cuda")]
1299 Accelerator::Cuda => {
1300 let dev = device.to_cuda()?;
1301 let client = <cubecl::cuda::CudaRuntime as Runtime>::client(&dev);
1302 Ok(Backend::Cuda(build_engine(
1303 &client,
1304 algorithm,
1305 params,
1306 width,
1307 height,
1308 output_format,
1309 )?))
1310 },
1311 #[cfg(feature = "rocm")]
1312 Accelerator::Rocm => {
1313 let dev = device.to_amd()?;
1314 let client = <cubecl::hip::HipRuntime as Runtime>::client(&dev);
1315 Ok(Backend::Rocm(build_engine(
1316 &client,
1317 algorithm,
1318 params,
1319 width,
1320 height,
1321 output_format,
1322 )?))
1323 },
1324 #[cfg(feature = "vulkan")]
1325 Accelerator::Vulkan => {
1326 let dev = device.to_wgpu()?;
1327 let client = <cubecl::wgpu::WgpuRuntime as Runtime>::client(&dev);
1328 Ok(Backend::Wgpu(build_engine(
1329 &client,
1330 algorithm,
1331 params,
1332 width,
1333 height,
1334 output_format,
1335 )?))
1336 },
1337 #[cfg(feature = "metal")]
1338 Accelerator::Metal => {
1339 let dev = device.to_wgpu()?;
1340 let client = <cubecl::wgpu::WgpuRuntime as Runtime>::client(&dev);
1341 Ok(Backend::Wgpu(build_engine(
1342 &client,
1343 algorithm,
1344 params,
1345 width,
1346 height,
1347 output_format,
1348 )?))
1349 },
1350 #[cfg(docsrs)]
1354 #[expect(
1355 unreachable_patterns,
1356 reason = "the arm only keeps the match exhaustive on docs.rs"
1357 )]
1358 _ => unreachable!(),
1359 }
1360}
1361
1362#[cfg(test)]
1363mod options_tests {
1364 use super::*;
1365
1366 fn hq(hq: HqParams) -> Algorithm {
1369 Algorithm::NlmeansHq(NlmeansHqOptions {
1370 hq,
1371 ..NlmeansHqOptions::default()
1372 })
1373 }
1374
1375 fn fast_tuned(tuning: NlmTuning) -> Algorithm {
1377 Algorithm::Nlmeans(NlmeansOptions {
1378 tuning,
1379 ..NlmeansOptions::default()
1380 })
1381 }
1382
1383 #[test]
1384 fn nl4d_default_lambda_ht_differs_between_luma_and_chroma() {
1385 let luma = nl4d_default_lambda_ht(ChannelMode::Luma);
1386 let chroma = nl4d_default_lambda_ht(ChannelMode::Chroma);
1387
1388 assert!((luma - 5.3).abs() < f32::EPSILON);
1389 assert!((chroma - 4.2).abs() < f32::EPSILON);
1390 assert!(
1391 (chroma - luma).abs() > f32::EPSILON,
1392 "the two planes should not resolve to the same default"
1393 );
1394 }
1395
1396 #[test]
1397 fn nl4d_default_lambda_ht_yuv_reads_the_luma_value() {
1398 let yuv = nl4d_default_lambda_ht(ChannelMode::Yuv);
1399 let luma = nl4d_default_lambda_ht(ChannelMode::Luma);
1400
1401 assert!((yuv - luma).abs() < f32::EPSILON);
1402 }
1403
1404 #[test]
1405 fn resolve_lambda_ht_unset_uses_the_per_plane_default() {
1406 let opts = Nl4dOptions::default();
1407
1408 let luma = resolve_lambda_ht(&opts, ChannelMode::Luma).expect("the default scale is in range");
1409 let chroma = resolve_lambda_ht(&opts, ChannelMode::Chroma).expect("the default scale is in range");
1410
1411 assert!((luma - 5.3).abs() < f32::EPSILON, "got {luma}");
1412 assert!((chroma - 4.2).abs() < f32::EPSILON, "got {chroma}");
1413 }
1414
1415 #[test]
1416 fn resolve_lambda_ht_explicit_value_overrides_every_plane() {
1417 let opts = Nl4dOptions {
1418 lambda_ht: Some(4.4),
1419 ..Nl4dOptions::default()
1420 };
1421
1422 for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
1423 let got = resolve_lambda_ht(&opts, channels).expect("the default scale is in range");
1424 assert!(
1425 (got - 4.4).abs() < f32::EPSILON,
1426 "channels {channels:?} got {got}"
1427 );
1428 }
1429 }
1430
1431 #[test]
1432 fn resolve_lambda_ht_default_scale_leaves_the_value_alone() {
1433 let opts = Nl4dOptions::default();
1434
1435 for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
1436 let got = resolve_lambda_ht(&opts, channels).expect("the default scale is in range");
1437 let want = nl4d_default_lambda_ht(channels);
1438 assert!(
1439 (got - want).abs() < f32::EPSILON,
1440 "channels {channels:?} got {got}"
1441 );
1442 }
1443 }
1444
1445 #[test]
1446 fn resolve_lambda_ht_scale_multiplies_the_per_plane_default() {
1447 let opts = Nl4dOptions {
1448 lambda_ht_scale: 1.1,
1449 ..Nl4dOptions::default()
1450 };
1451
1452 for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
1453 let got = resolve_lambda_ht(&opts, channels).expect("1.1 is in range");
1454 let want = nl4d_default_lambda_ht(channels) * 1.1;
1455 assert!(
1456 (got - want).abs() < 1e-5,
1457 "channels {channels:?} got {got}, want {want}"
1458 );
1459 }
1460 }
1461
1462 #[test]
1465 fn resolve_lambda_ht_scale_multiplies_an_explicit_value() {
1466 let opts = Nl4dOptions {
1467 lambda_ht: Some(5.0),
1468 lambda_ht_scale: 0.9,
1469 ..Nl4dOptions::default()
1470 };
1471
1472 let got = resolve_lambda_ht(&opts, ChannelMode::Luma).expect("0.9 is in range");
1473 assert!((got - 4.5).abs() < 1e-5, "got {got}");
1474 }
1475
1476 #[test]
1477 fn resolve_lambda_ht_rejects_an_out_of_range_scale() {
1478 for bad in [0.0, -1.0, 0.05, 10.5, f32::NAN, f32::INFINITY] {
1479 let opts = Nl4dOptions {
1480 lambda_ht_scale: bad,
1481 ..Nl4dOptions::default()
1482 };
1483 let err = resolve_lambda_ht(&opts, ChannelMode::Luma).unwrap_err();
1484 assert!(
1485 err.contains("lambda_ht_scale"),
1486 "lambda_ht_scale={bad} should be rejected, got {err}"
1487 );
1488 }
1489 }
1490
1491 #[test]
1492 fn the_default_algorithm_is_the_fast_nlmeans_path() {
1493 let opts = DenoiserOptions::builder().build();
1494 assert_eq!(opts.algorithm, Algorithm::Nlmeans(NlmeansOptions::default()));
1495 }
1496
1497 #[test]
1498 fn spatial_mode_maps_to_zero_temporal_radius() {
1499 let opts = DenoiserOptions::builder()
1500 .channel_mode(ChannelMode::Yuv)
1501 .mode(DenoisingMode::Spacial)
1502 .build();
1503 let params = opts.to_nlm_params();
1504
1505 assert_eq!(params.temporal_radius, 0);
1506 assert_eq!(params.channels, ChannelMode::Yuv);
1507 }
1508
1509 #[test]
1510 fn temporal_mode_propagates_radius() {
1511 let opts = DenoiserOptions::builder()
1512 .mode(DenoisingMode::Temporal { radius: 3 })
1513 .build();
1514 let params = opts.to_nlm_params();
1515
1516 assert_eq!(params.temporal_radius, 3);
1517 }
1518
1519 #[test]
1520 fn prefilter_passthrough() {
1521 let opts = DenoiserOptions::builder()
1522 .algorithm(Algorithm::Nlmeans(NlmeansOptions {
1523 prefilter: PrefilterMode::Bilateral {
1524 sigma_s: 3.0,
1525 sigma_r: 0.02,
1526 },
1527 ..NlmeansOptions::default()
1528 }))
1529 .build();
1530 let params = opts.to_nlm_params();
1531
1532 assert!(matches!(params.prefilter, PrefilterMode::Bilateral { .. }));
1533 }
1534
1535 #[test]
1536 fn hq_unset_prefilter_defaults_to_none() {
1537 let opts = DenoiserOptions::builder()
1538 .algorithm(hq(HqParams::default()))
1539 .build();
1540 let params = opts.to_nlm_params();
1541
1542 assert!(matches!(params.prefilter, PrefilterMode::None));
1543 }
1544
1545 #[test]
1546 fn fast_unset_prefilter_defaults_to_none() {
1547 let opts = DenoiserOptions::builder()
1548 .algorithm(Algorithm::Nlmeans(NlmeansOptions::default()))
1549 .build();
1550 let params = opts.to_nlm_params();
1551
1552 assert!(matches!(params.prefilter, PrefilterMode::None));
1553 }
1554
1555 #[test]
1556 fn hq_unset_strength_defaults_to_hq_default_strength() {
1557 let opts = DenoiserOptions::builder()
1559 .algorithm(hq(HqParams::default()))
1560 .build();
1561 let params = opts.to_nlm_params();
1562
1563 let expected = hq_default_strength(ChannelMode::Yuv, 0);
1564 assert!((params.strength - expected).abs() < f32::EPSILON);
1565 }
1566
1567 #[test]
1568 fn hq_no_auto_strength_falls_back_to_the_legacy_absolute_default() {
1569 let opts = DenoiserOptions::builder()
1576 .algorithm(hq(HqParams {
1577 auto_strength: false,
1578 ..HqParams::default()
1579 }))
1580 .build();
1581 let params = opts.to_nlm_params();
1582
1583 let expected = NlmParams::default().strength;
1584 assert!(
1585 (params.strength - expected).abs() < f32::EPSILON,
1586 "expected the legacy absolute default {expected}, got {}, which looks like the \
1587 auto-strength multiplier table leaking through",
1588 params.strength
1589 );
1590 }
1591
1592 #[test]
1593 fn hq_luma_r4_uses_measured_table_value() {
1594 let opts = DenoiserOptions::builder()
1595 .channel_mode(ChannelMode::Luma)
1596 .mode(DenoisingMode::Temporal { radius: 4 })
1597 .algorithm(hq(HqParams::default()))
1598 .build();
1599 let params = opts.to_nlm_params();
1600
1601 assert!((params.strength - 0.35).abs() < f32::EPSILON);
1602 }
1603
1604 #[test]
1605 fn hq_chroma_r4_uses_measured_table_value() {
1606 let opts = DenoiserOptions::builder()
1607 .channel_mode(ChannelMode::Chroma)
1608 .mode(DenoisingMode::Temporal { radius: 4 })
1609 .algorithm(hq(HqParams::default()))
1610 .build();
1611 let params = opts.to_nlm_params();
1612
1613 assert!((params.strength - 0.70).abs() < f32::EPSILON);
1614 }
1615
1616 #[test]
1617 fn hq_yuv_r8_uses_measured_table_value() {
1618 let opts = DenoiserOptions::builder()
1619 .channel_mode(ChannelMode::Yuv)
1620 .mode(DenoisingMode::Temporal { radius: 8 })
1621 .algorithm(hq(HqParams::default()))
1622 .build();
1623 let params = opts.to_nlm_params();
1624
1625 assert!((params.strength - 0.30).abs() < f32::EPSILON);
1626 }
1627
1628 #[test]
1629 fn hq_spacial_mode_uses_radius_zero_table_values() {
1630 for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
1631 let opts = DenoiserOptions::builder()
1632 .channel_mode(channels)
1633 .mode(DenoisingMode::Spacial)
1634 .algorithm(hq(HqParams::default()))
1635 .build();
1636 let params = opts.to_nlm_params();
1637
1638 let expected = hq_default_strength(channels, 0);
1639 assert!(
1640 (params.strength - expected).abs() < f32::EPSILON,
1641 "for channels {channels:?} expected {expected}, got {}",
1642 params.strength
1643 );
1644 }
1645 }
1646
1647 #[test]
1648 fn hq_explicit_strength_wins_over_the_table_for_every_plane() {
1649 for channels in [ChannelMode::Luma, ChannelMode::Chroma, ChannelMode::Yuv] {
1650 let opts = DenoiserOptions::builder()
1651 .channel_mode(channels)
1652 .mode(DenoisingMode::Temporal { radius: 4 })
1653 .algorithm(Algorithm::NlmeansHq(NlmeansHqOptions {
1654 nlm: NlmeansOptions {
1655 tuning: NlmTuning {
1656 strength: Some(0.99),
1657 ..NlmTuning::default()
1658 },
1659 ..NlmeansOptions::default()
1660 },
1661 hq: HqParams::default(),
1662 }))
1663 .build();
1664 let params = opts.to_nlm_params();
1665
1666 assert!(
1667 (params.strength - 0.99).abs() < f32::EPSILON,
1668 "for channels {channels:?} the explicit strength was overridden by the table"
1669 );
1670 }
1671 }
1672
1673 #[test]
1674 fn fast_unset_strength_defaults_to_legacy_default() {
1675 let opts = DenoiserOptions::builder()
1676 .algorithm(Algorithm::Nlmeans(NlmeansOptions::default()))
1677 .build();
1678 let params = opts.to_nlm_params();
1679
1680 assert!((params.strength - 1.2).abs() < f32::EPSILON);
1681 }
1682
1683 #[test]
1684 fn nl4d_options_default_matches_nl4d_params_default() {
1685 let opts = Nl4dOptions::default();
1686 let params = crate::nl4d::Nl4dParams::default();
1687
1688 assert_eq!(opts.refine, params.refine);
1689 assert_eq!(opts.spatial_radius, params.spatial_radius);
1690 assert!((opts.c_min - params.c_min).abs() < f32::EPSILON);
1691 assert_eq!(opts.confidence_variance, params.confidence_variance);
1692 assert_eq!(opts.lambda_ht, None);
1699 assert!((params.lambda_ht - nl4d_default_lambda_ht(ChannelMode::Yuv)).abs() < f32::EPSILON);
1700 }
1701
1702 #[test]
1706 fn nl4d_builds_the_front_ends_hq_params_from_its_own_fields() {
1707 let opts = DenoiserOptions::builder()
1708 .mode(DenoisingMode::Temporal { radius: 2 })
1709 .algorithm(Algorithm::Nl4d(Nl4dOptions {
1710 sigma: Some(0.02),
1711 sigma_scale: 1.3,
1712 thsad_scale: 0.8,
1713 ..Nl4dOptions::default()
1714 }))
1715 .build();
1716 let params = opts.to_nlm_params();
1717
1718 let hq = params.hq.expect("nl4d always runs the hq front end");
1719 assert_eq!(hq.sigma_override, Some(0.02));
1720 assert!((hq.sigma_scale - 1.3).abs() < f32::EPSILON);
1721 assert!((hq.thsad_scale - 0.8).abs() < f32::EPSILON);
1722 assert!(
1723 hq.temporal_confidence,
1724 "the grouping kernel reads the confidence scores, so this cannot be off"
1725 );
1726 }
1727
1728 #[test]
1731 fn nl4d_reads_its_temporal_radius_from_the_denoising_mode() {
1732 for radius in [1u32, 4, 8] {
1733 let opts = DenoiserOptions::builder()
1734 .mode(DenoisingMode::Temporal { radius })
1735 .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1736 .build();
1737
1738 assert_eq!(opts.to_nlm_params().temporal_radius, radius);
1739 }
1740 }
1741
1742 #[test]
1745 fn nl4d_never_builds_a_prefilter() {
1746 let opts = DenoiserOptions::builder()
1747 .mode(DenoisingMode::Temporal { radius: 2 })
1748 .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1749 .build();
1750
1751 assert!(matches!(opts.to_nlm_params().prefilter, PrefilterMode::None));
1752 }
1753
1754 #[test]
1757 fn nl4d_leaves_the_nlm_weighting_knobs_at_their_defaults() {
1758 let defaults = NlmParams::default();
1759 let opts = DenoiserOptions::builder()
1760 .channel_mode(ChannelMode::Luma)
1761 .mode(DenoisingMode::Temporal { radius: 4 })
1762 .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1763 .build();
1764 let params = opts.to_nlm_params();
1765
1766 assert!((params.strength - defaults.strength).abs() < f32::EPSILON);
1767 assert_eq!(params.search_radius, defaults.search_radius);
1768 assert_eq!(params.patch_radius, defaults.patch_radius);
1769 assert!((params.self_weight - defaults.self_weight).abs() < f32::EPSILON);
1770 }
1771
1772 #[test]
1775 fn nl4d_motion_search_becomes_an_active_mvtools_mode() {
1776 let opts = DenoiserOptions::builder()
1777 .mode(DenoisingMode::Temporal { radius: 2 })
1778 .algorithm(Algorithm::Nl4d(Nl4dOptions {
1779 motion: MotionSearch {
1780 blksize: 32,
1781 overlap: 16,
1782 search_radius: 6,
1783 pyramid_levels: 1,
1784 estimation: MotionEstimation::Direct,
1785 },
1786 ..Nl4dOptions::default()
1787 }))
1788 .build();
1789 let params = opts.to_nlm_params();
1790
1791 assert!(matches!(
1792 params.motion_compensation,
1793 MotionCompensationMode::Mvtools {
1794 blksize: 32,
1795 overlap: 16,
1796 search_radius: 6,
1797 pyramid_levels: 1,
1798 estimation: MotionEstimation::Direct,
1799 }
1800 ));
1801 }
1802
1803 #[test]
1804 fn nl4d_motion_search_defaults_match_the_front_ends_own_defaults() {
1805 let opts = DenoiserOptions::builder()
1806 .mode(DenoisingMode::Temporal { radius: 2 })
1807 .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1808 .build();
1809 let params = opts.to_nlm_params();
1810
1811 assert_eq!(
1812 params.motion_compensation,
1813 crate::nl4d::Nl4dParams::default().nlm.motion_compensation
1814 );
1815 }
1816
1817 #[test]
1818 fn motion_compensation_passthrough() {
1819 let opts = DenoiserOptions::builder()
1820 .mode(DenoisingMode::Temporal { radius: 1 })
1821 .algorithm(Algorithm::Nlmeans(NlmeansOptions {
1822 motion_compensation: MotionCompensationMode::Mvtools {
1823 blksize: 16,
1824 overlap: 8,
1825 search_radius: 4,
1826 pyramid_levels: 2,
1827 estimation: MotionEstimation::Direct,
1828 },
1829 ..NlmeansOptions::default()
1830 }))
1831 .build();
1832 let params = opts.to_nlm_params();
1833
1834 assert!(matches!(
1835 params.motion_compensation,
1836 MotionCompensationMode::Mvtools {
1837 blksize: 16,
1838 overlap: 8,
1839 search_radius: 4,
1840 pyramid_levels: 2,
1841 ..
1842 }
1843 ));
1844 }
1845
1846 #[test]
1847 fn motion_compensation_defaults_to_none() {
1848 let opts = DenoiserOptions::builder().build();
1849 let params = opts.to_nlm_params();
1850 assert!(matches!(params.motion_compensation, MotionCompensationMode::None));
1851 }
1852
1853 #[test]
1854 fn nlm_tuning_overrides_individual_fields() {
1855 let defaults = NlmParams::default();
1856 let opts = DenoiserOptions::builder()
1857 .algorithm(fast_tuned(NlmTuning {
1858 search_radius: Some(7),
1859 patch_radius: None,
1860 strength: Some(2.5),
1861 self_weight: None,
1862 }))
1863 .build();
1864 let params = opts.to_nlm_params();
1865
1866 assert_eq!(params.search_radius, 7);
1867 assert_eq!(params.patch_radius, defaults.patch_radius);
1868 assert!((params.strength - 2.5).abs() < f32::EPSILON);
1869 assert!((params.self_weight - defaults.self_weight).abs() < f32::EPSILON);
1870 }
1871}
1872
1873#[cfg(all(test, feature = "vulkan"))]
1874mod tests {
1875 use super::*;
1876
1877 fn opts(mode: DenoisingMode) -> DenoiserOptions {
1878 DenoiserOptions::builder()
1879 .channel_mode(ChannelMode::Luma)
1880 .mode(mode)
1881 .build()
1882 }
1883
1884 fn frame(w: u32, h: u32) -> Vec<f32> {
1885 vec![0.5f32; (w * h) as usize]
1886 }
1887
1888 fn f32_out(out: FrameOutput) -> Vec<f32> {
1889 out.into_f32().expect("f32 output")
1890 }
1891
1892 #[test]
1893 fn spatial_denoise_roundtrip() {
1894 let mut d = Denoiser::create(
1895 &[Accelerator::Vulkan],
1896 &Device::Default,
1897 16,
1898 16,
1899 opts(DenoisingMode::Spacial),
1900 )
1901 .expect("denoiser construction failed");
1902 assert_eq!(d.selected_accelerator(), Accelerator::Vulkan);
1903
1904 d.push_frame(&frame(16, 16)).expect("push failed");
1905 let out = f32_out(d.recv_frame().expect("recv failed").expect("no frame"));
1906 assert_eq!(out.len(), 16 * 16);
1907 }
1908
1909 #[test]
1910 fn nl4d_algorithm_round_trips_through_the_facade() {
1911 let opts = DenoiserOptions::builder()
1912 .channel_mode(ChannelMode::Luma)
1913 .mode(DenoisingMode::Temporal { radius: 2 })
1914 .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1915 .build();
1916 let mut d = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, opts)
1917 .expect("nl4d denoiser construction failed");
1918 assert_eq!(d.selected_accelerator(), Accelerator::Vulkan);
1919
1920 d.push_frame(&frame(16, 16)).expect("push failed");
1924 assert!(d.recv_frame().expect("recv failed").is_none());
1925
1926 let mut out = Vec::new();
1927 d.flush(|f| out.push(f32_out(f))).expect("flush failed");
1928 assert_eq!(out.len(), 1, "expected exactly one output for one pushed frame");
1929 assert_eq!(out[0].len(), 16 * 16);
1930 }
1931
1932 #[test]
1936 fn nl4d_rejects_a_spatial_denoising_mode() {
1937 let opts = DenoiserOptions::builder()
1938 .channel_mode(ChannelMode::Luma)
1939 .mode(DenoisingMode::Spacial)
1940 .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1941 .build();
1942 let result = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, opts);
1943
1944 match result {
1945 Err(DenoiserError::Other(e)) => assert!(
1946 e.to_string().contains("temporal window"),
1947 "unexpected error message: {e}"
1948 ),
1949 Err(other) => panic!("expected DenoiserError::Other, got {other:?}"),
1950 Ok(_) => panic!("expected a rejection, got Ok"),
1951 }
1952 }
1953
1954 #[test]
1957 fn window_span_is_symmetric_for_nlmeans() {
1958 let opts = DenoiserOptions::builder()
1959 .channel_mode(ChannelMode::Luma)
1960 .mode(DenoisingMode::Temporal { radius: 3 })
1961 .algorithm(Algorithm::Nlmeans(NlmeansOptions::default()))
1962 .build();
1963 let d = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, opts)
1964 .expect("denoiser construction failed");
1965
1966 let span = d.window_span();
1967 assert_eq!(span.behind, 3, "behind should equal the temporal radius");
1968 assert_eq!(span.ahead, 3, "ahead should equal the temporal radius");
1969 }
1970
1971 #[test]
1975 fn window_span_is_doubled_on_both_sides_for_nl4d() {
1976 let opts = DenoiserOptions::builder()
1977 .channel_mode(ChannelMode::Luma)
1978 .mode(DenoisingMode::Temporal { radius: 3 })
1979 .algorithm(Algorithm::Nl4d(Nl4dOptions::default()))
1980 .build();
1981 let d = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, opts)
1982 .expect("nl4d denoiser construction failed");
1983
1984 let span = d.window_span();
1985 assert_eq!(span.behind, 6, "behind should equal 2 * the temporal radius");
1986 assert_eq!(span.ahead, 6, "ahead should equal 2 * the temporal radius");
1987 }
1988
1989 #[test]
1990 fn invalid_params_surface_as_error() {
1991 let bad = DenoiserOptions::builder()
1992 .algorithm(Algorithm::Nlmeans(NlmeansOptions {
1993 tuning: NlmTuning {
1994 strength: Some(0.0),
1995 ..NlmTuning::default()
1996 },
1997 ..NlmeansOptions::default()
1998 }))
1999 .build();
2000 let result = Denoiser::create(&[Accelerator::Vulkan], &Device::Default, 16, 16, bad);
2001
2002 match result {
2003 Err(DenoiserError::Other(_)) => {},
2004 Err(other) => panic!("expected DenoiserError::Other, got {other:?}"),
2005 Ok(_) => panic!("expected validation error, got Ok"),
2006 }
2007 }
2008
2009 #[test]
2010 fn tiny_frame_dimensions_surface_as_error() {
2011 let result = Denoiser::create(
2012 &[Accelerator::Vulkan],
2013 &Device::Default,
2014 2,
2015 2,
2016 opts(DenoisingMode::Spacial),
2017 );
2018
2019 match result {
2020 Err(DenoiserError::Other(e)) => {
2021 assert!(
2022 e.to_string().contains("supported minimum"),
2023 "unexpected error message: {e}"
2024 );
2025 },
2026 Err(other) => panic!("expected DenoiserError::Other, got {other:?}"),
2027 Ok(_) => panic!("expected dimension validation error, got Ok"),
2028 }
2029 }
2030
2031 #[test]
2032 fn push_after_pending_returns_queue_full() {
2033 let mut d = Denoiser::create(
2034 &[Accelerator::Vulkan],
2035 &Device::Default,
2036 16,
2037 16,
2038 opts(DenoisingMode::Spacial),
2039 )
2040 .unwrap();
2041
2042 d.push_frame(&frame(16, 16)).unwrap();
2047 d.push_frame(&frame(16, 16)).unwrap();
2048 let err = d.push_frame(&frame(16, 16)).expect_err("expected QueueFull");
2049 assert!(matches!(err, DenoiserError::QueueFull));
2050
2051 let out = f32_out(d.recv_frame().unwrap().unwrap());
2052 assert_eq!(out.len(), 16 * 16);
2053
2054 d.push_frame(&frame(16, 16)).expect("push after drain failed");
2056 }
2057
2058 #[test]
2061 fn queue_full_does_not_poison() {
2062 let mut d = Denoiser::create(
2063 &[Accelerator::Vulkan],
2064 &Device::Default,
2065 16,
2066 16,
2067 opts(DenoisingMode::Spacial),
2068 )
2069 .unwrap();
2070
2071 d.push_frame(&frame(16, 16)).unwrap();
2072 d.push_frame(&frame(16, 16)).unwrap();
2073 let err = d.push_frame(&frame(16, 16)).expect_err("expected QueueFull");
2074 assert!(matches!(err, DenoiserError::QueueFull));
2075 assert!(!d.poisoned, "QueueFull must not poison the denoiser");
2076
2077 d.recv_frame().unwrap().expect("recv failed after QueueFull");
2078
2079 d.push_frame(&frame(16, 16))
2082 .expect("push after QueueFull drain should succeed, not poison");
2083 }
2084
2085 #[test]
2086 fn poisoned_denoiser_refuses_every_entry_point() {
2087 let mut d = Denoiser::create(
2088 &[Accelerator::Vulkan],
2089 &Device::Default,
2090 16,
2091 16,
2092 opts(DenoisingMode::Spacial),
2093 )
2094 .unwrap();
2095 d.poisoned = true;
2096
2097 assert!(matches!(
2098 d.push_frame(&frame(16, 16)),
2099 Err(DenoiserError::Poisoned)
2100 ));
2101 assert!(matches!(d.recv_frame(), Err(DenoiserError::Poisoned)));
2102 assert!(matches!(d.try_recv_frame(), Err(DenoiserError::Poisoned)));
2103 assert!(matches!(d.flush(|_| {}), Err(DenoiserError::Poisoned)));
2104 }
2105
2106 #[test]
2107 fn reset_stream_clears_poison() {
2108 let mut d = Denoiser::create(
2109 &[Accelerator::Vulkan],
2110 &Device::Default,
2111 16,
2112 16,
2113 opts(DenoisingMode::Spacial),
2114 )
2115 .unwrap();
2116 d.poisoned = true;
2117
2118 d.reset_stream();
2119 assert!(!d.poisoned, "reset_stream must clear the poison flag");
2120
2121 d.push_frame(&frame(16, 16))
2122 .expect("push after reset_stream should succeed");
2123 }
2124
2125 fn frame_filled(w: u32, h: u32, value: f32) -> Vec<f32> {
2126 vec![value; (w * h) as usize]
2127 }
2128
2129 fn push_n_with_drain(d: &mut Denoiser, n: usize, value: f32, out: &mut Vec<Vec<f32>>) {
2132 for _ in 0..n {
2133 loop {
2134 match d.push_frame(&frame_filled(16, 16, value)) {
2135 Ok(()) => break,
2136 Err(DenoiserError::QueueFull) => {
2137 let f = d
2138 .recv_frame()
2139 .expect("recv ok")
2140 .expect("queue full but recv yielded none");
2141 out.push(f32_out(f));
2142 },
2143 Err(e) => panic!("unexpected push error: {e:?}"),
2144 }
2145 }
2146 }
2147 }
2148
2149 #[test]
2150 fn flush_leaves_denoiser_reusable_spatial() {
2151 let mut d = Denoiser::create(
2152 &[Accelerator::Vulkan],
2153 &Device::Default,
2154 16,
2155 16,
2156 opts(DenoisingMode::Spacial),
2157 )
2158 .unwrap();
2159
2160 let mut batch_a = Vec::new();
2161 push_n_with_drain(&mut d, 5, 0.25, &mut batch_a);
2162 d.flush(|f| batch_a.push(f32_out(f))).expect("first flush failed");
2163 assert_eq!(batch_a.len(), 5);
2164
2165 assert!(d.recv_frame().unwrap().is_none());
2167
2168 let mut batch_b = Vec::new();
2169 push_n_with_drain(&mut d, 5, 0.75, &mut batch_b);
2170 d.flush(|f| batch_b.push(f32_out(f)))
2171 .expect("second flush failed");
2172 assert_eq!(batch_b.len(), 5);
2173
2174 for v in batch_b.iter().flatten() {
2175 assert!((v - 0.75).abs() < 0.1, "batch_b carried state from batch_a: {v}");
2176 }
2177 for v in batch_a.iter().flatten() {
2178 assert!((v - 0.25).abs() < 0.1, "batch_a value unexpectedly drifted: {v}");
2179 }
2180 }
2181
2182 #[test]
2183 fn flush_leaves_denoiser_reusable_temporal() {
2184 let mut d = Denoiser::create(
2185 &[Accelerator::Vulkan],
2186 &Device::Default,
2187 16,
2188 16,
2189 opts(DenoisingMode::Temporal { radius: 1 }),
2190 )
2191 .unwrap();
2192
2193 let mut batch_a = Vec::new();
2194 push_n_with_drain(&mut d, 5, 0.25, &mut batch_a);
2195 d.flush(|f| batch_a.push(f32_out(f))).expect("first flush failed");
2196 assert_eq!(batch_a.len(), 5, "expected 5 frames from first batch");
2197
2198 assert!(d.recv_frame().unwrap().is_none());
2203 d.push_frame(&frame_filled(16, 16, 0.75)).unwrap();
2204 assert!(
2205 d.recv_frame().unwrap().is_none(),
2206 "first push of new temporal stream should not produce output yet"
2207 );
2208
2209 let mut batch_b = Vec::new();
2211 push_n_with_drain(&mut d, 4, 0.75, &mut batch_b);
2212 d.flush(|f| batch_b.push(f32_out(f)))
2213 .expect("second flush failed");
2214 assert_eq!(batch_b.len(), 5, "expected 5 frames from second batch");
2215
2216 for v in batch_b.iter().flatten() {
2217 assert!((v - 0.75).abs() < 0.1, "batch_b carried state from batch_a: {v}");
2218 }
2219 }
2220
2221 #[test]
2222 fn flush_emits_exactly_n_outputs_for_small_n() {
2223 for n in 1..=5usize {
2228 let mut d = Denoiser::create(
2229 &[Accelerator::Vulkan],
2230 &Device::Default,
2231 16,
2232 16,
2233 opts(DenoisingMode::Temporal { radius: 2 }),
2234 )
2235 .unwrap();
2236
2237 let mut out = Vec::new();
2238 push_n_with_drain(&mut d, n, 0.5, &mut out);
2239 d.flush(|f| out.push(f32_out(f))).expect("flush failed");
2240 assert_eq!(
2241 out.len(),
2242 n,
2243 "expected {n} outputs for {n} pushes, got {}",
2244 out.len()
2245 );
2246 }
2247 }
2248}