From 96a349b9874c6cdaac525ba062a0f4f90c9bc3ed Mon Sep 17 00:00:00 2001 From: skal Date: Tue, 10 Feb 2026 10:27:44 +0100 Subject: feat: Add coordinate-aware CNN layer 0 for position-dependent stylization - Implement CoordConv2d custom layer accepting (x,y) patch center - Split layer 0 weights: rgba_weights (9x mat4x4) + coord_weights (mat2x4) - Add *_with_coord() functions to 3x3/5x5/7x7 convolution shaders - Update training script to generate coordinate grid and export split weights - Regenerate placeholder weights with new format Size impact: +32B coord weights + ~100B shader code = +132B total All 36 tests passing (100%) handoff(Claude): CNN coordinate awareness implemented, ready for training Co-Authored-By: Claude Sonnet 4.5 --- workspaces/main/shaders/cnn/cnn_conv3x3.wgsl | 27 ++++++++++++++++++++++ workspaces/main/shaders/cnn/cnn_conv5x5.wgsl | 27 ++++++++++++++++++++++ workspaces/main/shaders/cnn/cnn_conv7x7.wgsl | 27 ++++++++++++++++++++++ workspaces/main/shaders/cnn/cnn_layer.wgsl | 6 ++--- .../main/shaders/cnn/cnn_weights_generated.wgsl | 10 ++++++-- 5 files changed, 92 insertions(+), 5 deletions(-) (limited to 'workspaces') diff --git a/workspaces/main/shaders/cnn/cnn_conv3x3.wgsl b/workspaces/main/shaders/cnn/cnn_conv3x3.wgsl index 06ca73a..168c9e2 100644 --- a/workspaces/main/shaders/cnn/cnn_conv3x3.wgsl +++ b/workspaces/main/shaders/cnn/cnn_conv3x3.wgsl @@ -24,3 +24,30 @@ fn cnn_conv3x3( return sum; } + +fn cnn_conv3x3_with_coord( + tex: texture_2d, + samp: sampler, + uv: vec2, + resolution: vec2, + rgba_weights: array, 9>, + coord_weights: mat2x4, + bias: vec4 +) -> vec4 { + let step = 1.0 / resolution; + var sum = bias; + + sum += coord_weights * uv; + + var idx = 0; + for (var dy = -1; dy <= 1; dy++) { + for (var dx = -1; dx <= 1; dx++) { + let offset = vec2(f32(dx), f32(dy)) * step; + let rgba = textureSample(tex, samp, uv + offset); + sum += rgba_weights[idx] * rgba; + idx++; + } + } + + return sum; +} diff --git a/workspaces/main/shaders/cnn/cnn_conv5x5.wgsl b/workspaces/main/shaders/cnn/cnn_conv5x5.wgsl index 3d4a03a..bd9abfa 100644 --- a/workspaces/main/shaders/cnn/cnn_conv5x5.wgsl +++ b/workspaces/main/shaders/cnn/cnn_conv5x5.wgsl @@ -24,3 +24,30 @@ fn cnn_conv5x5( return sum; } + +fn cnn_conv5x5_with_coord( + tex: texture_2d, + samp: sampler, + uv: vec2, + resolution: vec2, + rgba_weights: array, 25>, + coord_weights: mat2x4, + bias: vec4 +) -> vec4 { + let step = 1.0 / resolution; + var sum = bias; + + sum += coord_weights * uv; + + var idx = 0; + for (var dy = -2; dy <= 2; dy++) { + for (var dx = -2; dx <= 2; dx++) { + let offset = vec2(f32(dx), f32(dy)) * step; + let rgba = textureSample(tex, samp, uv + offset); + sum += rgba_weights[idx] * rgba; + idx++; + } + } + + return sum; +} diff --git a/workspaces/main/shaders/cnn/cnn_conv7x7.wgsl b/workspaces/main/shaders/cnn/cnn_conv7x7.wgsl index ba28d64..e68d644 100644 --- a/workspaces/main/shaders/cnn/cnn_conv7x7.wgsl +++ b/workspaces/main/shaders/cnn/cnn_conv7x7.wgsl @@ -24,3 +24,30 @@ fn cnn_conv7x7( return sum; } + +fn cnn_conv7x7_with_coord( + tex: texture_2d, + samp: sampler, + uv: vec2, + resolution: vec2, + rgba_weights: array, 49>, + coord_weights: mat2x4, + bias: vec4 +) -> vec4 { + let step = 1.0 / resolution; + var sum = bias; + + sum += coord_weights * uv; + + var idx = 0; + for (var dy = -3; dy <= 3; dy++) { + for (var dx = -3; dx <= 3; dx++) { + let offset = vec2(f32(dx), f32(dy)) * step; + let rgba = textureSample(tex, samp, uv + offset); + sum += rgba_weights[idx] * rgba; + idx++; + } + } + + return sum; +} diff --git a/workspaces/main/shaders/cnn/cnn_layer.wgsl b/workspaces/main/shaders/cnn/cnn_layer.wgsl index e026ce8..b2bab26 100644 --- a/workspaces/main/shaders/cnn/cnn_layer.wgsl +++ b/workspaces/main/shaders/cnn/cnn_layer.wgsl @@ -29,10 +29,10 @@ struct CNNLayerParams { let uv = p.xy / uniforms.resolution; var result = vec4(0.0); - // Single layer for now (layer 0) + // Layer 0 uses coordinate-aware convolution if (params.layer_index == 0) { - result = cnn_conv3x3(txt, smplr, uv, uniforms.resolution, - weights_layer0, bias_layer0); + result = cnn_conv3x3_with_coord(txt, smplr, uv, uniforms.resolution, + rgba_weights_layer0, coord_weights_layer0, bias_layer0); result = cnn_tanh(result); } diff --git a/workspaces/main/shaders/cnn/cnn_weights_generated.wgsl b/workspaces/main/shaders/cnn/cnn_weights_generated.wgsl index 98c17ff..e0a7dc4 100644 --- a/workspaces/main/shaders/cnn/cnn_weights_generated.wgsl +++ b/workspaces/main/shaders/cnn/cnn_weights_generated.wgsl @@ -2,8 +2,8 @@ // 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, 9> = array( +// Layer 0: 3x3 convolution with coordinate awareness +const rgba_weights_layer0: array, 9> = array( mat4x4(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(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(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), @@ -14,4 +14,10 @@ const weights_layer0: array, 9> = array( mat4x4(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(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 coord_weights_layer0 = mat2x4( + 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0 +); + const bias_layer0 = vec4(0.0, 0.0, 0.0, 0.0); -- cgit v1.2.3