format .wgsl layer code (cosmetics)

1 files changed, 4 insertions, 7 deletions

diff --git a/training/train_cnn.py b/training/train_cnn.py
index 57b4da8..ea1c4d7 100755
--- a/training/train_cnn.py
+++ b/training/train_cnn.py
@@ -308,7 +308,7 @@ def generate_layer_shader(output_path, num_layers, kernel_sizes):
         f.write("    return vec4<f32>(pos[i], 0.0, 1.0);\n")
         f.write("}\n\n")
         f.write("@fragment fn fs_main(@builtin(position) p: vec4<f32>) -> @location(0) vec4<f32> {\n")
-        f.write("    // Match PyTorch linspace: pixel_idx / (size - 1), not pixel_center / size\n")
+        f.write("    // Match PyTorch linspace\n")
         f.write("    let uv = (p.xy - 0.5) / (uniforms.resolution - 1.0);\n")
         f.write("    let original_raw = textureSample(original_input, smplr, uv);\n")
         f.write("    let original = (original_raw - 0.5) * 2.0;  // Normalize to [-1,1]\n")
@@ -325,20 +325,17 @@ def generate_layer_shader(output_path, num_layers, kernel_sizes):
                 conv_fn_src = f"cnn_conv{ks}x{ks}_7to4_src"
                 f.write(f"    // Layer 0: 7→4 (RGBD output, normalizes [0,1] input)\n")
                 f.write(f"    if (params.layer_index == {layer_idx}) {{\n")
-                f.write(f"        result = {conv_fn_src}(txt, smplr, uv, uniforms.resolution,\n")
-                f.write(f"                                         weights_layer{layer_idx});\n")
+                f.write(f"        result = {conv_fn_src}(txt, smplr, uv, uniforms.resolution, weights_layer{layer_idx});\n")
                 f.write(f"        result = cnn_tanh(result);\n")
                 f.write(f"    }}\n")
             elif not is_final:
                 f.write(f"    else if (params.layer_index == {layer_idx}) {{\n")
-                f.write(f"        result = {conv_fn}(txt, smplr, uv, uniforms.resolution,\n")
-                f.write(f"                                   gray, weights_layer{layer_idx});\n")
+                f.write(f"        result = {conv_fn}(txt, smplr, uv, uniforms.resolution, gray, weights_layer{layer_idx});\n")
                 f.write(f"        result = cnn_tanh(result);  // Keep in [-1,1]\n")
                 f.write(f"    }}\n")
             else:
                 f.write(f"    else if (params.layer_index == {layer_idx}) {{\n")
-                f.write(f"        let gray_out = {conv_fn}(txt, smplr, uv, uniforms.resolution,\n")
-                f.write(f"                                         gray, weights_layer{layer_idx});\n")
+                f.write(f"        let gray_out = {conv_fn}(txt, smplr, uv, uniforms.resolution, gray, weights_layer{layer_idx});\n")
                 f.write(f"        // gray_out in [0,1] (clamped to match PyTorch training)\n")
                 f.write(f"        result = vec4<f32>(gray_out, gray_out, gray_out, 1.0);\n")
                 f.write(f"        return mix(original_raw, result, params.blend_amount); // [0,1]\n")