1 files changed, 63 insertions, 9 deletions

diff --git a/tools/cnn_test.cc b/tools/cnn_test.cc
index 3fad2ff..c504c3d 100644
--- a/tools/cnn_test.cc
+++ b/tools/cnn_test.cc
@@ -46,6 +46,8 @@ struct Args {
   int num_layers = 3; // Default to 3 layers
   bool debug_hex = false; // Print first 8 pixels as hex
   int cnn_version = 1; // 1=CNNEffect, 2=CNNv2Effect
+  const char* weights_path = nullptr; // Optional .bin weights file
+  bool cnn_version_explicit = false; // Track if --cnn-version was explicitly set
 };
 
 // Parse command-line arguments
@@ -87,10 +89,13 @@ static bool parse_args(int argc, char** argv, Args* args) {
       args->debug_hex = true;
     } else if (strcmp(argv[i], "--cnn-version") == 0 && i + 1 < argc) {
       args->cnn_version = atoi(argv[++i]);
+      args->cnn_version_explicit = true;
       if (args->cnn_version < 1 || args->cnn_version > 2) {
         fprintf(stderr, "Error: cnn-version must be 1 or 2\n");
         return false;
       }
+    } else if (strcmp(argv[i], "--weights") == 0 && i + 1 < argc) {
+      args->weights_path = argv[++i];
     } else if (strcmp(argv[i], "--help") == 0) {
       return false;
     } else {
@@ -99,6 +104,21 @@ static bool parse_args(int argc, char** argv, Args* args) {
     }
   }
 
+  // Force CNN v2 when --weights is specified
+  if (args->weights_path) {
+    if (args->cnn_version_explicit && args->cnn_version != 2) {
+      fprintf(stderr, "WARNING: --cnn-version %d ignored (--weights forces CNN v2)\n",
+              args->cnn_version);
+    }
+    args->cnn_version = 2;
+
+    // Warn if --layers was specified (binary file config takes precedence)
+    if (args->num_layers != 3) {  // 3 is the default
+      fprintf(stderr, "WARNING: --layers %d ignored (--weights loads layer config from .bin)\n",
+              args->num_layers);
+    }
+  }
+
   return true;
 }
 
@@ -108,10 +128,11 @@ static void print_usage(const char* prog) {
   fprintf(stderr, "\nOPTIONS:\n");
   fprintf(stderr, "  --blend F                Final blend amount (0.0-1.0, default: 1.0)\n");
   fprintf(stderr, "  --format ppm|png         Output format (default: png)\n");
-  fprintf(stderr, "  --layers N               Number of CNN layers (1-10, default: 3)\n");
+  fprintf(stderr, "  --layers N               Number of CNN layers (1-10, default: 3, ignored with --weights)\n");
   fprintf(stderr, "  --save-intermediates DIR Save intermediate layers to directory\n");
   fprintf(stderr, "  --debug-hex              Print first 8 pixels as hex (debug)\n");
-  fprintf(stderr, "  --cnn-version N          CNN version: 1 (default) or 2\n");
+  fprintf(stderr, "  --cnn-version N          CNN version: 1 (default) or 2 (ignored with --weights)\n");
+  fprintf(stderr, "  --weights PATH           Load weights from .bin (forces CNN v2, overrides layer config)\n");
   fprintf(stderr, "  --help                   Show this help\n");
 }
 
@@ -586,10 +607,38 @@ static bool process_cnn_v2(WGPUDevice device, WGPUQueue queue,
                            int width, int height, const Args& args) {
   printf("Using CNN v2 (storage buffer architecture)\n");
 
-  // Load weights
+  // Load weights (from file or asset system)
   size_t weights_size = 0;
-  const uint8_t* weights_data =
-      (const uint8_t*)GetAsset(AssetId::ASSET_WEIGHTS_CNN_V2, &weights_size);
+  const uint8_t* weights_data = nullptr;
+  std::vector<uint8_t> file_weights;  // For file-based loading
+
+  if (args.weights_path) {
+    // Load from file
+    printf("Loading weights from '%s'...\n", args.weights_path);
+    FILE* f = fopen(args.weights_path, "rb");
+    if (!f) {
+      fprintf(stderr, "Error: failed to open weights file '%s'\n", args.weights_path);
+      return false;
+    }
+
+    fseek(f, 0, SEEK_END);
+    weights_size = ftell(f);
+    fseek(f, 0, SEEK_SET);
+
+    file_weights.resize(weights_size);
+    size_t read = fread(file_weights.data(), 1, weights_size, f);
+    fclose(f);
+
+    if (read != weights_size) {
+      fprintf(stderr, "Error: failed to read weights file\n");
+      return false;
+    }
+
+    weights_data = file_weights.data();
+  } else {
+    // Load from asset system
+    weights_data = (const uint8_t*)GetAsset(AssetId::ASSET_WEIGHTS_CNN_V2, &weights_size);
+  }
 
   if (!weights_data || weights_size < 20) {
     fprintf(stderr, "Error: CNN v2 weights not available\n");
@@ -635,15 +684,20 @@ static bool process_cnn_v2(WGPUDevice device, WGPUQueue queue,
            info.out_channels, info.weight_count);
   }
 
-  // Create weights storage buffer
+  // Create weights storage buffer (skip header + layer info, upload only weights)
+  size_t header_size = 20;  // 5 u32
+  size_t layer_info_size = 20 * layer_info.size();  // 5 u32 per layer
+  size_t weights_offset = header_size + layer_info_size;
+  size_t weights_only_size = weights_size - weights_offset;
+
   WGPUBufferDescriptor weights_buffer_desc = {};
-  weights_buffer_desc.size = weights_size;
+  weights_buffer_desc.size = weights_only_size;
   weights_buffer_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst;
   weights_buffer_desc.mappedAtCreation = false;
 
   WGPUBuffer weights_buffer =
       wgpuDeviceCreateBuffer(device, &weights_buffer_desc);
-  wgpuQueueWriteBuffer(queue, weights_buffer, 0, weights_data, weights_size);
+  wgpuQueueWriteBuffer(queue, weights_buffer, 0, weights_data + weights_offset, weights_only_size);
 
   // Create input view
   const WGPUTextureViewDescriptor view_desc = {
@@ -1002,7 +1056,7 @@ static bool process_cnn_v2(WGPUDevice device, WGPUQueue queue,
 
     layer_bg_entries[3].binding = 3;
     layer_bg_entries[3].buffer = weights_buffer;
-    layer_bg_entries[3].size = weights_size;
+    layer_bg_entries[3].size = weights_only_size;
 
     layer_bg_entries[4].binding = 4;
     layer_bg_entries[4].buffer = layer_params_buffers[i];