From 0566bf11b375ddd220289f1e044d2dc213fd46bb Mon Sep 17 00:00:00 2001
From: Lukas Stockner <lukas.stockner@freenet.de>
Date: Tue, 4 Jun 2024 01:22:44 +0200
Subject: [PATCH] Cycles: Add debugging option to load custom OIDN weights

Pull Request: https://projects.blender.org/blender/blender/pulls/122591
---
 intern/cycles/integrator/denoiser_oidn.cpp     | 13 +++++++++++++
 intern/cycles/integrator/denoiser_oidn_gpu.cpp | 12 ++++++++++++
 intern/cycles/integrator/denoiser_oidn_gpu.h   |  2 ++
 3 files changed, 27 insertions(+)

diff --git a/intern/cycles/integrator/denoiser_oidn.cpp b/intern/cycles/integrator/denoiser_oidn.cpp
index 7fc2764e464..13f1e1c6b4d 100644
--- a/intern/cycles/integrator/denoiser_oidn.cpp
+++ b/intern/cycles/integrator/denoiser_oidn.cpp
@@ -13,6 +13,7 @@
 #include "util/array.h"
 #include "util/log.h"
 #include "util/openimagedenoise.h"
+#include "util/path.h"
 
 #include "kernel/device/cpu/compat.h"
 #include "kernel/device/cpu/kernel.h"
@@ -121,6 +122,13 @@ class OIDNDenoiseContext {
     if (denoise_params_.use_pass_normal) {
       oidn_normal_pass_ = OIDNPass(buffer_params_, "normal", PASS_DENOISING_NORMAL);
     }
+
+    const char *custom_weight_path = getenv("CYCLES_OIDN_CUSTOM_WEIGHTS");
+    if (custom_weight_path) {
+      if (!path_read_binary(custom_weight_path, custom_weights)) {
+        fprintf(stderr, "Cycles: Failed to load custom OIDN weights!");
+      }
+    }
   }
 
   bool need_denoising() const
@@ -174,6 +182,9 @@ class OIDNDenoiseContext {
     oidn_filter.setProgressMonitorFunction(oidn_progress_monitor_function, denoiser_);
     oidn_filter.set("hdr", true);
     oidn_filter.set("srgb", false);
+    if (custom_weights.size()) {
+      oidn_filter.setData("weights", custom_weights.data(), custom_weights.size());
+    }
     set_quality(oidn_filter);
 
     if (denoise_params_.prefilter == DENOISER_PREFILTER_NONE ||
@@ -556,6 +567,8 @@ class OIDNDenoiseContext {
   bool allow_inplace_modification_ = false;
   int pass_sample_count_ = PASS_UNUSED;
 
+  vector<uint8_t> custom_weights;
+
   /* Optional albedo and normal passes, reused by denoising of different pass types. */
   OIDNPass oidn_albedo_pass_;
   OIDNPass oidn_normal_pass_;
diff --git a/intern/cycles/integrator/denoiser_oidn_gpu.cpp b/intern/cycles/integrator/denoiser_oidn_gpu.cpp
index 262ce465205..4b756913363 100644
--- a/intern/cycles/integrator/denoiser_oidn_gpu.cpp
+++ b/intern/cycles/integrator/denoiser_oidn_gpu.cpp
@@ -15,6 +15,7 @@
 #  include "session/buffers.h"
 #  include "util/array.h"
 #  include "util/log.h"
+#  include "util/path.h"
 
 #  include "kernel/device/cpu/compat.h"
 #  include "kernel/device/cpu/kernel.h"
@@ -332,6 +333,17 @@ bool OIDNDenoiserGPU::denoise_create_if_needed(DenoiseContext &context)
   oidnSetFilterBool(oidn_filter_, "hdr", true);
   oidnSetFilterBool(oidn_filter_, "srgb", false);
 
+  const char *custom_weight_path = getenv("CYCLES_OIDN_CUSTOM_WEIGHTS");
+  if (custom_weight_path) {
+    if (path_read_binary(custom_weight_path, custom_weights)) {
+      oidnSetSharedFilterData(
+          oidn_filter_, "weights", custom_weights.data(), custom_weights.size());
+    }
+    else {
+      fprintf(stderr, "Cycles: Failed to load custom OIDN weights!");
+    }
+  }
+
   if (context.use_pass_albedo) {
     albedo_filter_ = create_filter();
     if (albedo_filter_ == nullptr) {
diff --git a/intern/cycles/integrator/denoiser_oidn_gpu.h b/intern/cycles/integrator/denoiser_oidn_gpu.h
index c3ae9f3affe..abcd9d56952 100644
--- a/intern/cycles/integrator/denoiser_oidn_gpu.h
+++ b/intern/cycles/integrator/denoiser_oidn_gpu.h
@@ -75,6 +75,8 @@ class OIDNDenoiserGPU : public DenoiserGPU {
   bool is_configured_ = false;
   int2 configured_size_ = make_int2(0, 0);
 
+  vector<uint8_t> custom_weights;
+
   bool use_pass_albedo_ = false;
   bool use_pass_normal_ = false;
   DenoiserQuality quality_ = DENOISER_QUALITY_HIGH;