5 files changed, 92 insertions, 5 deletions
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<f32>,
+  samp: sampler,
+  uv: vec2<f32>,
+  resolution: vec2<f32>,
+  rgba_weights: array<mat4x4<f32>, 9>,
+  coord_weights: mat2x4<f32>,
+  bias: vec4<f32>
+) -> vec4<f32> {
+  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>(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<f32>,
+  samp: sampler,
+  uv: vec2<f32>,
+  resolution: vec2<f32>,
+  rgba_weights: array<mat4x4<f32>, 25>,
+  coord_weights: mat2x4<f32>,
+  bias: vec4<f32>
+) -> vec4<f32> {
+  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>(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<f32>,
+  samp: sampler,
+  uv: vec2<f32>,
+  resolution: vec2<f32>,
+  rgba_weights: array<mat4x4<f32>, 49>,
+  coord_weights: mat2x4<f32>,
+  bias: vec4<f32>
+) -> vec4<f32> {
+  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>(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<f32>(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<mat4x4<f32>, 9> = array(
+// Layer 0: 3x3 convolution with coordinate awareness
+const rgba_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),
@@ -14,4 +14,10 @@ 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)
 );
+
+const coord_weights_layer0 = mat2x4<f32>(
+  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);