diff options
| -rw-r--r-- | CMakeLists.txt | 2 | ||||
| -rw-r--r-- | src/gpu/demo_effects.h | 2 | ||||
| -rw-r--r-- | src/gpu/effects/cnn_effect.cc | 36 | ||||
| -rw-r--r-- | src/gpu/effects/cnn_effect.h | 29 | ||||
| -rw-r--r-- | src/gpu/effects/shaders.cc | 11 | ||||
| -rw-r--r-- | src/gpu/effects/shaders.h | 1 | ||||
| -rw-r--r-- | src/tests/gpu/test_demo_effects.cc | 2 | ||||
| -rw-r--r-- | workspaces/main/assets.txt | 6 | ||||
| -rw-r--r-- | workspaces/main/shaders/cnn/cnn_activation.wgsl | 18 | ||||
| -rw-r--r-- | workspaces/main/shaders/cnn/cnn_conv3x3.wgsl | 26 | ||||
| -rw-r--r-- | workspaces/main/shaders/cnn/cnn_conv5x5.wgsl | 26 | ||||
| -rw-r--r-- | workspaces/main/shaders/cnn/cnn_conv7x7.wgsl | 26 | ||||
| -rw-r--r-- | workspaces/main/shaders/cnn/cnn_layer.wgsl | 46 | ||||
| -rw-r--r-- | workspaces/main/shaders/cnn/cnn_weights_generated.wgsl | 17 |
14 files changed, 248 insertions, 0 deletions
diff --git a/CMakeLists.txt b/CMakeLists.txt index ee3dfc3..97b371a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -151,6 +151,7 @@ if (DEMO_HEADLESS) src/gpu/effects/solarize_effect.cc src/gpu/effects/chroma_aberration_effect.cc src/gpu/effects/vignette_effect.cc + src/gpu/effects/cnn_effect.cc src/gpu/effects/post_process_helper.cc src/gpu/effects/shaders.cc src/gpu/effects/hybrid_3d_effect.cc @@ -180,6 +181,7 @@ else() src/gpu/effects/solarize_effect.cc src/gpu/effects/chroma_aberration_effect.cc src/gpu/effects/vignette_effect.cc + src/gpu/effects/cnn_effect.cc src/gpu/effects/post_process_helper.cc src/gpu/effects/shaders.cc src/gpu/effects/hybrid_3d_effect.cc diff --git a/src/gpu/demo_effects.h b/src/gpu/demo_effects.h index ff7e017..6e6cb73 100644 --- a/src/gpu/demo_effects.h +++ b/src/gpu/demo_effects.h @@ -184,6 +184,8 @@ class DistortEffect : public PostProcessEffect { // (included above) FlashEffect now defined in gpu/effects/flash_effect.h // (included above) +class CNNEffect; + // Auto-generated functions void LoadTimeline(MainSequence& main_seq, const GpuContext& ctx); float GetDemoDuration(); // Returns demo end time in seconds, or -1 if not diff --git a/src/gpu/effects/cnn_effect.cc b/src/gpu/effects/cnn_effect.cc new file mode 100644 index 0000000..25db0c2 --- /dev/null +++ b/src/gpu/effects/cnn_effect.cc @@ -0,0 +1,36 @@ +// CNN post-processing effect implementation +// Neural network-based stylization with modular WGSL + +#include "gpu/effects/cnn_effect.h" +#include "gpu/effects/post_process_helper.h" +#include "gpu/effects/shaders.h" + +CNNEffect::CNNEffect(const GpuContext& ctx, int num_layers) + : PostProcessEffect(ctx), num_layers_(num_layers), input_view_(nullptr), + bind_group_(nullptr) { + pipeline_ = create_post_process_pipeline(ctx_.device, ctx_.format, + cnn_layer_shader_wgsl); +} + +void CNNEffect::init(MainSequence* demo) { + PostProcessEffect::init(demo); + params_buffer_.init(ctx_.device); + + CNNLayerParams params = {0, 1, {0.0f, 0.0f}}; + params_buffer_.update(ctx_.queue, params); +} + +void CNNEffect::render(WGPURenderPassEncoder pass, float time, float beat, + float intensity, float aspect_ratio) { + if (!bind_group_) return; + + wgpuRenderPassEncoderSetPipeline(pass, pipeline_); + wgpuRenderPassEncoderSetBindGroup(pass, 0, bind_group_, 0, nullptr); + wgpuRenderPassEncoderDraw(pass, 3, 1, 0, 0); +} + +void CNNEffect::update_bind_group(WGPUTextureView input_view) { + input_view_ = input_view; + pp_update_bind_group(ctx_.device, pipeline_, &bind_group_, + input_view_, uniforms_.get(), params_buffer_.get()); +} diff --git a/src/gpu/effects/cnn_effect.h b/src/gpu/effects/cnn_effect.h new file mode 100644 index 0000000..9cc4935 --- /dev/null +++ b/src/gpu/effects/cnn_effect.h @@ -0,0 +1,29 @@ +// CNN post-processing effect header +// Multi-layer neural network stylization + +#pragma once +#include "gpu/effect.h" +#include "gpu/uniform_helper.h" + +struct CNNLayerParams { + int layer_index; + int use_residual; + float _pad[2]; +}; +static_assert(sizeof(CNNLayerParams) == 16); + +class CNNEffect : public PostProcessEffect { + public: + explicit CNNEffect(const GpuContext& ctx, int num_layers = 1); + + void init(MainSequence* demo) override; + void render(WGPURenderPassEncoder pass, float time, float beat, + float intensity, float aspect_ratio) override; + void update_bind_group(WGPUTextureView input_view) override; + + private: + int num_layers_; + WGPUTextureView input_view_; + UniformBuffer<CNNLayerParams> params_buffer_; + WGPUBindGroup bind_group_; +}; diff --git a/src/gpu/effects/shaders.cc b/src/gpu/effects/shaders.cc index 625c5b6..6559bf5 100644 --- a/src/gpu/effects/shaders.cc +++ b/src/gpu/effects/shaders.cc @@ -51,6 +51,13 @@ void InitShaderComposer() { register_if_exists("ray_box", AssetId::ASSET_SHADER_RAY_BOX); register_if_exists("ray_triangle", AssetId::ASSET_SHADER_RAY_TRIANGLE); + + register_if_exists("cnn_activation", AssetId::ASSET_SHADER_CNN_ACTIVATION); + register_if_exists("cnn_conv3x3", AssetId::ASSET_SHADER_CNN_CONV3X3); + register_if_exists("cnn_conv5x5", AssetId::ASSET_SHADER_CNN_CONV5X5); + register_if_exists("cnn_conv7x7", AssetId::ASSET_SHADER_CNN_CONV7X7); + register_if_exists("cnn_weights_generated", + AssetId::ASSET_SHADER_CNN_WEIGHTS); } // Helper to get asset string or empty string @@ -99,6 +106,10 @@ const char* chroma_aberration_shader_wgsl = SafeGetAsset(AssetId::ASSET_SHADER_CHROMA_ABERRATION); +const char* cnn_layer_shader_wgsl = + + SafeGetAsset(AssetId::ASSET_SHADER_CNN_LAYER); + const char* gen_noise_compute_wgsl = SafeGetAsset(AssetId::ASSET_SHADER_COMPUTE_GEN_NOISE); diff --git a/src/gpu/effects/shaders.h b/src/gpu/effects/shaders.h index 68b8834..7acc2a6 100644 --- a/src/gpu/effects/shaders.h +++ b/src/gpu/effects/shaders.h @@ -18,6 +18,7 @@ extern const char* solarize_shader_wgsl; extern const char* distort_shader_wgsl; extern const char* chroma_aberration_shader_wgsl; extern const char* vignette_shader_wgsl; +extern const char* cnn_layer_shader_wgsl; extern const char* gen_noise_compute_wgsl; extern const char* gen_perlin_compute_wgsl; extern const char* gen_grid_compute_wgsl; diff --git a/src/tests/gpu/test_demo_effects.cc b/src/tests/gpu/test_demo_effects.cc index 8a7d8af..9281413 100644 --- a/src/tests/gpu/test_demo_effects.cc +++ b/src/tests/gpu/test_demo_effects.cc @@ -13,6 +13,7 @@ #include "../common/effect_test_helpers.h" #include "gpu/demo_effects.h" #include "gpu/effect.h" +#include "gpu/effects/cnn_effect.h" #include "../common/webgpu_test_fixture.h" #include <cassert> #include <cstdio> @@ -88,6 +89,7 @@ static void test_post_process_effects() { {"ThemeModulationEffect", std::make_shared<ThemeModulationEffect>(fixture.ctx())}, {"VignetteEffect", std::make_shared<VignetteEffect>(fixture.ctx())}, + {"CNNEffect", std::make_shared<CNNEffect>(fixture.ctx(), 1)}, }; int passed = 0; diff --git a/workspaces/main/assets.txt b/workspaces/main/assets.txt index ca77e21..53c8b3e 100644 --- a/workspaces/main/assets.txt +++ b/workspaces/main/assets.txt @@ -36,6 +36,12 @@ SHADER_PASSTHROUGH, NONE, shaders/passthrough.wgsl, "Passthrough Shader" SHADER_ELLIPSE, NONE, shaders/ellipse.wgsl, "Ellipse Shader" SHADER_PARTICLE_SPRAY_COMPUTE, NONE, shaders/particle_spray_compute.wgsl, "Particle Spray Compute" SHADER_GAUSSIAN_BLUR, NONE, shaders/gaussian_blur.wgsl, "Gaussian Blur Shader" +SHADER_CNN_ACTIVATION, NONE, shaders/cnn/cnn_activation.wgsl, "CNN Activation Functions" +SHADER_CNN_CONV3X3, NONE, shaders/cnn/cnn_conv3x3.wgsl, "CNN 3x3 Convolution" +SHADER_CNN_CONV5X5, NONE, shaders/cnn/cnn_conv5x5.wgsl, "CNN 5x5 Convolution" +SHADER_CNN_CONV7X7, NONE, shaders/cnn/cnn_conv7x7.wgsl, "CNN 7x7 Convolution" +SHADER_CNN_WEIGHTS, NONE, shaders/cnn/cnn_weights_generated.wgsl, "CNN Weights (Generated)" +SHADER_CNN_LAYER, NONE, shaders/cnn/cnn_layer.wgsl, "CNN Layer Shader" SHADER_SOLARIZE, NONE, shaders/solarize.wgsl, "Solarize Shader" SHADER_DISTORT, NONE, shaders/distort.wgsl, "Distort Shader" SHADER_CHROMA_ABERRATION, NONE, shaders/chroma_aberration.wgsl, "Chroma Aberration Shader" diff --git a/workspaces/main/shaders/cnn/cnn_activation.wgsl b/workspaces/main/shaders/cnn/cnn_activation.wgsl new file mode 100644 index 0000000..4fe771e --- /dev/null +++ b/workspaces/main/shaders/cnn/cnn_activation.wgsl @@ -0,0 +1,18 @@ +// CNN activation functions +// 4 functions: tanh, ReLU, sigmoid, leaky_relu + +fn cnn_tanh(x: vec4<f32>) -> vec4<f32> { + return tanh(x); +} + +fn cnn_relu(x: vec4<f32>) -> vec4<f32> { + return max(vec4<f32>(0.0), x); +} + +fn cnn_sigmoid(x: vec4<f32>) -> vec4<f32> { + return 1.0 / (1.0 + exp(-x)); +} + +fn cnn_leaky_relu(x: vec4<f32>, alpha: f32) -> vec4<f32> { + return max(alpha * x, x); +} diff --git a/workspaces/main/shaders/cnn/cnn_conv3x3.wgsl b/workspaces/main/shaders/cnn/cnn_conv3x3.wgsl new file mode 100644 index 0000000..06ca73a --- /dev/null +++ b/workspaces/main/shaders/cnn/cnn_conv3x3.wgsl @@ -0,0 +1,26 @@ +// 3x3 convolution with weight indexing +// Samples 9 pixels, applies mat4 weights per sample + +fn cnn_conv3x3( + tex: texture_2d<f32>, + samp: sampler, + uv: vec2<f32>, + resolution: vec2<f32>, + weights: array<mat4x4<f32>, 9>, + bias: vec4<f32> +) -> vec4<f32> { + let step = 1.0 / resolution; + var sum = bias; + var idx = 0; + + for (var dy = -1; dy <= 1; dy++) { + for (var dx = -1; dx <= 1; dx++) { + let offset = vec2<f32>(f32(dx), f32(dy)) * step; + let sample = textureSample(tex, samp, uv + offset); + sum += weights[idx] * sample; + idx++; + } + } + + return sum; +} diff --git a/workspaces/main/shaders/cnn/cnn_conv5x5.wgsl b/workspaces/main/shaders/cnn/cnn_conv5x5.wgsl new file mode 100644 index 0000000..3d4a03a --- /dev/null +++ b/workspaces/main/shaders/cnn/cnn_conv5x5.wgsl @@ -0,0 +1,26 @@ +// 5x5 convolution with 25 samples +// Applies mat4 weights per sample + +fn cnn_conv5x5( + tex: texture_2d<f32>, + samp: sampler, + uv: vec2<f32>, + resolution: vec2<f32>, + weights: array<mat4x4<f32>, 25>, + bias: vec4<f32> +) -> vec4<f32> { + let step = 1.0 / resolution; + var sum = bias; + var idx = 0; + + for (var dy = -2; dy <= 2; dy++) { + for (var dx = -2; dx <= 2; dx++) { + let offset = vec2<f32>(f32(dx), f32(dy)) * step; + let sample = textureSample(tex, samp, uv + offset); + sum += weights[idx] * sample; + idx++; + } + } + + return sum; +} diff --git a/workspaces/main/shaders/cnn/cnn_conv7x7.wgsl b/workspaces/main/shaders/cnn/cnn_conv7x7.wgsl new file mode 100644 index 0000000..ba28d64 --- /dev/null +++ b/workspaces/main/shaders/cnn/cnn_conv7x7.wgsl @@ -0,0 +1,26 @@ +// 7x7 convolution with 49 samples +// Applies mat4 weights per sample + +fn cnn_conv7x7( + tex: texture_2d<f32>, + samp: sampler, + uv: vec2<f32>, + resolution: vec2<f32>, + weights: array<mat4x4<f32>, 49>, + bias: vec4<f32> +) -> vec4<f32> { + let step = 1.0 / resolution; + var sum = bias; + var idx = 0; + + for (var dy = -3; dy <= 3; dy++) { + for (var dx = -3; dx <= 3; dx++) { + let offset = vec2<f32>(f32(dx), f32(dy)) * step; + let sample = textureSample(tex, samp, uv + offset); + sum += weights[idx] * sample; + idx++; + } + } + + return sum; +} diff --git a/workspaces/main/shaders/cnn/cnn_layer.wgsl b/workspaces/main/shaders/cnn/cnn_layer.wgsl new file mode 100644 index 0000000..e026ce8 --- /dev/null +++ b/workspaces/main/shaders/cnn/cnn_layer.wgsl @@ -0,0 +1,46 @@ +// CNN layer shader - uses modular convolution snippets +// Supports multi-pass rendering with residual connections + +@group(0) @binding(0) var smplr: sampler; +@group(0) @binding(1) var txt: texture_2d<f32>; + +#include "common_uniforms" +#include "cnn_activation" +#include "cnn_conv3x3" +#include "cnn_weights_generated" + +struct CNNLayerParams { + layer_index: i32, + use_residual: i32, + _pad: vec2<f32>, +}; + +@group(0) @binding(2) var<uniform> uniforms: CommonUniforms; +@group(0) @binding(3) var<uniform> params: CNNLayerParams; + +@vertex fn vs_main(@builtin(vertex_index) i: u32) -> @builtin(position) vec4<f32> { + var pos = array<vec2<f32>, 3>( + vec2<f32>(-1.0, -1.0), vec2<f32>(3.0, -1.0), vec2<f32>(-1.0, 3.0) + ); + return vec4<f32>(pos[i], 0.0, 1.0); +} + +@fragment fn fs_main(@builtin(position) p: vec4<f32>) -> @location(0) vec4<f32> { + let uv = p.xy / uniforms.resolution; + var result = vec4<f32>(0.0); + + // Single layer for now (layer 0) + if (params.layer_index == 0) { + result = cnn_conv3x3(txt, smplr, uv, uniforms.resolution, + weights_layer0, bias_layer0); + result = cnn_tanh(result); + } + + // Residual connection + if (params.use_residual != 0) { + let input = textureSample(txt, smplr, uv); + result = input + result * 0.3; + } + + return result; +} diff --git a/workspaces/main/shaders/cnn/cnn_weights_generated.wgsl b/workspaces/main/shaders/cnn/cnn_weights_generated.wgsl new file mode 100644 index 0000000..98c17ff --- /dev/null +++ b/workspaces/main/shaders/cnn/cnn_weights_generated.wgsl @@ -0,0 +1,17 @@ +// Generated CNN weights and biases +// DO NOT EDIT MANUALLY - regenerate with scripts/train_cnn.py + +// Placeholder identity-like weights for initial testing +// Layer 0: 3x3 convolution +const weights_layer0: array<mat4x4<f32>, 9> = array( + mat4x4<f32>(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + mat4x4<f32>(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + mat4x4<f32>(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + mat4x4<f32>(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + mat4x4<f32>(1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0), + mat4x4<f32>(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + mat4x4<f32>(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + mat4x4<f32>(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), + mat4x4<f32>(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0) +); +const bias_layer0 = vec4<f32>(0.0, 0.0, 0.0, 0.0); |
