Compositor: Implement Denoise for new CPU compositor

Reference #125968.
This commit is contained in:
Omar Emara
2024-11-20 15:55:33 +02:00
parent 5e806db594
commit 11a5e362d4

View File

@@ -19,6 +19,7 @@
#include "DNA_node_types.h"
#include "COM_node_operation.hh"
#include "COM_utilities.hh"
#include "node_composite_util.hh"
@@ -93,17 +94,6 @@ class DenoiseOperation : public NodeOperation {
void execute() override
{
/* Not yet supported on CPU. */
if (!context().use_gpu()) {
for (const bNodeSocket *output : this->node()->output_sockets()) {
Result &output_result = get_result(output->identifier);
if (output_result.should_compute()) {
output_result.allocate_invalid();
}
}
return;
}
Result &input_image = get_input("Image");
Result &output_image = get_result("Image");
@@ -112,6 +102,8 @@ class DenoiseOperation : public NodeOperation {
return;
}
output_image.allocate_texture(input_image.domain());
#ifdef WITH_OPENIMAGEDENOISE
oidn::DeviceRef device = oidn::newDevice(oidn::DeviceType::CPU);
device.commit();
@@ -121,13 +113,22 @@ class DenoiseOperation : public NodeOperation {
const int pixel_stride = sizeof(float) * 4;
const eGPUDataFormat data_format = GPU_DATA_FLOAT;
/* Download the input texture and set it as both the input and output of the filter to denoise
* it in-place. */
GPU_memory_barrier(GPU_BARRIER_TEXTURE_UPDATE);
float *color = static_cast<float *>(GPU_texture_read(input_image, data_format, 0));
float *input_color = nullptr;
float *output_color = nullptr;
if (this->context().use_gpu()) {
/* Download the input texture and set it as both the input and output of the filter to
* denoise it in-place. */
GPU_memory_barrier(GPU_BARRIER_TEXTURE_UPDATE);
input_color = static_cast<float *>(GPU_texture_read(input_image, data_format, 0));
output_color = input_color;
}
else {
input_color = input_image.float_texture();
output_color = output_image.float_texture();
}
oidn::FilterRef filter = device.newFilter("RT");
filter.setImage("color", color, oidn::Format::Float3, width, height, 0, pixel_stride);
filter.setImage("output", color, oidn::Format::Float3, width, height, 0, pixel_stride);
filter.setImage("color", input_color, oidn::Format::Float3, width, height, 0, pixel_stride);
filter.setImage("output", output_color, oidn::Format::Float3, width, height, 0, pixel_stride);
filter.set("hdr", use_hdr());
filter.set("cleanAux", auxiliary_passes_are_clean());
filter.setProgressMonitorFunction(oidn_progress_monitor_function, &context());
@@ -137,7 +138,12 @@ class DenoiseOperation : public NodeOperation {
float *albedo = nullptr;
Result &input_albedo = get_input("Albedo");
if (!input_albedo.is_single_value()) {
albedo = static_cast<float *>(GPU_texture_read(input_albedo, data_format, 0));
if (this->context().use_gpu()) {
albedo = static_cast<float *>(GPU_texture_read(input_albedo, data_format, 0));
}
else {
albedo = input_albedo.float_texture();
}
if (should_denoise_auxiliary_passes()) {
oidn::FilterRef albedoFilter = device.newFilter("RT");
@@ -160,7 +166,12 @@ class DenoiseOperation : public NodeOperation {
float *normal = nullptr;
Result &input_normal = get_input("Normal");
if (albedo && !input_normal.is_single_value()) {
normal = static_cast<float *>(GPU_texture_read(input_normal, data_format, 0));
if (this->context().use_gpu()) {
normal = static_cast<float *>(GPU_texture_read(input_normal, data_format, 0));
}
else {
normal = input_normal.float_texture();
}
if (should_denoise_auxiliary_passes()) {
oidn::FilterRef normalFilter = device.newFilter("RT");
@@ -179,15 +190,28 @@ class DenoiseOperation : public NodeOperation {
filter.commit();
filter.execute();
output_image.allocate_texture(input_image.domain());
GPU_texture_update(output_image, data_format, color);
MEM_freeN(color);
if (albedo) {
MEM_freeN(albedo);
if (this->context().use_gpu()) {
GPU_texture_update(output_image, data_format, output_color);
}
if (normal) {
MEM_freeN(normal);
else {
/* OIDN already wrote to the output directly, however, OIDN skips the alpha channel, so we
* need to restore it. */
parallel_for(int2(width, height), [&](const int2 texel) {
const float alpha = input_image.load_pixel(texel).w;
output_image.store_pixel(texel, float4(output_image.load_pixel(texel).xyz(), alpha));
});
}
/* Buffers for the CPU case are owned by the inputs, while for GPU, they are temporally read
* from the GPU texture, so they need to be freed if they were read. */
if (this->context().use_gpu()) {
MEM_freeN(input_color);
if (albedo) {
MEM_freeN(albedo);
}
if (normal) {
MEM_freeN(normal);
}
}
#endif
}