Upload model files

Cargo.toml

@@ -1,37 +0,0 @@
-[package]
-name = "flowstate-r1-rs"
-version = "0.1.0"
-edition = "2021"
-
-[[bin]]
-name = "flowstate-r1-rs"
-path = "src/main.rs"
-
-[dependencies]
-pyo3 = { version = "0.29", features = ["extension-module"], optional = true }
-reqwest    = { version = "0.11", features = ["stream", "json"] }
-futures-util = "0.3"
-tokio      = { version = "1", features = ["full"] }
-safetensors = "0.4"
-serde      = { version = "1", features = ["derive"] }
-serde_json = "1"
-clap       = { version = "4", features = ["derive", "env"] }
-anyhow     = "1"
-byteorder  = "1"
-indicatif  = "0.17"
-sha2 = "0.10"
-base64 = "0.22"
-tokio-util = { version = "0.7", features = ["io"] }
-candle-core = { version = "0.8", features = ["accelerate"] }
-candle-nn   = { version = "0.8", features = ["accelerate"] }
-
-[profile.release]
-opt-level = 3
-
-
-[features]
-python = ["dep:pyo3"]
-
-[lib]
-path = "src/lib.rs"
-crate-type = ["cdylib", "rlib"]

LICENSE

@@ -1,21 +0,0 @@
-MIT License
-
-Copyright (c) 2024 Andrew Mayes
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.

flowstate-r1-rs/Cargo.lock

@@ -481,6 +481,7 @@ dependencies = [
  "clap",
  "futures-util",
  "indicatif",
+ "pyo3",
  "reqwest",
  "safetensors",
  "serde",
@@ -1552,6 +1553,63 @@ dependencies = [
  "version_check",
 ]
 
+[[package]]
+name = "pyo3"
+version = "0.29.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c"
+dependencies = [
+ "libc",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.29.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078"
+dependencies = [
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.29.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.29.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.29.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "quote"
 version = "1.0.45"
@@ -2013,6 +2071,12 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "target-lexicon"
+version = "0.13.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
+
 [[package]]
 name = "tempfile"
 version = "3.27.0"

flowstate-r1-rs/Cargo.toml

@@ -8,6 +8,7 @@ name = "flowstate-r1-rs"
 path = "src/main.rs"
 
 [dependencies]
+pyo3 = { version = "0.29", features = ["extension-module"], optional = true }
 reqwest    = { version = "0.11", features = ["stream", "json"] }
 futures-util = "0.3"
 tokio      = { version = "1", features = ["full"] }
@@ -26,3 +27,11 @@ candle-nn   = { version = "0.8", features = ["accelerate"] }
 
 [profile.release]
 opt-level = 3
+
+
+[features]
+python = ["dep:pyo3"]
+
+[lib]
+path = "src/lib.rs"
+crate-type = ["cdylib", "rlib"]

flowstate-r1-rs/pyproject.toml

@@ -0,0 +1,17 @@
+[build-system]
+requires = ["maturin>=1,<2"]
+build-backend = "maturin"
+
+[project]
+name = "flowstate-r1-rs"
+version = "0.1.0"
+requires-python = ">=3.8"
+
+[tool.maturin]
+features = ["python"]
+module-name = "flowstate_r1_rs"
+
+[dependency-groups]
+dev = [
+    "maturin>=1.14.0",
+]

flowstate-r1-rs/src/infer/mod.rs

@@ -1,8 +1,11 @@
+use std::collections::HashMap;
+use std::sync::Mutex;
 use std::path::Path;
 
 use anyhow::Context;
-use candle_core::{DType, Device, Tensor};
+use candle_core::{DType, Device, Tensor, D};
 use candle_core::quantized::gguf_file;
+use candle_nn::ops;
 
 // ---------------------------------------------------------------------------
 // Config
@@ -29,30 +32,33 @@ pub struct InferConfig {
 struct S5Weights {
     log_lambda_real: Vec<f32>,  // [state_dim]
     lambda_imag:     Vec<f32>,  // [state_dim]
-    b_r: Tensor,                // [state_dim, embed_dim]  (row-major, HF order)
+    b_r: Tensor,                // [state_dim, embed_dim]
     b_i: Tensor,                // [state_dim, embed_dim]
-    c_r: Tensor,                // [embed_dim, state_dim]  (Python C_tilde shape)
+    c_r: Tensor,                // [embed_dim, state_dim]
     c_i: Tensor,                // [embed_dim, state_dim]
-    d:   Vec<f32>,              // [embed_dim]
+    d:   Tensor,                // [embed_dim]
     log_delta: Vec<f32>,        // [state_dim]
 }
 
 struct BlockWeights {
     ssm:        S5Weights,
-    out_weight: Tensor,     // [embed_dim, embed_dim]
-    out_bias:   Vec<f32>,   // [embed_dim]
-    norm_weight: Vec<f32>,  // [embed_dim]
-    norm_bias:   Vec<f32>,  // [embed_dim]
+    out_weight: Tensor,   // [embed_dim, embed_dim]
+    out_bias:   Tensor,   // [embed_dim]
+    norm_weight: Tensor,  // [embed_dim]
+    norm_bias:   Tensor,  // [embed_dim]
+    // cached per scale_factor: (A_bar_real, A_bar_imag, B_bar_real_t, B_bar_imag_t)
+    disc_cache: Mutex<HashMap<u32, (Vec<f32>, Vec<f32>, Tensor, Tensor)>>,
 }
 
 pub struct FlowStateModel {
     device:     Device,
     config:     InferConfig,
-    embed_w:    Tensor,     // [embed_dim, n_inputs]
-    embed_b:    Vec<f32>,   // [embed_dim]
+    embed_w:    Tensor,   // [embed_dim, n_inputs]
+    embed_b:    Tensor,   // [embed_dim]
     blocks:     Vec<BlockWeights>,
-    decoder_w:  Tensor,     // [n_quantiles * decoder_dim, embed_dim]
-    decoder_b:  Vec<f32>,   // [n_quantiles * decoder_dim]
+    decoder_w:  Tensor,   // [n_quantiles * decoder_dim, embed_dim]
+    decoder_b:  Tensor,   // [n_quantiles * decoder_dim]
+    legendre_cache: Mutex<HashMap<usize, Tensor>>,
 }
 
 // ---------------------------------------------------------------------------
@@ -91,11 +97,9 @@ impl FlowStateModel {
             .with_context(|| format!("open {}", gguf_path.display()))?;
         let content = gguf_file::Content::read(&mut file).context("read GGUF header")?;
 
-        // Embedding
         let embed_w = load_matrix(&content, &mut file, "embed.weight", &device)?;
-        let embed_b = load_f32_vec(&content, &mut file, "embed.bias", &device)?;
+        let embed_b = load_matrix(&content, &mut file, "embed.bias", &device)?;
 
-        // Encoder blocks
         let mut blocks = Vec::with_capacity(config.num_layers);
         for n in 0..config.num_layers {
             let ssm = S5Weights {
@@ -105,23 +109,32 @@ impl FlowStateModel {
                 b_i: load_matrix(&content, &mut file, &format!("blk.{n}.ssm.b_i"), &device)?,
                 c_r: load_matrix(&content, &mut file, &format!("blk.{n}.ssm.c_r"), &device)?,
                 c_i: load_matrix(&content, &mut file, &format!("blk.{n}.ssm.c_i"), &device)?,
-                d:          load_f32_vec(&content, &mut file, &format!("blk.{n}.ssm.d"),         &device)?,
+                d:          load_matrix(&content, &mut file, &format!("blk.{n}.ssm.d"),         &device)?,
                 log_delta:  load_f32_vec(&content, &mut file, &format!("blk.{n}.ssm.log_delta"), &device)?,
             };
             blocks.push(BlockWeights {
                 ssm,
-                out_weight:  load_matrix(&content, &mut file, &format!("blk.{n}.out.weight"), &device)?,
-                out_bias:    load_f32_vec(&content, &mut file, &format!("blk.{n}.out.bias"),  &device)?,
-                norm_weight: load_f32_vec(&content, &mut file, &format!("blk.{n}.norm.weight"), &device)?,
-                norm_bias:   load_f32_vec(&content, &mut file, &format!("blk.{n}.norm.bias"),   &device)?,
+                out_weight:  load_matrix(&content, &mut file, &format!("blk.{n}.out.weight"),  &device)?,
+                out_bias:    load_matrix(&content, &mut file, &format!("blk.{n}.out.bias"),    &device)?,
+                norm_weight: load_matrix(&content, &mut file, &format!("blk.{n}.norm.weight"), &device)?,
+                norm_bias:   load_matrix(&content, &mut file, &format!("blk.{n}.norm.bias"),   &device)?,
+                disc_cache:  Mutex::new(HashMap::new()),
             });
         }
 
-        // Decoder
         let decoder_w = load_matrix(&content, &mut file, "decoder.weight", &device)?;
-        let decoder_b = load_f32_vec(&content, &mut file, "decoder.bias", &device)?;
-
-        Ok(Self { device, config, embed_w, embed_b, blocks, decoder_w, decoder_b })
+        let decoder_b = load_matrix(&content, &mut file, "decoder.bias", &device)?;
+
+        Ok(Self {
+            device,
+            config,
+            embed_w,
+            embed_b,
+            blocks,
+            decoder_w,
+            decoder_b,
+            legendre_cache: Mutex::new(HashMap::new()),
+        })
     }
 
     // -----------------------------------------------------------------------
@@ -160,29 +173,26 @@ impl FlowStateModel {
 
         // 4. Embedding: [seq_len, n_inputs] × embed_w^T + embed_b → [seq_len, embed_dim]
         let input_t = Tensor::from_vec(input_data, (seq_len, cfg.n_inputs), &self.device)?;
-        let mut hidden = linear(&input_t, &self.embed_w, &self.embed_b, &self.device)?;
-        // hidden: [seq_len, embed_dim]
+        let mut hidden = linear(&input_t, &self.embed_w, &self.embed_b)?;
 
         // 5. Scale factor for discretization: decoder_patch_len / prediction_length
         let scale_factor = cfg.decoder_patch_len as f32 / prediction_length as f32;
 
-        // 6. Encoder: 6 S5 layers
+        // 6. Encoder: S5 layers
         let num_layers = self.blocks.len();
         for (i, block) in self.blocks.iter().enumerate() {
             let is_last = i == num_layers - 1;
-            hidden = self.apply_s5_layer(hidden, block, scale_factor, is_last, seq_len)?;
+            hidden = self.apply_s5_layer(hidden, block, scale_factor, is_last)?;
         }
         // After last layer: hidden is [1, embed_dim]
 
-        // 7. Decoder: linear → [1, n_quantiles * decoder_dim]
-        let coeffs = linear(&hidden, &self.decoder_w, &self.decoder_b, &self.device)?;
+        // 7. Decoder: linear → [n_q, decoder_dim]
         let n_q = cfg.quantiles.len();
-        let coeffs_data: Vec<f32> = coeffs.flatten_all()?.to_vec1()?;
-        // coeffs_data: [n_quantiles * decoder_dim]
+        let coeffs = linear(&hidden, &self.decoder_w, &self.decoder_b)?
+            .reshape((n_q, cfg.decoder_dim))?;
 
-        // Diagnostic: dump all coefficients when FLOWSTATE_DEBUG is set
         if std::env::var("FLOWSTATE_DEBUG").is_ok() {
-            // Print all coefficients as CSV lines to stderr: COEFF,{qi},{d},{val}
+            let coeffs_data: Vec<f32> = coeffs.flatten_all()?.to_vec1()?;
             for qi in 0..n_q {
                 for d in 0..cfg.decoder_dim {
                     eprintln!("COEFF,{qi},{d},{:.8}", coeffs_data[qi * cfg.decoder_dim + d]);
@@ -190,30 +200,28 @@ impl FlowStateModel {
             }
         }
 
-        // 8. Compute Legendre basis [prediction_length, decoder_dim]
-        // Python: t[i] = range[0] + i * scale * (range[1]-range[0]) / decoder_patch_len  (i=1..n)
-        let basis = legendre_basis(prediction_length, cfg.decoder_dim, cfg.basis_range,
-                                   scale_factor, cfg.decoder_patch_len);
-        // basis: [prediction_length][decoder_dim], scaled by 1/4
-
-        // 9. Matrix multiply: [n_q, decoder_dim] × [prediction_length, decoder_dim]^T
-        //    = [n_q, prediction_length]
-        let mut output = vec![vec![0.0f32; prediction_length]; n_q];
-        for q in 0..n_q {
-            let coeff_offset = q * cfg.decoder_dim;
-            for p in 0..prediction_length {
-                let mut sum = 0.0f32;
-                for d in 0..cfg.decoder_dim {
-                    sum += coeffs_data[coeff_offset + d] * basis[p][d];
-                }
-                output[q][p] = sum;
+        // 8. Legendre basis [prediction_length, decoder_dim] — cached per prediction_length
+        let basis = {
+            let mut cache = self.legendre_cache.lock().unwrap();
+            if !cache.contains_key(&prediction_length) {
+                let raw = legendre_basis(prediction_length, cfg.decoder_dim, cfg.basis_range,
+                                        scale_factor, cfg.decoder_patch_len);
+                let flat: Vec<f32> = raw.into_iter().flatten().collect();
+                cache.insert(prediction_length,
+                    Tensor::from_vec(flat, (prediction_length, cfg.decoder_dim), &self.device)?);
             }
-        }
+            cache[&prediction_length].clone()
+        };
 
-        // 10. Denormalize: output * final_std + final_mean
+        // 9. [n_q, decoder_dim] @ [decoder_dim, prediction_length] → [n_q, prediction_length]
+        let out_t = coeffs.matmul(&basis.t()?)?;
+
+        // 10. Denormalize and convert to Vec<Vec<f32>>
+        let out_raw: Vec<f32> = out_t.flatten_all()?.to_vec1()?;
+        let mut output = vec![vec![0.0f32; prediction_length]; n_q];
         for q in 0..n_q {
             for p in 0..prediction_length {
-                output[q][p] = output[q][p] * final_std + final_mean;
+                output[q][p] = out_raw[q * prediction_length + p] * final_std + final_mean;
             }
         }
 
@@ -226,11 +234,10 @@ impl FlowStateModel {
 
     fn apply_s5_layer(
         &self,
-        x: Tensor,      // [seq_len, embed_dim] for non-last; [seq_len, embed_dim] → output [1, embed_dim] for last
+        x: Tensor,      // [seq_len, embed_dim]
         block: &BlockWeights,
         scale_factor: f32,
         is_last: bool,
-        _seq_len: usize,
     ) -> anyhow::Result<Tensor> {
         let cfg = &self.config;
         let state_dim = cfg.state_dim;
@@ -240,81 +247,76 @@ impl FlowStateModel {
 
         // Save skip connection (trimmed for last layer)
         let skip = if is_last {
-            x.narrow(0, seq_len - 1, 1)?  // last row: [1, embed_dim]
+            x.narrow(0, seq_len - 1, 1)?  // [1, embed_dim]
         } else {
             x.clone()
         };
 
-        // ---- SSM forward ----
-
-        // Discretize: compute A_bar and B_bar from SSM parameters
-        let ssm = &block.ssm;
-        let (a_bar_r, a_bar_i, b_bar_r_data, b_bar_i_data) =
-            discretize(ssm, scale_factor, state_dim, embed_dim);
+        // ---- Get or compute discretized SSM matrices (cached per scale_factor) ----
+        let (a_bar_r, a_bar_i, b_bar_r_t, b_bar_i_t) = {
+            let key = scale_factor.to_bits();
+            let mut cache = block.disc_cache.lock().unwrap();
+            if !cache.contains_key(&key) {
+                let result = discretize(&block.ssm, scale_factor, state_dim, embed_dim, &self.device)?;
+                cache.insert(key, result);
+            }
+            let (ar, ai, brt, bit) = &cache[&key];
+            (ar.clone(), ai.clone(), brt.clone(), bit.clone())
+        };
 
         // B @ x for all timesteps at once: x [seq_len, embed_dim] × B^T [embed_dim, state_dim]
-        // B_tilde stored as [state_dim, embed_dim], so B @ x = x @ B^T
-        let b_bar_r_t = Tensor::from_vec(b_bar_r_data, (state_dim, embed_dim), &self.device)?;
-        let b_bar_i_t = Tensor::from_vec(b_bar_i_data, (state_dim, embed_dim), &self.device)?;
-
-        // x @ B^T → [seq_len, state_dim]
-        let bu_r = x.matmul(&b_bar_r_t.t()?)?;
+        let bu_r = x.matmul(&b_bar_r_t.t()?)?;  // [seq_len, state_dim]
         let bu_i = x.matmul(&b_bar_i_t.t()?)?;
         let bu_r_data: Vec<f32> = bu_r.flatten_all()?.to_vec1()?;
         let bu_i_data: Vec<f32> = bu_i.flatten_all()?.to_vec1()?;
 
-        // Sequential SSM scan
-        let mut h_r = vec![0.0f32; state_dim]; // real part of hidden state
-        let mut h_i = vec![0.0f32; state_dim]; // imag part
-        let mut all_h_r = vec![0.0f32; seq_len * state_dim];
-        let mut all_h_i = vec![0.0f32; seq_len * state_dim];
-
-        for t in 0..seq_len {
-            for s in 0..state_dim {
-                let new_r = a_bar_r[s] * h_r[s] - a_bar_i[s] * h_i[s] + bu_r_data[t * state_dim + s];
-                let new_i = a_bar_r[s] * h_i[s] + a_bar_i[s] * h_r[s] + bu_i_data[t * state_dim + s];
-                h_r[s] = new_r;
-                h_i[s] = new_i;
+        // Sequential SSM scan (inherently sequential — recurrence prevents parallelism).
+        // For the last block only the final hidden state is needed, so skip the history buffer.
+        let (h_r_t, h_i_t) = if is_last {
+            let mut h_r = vec![0.0f32; state_dim];
+            let mut h_i = vec![0.0f32; state_dim];
+            for t in 0..seq_len {
+                for s in 0..state_dim {
+                    let new_r = a_bar_r[s] * h_r[s] - a_bar_i[s] * h_i[s] + bu_r_data[t * state_dim + s];
+                    let new_i = a_bar_r[s] * h_i[s] + a_bar_i[s] * h_r[s] + bu_i_data[t * state_dim + s];
+                    h_r[s] = new_r;
+                    h_i[s] = new_i;
+                }
             }
-            all_h_r[t * state_dim..(t + 1) * state_dim].copy_from_slice(&h_r);
-            all_h_i[t * state_dim..(t + 1) * state_dim].copy_from_slice(&h_i);
-        }
-
-        // For last layer, keep only the final timestep
-        let (h_r_t, h_i_t, out_seq_len) = if is_last {
-            let hr = Tensor::from_vec(all_h_r[(seq_len - 1) * state_dim..].to_vec(), (1, state_dim), &self.device)?;
-            let hi = Tensor::from_vec(all_h_i[(seq_len - 1) * state_dim..].to_vec(), (1, state_dim), &self.device)?;
-            (hr, hi, 1usize)
+            let hr = Tensor::from_vec(h_r, (1, state_dim), &self.device)?;
+            let hi = Tensor::from_vec(h_i, (1, state_dim), &self.device)?;
+            (hr, hi)
         } else {
+            let mut h_r = vec![0.0f32; state_dim];
+            let mut h_i = vec![0.0f32; state_dim];
+            let mut all_h_r = vec![0.0f32; seq_len * state_dim];
+            let mut all_h_i = vec![0.0f32; seq_len * state_dim];
+            for t in 0..seq_len {
+                for s in 0..state_dim {
+                    let new_r = a_bar_r[s] * h_r[s] - a_bar_i[s] * h_i[s] + bu_r_data[t * state_dim + s];
+                    let new_i = a_bar_r[s] * h_i[s] + a_bar_i[s] * h_r[s] + bu_i_data[t * state_dim + s];
+                    h_r[s] = new_r;
+                    h_i[s] = new_i;
+                }
+                all_h_r[t * state_dim..(t + 1) * state_dim].copy_from_slice(&h_r);
+                all_h_i[t * state_dim..(t + 1) * state_dim].copy_from_slice(&h_i);
+            }
             let hr = Tensor::from_vec(all_h_r, (seq_len, state_dim), &self.device)?;
             let hi = Tensor::from_vec(all_h_i, (seq_len, state_dim), &self.device)?;
-            (hr, hi, seq_len)
+            (hr, hi)
         };
 
-        // C @ h: y_real = C_r @ h_r - C_i @ h_i
-        // C_tilde_r stored as [embed_dim, state_dim]; h_r as [out_seq_len, state_dim]
-        // y = h @ C^T → [out_seq_len, embed_dim]
-        let y_from_cr = h_r_t.matmul(&block.ssm.c_r.t()?)?;  // [out_seq_len, embed_dim]
-        let y_from_ci = h_i_t.matmul(&block.ssm.c_i.t()?)?;  // [out_seq_len, embed_dim]
-        let mut y_data: Vec<f32> = (y_from_cr - y_from_ci)?.flatten_all()?.to_vec1()?;
+        // C @ h: y_real = C_r @ h_r - C_i @ h_i  → [out_seq_len, embed_dim]
+        let y_from_cr = h_r_t.matmul(&block.ssm.c_r.t()?)?;
+        let y_from_ci = h_i_t.matmul(&block.ssm.c_i.t()?)?;
+        let y_raw = (y_from_cr - y_from_ci)?;
 
-        // Add D skip: y += D * x_at_positions
-        let x_positions = if is_last {
-            x.narrow(0, seq_len - 1, 1)?.flatten_all()?.to_vec1::<f32>()?
-        } else {
-            x.flatten_all()?.to_vec1::<f32>()?
-        };
-        let d = &block.ssm.d;
-        for t in 0..out_seq_len {
-            for e in 0..embed_dim {
-                y_data[t * embed_dim + e] += d[e] * x_positions[t * embed_dim + e];
-            }
-        }
+        // D skip: y += D * x_at_positions  (skip is the correct slice in both cases)
+        let y_t = (y_raw + skip.broadcast_mul(&block.ssm.d)?)?;
 
         // ---- MLP: selu(y) * sigmoid(out_linear(selu(y))) ----
-        let y_t = Tensor::from_vec(y_data, (out_seq_len, embed_dim), &self.device)?;
         let y_selu = selu_tensor(&y_t)?;
-        let gate_pre = linear(&y_selu, &block.out_weight, &block.out_bias, &self.device)?;
+        let gate_pre = linear(&y_selu, &block.out_weight, &block.out_bias)?;
         let gate = sigmoid_tensor(&gate_pre)?;
         let y_gated = y_selu.mul(&gate)?;
 
@@ -322,9 +324,7 @@ impl FlowStateModel {
         let y_normed = layer_norm(&y_gated, &block.norm_weight, &block.norm_bias, self.config.eps)?;
 
         // ---- Residual ----
-        let out = (y_normed + skip)?;
-
-        Ok(out)
+        Ok((y_normed + skip)?)
     }
 }
 
@@ -347,7 +347,6 @@ fn causal_revin_norm(x: &[f32], eps: f32) -> (Vec<f32>, f32, f32) {
         cum_sum += x[t];
         let mean_t = cum_sum / count;
 
-        // Python: cumsum((x[s] - mean_s)^2) / count — each step uses its own running mean
         cum_sq_diff += (x[t] - mean_t) * (x[t] - mean_t);
         let var_t = (cum_sq_diff / count).max(0.0);
         let std_t = (var_t + eps).sqrt();
@@ -367,23 +366,17 @@ fn causal_revin_norm(x: &[f32], eps: f32) -> (Vec<f32>, f32, f32) {
 // SSM discretization
 // ---------------------------------------------------------------------------
 
-/// Returns (A_bar_real, A_bar_imag, B_bar_real_flat, B_bar_imag_flat).
-/// B_bar shape: [state_dim, embed_dim] flattened row-major.
+/// Returns (A_bar_real, A_bar_imag, B_bar_real_tensor, B_bar_imag_tensor).
+/// B_bar tensors have shape [state_dim, embed_dim].
 fn discretize(
     ssm: &S5Weights,
     scale_factor: f32,
     state_dim: usize,
     embed_dim: usize,
-) -> (Vec<f32>, Vec<f32>, Vec<f32>, Vec<f32>) {
-    // lambda = -exp(log_lambda_real) + j * lambda_imag
-    // delta  = exp(scale_factor * log_delta)
-    // A_bar  = exp(lambda * delta) = exp((lambda_r + j*lambda_i) * delta)
-    //        = exp(lambda_r * delta) * (cos(lambda_i * delta) + j*sin(lambda_i * delta))
-
+    device: &Device,
+) -> anyhow::Result<(Vec<f32>, Vec<f32>, Tensor, Tensor)> {
     let mut a_r = vec![0.0f32; state_dim];
     let mut a_i = vec![0.0f32; state_dim];
-
-    // (A_bar - 1) / lambda per element (complex division)
     let mut coeff_r = vec![0.0f32; state_dim];
     let mut coeff_i = vec![0.0f32; state_dim];
 
@@ -392,16 +385,12 @@ fn discretize(
         let lam_i = ssm.lambda_imag[s];
         let delta = (scale_factor * ssm.log_delta[s]).exp();
 
-        // A_bar = exp(lambda * delta)
         let exp_r = lam_r * delta;
         let exp_i = lam_i * delta;
         let mag = exp_r.exp();
         a_r[s] = mag * exp_i.cos();
         a_i[s] = mag * exp_i.sin();
 
-        // (A_bar - 1) / lambda
-        // numerator: (a_r - 1) + j*a_i
-        // denominator: lam_r + j*lam_i
         let num_r = a_r[s] - 1.0;
         let num_i = a_i[s];
         let denom_sq = lam_r * lam_r + lam_i * lam_i;
@@ -409,31 +398,14 @@ fn discretize(
             coeff_r[s] = (num_r * lam_r + num_i * lam_i) / denom_sq;
             coeff_i[s] = (num_i * lam_r - num_r * lam_i) / denom_sq;
         } else {
-            coeff_r[s] = delta; // limit as lambda → 0
+            coeff_r[s] = delta;
             coeff_i[s] = 0.0;
         }
     }
 
-    // B_bar = coeff * B_tilde (element-wise complex multiply, then multiply B matrix)
-    // B_tilde stored as [state_dim, embed_dim] in GGUF (transposed from Python HF which is also [state_dim, emb])
-    // After load_matrix, candle reverses dims: stored [embed_dim, state_dim] in GGUF → loaded as [state_dim, embed_dim]
-    // Actually GGUF reversal: Python [state_dim, embed_dim] → GGUF stores reversed → candle sees [embed_dim, state_dim]
-    // But we treat it as [state_dim, embed_dim] for the einsum B @ u = u @ B^T
-    // Let's get B_tilde data and compute B_bar manually
-    // The GGUF shape is reversed from Python. Python B_tilde: [512, 512].
-    // So GGUF/candle sees it as [512, 512] (square, so same either way).
-
-    // Get B_tilde data as flat arrays [state_dim * embed_dim]
-    // b_r and b_i are loaded as candle tensors; we need them as flat vecs for the per-state multiply
-    // They were stored with reversed shape but 512×512 is symmetric in shape
-    // We extract row by row: B_tilde[s, :] = row s of the [state_dim, embed_dim] matrix
-
-    // Use flat data from the tensors
     let b_r_data = get_tensor_data_row_major(&ssm.b_r, state_dim, embed_dim);
     let b_i_data = get_tensor_data_row_major(&ssm.b_i, state_dim, embed_dim);
 
-    // B_bar[s, e] = coeff_r[s]*b_r[s,e] - coeff_i[s]*b_i[s,e]  (real)
-    //             + coeff_r[s]*b_i[s,e] + coeff_i[s]*b_r[s,e]  (imag)
     let mut b_bar_r = vec![0.0f32; state_dim * embed_dim];
     let mut b_bar_i = vec![0.0f32; state_dim * embed_dim];
 
@@ -446,11 +418,13 @@ fn discretize(
         }
     }
 
-    (a_r, a_i, b_bar_r, b_bar_i)
+    let b_bar_r_t = Tensor::from_vec(b_bar_r, (state_dim, embed_dim), device)?;
+    let b_bar_i_t = Tensor::from_vec(b_bar_i, (state_dim, embed_dim), device)?;
+
+    Ok((a_r, a_i, b_bar_r_t, b_bar_i_t))
 }
 
 /// Extract tensor data in row-major order as Vec<f32>.
-/// Handles candle's internal memory layout for F32 tensors.
 fn get_tensor_data_row_major(t: &Tensor, rows: usize, cols: usize) -> Vec<f32> {
     t.to_dtype(DType::F32)
         .and_then(|t| t.reshape((rows, cols)))
@@ -471,14 +445,13 @@ pub fn dump_legendre_basis(n_points: usize, degree: usize, range: [f32; 2],
 
 /// Compute Legendre polynomial basis matrix.
 /// Returns [n_points][degree+1] scaled by 1/4 (as in get_kernel).
-/// Python formula: t[i] = range[0] + i * scale * (range[1]-range[0]) / pred_dist  (i=1..n_points)
 fn legendre_basis(n_points: usize, degree: usize, range: [f32; 2],
                   scale: f32, pred_dist: usize) -> Vec<Vec<f32>> {
     let dt = scale * (range[1] - range[0]) / pred_dist as f32;
     let t: Vec<f32> = (1..=n_points).map(|i| range[0] + i as f32 * dt).collect();
 
-    // Evaluate Legendre polynomials at each t: [n_points, degree+1]
-    // P_0 = 1, P_1 = x, P_k = ((2k-1)*x*P_{k-1} - (k-1)*P_{k-2}) / k
+    // Compute degree Legendre polynomials (P0..P_{degree-1}) per point.
+    // degree+1 scratch columns needed during recurrence, then truncated to degree.
     let mut basis = vec![vec![0.0f32; degree + 1]; n_points];
     for (p, &x) in t.iter().enumerate() {
         basis[p][0] = 1.0;
@@ -490,10 +463,10 @@ fn legendre_basis(n_points: usize, degree: usize, range: [f32; 2],
             basis[p][k + 1] =
                 ((2.0 * kf + 1.0) * x * basis[p][k] - kf * basis[p][k - 1]) / (kf + 1.0);
         }
-        // Scale by 1/4 (from get_kernel: return f.float() / 4.0)
-        for d in 0..=degree {
+        for d in 0..degree {
             basis[p][d] /= 4.0;
         }
+        basis[p].truncate(degree);
     }
 
     basis
@@ -503,55 +476,33 @@ fn legendre_basis(n_points: usize, degree: usize, range: [f32; 2],
 // Neural network primitives
 // ---------------------------------------------------------------------------
 
-/// y = x @ w^T + b  (w: [out, in])
-fn linear(x: &Tensor, w: &Tensor, b: &[f32], device: &Device) -> anyhow::Result<Tensor> {
-    let out_size = b.len();
-    let y = x.matmul(&w.t()?)?;
-    let b_t = Tensor::from_vec(b.to_vec(), (1, out_size), device)?.broadcast_as(y.shape())?;
-    Ok((y + b_t)?)
+/// y = x @ w^T + b  (w: [out, in], b: [out])
+fn linear(x: &Tensor, w: &Tensor, b: &Tensor) -> anyhow::Result<Tensor> {
+    Ok(x.matmul(&w.t()?)?.broadcast_add(b)?)
 }
 
-/// SELU activation: selu(x) = scale * (x if x>0 else alpha*(exp(x)-1))
+/// SELU activation using candle ops (no Vec roundtrip).
 fn selu_tensor(x: &Tensor) -> anyhow::Result<Tensor> {
     const SCALE: f64 = 1.0507009873554804934193349852946;
     const ALPHA: f64 = 1.6732632423543772848170429916717;
-    let data: Vec<f32> = x.flatten_all()?.to_vec1()?;
-    let selu_data: Vec<f32> = data.iter().map(|&v| {
-        if v > 0.0 {
-            (SCALE * v as f64) as f32
-        } else {
-            (SCALE * ALPHA * (v.exp() as f64 - 1.0)) as f32
-        }
-    }).collect();
-    Ok(Tensor::from_vec(selu_data, x.shape(), x.device())?)
+    const ALPHA_SCALE: f64 = SCALE * ALPHA;
+    let pos = x.relu()?;
+    let neg = (x - &pos)?;            // min(x, 0)
+    let selu_pos = (pos * SCALE)?;
+    let selu_neg = ((neg.exp()? - 1.0)? * ALPHA_SCALE)?;
+    Ok((selu_pos + selu_neg)?)
 }
 
-/// Sigmoid: 1 / (1 + exp(-x))
 fn sigmoid_tensor(x: &Tensor) -> anyhow::Result<Tensor> {
-    let data: Vec<f32> = x.flatten_all()?.to_vec1()?;
-    let sig_data: Vec<f32> = data.iter().map(|&v| 1.0 / (1.0 + (-v).exp())).collect();
-    Ok(Tensor::from_vec(sig_data, x.shape(), x.device())?)
+    Ok(ops::sigmoid(x)?)
 }
 
-/// LayerNorm: (x - mean) / std * weight + bias (over last dimension)
-fn layer_norm(x: &Tensor, weight: &[f32], bias: &[f32], eps: f32) -> anyhow::Result<Tensor> {
-    let shape = x.shape().clone();
-    let last_dim = *shape.dims().last().context("empty tensor")?;
-    let n = shape.elem_count();
-    let n_rows = n / last_dim;
-
-    let data: Vec<f32> = x.flatten_all()?.to_vec1()?;
-    let mut out = vec![0.0f32; n];
-
-    for r in 0..n_rows {
-        let row = &data[r * last_dim..(r + 1) * last_dim];
-        let mean: f32 = row.iter().sum::<f32>() / last_dim as f32;
-        let var: f32 = row.iter().map(|&v| (v - mean).powi(2)).sum::<f32>() / last_dim as f32;
-        let std_inv = 1.0 / (var + eps).sqrt();
-        for i in 0..last_dim {
-            out[r * last_dim + i] = (row[i] - mean) * std_inv * weight[i] + bias[i];
-        }
-    }
-
-    Ok(Tensor::from_vec(out, shape, x.device())?)
+/// LayerNorm using candle ops (no Vec roundtrip).
+fn layer_norm(x: &Tensor, weight: &Tensor, bias: &Tensor, eps: f32) -> anyhow::Result<Tensor> {
+    let mean = x.mean_keepdim(D::Minus1)?;
+    let x_c = x.broadcast_sub(&mean)?;
+    let var = x_c.sqr()?.mean_keepdim(D::Minus1)?;
+    let std = (var + eps as f64)?.sqrt()?;
+    let normed = x_c.broadcast_div(&std)?;
+    Ok(normed.broadcast_mul(weight)?.broadcast_add(bias)?)
 }

flowstate-r1-rs/src/main.rs

@@ -184,7 +184,7 @@ async fn main() -> anyhow::Result<()> {
             let mut buf = String::new();
             std::io::stdin().read_to_string(&mut buf).context("read stdin")?;
             let req: serde_json::Value = serde_json::from_str(&buf).context("parse JSON input")?;
-            let ctx: Vec<f32> = serde_json::from_value(req["context"].clone()).context("context must be a JSON array of numbers")?;
+            let contexts = parse_contexts(req["context"].clone())?;
             let horizon: usize = req["horizon"].as_u64().context("horizon must be a positive integer")? as usize;
 
             let config_str = std::fs::read_to_string(&config)
@@ -205,47 +205,68 @@ async fn main() -> anyhow::Result<()> {
                 eps:               1e-5,
             };
 
-            anyhow::ensure!(!ctx.is_empty(), "No context values provided");
-
-            eprintln!("Loading model from {} …", gguf.display());
-            let model = FlowStateModel::load(&gguf, infer_config.clone())
-                .context("load model")?;
-
-            eprintln!("Running forecast ({} context steps → {horizon} future steps) …", ctx.len());
-            let quantile_mat = model.forecast(&ctx, horizon)
-                .context("forecast")?;
-            // quantile_mat: [n_quantiles][horizon]
-
-            let quantile_levels = &infer_config.quantiles;
-
-            // Find median index (q0.5), fall back to middle
+            let quantile_levels = infer_config.quantiles.clone();
             let median_idx = quantile_levels
                 .iter()
                 .position(|&q| (q - 0.5).abs() < 1e-6)
                 .unwrap_or(quantile_levels.len() / 2);
 
-            let point = quantile_mat.get(median_idx).cloned().unwrap_or_default();
+            eprintln!("Loading model from {} …", gguf.display());
+            let model = FlowStateModel::load(&gguf, infer_config)
+                .context("load model")?;
 
-            let mut quantiles = BTreeMap::new();
-            for (i, &level) in quantile_levels.iter().enumerate() {
-                if let Some(q) = quantile_mat.get(i) {
-                    quantiles.insert(format!("{level:.2}"), q.clone());
+            let mut fc_choices = Vec::new();
+            for ctx in &contexts {
+                anyhow::ensure!(!ctx.is_empty(), "context series must not be empty");
+                eprintln!("Running forecast ({} context steps → {horizon} future steps) …", ctx.len());
+                let quantile_mat = model.forecast(ctx, horizon).context("forecast")?;
+
+                let point = quantile_mat.get(median_idx).cloned().unwrap_or_default();
+                let mut quantiles = BTreeMap::new();
+                for (i, &level) in quantile_levels.iter().enumerate() {
+                    if let Some(q) = quantile_mat.get(i) {
+                        quantiles.insert(format!("{level:.2}"), q.clone());
+                    }
                 }
+                fc_choices.push((point, quantiles));
             }
-
-            println!("{}", forecast_json("flowstate-r1", ctx.len(), horizon, point, quantiles)?);
+            let total_ctx: usize = contexts.iter().map(|c| c.len()).sum();
+            println!("{}", forecast_json("flowstate-r1", total_ctx, horizon, fc_choices)?);
         }
     }
 
     Ok(())
 }
 
+fn parse_contexts(val: serde_json::Value) -> anyhow::Result<Vec<Vec<f32>>> {
+    match val {
+        serde_json::Value::Array(arr) if arr.is_empty() => {
+            anyhow::bail!("context must be a non-empty array")
+        }
+        serde_json::Value::Array(arr) => {
+            if arr.first().map(|v| v.is_array()).unwrap_or(false) {
+                arr.into_iter()
+                    .enumerate()
+                    .map(|(i, v)| {
+                        serde_json::from_value::<Vec<f32>>(v)
+                            .with_context(|| format!("context[{i}] must be an array of numbers"))
+                    })
+                    .collect()
+            } else {
+                let ctx = serde_json::from_value::<Vec<f32>>(serde_json::Value::Array(arr))
+                    .context("context must be a JSON array of numbers")?;
+                Ok(vec![ctx])
+            }
+        }
+        _ => anyhow::bail!("context must be a JSON array"),
+    }
+}
+
 fn forecast_json(
     model_name: &str,
     context_length: usize,
     forecast_length: usize,
-    point: Vec<f32>,
-    quantiles: BTreeMap<String, Vec<f32>>,
+    fc_choices: Vec<(Vec<f32>, BTreeMap<String, Vec<f32>>)>,
 ) -> anyhow::Result<String> {
     #[derive(Serialize)]
     struct ForecastResponse {
@@ -283,11 +304,11 @@ fn forecast_json(
         object: "forecast",
         created,
         model: model_name.to_string(),
-        choices: vec![Choice {
-            index: 0,
+        choices: fc_choices.into_iter().enumerate().map(|(i, (point, quantiles))| Choice {
+            index: i,
             forecast: ForecastOutput { point, quantiles },
             finish_reason: "stop",
-        }],
+        }).collect(),
         usage: Usage { context_length, forecast_length },
     };
 

flowstate-r1-rs/uv.lock

@@ -0,0 +1,96 @@
+version = 1
+revision = 3
+requires-python = ">=3.8"
+
+[[package]]
+name = "flowstate-r1-rs"
+version = "0.1.0"
+source = { editable = "." }
+
+[package.dev-dependencies]
+dev = [
+    { name = "maturin" },
+]
+
+[package.metadata]
+
+[package.metadata.requires-dev]
+dev = [{ name = "maturin", specifier = ">=1.14.0" }]
+
+[[package]]
+name = "maturin"
+version = "1.14.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "tomli", marker = "python_full_version < '3.11'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a7/d0/b7c8b7778cc44df3efbc96eb23acaa995e06ea1a60eb9b02f29858fcbd08/maturin-1.14.0.tar.gz", hash = "sha256:f7f82a6aca4a6c402bf00b99200be199d4874d04b9b9e74e825726a3478bba7f", size = 367010, upload-time = "2026-06-12T00:13:30.811Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/88/51/49367dcd8f6ec139e69ef0c695c8ff5075223673382101812b4affa53216/maturin-1.14.0-py3-none-linux_armv6l.whl", hash = "sha256:019ea3ec7e71f4c9759a367d4d21022ed5a3a621a2ce123abf3fb114ab3711ca", size = 10204135, upload-time = "2026-06-12T00:13:34.308Z" },
+    { url = "https://files.pythonhosted.org/packages/dd/2a/487ce56c838d25e0ce64350e75ec4e3dc89544c0a6233221c229d6aa1a84/maturin-1.14.0-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:6948a10f5f3470b791f79319be51debdd8bfd1778b36f2409f98e1314bc3859b", size = 19736800, upload-time = "2026-06-12T00:13:40.456Z" },
+    { url = "https://files.pythonhosted.org/packages/a8/a5/12f2efc18f419edce3282a93629cba16278bb502135dac95cd04ef7c2eae/maturin-1.14.0-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:1506e86b1e273a98074a62e281b13f27ac96f8cdef85f7f98d3e3589a9387a23", size = 10201144, upload-time = "2026-06-12T00:13:26.842Z" },
+    { url = "https://files.pythonhosted.org/packages/bf/95/3789e72273fd8bc80c33a11c787634b3251c4989d7a7203a92438836d4ff/maturin-1.14.0-py3-none-manylinux_2_12_i686.manylinux2010_i686.musllinux_1_1_i686.whl", hash = "sha256:df10ce4f7ba97fd3423f624f39b94c888ae3e5b470642a91918e1ccec81282fd", size = 10182394, upload-time = "2026-06-12T00:13:13.693Z" },
+    { url = "https://files.pythonhosted.org/packages/40/79/15957eb4e055597f217e6310963a9c1371372e63c5b4a3e30803365addd2/maturin-1.14.0-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl", hash = "sha256:75bcd4468a7fe597652cc2980c6bb16ce4bb8c411e3eb85dac2c4418cef0e95a", size = 10616603, upload-time = "2026-06-12T00:13:22.795Z" },
+    { url = "https://files.pythonhosted.org/packages/3e/4b/d1822f88cd5e855640f0e10ee00c39b9be614c1ef2f827e9792332d94b9f/maturin-1.14.0-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:2d123337e817f8dfe23755d6760139c01104137bb63e9e20c289c547e25ec857", size = 10075309, upload-time = "2026-06-12T00:13:38.274Z" },
+    { url = "https://files.pythonhosted.org/packages/c0/82/c1b160d2163e8784489285e82a5c811fdcef3e0704e35b34c1cfe1828de3/maturin-1.14.0-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:107f84110d890090a01bb1ecd01761fdfae925c23c659ba492c9b83dd179eab4", size = 10024058, upload-time = "2026-06-12T00:13:16.49Z" },
+    { url = "https://files.pythonhosted.org/packages/0c/e8/88a9d1872997d4535af10ebe79f550e834880bf613cf8e50b50d2d938e3b/maturin-1.14.0-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.musllinux_1_1_ppc64le.whl", hash = "sha256:9a84277aa907961cd47ad26fef1539e79efa30611972eaf7499606e773e991b2", size = 13302073, upload-time = "2026-06-12T00:13:29.027Z" },
+    { url = "https://files.pythonhosted.org/packages/4a/13/3f6d28bb7b744558b9bc78c995c1855d7e5ff21ad475f46d9de5c3dab039/maturin-1.14.0-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:095714b2a904927e3c868a1c5d078257ff0443c5049f7623777352966768306e", size = 10863616, upload-time = "2026-06-12T00:13:32.191Z" },
+    { url = "https://files.pythonhosted.org/packages/24/06/39352d2b402efa3a7dd01d4ed197b301ea35eec10208ba2b8c649101f4df/maturin-1.14.0-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:20229d332f87166b930e4ca07cdbee8a1726f2eea87a337610aa25bba3ddf4b4", size = 10399943, upload-time = "2026-06-12T00:13:36.273Z" },
+    { url = "https://files.pythonhosted.org/packages/58/77/641504541336240fef3836b2d15a785eaeb33c941fb118513c267dd70840/maturin-1.14.0-py3-none-win32.whl", hash = "sha256:4ba1e3c3f33609f461d587b7549104c81a15fd6d42ba63a73cea9376a1e9876e", size = 8905117, upload-time = "2026-06-12T00:13:18.38Z" },
+    { url = "https://files.pythonhosted.org/packages/02/4a/ca247a0c43069b2f48cf783c5b13c3a9eb92c8f596dc7fbdb9f75fea4414/maturin-1.14.0-py3-none-win_amd64.whl", hash = "sha256:cb09a313f097adeb4dda0082277871a28d1bd26615dbadab42e6234b6df6fe69", size = 10309099, upload-time = "2026-06-12T00:13:20.523Z" },
+    { url = "https://files.pythonhosted.org/packages/8b/a4/f14a3f6086cc3caaa90d12e832e4aa41de771c310041959f0d35dd4efe17/maturin-1.14.0-py3-none-win_arm64.whl", hash = "sha256:8c1a8188195f5b6ce1aab99ae2d92e342900298f901456b43ca028947fd3b288", size = 9719100, upload-time = "2026-06-12T00:13:24.741Z" },
+]
+
+[[package]]
+name = "tomli"
+version = "2.4.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/22/de/48c59722572767841493b26183a0d1cc411d54fd759c5607c4590b6563a6/tomli-2.4.1.tar.gz", hash = "sha256:7c7e1a961a0b2f2472c1ac5b69affa0ae1132c39adcb67aba98568702b9cc23f", size = 17543, upload-time = "2026-03-25T20:22:03.828Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f4/11/db3d5885d8528263d8adc260bb2d28ebf1270b96e98f0e0268d32b8d9900/tomli-2.4.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:f8f0fc26ec2cc2b965b7a3b87cd19c5c6b8c5e5f436b984e85f486d652285c30", size = 154704, upload-time = "2026-03-25T20:21:10.473Z" },
+    { url = "https://files.pythonhosted.org/packages/6d/f7/675db52c7e46064a9aa928885a9b20f4124ecb9bc2e1ce74c9106648d202/tomli-2.4.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:4ab97e64ccda8756376892c53a72bd1f964e519c77236368527f758fbc36a53a", size = 149454, upload-time = "2026-03-25T20:21:12.036Z" },
+    { url = "https://files.pythonhosted.org/packages/61/71/81c50943cf953efa35bce7646caab3cf457a7d8c030b27cfb40d7235f9ee/tomli-2.4.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:96481a5786729fd470164b47cdb3e0e58062a496f455ee41b4403be77cb5a076", size = 237561, upload-time = "2026-03-25T20:21:13.098Z" },
+    { url = "https://files.pythonhosted.org/packages/48/c1/f41d9cb618acccca7df82aaf682f9b49013c9397212cb9f53219e3abac37/tomli-2.4.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5a881ab208c0baf688221f8cecc5401bd291d67e38a1ac884d6736cbcd8247e9", size = 243824, upload-time = "2026-03-25T20:21:14.569Z" },
+    { url = "https://files.pythonhosted.org/packages/22/e4/5a816ecdd1f8ca51fb756ef684b90f2780afc52fc67f987e3c61d800a46d/tomli-2.4.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:47149d5bd38761ac8be13a84864bf0b7b70bc051806bc3669ab1cbc56216b23c", size = 242227, upload-time = "2026-03-25T20:21:15.712Z" },
+    { url = "https://files.pythonhosted.org/packages/6b/49/2b2a0ef529aa6eec245d25f0c703e020a73955ad7edf73e7f54ddc608aa5/tomli-2.4.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:ec9bfaf3ad2df51ace80688143a6a4ebc09a248f6ff781a9945e51937008fcbc", size = 247859, upload-time = "2026-03-25T20:21:17.001Z" },
+    { url = "https://files.pythonhosted.org/packages/83/bd/6c1a630eaca337e1e78c5903104f831bda934c426f9231429396ce3c3467/tomli-2.4.1-cp311-cp311-win32.whl", hash = "sha256:ff2983983d34813c1aeb0fa89091e76c3a22889ee83ab27c5eeb45100560c049", size = 97204, upload-time = "2026-03-25T20:21:18.079Z" },
+    { url = "https://files.pythonhosted.org/packages/42/59/71461df1a885647e10b6bb7802d0b8e66480c61f3f43079e0dcd315b3954/tomli-2.4.1-cp311-cp311-win_amd64.whl", hash = "sha256:5ee18d9ebdb417e384b58fe414e8d6af9f4e7a0ae761519fb50f721de398dd4e", size = 108084, upload-time = "2026-03-25T20:21:18.978Z" },
+    { url = "https://files.pythonhosted.org/packages/b8/83/dceca96142499c069475b790e7913b1044c1a4337e700751f48ed723f883/tomli-2.4.1-cp311-cp311-win_arm64.whl", hash = "sha256:c2541745709bad0264b7d4705ad453b76ccd191e64aa6f0fc66b69a293a45ece", size = 95285, upload-time = "2026-03-25T20:21:20.309Z" },
+    { url = "https://files.pythonhosted.org/packages/c1/ba/42f134a3fe2b370f555f44b1d72feebb94debcab01676bf918d0cb70e9aa/tomli-2.4.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:c742f741d58a28940ce01d58f0ab2ea3ced8b12402f162f4d534dfe18ba1cd6a", size = 155924, upload-time = "2026-03-25T20:21:21.626Z" },
+    { url = "https://files.pythonhosted.org/packages/dc/c7/62d7a17c26487ade21c5422b646110f2162f1fcc95980ef7f63e73c68f14/tomli-2.4.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:7f86fd587c4ed9dd76f318225e7d9b29cfc5a9d43de44e5754db8d1128487085", size = 150018, upload-time = "2026-03-25T20:21:23.002Z" },
+    { url = "https://files.pythonhosted.org/packages/5c/05/79d13d7c15f13bdef410bdd49a6485b1c37d28968314eabee452c22a7fda/tomli-2.4.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ff18e6a727ee0ab0388507b89d1bc6a22b138d1e2fa56d1ad494586d61d2eae9", size = 244948, upload-time = "2026-03-25T20:21:24.04Z" },
+    { url = "https://files.pythonhosted.org/packages/10/90/d62ce007a1c80d0b2c93e02cab211224756240884751b94ca72df8a875ca/tomli-2.4.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:136443dbd7e1dee43c68ac2694fde36b2849865fa258d39bf822c10e8068eac5", size = 253341, upload-time = "2026-03-25T20:21:25.177Z" },
+    { url = "https://files.pythonhosted.org/packages/1a/7e/caf6496d60152ad4ed09282c1885cca4eea150bfd007da84aea07bcc0a3e/tomli-2.4.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:5e262d41726bc187e69af7825504c933b6794dc3fbd5945e41a79bb14c31f585", size = 248159, upload-time = "2026-03-25T20:21:26.364Z" },
+    { url = "https://files.pythonhosted.org/packages/99/e7/c6f69c3120de34bbd882c6fba7975f3d7a746e9218e56ab46a1bc4b42552/tomli-2.4.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:5cb41aa38891e073ee49d55fbc7839cfdb2bc0e600add13874d048c94aadddd1", size = 253290, upload-time = "2026-03-25T20:21:27.46Z" },
+    { url = "https://files.pythonhosted.org/packages/d6/2f/4a3c322f22c5c66c4b836ec58211641a4067364f5dcdd7b974b4c5da300c/tomli-2.4.1-cp312-cp312-win32.whl", hash = "sha256:da25dc3563bff5965356133435b757a795a17b17d01dbc0f42fb32447ddfd917", size = 98141, upload-time = "2026-03-25T20:21:28.492Z" },
+    { url = "https://files.pythonhosted.org/packages/24/22/4daacd05391b92c55759d55eaee21e1dfaea86ce5c571f10083360adf534/tomli-2.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:52c8ef851d9a240f11a88c003eacb03c31fc1c9c4ec64a99a0f922b93874fda9", size = 108847, upload-time = "2026-03-25T20:21:29.386Z" },
+    { url = "https://files.pythonhosted.org/packages/68/fd/70e768887666ddd9e9f5d85129e84910f2db2796f9096aa02b721a53098d/tomli-2.4.1-cp312-cp312-win_arm64.whl", hash = "sha256:f758f1b9299d059cc3f6546ae2af89670cb1c4d48ea29c3cacc4fe7de3058257", size = 95088, upload-time = "2026-03-25T20:21:30.677Z" },
+    { url = "https://files.pythonhosted.org/packages/07/06/b823a7e818c756d9a7123ba2cda7d07bc2dd32835648d1a7b7b7a05d848d/tomli-2.4.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:36d2bd2ad5fb9eaddba5226aa02c8ec3fa4f192631e347b3ed28186d43be6b54", size = 155866, upload-time = "2026-03-25T20:21:31.65Z" },
+    { url = "https://files.pythonhosted.org/packages/14/6f/12645cf7f08e1a20c7eb8c297c6f11d31c1b50f316a7e7e1e1de6e2e7b7e/tomli-2.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:eb0dc4e38e6a1fd579e5d50369aa2e10acfc9cace504579b2faabb478e76941a", size = 149887, upload-time = "2026-03-25T20:21:33.028Z" },
+    { url = "https://files.pythonhosted.org/packages/5c/e0/90637574e5e7212c09099c67ad349b04ec4d6020324539297b634a0192b0/tomli-2.4.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7f2c7f2b9ca6bdeef8f0fa897f8e05085923eb091721675170254cbc5b02897", size = 243704, upload-time = "2026-03-25T20:21:34.51Z" },
+    { url = "https://files.pythonhosted.org/packages/10/8f/d3ddb16c5a4befdf31a23307f72828686ab2096f068eaf56631e136c1fdd/tomli-2.4.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f3c6818a1a86dd6dca7ddcaaf76947d5ba31aecc28cb1b67009a5877c9a64f3f", size = 251628, upload-time = "2026-03-25T20:21:36.012Z" },
+    { url = "https://files.pythonhosted.org/packages/e3/f1/dbeeb9116715abee2485bf0a12d07a8f31af94d71608c171c45f64c0469d/tomli-2.4.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d312ef37c91508b0ab2cee7da26ec0b3ed2f03ce12bd87a588d771ae15dcf82d", size = 247180, upload-time = "2026-03-25T20:21:37.136Z" },
+    { url = "https://files.pythonhosted.org/packages/d3/74/16336ffd19ed4da28a70959f92f506233bd7cfc2332b20bdb01591e8b1d1/tomli-2.4.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:51529d40e3ca50046d7606fa99ce3956a617f9b36380da3b7f0dd3dd28e68cb5", size = 251674, upload-time = "2026-03-25T20:21:38.298Z" },
+    { url = "https://files.pythonhosted.org/packages/16/f9/229fa3434c590ddf6c0aa9af64d3af4b752540686cace29e6281e3458469/tomli-2.4.1-cp313-cp313-win32.whl", hash = "sha256:2190f2e9dd7508d2a90ded5ed369255980a1bcdd58e52f7fe24b8162bf9fedbd", size = 97976, upload-time = "2026-03-25T20:21:39.316Z" },
+    { url = "https://files.pythonhosted.org/packages/6a/1e/71dfd96bcc1c775420cb8befe7a9d35f2e5b1309798f009dca17b7708c1e/tomli-2.4.1-cp313-cp313-win_amd64.whl", hash = "sha256:8d65a2fbf9d2f8352685bc1364177ee3923d6baf5e7f43ea4959d7d8bc326a36", size = 108755, upload-time = "2026-03-25T20:21:40.248Z" },
+    { url = "https://files.pythonhosted.org/packages/83/7a/d34f422a021d62420b78f5c538e5b102f62bea616d1d75a13f0a88acb04a/tomli-2.4.1-cp313-cp313-win_arm64.whl", hash = "sha256:4b605484e43cdc43f0954ddae319fb75f04cc10dd80d830540060ee7cd0243cd", size = 95265, upload-time = "2026-03-25T20:21:41.219Z" },
+    { url = "https://files.pythonhosted.org/packages/3c/fb/9a5c8d27dbab540869f7c1f8eb0abb3244189ce780ba9cd73f3770662072/tomli-2.4.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:fd0409a3653af6c147209d267a0e4243f0ae46b011aa978b1080359fddc9b6cf", size = 155726, upload-time = "2026-03-25T20:21:42.23Z" },
+    { url = "https://files.pythonhosted.org/packages/62/05/d2f816630cc771ad836af54f5001f47a6f611d2d39535364f148b6a92d6b/tomli-2.4.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:a120733b01c45e9a0c34aeef92bf0cf1d56cfe81ed9d47d562f9ed591a9828ac", size = 149859, upload-time = "2026-03-25T20:21:43.386Z" },
+    { url = "https://files.pythonhosted.org/packages/ce/48/66341bdb858ad9bd0ceab5a86f90eddab127cf8b046418009f2125630ecb/tomli-2.4.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:559db847dc486944896521f68d8190be1c9e719fced785720d2216fe7022b662", size = 244713, upload-time = "2026-03-25T20:21:44.474Z" },
+    { url = "https://files.pythonhosted.org/packages/df/6d/c5fad00d82b3c7a3ab6189bd4b10e60466f22cfe8a08a9394185c8a8111c/tomli-2.4.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:01f520d4f53ef97964a240a035ec2a869fe1a37dde002b57ebc4417a27ccd853", size = 252084, upload-time = "2026-03-25T20:21:45.62Z" },
+    { url = "https://files.pythonhosted.org/packages/00/71/3a69e86f3eafe8c7a59d008d245888051005bd657760e96d5fbfb0b740c2/tomli-2.4.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7f94b27a62cfad8496c8d2513e1a222dd446f095fca8987fceef261225538a15", size = 247973, upload-time = "2026-03-25T20:21:46.937Z" },
+    { url = "https://files.pythonhosted.org/packages/67/50/361e986652847fec4bd5e4a0208752fbe64689c603c7ae5ea7cb16b1c0ca/tomli-2.4.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:ede3e6487c5ef5d28634ba3f31f989030ad6af71edfb0055cbbd14189ff240ba", size = 256223, upload-time = "2026-03-25T20:21:48.467Z" },
+    { url = "https://files.pythonhosted.org/packages/8c/9a/b4173689a9203472e5467217e0154b00e260621caa227b6fa01feab16998/tomli-2.4.1-cp314-cp314-win32.whl", hash = "sha256:3d48a93ee1c9b79c04bb38772ee1b64dcf18ff43085896ea460ca8dec96f35f6", size = 98973, upload-time = "2026-03-25T20:21:49.526Z" },
+    { url = "https://files.pythonhosted.org/packages/14/58/640ac93bf230cd27d002462c9af0d837779f8773bc03dee06b5835208214/tomli-2.4.1-cp314-cp314-win_amd64.whl", hash = "sha256:88dceee75c2c63af144e456745e10101eb67361050196b0b6af5d717254dddf7", size = 109082, upload-time = "2026-03-25T20:21:50.506Z" },
+    { url = "https://files.pythonhosted.org/packages/d5/2f/702d5e05b227401c1068f0d386d79a589bb12bf64c3d2c72ce0631e3bc49/tomli-2.4.1-cp314-cp314-win_arm64.whl", hash = "sha256:b8c198f8c1805dc42708689ed6864951fd2494f924149d3e4bce7710f8eb5232", size = 96490, upload-time = "2026-03-25T20:21:51.474Z" },
+    { url = "https://files.pythonhosted.org/packages/45/4b/b877b05c8ba62927d9865dd980e34a755de541eb65fffba52b4cc495d4d2/tomli-2.4.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:d4d8fe59808a54658fcc0160ecfb1b30f9089906c50b23bcb4c69eddc19ec2b4", size = 164263, upload-time = "2026-03-25T20:21:52.543Z" },
+    { url = "https://files.pythonhosted.org/packages/24/79/6ab420d37a270b89f7195dec5448f79400d9e9c1826df982f3f8e97b24fd/tomli-2.4.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7008df2e7655c495dd12d2a4ad038ff878d4ca4b81fccaf82b714e07eae4402c", size = 160736, upload-time = "2026-03-25T20:21:53.674Z" },
+    { url = "https://files.pythonhosted.org/packages/02/e0/3630057d8eb170310785723ed5adcdfb7d50cb7e6455f85ba8a3deed642b/tomli-2.4.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1d8591993e228b0c930c4bb0db464bdad97b3289fb981255d6c9a41aedc84b2d", size = 270717, upload-time = "2026-03-25T20:21:55.129Z" },
+    { url = "https://files.pythonhosted.org/packages/7a/b4/1613716072e544d1a7891f548d8f9ec6ce2faf42ca65acae01d76ea06bb0/tomli-2.4.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:734e20b57ba95624ecf1841e72b53f6e186355e216e5412de414e3c51e5e3c41", size = 278461, upload-time = "2026-03-25T20:21:56.228Z" },
+    { url = "https://files.pythonhosted.org/packages/05/38/30f541baf6a3f6df77b3df16b01ba319221389e2da59427e221ef417ac0c/tomli-2.4.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8a650c2dbafa08d42e51ba0b62740dae4ecb9338eefa093aa5c78ceb546fcd5c", size = 274855, upload-time = "2026-03-25T20:21:57.653Z" },
+    { url = "https://files.pythonhosted.org/packages/77/a3/ec9dd4fd2c38e98de34223b995a3b34813e6bdadf86c75314c928350ed14/tomli-2.4.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:504aa796fe0569bb43171066009ead363de03675276d2d121ac1a4572397870f", size = 283144, upload-time = "2026-03-25T20:21:59.089Z" },
+    { url = "https://files.pythonhosted.org/packages/ef/be/605a6261cac79fba2ec0c9827e986e00323a1945700969b8ee0b30d85453/tomli-2.4.1-cp314-cp314t-win32.whl", hash = "sha256:b1d22e6e9387bf4739fbe23bfa80e93f6b0373a7f1b96c6227c32bef95a4d7a8", size = 108683, upload-time = "2026-03-25T20:22:00.214Z" },
+    { url = "https://files.pythonhosted.org/packages/12/64/da524626d3b9cc40c168a13da8335fe1c51be12c0a63685cc6db7308daae/tomli-2.4.1-cp314-cp314t-win_amd64.whl", hash = "sha256:2c1c351919aca02858f740c6d33adea0c5deea37f9ecca1cc1ef9e884a619d26", size = 121196, upload-time = "2026-03-25T20:22:01.169Z" },
+    { url = "https://files.pythonhosted.org/packages/5a/cd/e80b62269fc78fc36c9af5a6b89c835baa8af28ff5ad28c7028d60860320/tomli-2.4.1-cp314-cp314t-win_arm64.whl", hash = "sha256:eab21f45c7f66c13f2a9e0e1535309cee140182a9cdae1e041d02e47291e8396", size = 100393, upload-time = "2026-03-25T20:22:02.137Z" },
+    { url = "https://files.pythonhosted.org/packages/7b/61/cceae43728b7de99d9b847560c262873a1f6c98202171fd5ed62640b494b/tomli-2.4.1-py3-none-any.whl", hash = "sha256:0d85819802132122da43cb86656f8d1f8c6587d54ae7dcaf30e90533028b49fe", size = 14583, upload-time = "2026-03-25T20:22:03.012Z" },
+]

pyproject.toml

@@ -1,12 +1,12 @@
-[build-system]
-requires = ["maturin>=1,<2"]
-build-backend = "maturin"
-
 [project]
-name = "flowstate-r1-rs"
+name = "flowstate-r1-env"
 version = "0.1.0"
-requires-python = ">=3.8"
+requires-python = ">=3.11"
+dependencies = [
+    "granite-tsfm",
+    "numpy",
+    "torch",
+]
 
-[tool.maturin]
-features = ["python"]
-module-name = "flowstate_r1_rs"
+[tool.uv]
+package = false

scripts/compare_python.py

@@ -1,138 +0,0 @@
-#!/usr/bin/env python3
-"""
-Numerical comparison: Rust flowstate-r1-rs vs Python FlowStateForPrediction.
-
-Both engines use the same local safetensors checkpoint (models/).
-Input is a fixed 256-step sine+trend series; prediction_length=24.
-
-Usage:
-    python3 scripts/compare_python.py \
-        --model-dir models \
-        --gguf gguf/flowstate-r1-f32.gguf \
-        --config models/config.json \
-        --horizon 24
-"""
-
-import argparse
-import json
-import math
-import os
-import subprocess
-import sys
-
-import numpy as np
-
-
-def make_test_series(n: int = 256) -> np.ndarray:
-    t = np.arange(n, dtype=np.float32)
-    return np.sin(2 * math.pi * t / 48) + 0.02 * t
-
-
-def run_python(model_dir: str, series: np.ndarray, horizon: int) -> np.ndarray:
-    try:
-        import torch
-        from tsfm_public import FlowStateForPrediction
-    except ImportError:
-        print("tsfm_public not installed. Install with:")
-        print("  pip install git+https://github.com/ibm-granite/granite-tsfm.git")
-        sys.exit(1)
-
-    predictor = FlowStateForPrediction.from_pretrained(model_dir, torch_dtype=torch.float32)
-    predictor.model.eval()
-
-    with open(os.path.join(model_dir, "config.json")) as f:
-        cfg = json.load(f)
-    decoder_patch_len = cfg["decoder_patch_len"]
-    scale_factor = decoder_patch_len / horizon
-
-    with torch.no_grad():
-        ts = torch.tensor(series, dtype=torch.float32).unsqueeze(-1).unsqueeze(1)
-        preds = predictor(
-            ts,
-            scale_factor=scale_factor,
-            prediction_length=horizon,
-            batch_first=False,
-        )
-
-    if hasattr(preds, "quantile_outputs") and preds.quantile_outputs is not None:
-        arr = preds.quantile_outputs[0, :, :, 0].detach().numpy()  # [n_q, pred_len]
-    elif hasattr(preds, "last_hidden_state"):
-        arr = preds.last_hidden_state.squeeze().detach().numpy()
-    else:
-        arr = preds.squeeze().detach().numpy()
-
-    return arr.astype(np.float32)
-
-
-def run_rust(binary: str, gguf: str, config: str, series: np.ndarray,
-             horizon: int) -> np.ndarray:
-    request = json.dumps({"context": series.tolist(), "horizon": horizon})
-    result = subprocess.run(
-        [binary, "infer", "--gguf", gguf, "--config", config],
-        input=request, capture_output=True, text=True, check=True,
-    )
-    fc = json.loads(result.stdout)["choices"][0]["forecast"]
-    quants = fc.get("quantiles", {})
-    keys = sorted(quants.keys(), key=float)
-    rows = [np.array(quants[k], dtype=np.float32) for k in keys]
-    if not rows:
-        return np.array(fc["point"], dtype=np.float32)[None, :]
-    return np.stack(rows)  # (n_quantiles, horizon)
-
-
-def report(py_preds: np.ndarray, rs_preds: np.ndarray, quantile_labels) -> None:
-    diff = np.abs(py_preds - rs_preds)
-    print(f"\n{'Quantile':<12} {'MaxAbsErr':>12} {'MeanAbsErr':>12} {'MedianAbsErr':>14}")
-    print("-" * 54)
-    for i, q in enumerate(quantile_labels):
-        d = diff[i]
-        print(f"{q:<12} {d.max():>12.6f} {d.mean():>12.6f} {np.median(d):>14.6f}")
-    print("-" * 54)
-    print(f"{'ALL':12} {diff.max():>12.6f} {diff.mean():>12.6f} {np.median(diff):>14.6f}")
-
-    for idx in [0, len(quantile_labels) // 2, len(quantile_labels) - 1]:
-        q = quantile_labels[idx]
-        py = py_preds[idx, :8]
-        rs = rs_preds[idx, :8]
-        print(f"\nSample [{q}] first 8 steps:")
-        print(f"  Python: {' '.join(f'{v:.4f}' for v in py)}")
-        print(f"    Rust: {' '.join(f'{v:.4f}' for v in rs)}")
-        print(f"    Diff: {' '.join(f'{abs(a-b):.4f}' for a, b in zip(py, rs))}")
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model-dir", default="models")
-    parser.add_argument("--gguf", default="gguf/flowstate-r1-f32.gguf")
-    parser.add_argument("--config", default="models/config.json")
-    parser.add_argument("--horizon", type=int, default=24)
-    parser.add_argument("--binary", default="./target/release/flowstate-r1-rs")
-    args = parser.parse_args()
-
-    with open(args.config) as f:
-        cfg = json.load(f)
-    quantiles = cfg["quantiles"]
-    quantile_labels = [f"q{q:.2f}" for q in quantiles]
-
-    series = make_test_series(256)
-
-    print(f"Input: {len(series)}-step sine+trend  |  Horizon: {args.horizon}")
-    print(f"Quantiles: {quantile_labels}")
-
-    print("\n--- Running Python (HuggingFace safetensors) ---")
-    py_preds = run_python(args.model_dir, series, args.horizon)
-    print(f"Python output shape: {py_preds.shape}")
-
-    print("\n--- Running Rust (GGUF) ---")
-    rs_preds = run_rust(args.binary, args.gguf, args.config, series, args.horizon)
-    print(f"Rust output shape: {rs_preds.shape}")
-
-    if py_preds.shape != rs_preds.shape:
-        print(f"Shape mismatch: Python {py_preds.shape} vs Rust {rs_preds.shape}")
-        sys.exit(1)
-
-    report(py_preds, rs_preds, quantile_labels)
-
-
-if __name__ == "__main__":
-    main()

scripts/convert_all.sh

@@ -1,30 +0,0 @@
-#!/usr/bin/env bash
-# Download ibm-granite/granite-timeseries-flowstate-r1 and convert to GGUF in all dtypes.
-set -euo pipefail
-cd "$(dirname "$0")/.."
-
-BIN=./target/release/flowstate-r1-rs
-MODEL=${1:-ibm-granite/granite-timeseries-flowstate-r1}
-MODEL_DIR=models
-TOKEN_ARG=""
-if [[ -n "${HF_TOKEN:-}" ]]; then
-  TOKEN_ARG="--token $HF_TOKEN"
-fi
-
-if [[ ! -x "$BIN" ]]; then
-  echo "Binary not found — run: cargo build --release"
-  exit 1
-fi
-
-mkdir -p gguf
-
-for DTYPE in f32 f16 q8; do
-  OUT="gguf/flowstate-r1-${DTYPE}.gguf"
-  echo "=== Converting ${MODEL} → ${OUT} (dtype=${DTYPE}) ==="
-  $BIN convert --model "$MODEL" --dtype "$DTYPE" \
-    --model-dir "$MODEL_DIR" --output "$OUT" $TOKEN_ARG
-  ls -lh "$OUT"
-done
-
-echo "Done. Files:"
-ls -lh gguf/flowstate-r1-*.gguf

scripts/upload.sh

@@ -1,20 +0,0 @@
-#!/usr/bin/env bash
-# Upload source + GGUF files to HuggingFace Hub via the built-in upload subcommand.
-#
-# Usage:
-#   HF_TOKEN=hf_... ./scripts/upload.sh [--repo owner/repo-name]
-#
-set -euo pipefail
-cd "$(dirname "$0")/.."
-
-BIN=./target/release/flowstate-r1-rs
-
-if [[ ! -f "$BIN" ]]; then
-  echo "Binary not found — building release …"
-  cargo build --release
-fi
-
-exec "$BIN" upload \
-  --repo "amaye15/flowstate-r1-gguf" \
-  ${HF_TOKEN:+--token "$HF_TOKEN"} \
-  "$@"

src/config.rs

@@ -1,51 +0,0 @@
-use serde::Deserialize;
-
-/// Top-level FlowState `config.json`.
-#[derive(Debug, Deserialize, Clone)]
-pub struct FlowStateConfig {
-    pub context_length: u32,
-    pub decoder_dim: u32,
-    pub decoder_patch_len: u32,
-    pub decoder_type: String,
-    pub embedding_feature_dim: u32,
-    pub encoder_num_hippo_blocks: u32,
-    pub encoder_num_layers: u32,
-    pub encoder_state_dim: u32,
-    pub quantiles: Vec<f32>,
-    #[serde(default = "default_bool_true")]
-    pub with_missing: bool,
-    #[serde(default = "default_bool_true")]
-    pub use_freq: bool,
-    #[serde(default = "default_bool_true")]
-    pub init_processing: bool,
-    #[serde(default = "default_u32_2048")]
-    pub min_context: u32,
-}
-
-impl FlowStateConfig {
-    pub fn from_json(s: &str) -> anyhow::Result<Self> {
-        Ok(serde_json::from_str(s)?)
-    }
-
-    pub fn n_quantiles(&self) -> u32 {
-        self.quantiles.len() as u32
-    }
-
-    /// Number of input channels (value + missing mask).
-    pub fn n_inputs(&self) -> u32 {
-        if self.with_missing { 2 } else { 1 }
-    }
-
-    /// Legendre basis range for "legs" / "hlegs" decoder.
-    pub fn basis_range(&self) -> [f32; 2] {
-        let dt = self.decoder_type.to_lowercase();
-        if dt == "hlegs" {
-            [0.0, 0.95]
-        } else {
-            [-1.0, 0.95]
-        }
-    }
-}
-
-fn default_bool_true() -> bool { true }
-fn default_u32_2048() -> u32 { 2048 }

src/convert.rs

@@ -1,265 +0,0 @@
-use std::fs::File;
-use std::io::BufWriter;
-use std::path::Path;
-
-use anyhow::Context;
-use indicatif::{ProgressBar, ProgressStyle};
-use safetensors::SafeTensors;
-use safetensors::Dtype as StDtype;
-
-use crate::config::FlowStateConfig;
-use crate::download::ModelFiles;
-use crate::gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
-use crate::tensor_map::map_tensor_name;
-
-pub struct ConvertOptions {
-    pub output_dtype: GGMLType,
-}
-
-pub fn convert(
-    model_id: &str,
-    files: &ModelFiles,
-    config: &FlowStateConfig,
-    opts: &ConvertOptions,
-    output_path: &Path,
-) -> anyhow::Result<()> {
-    let mut writer = GGUFWriter::new();
-    write_metadata(&mut writer, model_id, config);
-
-    let shard_bytes = load_shard_bytes(&files.safetensors_shards)?;
-    let shard_views: Vec<SafeTensors> = shard_bytes
-        .iter()
-        .map(|b| SafeTensors::deserialize(b).context("deserialize shard"))
-        .collect::<anyhow::Result<_>>()?;
-
-    let total_tensors: usize = shard_views.iter().map(|s| s.len()).sum();
-    println!("Found {} tensors across {} shard(s).", total_tensors, shard_views.len());
-
-    let pb = ProgressBar::new(total_tensors as u64);
-    pb.set_style(
-        ProgressStyle::with_template(
-            "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} {msg}",
-        )
-        .unwrap()
-        .progress_chars("=>-"),
-    );
-
-    let mut mapped = 0usize;
-    let mut skipped: Vec<String> = Vec::new();
-    let mut fallback_count = 0usize;
-
-    for shard in &shard_views {
-        for (hf_name, tensor_view) in shard.tensors() {
-            pb.set_message(hf_name.to_string());
-
-            let gguf_name = match map_tensor_name(&hf_name) {
-                Some(n) => n,
-                None => {
-                    skipped.push(hf_name.to_string());
-                    pb.inc(1);
-                    continue;
-                }
-            };
-
-            let src_dtype = ggml_type_from_st(tensor_view.dtype())
-                .with_context(|| format!("tensor {hf_name}: unsupported dtype {:?}", tensor_view.dtype()))?;
-
-            let raw_data = tensor_view.data();
-            let py_shape = tensor_view.shape();
-            let n_elems: usize = py_shape.iter().product();
-            let innermost = py_shape.last().copied().unwrap_or(1);
-
-            let (dst_dtype, gguf_shape, tensor_data) =
-                if opts.output_dtype == GGMLType::Q8_0 && (innermost % 32 != 0 || n_elems % 32 != 0) {
-                    fallback_count += 1;
-                    let data = cast_data(raw_data, src_dtype, GGMLType::F32)
-                        .with_context(|| format!("tensor {hf_name}: cast failed"))?;
-                    let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
-                    (GGMLType::F32, gs, data)
-                } else {
-                    let dst = opts.output_dtype;
-                    let data = cast_data(raw_data, src_dtype, dst)
-                        .with_context(|| format!("tensor {hf_name}: cast failed"))?;
-                    let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
-                    (dst, gs, data)
-                };
-
-            writer.add_tensor(gguf_name, gguf_shape, dst_dtype, tensor_data);
-            mapped += 1;
-            pb.inc(1);
-        }
-    }
-
-    pb.finish_with_message("tensors processed");
-
-    if !skipped.is_empty() {
-        eprintln!("\nWarning: {} tensor(s) skipped (not mapped):", skipped.len());
-        for name in &skipped {
-            eprintln!("  {name}");
-        }
-    }
-    if fallback_count > 0 {
-        eprintln!("\nNote: {fallback_count} tensor(s) fell back to F32 (too small for Q8_0 blocks).");
-    }
-
-    println!("Writing {mapped} tensors to {} …", output_path.display());
-    let out_file = File::create(output_path)
-        .with_context(|| format!("create {}", output_path.display()))?;
-    let mut buf_writer = BufWriter::new(out_file);
-    writer.write_to(&mut buf_writer)?;
-    println!("Done.");
-    Ok(())
-}
-
-fn write_metadata(writer: &mut GGUFWriter, model_id: &str, config: &FlowStateConfig) {
-    writer.add_metadata("general.architecture", GGUFMetaValue::String("flowstate".into()));
-    writer.add_metadata("general.name",         GGUFMetaValue::String(model_id.into()));
-
-    writer.add_metadata("flowstate.block_count",        GGUFMetaValue::Uint32(config.encoder_num_layers));
-    writer.add_metadata("flowstate.embedding_length",   GGUFMetaValue::Uint32(config.embedding_feature_dim));
-    writer.add_metadata("flowstate.state_dim",          GGUFMetaValue::Uint32(config.encoder_state_dim));
-    writer.add_metadata("flowstate.num_hippo_blocks",   GGUFMetaValue::Uint32(config.encoder_num_hippo_blocks));
-    writer.add_metadata("flowstate.context_length",     GGUFMetaValue::Uint32(config.context_length));
-    writer.add_metadata("flowstate.min_context",        GGUFMetaValue::Uint32(config.min_context));
-    writer.add_metadata("flowstate.decoder_dim",        GGUFMetaValue::Uint32(config.decoder_dim));
-    writer.add_metadata("flowstate.decoder_patch_len",  GGUFMetaValue::Uint32(config.decoder_patch_len));
-    writer.add_metadata("flowstate.decoder_type",       GGUFMetaValue::String(config.decoder_type.clone()));
-    writer.add_metadata("flowstate.quantile_count",     GGUFMetaValue::Uint32(config.n_quantiles()));
-    writer.add_metadata("flowstate.quantiles",          GGUFMetaValue::ArrayFloat32(config.quantiles.clone()));
-    writer.add_metadata("flowstate.with_missing",       GGUFMetaValue::Bool(config.with_missing));
-
-    let range = config.basis_range();
-    writer.add_metadata("flowstate.basis_range_low",    GGUFMetaValue::Float32(range[0]));
-    writer.add_metadata("flowstate.basis_range_high",   GGUFMetaValue::Float32(range[1]));
-}
-
-fn ggml_type_from_st(dtype: StDtype) -> anyhow::Result<GGMLType> {
-    match dtype {
-        StDtype::F32  => Ok(GGMLType::F32),
-        StDtype::F16  => Ok(GGMLType::F16),
-        StDtype::BF16 => Ok(GGMLType::BF16),
-        other => anyhow::bail!("unsupported safetensors dtype: {other:?}"),
-    }
-}
-
-fn load_shard_bytes(shards: &[std::path::PathBuf]) -> anyhow::Result<Vec<Vec<u8>>> {
-    shards
-        .iter()
-        .map(|p| std::fs::read(p).with_context(|| format!("read shard {}", p.display())))
-        .collect()
-}
-
-fn cast_data(data: &[u8], src: GGMLType, dst: GGMLType) -> anyhow::Result<Vec<u8>> {
-    if src == dst { return Ok(data.to_vec()); }
-    if dst == GGMLType::Q8_0 {
-        let f32_values = decode_to_f32(data, src)?;
-        return quantize_q8_0(&f32_values);
-    }
-    match (src, dst) {
-        (GGMLType::F32, GGMLType::F16) => {
-            let f32_values = parse_f32_le(data)?;
-            let mut out = Vec::with_capacity(f32_values.len() * 2);
-            for v in f32_values { let bits = f32_to_f16_bits(v); out.extend_from_slice(&bits.to_le_bytes()); }
-            Ok(out)
-        }
-        (GGMLType::BF16, GGMLType::F32) => {
-            let mut out = Vec::with_capacity(data.len() * 2);
-            for chunk in data.chunks_exact(2) {
-                let bf16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
-                out.extend_from_slice(&((bf16_bits as u32) << 16).to_le_bytes());
-            }
-            Ok(out)
-        }
-        (GGMLType::BF16, GGMLType::F16) => {
-            let mut out = Vec::with_capacity(data.len());
-            for chunk in data.chunks_exact(2) {
-                let bf16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
-                let f32_val = f32::from_bits((bf16_bits as u32) << 16);
-                out.extend_from_slice(&f32_to_f16_bits(f32_val).to_le_bytes());
-            }
-            Ok(out)
-        }
-        (GGMLType::F16, GGMLType::F32) => {
-            let mut out = Vec::with_capacity(data.len() * 2);
-            for chunk in data.chunks_exact(2) {
-                let f16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
-                out.extend_from_slice(&f16_to_f32(f16_bits).to_bits().to_le_bytes());
-            }
-            Ok(out)
-        }
-        _ => anyhow::bail!("unsupported cast: {src:?} → {dst:?}"),
-    }
-}
-
-fn decode_to_f32(data: &[u8], src: GGMLType) -> anyhow::Result<Vec<f32>> {
-    match src {
-        GGMLType::F32  => parse_f32_le(data),
-        GGMLType::F16  => data.chunks_exact(2).map(|c| Ok(f16_to_f32(u16::from_le_bytes([c[0], c[1]])))).collect(),
-        GGMLType::BF16 => data.chunks_exact(2).map(|c| {
-            let bits = u16::from_le_bytes([c[0], c[1]]);
-            Ok(f32::from_bits((bits as u32) << 16))
-        }).collect(),
-        GGMLType::Q8_0 => anyhow::bail!("Q8_0 → Q8_0 re-quantization not supported"),
-    }
-}
-
-fn quantize_q8_0(values: &[f32]) -> anyhow::Result<Vec<u8>> {
-    const BLOCK: usize = 32;
-    if values.len() % BLOCK != 0 {
-        anyhow::bail!("Q8_0 requires element count divisible by {BLOCK}, got {}", values.len());
-    }
-    let n_blocks = values.len() / BLOCK;
-    let mut out = vec![0u8; n_blocks * 34];
-    for b in 0..n_blocks {
-        let blk = &values[b * BLOCK..(b + 1) * BLOCK];
-        let amax = blk.iter().copied().map(f32::abs).fold(0.0f32, f32::max);
-        let d = if amax == 0.0 { 0.0f32 } else { amax / 127.0 };
-        let d_inv = if d == 0.0 { 0.0f32 } else { 1.0 / d };
-        let base = b * 34;
-        out[base..base + 2].copy_from_slice(&f32_to_f16_bits(d).to_le_bytes());
-        for i in 0..BLOCK {
-            out[base + 2 + i] = (blk[i] * d_inv).round().clamp(-127.0, 127.0) as i8 as u8;
-        }
-    }
-    Ok(out)
-}
-
-fn parse_f32_le(data: &[u8]) -> anyhow::Result<Vec<f32>> {
-    if data.len() % 4 != 0 { anyhow::bail!("f32 data length not divisible by 4"); }
-    Ok(data.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect())
-}
-
-fn f32_to_f16_bits(v: f32) -> u16 {
-    let bits = v.to_bits();
-    let sign = ((bits >> 16) & 0x8000) as u16;
-    let exp = ((bits >> 23) & 0xFF) as i32;
-    let mantissa = bits & 0x007F_FFFF;
-    if exp == 0xFF { return sign | 0x7C00 | if mantissa != 0 { 0x0200 } else { 0 }; }
-    let new_exp = exp - 127 + 15;
-    if new_exp >= 31 { return sign | 0x7C00; }
-    if new_exp <= 0 {
-        if new_exp < -10 { return sign; }
-        let m = (mantissa | 0x0080_0000) >> (1 - new_exp);
-        return sign | (m >> 13) as u16;
-    }
-    sign | ((new_exp as u16) << 10) | (mantissa >> 13) as u16
-}
-
-fn f16_to_f32(bits: u16) -> f32 {
-    let sign = ((bits & 0x8000) as u32) << 16;
-    let exp = ((bits >> 10) & 0x1F) as i32;
-    let mantissa = (bits & 0x03FF) as u32;
-    let f32_bits = if exp == 0 {
-        if mantissa == 0 { sign }
-        else {
-            let mut m = mantissa; let mut e = 0i32;
-            while m & 0x0400 == 0 { m <<= 1; e += 1; }
-            sign | ((127 - 15 - e + 1) as u32) << 23 | (m & 0x03FF) << 13
-        }
-    } else if exp == 31 {
-        sign | 0x7F80_0000 | (mantissa << 13)
-    } else {
-        sign | ((exp + 127 - 15) as u32) << 23 | (mantissa << 13)
-    };
-    f32::from_bits(f32_bits)
-}

src/download.rs

@@ -1,182 +0,0 @@
-use std::path::{Path, PathBuf};
-
-use anyhow::Context;
-use futures_util::StreamExt;
-use serde::Deserialize;
-use tokio::io::AsyncWriteExt;
-
-const HF_BASE: &str = "https://huggingface.co";
-
-/// Paths to all local files needed for conversion.
-pub struct ModelFiles {
-    pub config_json: PathBuf,
-    /// Ordered list of safetensors shard paths, already downloaded locally.
-    pub safetensors_shards: Vec<PathBuf>,
-}
-
-/// Download (or locate from cache) all model files for `repo_id` into `model_dir`.
-pub async fn download_model(
-    repo_id: &str,
-    hf_token: Option<&str>,
-    model_dir: &Path,
-) -> anyhow::Result<ModelFiles> {
-    let client = build_client(hf_token)?;
-    std::fs::create_dir_all(model_dir).context("create model dir")?;
-
-    println!("Fetching config.json …");
-    let config_json = fetch_file(&client, repo_id, "config.json", model_dir).await?;
-
-    // Detect sharded model by fetching the index file.
-    let shards =
-        match fetch_file(&client, repo_id, "model.safetensors.index.json", model_dir).await {
-            Ok(index_path) => {
-                println!("Found sharded model — reading index …");
-                resolve_shards(&client, repo_id, &index_path, model_dir).await?
-            }
-            Err(_) => {
-                println!("Fetching model.safetensors …");
-                let shard =
-                    fetch_file(&client, repo_id, "model.safetensors", model_dir).await?;
-                vec![shard]
-            }
-        };
-
-    Ok(ModelFiles {
-        config_json,
-        safetensors_shards: shards,
-    })
-}
-
-fn build_client(hf_token: Option<&str>) -> anyhow::Result<reqwest::Client> {
-    let mut headers = reqwest::header::HeaderMap::new();
-    headers.insert(
-        reqwest::header::USER_AGENT,
-        "chronos-rs/0.1".parse().unwrap(),
-    );
-    if let Some(token) = hf_token {
-        headers.insert(
-            reqwest::header::AUTHORIZATION,
-            format!("Bearer {token}").parse().context("invalid HF token")?,
-        );
-    }
-    Ok(reqwest::Client::builder()
-        .default_headers(headers)
-        .redirect(reqwest::redirect::Policy::limited(10))
-        .build()?)
-}
-
-/// Download `filename` from `repo_id` into `dest_dir`, resuming if a partial
-/// `.tmp` file already exists. Returns the local path of the completed file.
-async fn fetch_file(
-    client: &reqwest::Client,
-    repo_id: &str,
-    filename: &str,
-    dest_dir: &Path,
-) -> anyhow::Result<PathBuf> {
-    let dest = dest_dir.join(filename.replace('/', "_"));
-    if dest.exists() {
-        println!("  (cached) {filename}");
-        return Ok(dest);
-    }
-
-    let dest_tmp = dest.with_extension("tmp");
-    let already = if dest_tmp.exists() {
-        dest_tmp.metadata()?.len()
-    } else {
-        0
-    };
-
-    let url = format!("{HF_BASE}/{repo_id}/resolve/main/{filename}");
-
-    let mut req = client.get(&url);
-    if already > 0 {
-        req = req.header(reqwest::header::RANGE, format!("bytes={already}-"));
-        println!("  Resuming {filename} from {} MB …", already / 1_000_000);
-    }
-
-    let response = req
-        .send()
-        .await
-        .with_context(|| format!("GET {url}"))?;
-
-    let status = response.status();
-    if !status.is_success() {
-        anyhow::bail!("HTTP {status} fetching {filename} from {repo_id}");
-    }
-
-    let (file, resume_offset) = if status == reqwest::StatusCode::PARTIAL_CONTENT {
-        let f = tokio::fs::OpenOptions::new()
-            .append(true)
-            .open(&dest_tmp)
-            .await
-            .with_context(|| format!("open tmp {}", dest_tmp.display()))?;
-        (f, already)
-    } else {
-        let f = tokio::fs::File::create(&dest_tmp)
-            .await
-            .with_context(|| format!("create tmp {}", dest_tmp.display()))?;
-        (f, 0)
-    };
-
-    let total = response
-        .content_length()
-        .map(|n| n + resume_offset)
-        .unwrap_or(0);
-
-    let pb = indicatif::ProgressBar::new(total);
-    pb.set_style(
-        indicatif::ProgressStyle::with_template(
-            "  {msg} [{bar:40}] {bytes}/{total_bytes} ({bytes_per_sec}, eta {eta})",
-        )
-        .unwrap()
-        .progress_chars("=>-"),
-    );
-    pb.set_message(filename.to_string());
-    pb.set_position(resume_offset);
-
-    {
-        let mut file = file;
-        let mut stream = response.bytes_stream();
-        while let Some(chunk) = stream.next().await {
-            let chunk = chunk.with_context(|| format!("stream chunk of {filename}"))?;
-            pb.inc(chunk.len() as u64);
-            file.write_all(&chunk)
-                .await
-                .with_context(|| format!("write chunk to {}", dest_tmp.display()))?;
-        }
-    }
-
-    pb.finish_and_clear();
-    std::fs::rename(&dest_tmp, &dest)
-        .with_context(|| format!("rename tmp → {}", dest.display()))?;
-
-    Ok(dest)
-}
-
-/// Parse the shard index JSON and download every unique shard.
-async fn resolve_shards(
-    client: &reqwest::Client,
-    repo_id: &str,
-    index_path: &Path,
-    cache_dir: &Path,
-) -> anyhow::Result<Vec<PathBuf>> {
-    #[derive(Deserialize)]
-    struct Index {
-        weight_map: std::collections::HashMap<String, String>,
-    }
-
-    let raw = std::fs::read_to_string(index_path).context("read index json")?;
-    let index: Index = serde_json::from_str(&raw).context("parse index json")?;
-
-    let mut shard_names: Vec<String> = index.weight_map.into_values().collect();
-    shard_names.sort();
-    shard_names.dedup();
-
-    let mut paths = Vec::with_capacity(shard_names.len());
-    for name in &shard_names {
-        println!("  Fetching {name} …");
-        let p = fetch_file(client, repo_id, name, cache_dir).await?;
-        paths.push(p);
-    }
-    Ok(paths)
-}

src/gguf/mod.rs

@@ -1,5 +0,0 @@
-pub mod types;
-pub mod writer;
-
-pub use types::{GGMLType, GGUFMetaValue};
-pub use writer::GGUFWriter;

src/gguf/types.rs

@@ -1,86 +0,0 @@
-/// GGML tensor data types used in GGUF files.
-/// Values match the ggml_type enum in ggml.h.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-#[repr(u32)]
-pub enum GGMLType {
-    F32 = 0,
-    F16 = 1,
-    /// Q8_0: blocks of 32 × i8 with a shared f16 scale (34 bytes/block).
-    Q8_0 = 8,
-    BF16 = 30,
-}
-
-impl GGMLType {
-    /// (block_elems, bytes_per_block) for block-quantized types; None for float types.
-    /// Q8_0: 32 × i8 values + 1 × f16 scale = 34 bytes/block.
-    #[allow(dead_code)]
-    pub fn block_shape(self) -> Option<(usize, usize)> {
-        match self {
-            GGMLType::Q8_0 => Some((32, 34)),
-            _ => None,
-        }
-    }
-}
-
-/// GGUF metadata value types (gguf_metadata_value_type).
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-#[repr(u32)]
-#[allow(dead_code)]
-pub enum GGUFValueType {
-    Uint8 = 0,
-    Int8 = 1,
-    Uint16 = 2,
-    Int16 = 3,
-    Uint32 = 4,
-    Int32 = 5,
-    Float32 = 6,
-    Bool = 7,
-    String = 8,
-    Array = 9,
-    Uint64 = 10,
-    Int64 = 11,
-    Float64 = 12,
-}
-
-/// A typed metadata value for a GGUF key-value pair.
-#[derive(Debug, Clone)]
-#[allow(dead_code)]
-pub enum GGUFMetaValue {
-    Uint8(u8),
-    Int8(i8),
-    Uint16(u16),
-    Int16(i16),
-    Uint32(u32),
-    Int32(i32),
-    Float32(f32),
-    Bool(bool),
-    String(String),
-    Uint64(u64),
-    Int64(i64),
-    Float64(f64),
-    ArrayUint32(Vec<u32>),
-    ArrayString(Vec<String>),
-    ArrayFloat32(Vec<f32>),
-}
-
-impl GGUFMetaValue {
-    pub fn value_type(&self) -> GGUFValueType {
-        match self {
-            GGUFMetaValue::Uint8(_) => GGUFValueType::Uint8,
-            GGUFMetaValue::Int8(_) => GGUFValueType::Int8,
-            GGUFMetaValue::Uint16(_) => GGUFValueType::Uint16,
-            GGUFMetaValue::Int16(_) => GGUFValueType::Int16,
-            GGUFMetaValue::Uint32(_) => GGUFValueType::Uint32,
-            GGUFMetaValue::Int32(_) => GGUFValueType::Int32,
-            GGUFMetaValue::Float32(_) => GGUFValueType::Float32,
-            GGUFMetaValue::Bool(_) => GGUFValueType::Bool,
-            GGUFMetaValue::String(_) => GGUFValueType::String,
-            GGUFMetaValue::Uint64(_) => GGUFValueType::Uint64,
-            GGUFMetaValue::Int64(_) => GGUFValueType::Int64,
-            GGUFMetaValue::Float64(_) => GGUFValueType::Float64,
-            GGUFMetaValue::ArrayUint32(_)
-            | GGUFMetaValue::ArrayString(_)
-            | GGUFMetaValue::ArrayFloat32(_) => GGUFValueType::Array,
-        }
-    }
-}

src/gguf/writer.rs

@@ -1,161 +0,0 @@
-use std::io::{self, Seek, Write};
-
-use byteorder::{LittleEndian, WriteBytesExt};
-
-use super::types::{GGMLType, GGUFMetaValue, GGUFValueType};
-
-const GGUF_MAGIC: &[u8; 4] = b"GGUF";
-const GGUF_VERSION: u32 = 3;
-const ALIGNMENT: u64 = 32;
-
-/// Describes one tensor's position in the GGUF data section.
-#[derive(Debug)]
-struct TensorInfo {
-    name: String,
-    shape: Vec<u64>,
-    dtype: GGMLType,
-    /// Byte offset from the start of the tensor data block.
-    offset: u64,
-    /// Raw tensor bytes (row-major, little-endian).
-    data: Vec<u8>,
-}
-
-/// Streaming GGUF v3 writer.
-///
-/// Call [`add_metadata`] for every key-value pair, then [`add_tensor`] for
-/// every tensor, then [`write_to`] to flush the complete file.
-pub struct GGUFWriter {
-    metadata: Vec<(String, GGUFMetaValue)>,
-    tensors: Vec<TensorInfo>,
-    /// Running data-block offset for the next tensor.
-    data_offset: u64,
-}
-
-impl GGUFWriter {
-    pub fn new() -> Self {
-        Self {
-            metadata: Vec::new(),
-            tensors: Vec::new(),
-            data_offset: 0,
-        }
-    }
-
-    pub fn add_metadata(&mut self, key: impl Into<String>, value: GGUFMetaValue) {
-        self.metadata.push((key.into(), value));
-    }
-
-    /// Buffer a tensor. `data` must already be in the target dtype byte layout.
-    pub fn add_tensor(
-        &mut self,
-        name: impl Into<String>,
-        shape: Vec<u64>,
-        dtype: GGMLType,
-        data: Vec<u8>,
-    ) {
-        let offset = self.data_offset;
-        let padded_len = round_up(data.len() as u64, ALIGNMENT);
-        self.data_offset += padded_len;
-        self.tensors.push(TensorInfo {
-            name: name.into(),
-            shape,
-            dtype,
-            offset,
-            data,
-        });
-    }
-
-    /// Serialize the complete GGUF file to `writer`.
-    pub fn write_to<W: Write + Seek>(&self, writer: &mut W) -> anyhow::Result<()> {
-        // --- header ---
-        writer.write_all(GGUF_MAGIC)?;
-        writer.write_u32::<LittleEndian>(GGUF_VERSION)?;
-        writer.write_u64::<LittleEndian>(self.tensors.len() as u64)?;
-        writer.write_u64::<LittleEndian>(self.metadata.len() as u64)?;
-
-        // --- metadata key-value pairs ---
-        for (key, value) in &self.metadata {
-            write_string(writer, key)?;
-            writer.write_u32::<LittleEndian>(value.value_type() as u32)?;
-            write_value(writer, value)?;
-        }
-
-        // --- tensor info ---
-        for t in &self.tensors {
-            write_string(writer, &t.name)?;
-            writer.write_u32::<LittleEndian>(t.shape.len() as u32)?;
-            for &dim in &t.shape {
-                writer.write_u64::<LittleEndian>(dim)?;
-            }
-            writer.write_u32::<LittleEndian>(t.dtype as u32)?;
-            writer.write_u64::<LittleEndian>(t.offset)?;
-        }
-
-        // --- align to ALIGNMENT before tensor data ---
-        let pos = writer.stream_position()?;
-        let aligned = round_up(pos, ALIGNMENT);
-        if aligned > pos {
-            let pad = vec![0u8; (aligned - pos) as usize];
-            writer.write_all(&pad)?;
-        }
-
-        // --- tensor data (each padded to ALIGNMENT) ---
-        for t in &self.tensors {
-            writer.write_all(&t.data)?;
-            let remainder = t.data.len() as u64 % ALIGNMENT;
-            if remainder != 0 {
-                let pad = vec![0u8; (ALIGNMENT - remainder) as usize];
-                writer.write_all(&pad)?;
-            }
-        }
-
-        Ok(())
-    }
-}
-
-fn round_up(value: u64, align: u64) -> u64 {
-    (value + align - 1) / align * align
-}
-
-fn write_string<W: Write>(writer: &mut W, s: &str) -> io::Result<()> {
-    writer.write_u64::<LittleEndian>(s.len() as u64)?;
-    writer.write_all(s.as_bytes())
-}
-
-fn write_value<W: Write>(writer: &mut W, value: &GGUFMetaValue) -> anyhow::Result<()> {
-    match value {
-        GGUFMetaValue::Uint8(v) => writer.write_u8(*v)?,
-        GGUFMetaValue::Int8(v) => writer.write_i8(*v)?,
-        GGUFMetaValue::Uint16(v) => writer.write_u16::<LittleEndian>(*v)?,
-        GGUFMetaValue::Int16(v) => writer.write_i16::<LittleEndian>(*v)?,
-        GGUFMetaValue::Uint32(v) => writer.write_u32::<LittleEndian>(*v)?,
-        GGUFMetaValue::Int32(v) => writer.write_i32::<LittleEndian>(*v)?,
-        GGUFMetaValue::Float32(v) => writer.write_f32::<LittleEndian>(*v)?,
-        GGUFMetaValue::Bool(v) => writer.write_u8(*v as u8)?,
-        GGUFMetaValue::String(v) => write_string(writer, v)?,
-        GGUFMetaValue::Uint64(v) => writer.write_u64::<LittleEndian>(*v)?,
-        GGUFMetaValue::Int64(v) => writer.write_i64::<LittleEndian>(*v)?,
-        GGUFMetaValue::Float64(v) => writer.write_f64::<LittleEndian>(*v)?,
-        GGUFMetaValue::ArrayUint32(arr) => {
-            writer.write_u32::<LittleEndian>(GGUFValueType::Uint32 as u32)?;
-            writer.write_u64::<LittleEndian>(arr.len() as u64)?;
-            for v in arr {
-                writer.write_u32::<LittleEndian>(*v)?;
-            }
-        }
-        GGUFMetaValue::ArrayString(arr) => {
-            writer.write_u32::<LittleEndian>(GGUFValueType::String as u32)?;
-            writer.write_u64::<LittleEndian>(arr.len() as u64)?;
-            for s in arr {
-                write_string(writer, s)?;
-            }
-        }
-        GGUFMetaValue::ArrayFloat32(arr) => {
-            writer.write_u32::<LittleEndian>(GGUFValueType::Float32 as u32)?;
-            writer.write_u64::<LittleEndian>(arr.len() as u64)?;
-            for v in arr {
-                writer.write_f32::<LittleEndian>(*v)?;
-            }
-        }
-    }
-    Ok(())
-}

src/infer/mod.rs

@@ -1,557 +0,0 @@
-use std::path::Path;
-
-use anyhow::Context;
-use candle_core::{DType, Device, Tensor};
-use candle_core::quantized::gguf_file;
-
-// ---------------------------------------------------------------------------
-// Config
-// ---------------------------------------------------------------------------
-
-#[derive(Debug, Clone)]
-pub struct InferConfig {
-    pub num_layers: usize,
-    pub embed_dim: usize,
-    pub state_dim: usize,
-    pub n_inputs: usize,        // 2 for with_missing (value + mask)
-    pub decoder_dim: usize,
-    pub decoder_patch_len: usize,
-    pub quantiles: Vec<f32>,
-    pub basis_range: [f32; 2],  // e.g. [-1.0, 0.95] for "legs"
-    pub context_length: usize,
-    pub eps: f32,
-}
-
-// ---------------------------------------------------------------------------
-// Weight structs
-// ---------------------------------------------------------------------------
-
-struct S5Weights {
-    log_lambda_real: Vec<f32>,  // [state_dim]
-    lambda_imag:     Vec<f32>,  // [state_dim]
-    b_r: Tensor,                // [state_dim, embed_dim]  (row-major, HF order)
-    b_i: Tensor,                // [state_dim, embed_dim]
-    c_r: Tensor,                // [embed_dim, state_dim]  (Python C_tilde shape)
-    c_i: Tensor,                // [embed_dim, state_dim]
-    d:   Vec<f32>,              // [embed_dim]
-    log_delta: Vec<f32>,        // [state_dim]
-}
-
-struct BlockWeights {
-    ssm:        S5Weights,
-    out_weight: Tensor,     // [embed_dim, embed_dim]
-    out_bias:   Vec<f32>,   // [embed_dim]
-    norm_weight: Vec<f32>,  // [embed_dim]
-    norm_bias:   Vec<f32>,  // [embed_dim]
-}
-
-pub struct FlowStateModel {
-    device:     Device,
-    config:     InferConfig,
-    embed_w:    Tensor,     // [embed_dim, n_inputs]
-    embed_b:    Vec<f32>,   // [embed_dim]
-    blocks:     Vec<BlockWeights>,
-    decoder_w:  Tensor,     // [n_quantiles * decoder_dim, embed_dim]
-    decoder_b:  Vec<f32>,   // [n_quantiles * decoder_dim]
-}
-
-// ---------------------------------------------------------------------------
-// Loading
-// ---------------------------------------------------------------------------
-
-fn load_f32_vec(
-    content: &gguf_file::Content,
-    reader: &mut (impl std::io::Read + std::io::Seek),
-    name: &str,
-    device: &Device,
-) -> anyhow::Result<Vec<f32>> {
-    let t = content.tensor(reader, name, device)
-        .with_context(|| format!("missing tensor: {name}"))?;
-    let t = t.dequantize(device)?;
-    let t = t.to_dtype(DType::F32)?;
-    Ok(t.flatten_all()?.to_vec1()?)
-}
-
-fn load_matrix(
-    content: &gguf_file::Content,
-    reader: &mut (impl std::io::Read + std::io::Seek),
-    name: &str,
-    device: &Device,
-) -> anyhow::Result<Tensor> {
-    let t = content.tensor(reader, name, device)
-        .with_context(|| format!("missing tensor: {name}"))?;
-    let t = t.dequantize(device)?;
-    Ok(t.to_dtype(DType::F32)?)
-}
-
-impl FlowStateModel {
-    pub fn load(gguf_path: &Path, config: InferConfig) -> anyhow::Result<Self> {
-        let device = Device::Cpu;
-        let mut file = std::fs::File::open(gguf_path)
-            .with_context(|| format!("open {}", gguf_path.display()))?;
-        let content = gguf_file::Content::read(&mut file).context("read GGUF header")?;
-
-        // Embedding
-        let embed_w = load_matrix(&content, &mut file, "embed.weight", &device)?;
-        let embed_b = load_f32_vec(&content, &mut file, "embed.bias", &device)?;
-
-        // Encoder blocks
-        let mut blocks = Vec::with_capacity(config.num_layers);
-        for n in 0..config.num_layers {
-            let ssm = S5Weights {
-                log_lambda_real: load_f32_vec(&content, &mut file, &format!("blk.{n}.ssm.log_lambda_real"), &device)?,
-                lambda_imag:     load_f32_vec(&content, &mut file, &format!("blk.{n}.ssm.lambda_imag"),     &device)?,
-                b_r: load_matrix(&content, &mut file, &format!("blk.{n}.ssm.b_r"), &device)?,
-                b_i: load_matrix(&content, &mut file, &format!("blk.{n}.ssm.b_i"), &device)?,
-                c_r: load_matrix(&content, &mut file, &format!("blk.{n}.ssm.c_r"), &device)?,
-                c_i: load_matrix(&content, &mut file, &format!("blk.{n}.ssm.c_i"), &device)?,
-                d:          load_f32_vec(&content, &mut file, &format!("blk.{n}.ssm.d"),         &device)?,
-                log_delta:  load_f32_vec(&content, &mut file, &format!("blk.{n}.ssm.log_delta"), &device)?,
-            };
-            blocks.push(BlockWeights {
-                ssm,
-                out_weight:  load_matrix(&content, &mut file, &format!("blk.{n}.out.weight"), &device)?,
-                out_bias:    load_f32_vec(&content, &mut file, &format!("blk.{n}.out.bias"),  &device)?,
-                norm_weight: load_f32_vec(&content, &mut file, &format!("blk.{n}.norm.weight"), &device)?,
-                norm_bias:   load_f32_vec(&content, &mut file, &format!("blk.{n}.norm.bias"),   &device)?,
-            });
-        }
-
-        // Decoder
-        let decoder_w = load_matrix(&content, &mut file, "decoder.weight", &device)?;
-        let decoder_b = load_f32_vec(&content, &mut file, "decoder.bias", &device)?;
-
-        Ok(Self { device, config, embed_w, embed_b, blocks, decoder_w, decoder_b })
-    }
-
-    // -----------------------------------------------------------------------
-    // Public inference entry point
-    // -----------------------------------------------------------------------
-
-    /// Forecast `prediction_length` steps from a univariate context series.
-    /// Returns `Vec<Vec<f32>>` of shape `[n_quantiles][prediction_length]`.
-    pub fn forecast(&self, context: &[f32], prediction_length: usize) -> anyhow::Result<Vec<Vec<f32>>> {
-        let cfg = &self.config;
-
-        // 1. Pad or trim context to match model requirements
-        let ctx_len = context.len().min(cfg.context_length);
-        let start = context.len().saturating_sub(ctx_len);
-        let context = &context[start..];
-        let seq_len = context.len();
-
-        // 2. Causal RevIN: compute prefix statistics
-        let (normed_values, final_mean, final_std) = causal_revin_norm(context, cfg.eps);
-        if std::env::var("FLOWSTATE_DEBUG").is_ok() {
-            eprintln!("RevIN final_mean={:.8} final_std={:.8}", final_mean, final_std);
-            eprintln!("normed[0..4]: {:.6} {:.6} {:.6} {:.6}",
-                      normed_values[0], normed_values[1], normed_values[2], normed_values[3]);
-            eprintln!("normed[252..256]: {:.6} {:.6} {:.6} {:.6}",
-                      normed_values[252], normed_values[253], normed_values[254], normed_values[255]);
-        }
-
-        // 3. Build input tensor [seq_len, n_inputs] with mask channel = 0 (no missing)
-        let mut input_data = vec![0.0f32; seq_len * cfg.n_inputs];
-        for t in 0..seq_len {
-            input_data[t * cfg.n_inputs] = normed_values[t];
-            if cfg.n_inputs > 1 {
-                input_data[t * cfg.n_inputs + 1] = 0.0; // mask = 0 means known
-            }
-        }
-
-        // 4. Embedding: [seq_len, n_inputs] × embed_w^T + embed_b → [seq_len, embed_dim]
-        let input_t = Tensor::from_vec(input_data, (seq_len, cfg.n_inputs), &self.device)?;
-        let mut hidden = linear(&input_t, &self.embed_w, &self.embed_b, &self.device)?;
-        // hidden: [seq_len, embed_dim]
-
-        // 5. Scale factor for discretization: decoder_patch_len / prediction_length
-        let scale_factor = cfg.decoder_patch_len as f32 / prediction_length as f32;
-
-        // 6. Encoder: 6 S5 layers
-        let num_layers = self.blocks.len();
-        for (i, block) in self.blocks.iter().enumerate() {
-            let is_last = i == num_layers - 1;
-            hidden = self.apply_s5_layer(hidden, block, scale_factor, is_last, seq_len)?;
-        }
-        // After last layer: hidden is [1, embed_dim]
-
-        // 7. Decoder: linear → [1, n_quantiles * decoder_dim]
-        let coeffs = linear(&hidden, &self.decoder_w, &self.decoder_b, &self.device)?;
-        let n_q = cfg.quantiles.len();
-        let coeffs_data: Vec<f32> = coeffs.flatten_all()?.to_vec1()?;
-        // coeffs_data: [n_quantiles * decoder_dim]
-
-        // Diagnostic: dump all coefficients when FLOWSTATE_DEBUG is set
-        if std::env::var("FLOWSTATE_DEBUG").is_ok() {
-            // Print all coefficients as CSV lines to stderr: COEFF,{qi},{d},{val}
-            for qi in 0..n_q {
-                for d in 0..cfg.decoder_dim {
-                    eprintln!("COEFF,{qi},{d},{:.8}", coeffs_data[qi * cfg.decoder_dim + d]);
-                }
-            }
-        }
-
-        // 8. Compute Legendre basis [prediction_length, decoder_dim]
-        // Python: t[i] = range[0] + i * scale * (range[1]-range[0]) / decoder_patch_len  (i=1..n)
-        let basis = legendre_basis(prediction_length, cfg.decoder_dim, cfg.basis_range,
-                                   scale_factor, cfg.decoder_patch_len);
-        // basis: [prediction_length][decoder_dim], scaled by 1/4
-
-        // 9. Matrix multiply: [n_q, decoder_dim] × [prediction_length, decoder_dim]^T
-        //    = [n_q, prediction_length]
-        let mut output = vec![vec![0.0f32; prediction_length]; n_q];
-        for q in 0..n_q {
-            let coeff_offset = q * cfg.decoder_dim;
-            for p in 0..prediction_length {
-                let mut sum = 0.0f32;
-                for d in 0..cfg.decoder_dim {
-                    sum += coeffs_data[coeff_offset + d] * basis[p][d];
-                }
-                output[q][p] = sum;
-            }
-        }
-
-        // 10. Denormalize: output * final_std + final_mean
-        for q in 0..n_q {
-            for p in 0..prediction_length {
-                output[q][p] = output[q][p] * final_std + final_mean;
-            }
-        }
-
-        Ok(output)
-    }
-
-    // -----------------------------------------------------------------------
-    // S5 layer (one encoder block)
-    // -----------------------------------------------------------------------
-
-    fn apply_s5_layer(
-        &self,
-        x: Tensor,      // [seq_len, embed_dim] for non-last; [seq_len, embed_dim] → output [1, embed_dim] for last
-        block: &BlockWeights,
-        scale_factor: f32,
-        is_last: bool,
-        _seq_len: usize,
-    ) -> anyhow::Result<Tensor> {
-        let cfg = &self.config;
-        let state_dim = cfg.state_dim;
-        let embed_dim = cfg.embed_dim;
-
-        let seq_len = x.dim(0)?;
-
-        // Save skip connection (trimmed for last layer)
-        let skip = if is_last {
-            x.narrow(0, seq_len - 1, 1)?  // last row: [1, embed_dim]
-        } else {
-            x.clone()
-        };
-
-        // ---- SSM forward ----
-
-        // Discretize: compute A_bar and B_bar from SSM parameters
-        let ssm = &block.ssm;
-        let (a_bar_r, a_bar_i, b_bar_r_data, b_bar_i_data) =
-            discretize(ssm, scale_factor, state_dim, embed_dim);
-
-        // B @ x for all timesteps at once: x [seq_len, embed_dim] × B^T [embed_dim, state_dim]
-        // B_tilde stored as [state_dim, embed_dim], so B @ x = x @ B^T
-        let b_bar_r_t = Tensor::from_vec(b_bar_r_data, (state_dim, embed_dim), &self.device)?;
-        let b_bar_i_t = Tensor::from_vec(b_bar_i_data, (state_dim, embed_dim), &self.device)?;
-
-        // x @ B^T → [seq_len, state_dim]
-        let bu_r = x.matmul(&b_bar_r_t.t()?)?;
-        let bu_i = x.matmul(&b_bar_i_t.t()?)?;
-        let bu_r_data: Vec<f32> = bu_r.flatten_all()?.to_vec1()?;
-        let bu_i_data: Vec<f32> = bu_i.flatten_all()?.to_vec1()?;
-
-        // Sequential SSM scan
-        let mut h_r = vec![0.0f32; state_dim]; // real part of hidden state
-        let mut h_i = vec![0.0f32; state_dim]; // imag part
-        let mut all_h_r = vec![0.0f32; seq_len * state_dim];
-        let mut all_h_i = vec![0.0f32; seq_len * state_dim];
-
-        for t in 0..seq_len {
-            for s in 0..state_dim {
-                let new_r = a_bar_r[s] * h_r[s] - a_bar_i[s] * h_i[s] + bu_r_data[t * state_dim + s];
-                let new_i = a_bar_r[s] * h_i[s] + a_bar_i[s] * h_r[s] + bu_i_data[t * state_dim + s];
-                h_r[s] = new_r;
-                h_i[s] = new_i;
-            }
-            all_h_r[t * state_dim..(t + 1) * state_dim].copy_from_slice(&h_r);
-            all_h_i[t * state_dim..(t + 1) * state_dim].copy_from_slice(&h_i);
-        }
-
-        // For last layer, keep only the final timestep
-        let (h_r_t, h_i_t, out_seq_len) = if is_last {
-            let hr = Tensor::from_vec(all_h_r[(seq_len - 1) * state_dim..].to_vec(), (1, state_dim), &self.device)?;
-            let hi = Tensor::from_vec(all_h_i[(seq_len - 1) * state_dim..].to_vec(), (1, state_dim), &self.device)?;
-            (hr, hi, 1usize)
-        } else {
-            let hr = Tensor::from_vec(all_h_r, (seq_len, state_dim), &self.device)?;
-            let hi = Tensor::from_vec(all_h_i, (seq_len, state_dim), &self.device)?;
-            (hr, hi, seq_len)
-        };
-
-        // C @ h: y_real = C_r @ h_r - C_i @ h_i
-        // C_tilde_r stored as [embed_dim, state_dim]; h_r as [out_seq_len, state_dim]
-        // y = h @ C^T → [out_seq_len, embed_dim]
-        let y_from_cr = h_r_t.matmul(&block.ssm.c_r.t()?)?;  // [out_seq_len, embed_dim]
-        let y_from_ci = h_i_t.matmul(&block.ssm.c_i.t()?)?;  // [out_seq_len, embed_dim]
-        let mut y_data: Vec<f32> = (y_from_cr - y_from_ci)?.flatten_all()?.to_vec1()?;
-
-        // Add D skip: y += D * x_at_positions
-        let x_positions = if is_last {
-            x.narrow(0, seq_len - 1, 1)?.flatten_all()?.to_vec1::<f32>()?
-        } else {
-            x.flatten_all()?.to_vec1::<f32>()?
-        };
-        let d = &block.ssm.d;
-        for t in 0..out_seq_len {
-            for e in 0..embed_dim {
-                y_data[t * embed_dim + e] += d[e] * x_positions[t * embed_dim + e];
-            }
-        }
-
-        // ---- MLP: selu(y) * sigmoid(out_linear(selu(y))) ----
-        let y_t = Tensor::from_vec(y_data, (out_seq_len, embed_dim), &self.device)?;
-        let y_selu = selu_tensor(&y_t)?;
-        let gate_pre = linear(&y_selu, &block.out_weight, &block.out_bias, &self.device)?;
-        let gate = sigmoid_tensor(&gate_pre)?;
-        let y_gated = y_selu.mul(&gate)?;
-
-        // ---- LayerNorm ----
-        let y_normed = layer_norm(&y_gated, &block.norm_weight, &block.norm_bias, self.config.eps)?;
-
-        // ---- Residual ----
-        let out = (y_normed + skip)?;
-
-        Ok(out)
-    }
-}
-
-// ---------------------------------------------------------------------------
-// Causal RevIN
-// ---------------------------------------------------------------------------
-
-/// Returns (normalized_values[seq_len], final_mean, final_std).
-/// Each position t is normalized by the cumulative mean/std of x[0..=t].
-fn causal_revin_norm(x: &[f32], eps: f32) -> (Vec<f32>, f32, f32) {
-    let n = x.len();
-    let mut normed = vec![0.0f32; n];
-    let mut cum_sum = 0.0f32;
-    let mut cum_sq_diff = 0.0f32;
-    let mut final_mean = 0.0f32;
-    let mut final_std = 1.0f32;
-
-    for t in 0..n {
-        let count = (t + 1) as f32;
-        cum_sum += x[t];
-        let mean_t = cum_sum / count;
-
-        // Python: cumsum((x[s] - mean_s)^2) / count — each step uses its own running mean
-        cum_sq_diff += (x[t] - mean_t) * (x[t] - mean_t);
-        let var_t = (cum_sq_diff / count).max(0.0);
-        let std_t = (var_t + eps).sqrt();
-
-        normed[t] = (x[t] - mean_t) / std_t;
-
-        if t == n - 1 {
-            final_mean = mean_t;
-            final_std = std_t;
-        }
-    }
-
-    (normed, final_mean, final_std)
-}
-
-// ---------------------------------------------------------------------------
-// SSM discretization
-// ---------------------------------------------------------------------------
-
-/// Returns (A_bar_real, A_bar_imag, B_bar_real_flat, B_bar_imag_flat).
-/// B_bar shape: [state_dim, embed_dim] flattened row-major.
-fn discretize(
-    ssm: &S5Weights,
-    scale_factor: f32,
-    state_dim: usize,
-    embed_dim: usize,
-) -> (Vec<f32>, Vec<f32>, Vec<f32>, Vec<f32>) {
-    // lambda = -exp(log_lambda_real) + j * lambda_imag
-    // delta  = exp(scale_factor * log_delta)
-    // A_bar  = exp(lambda * delta) = exp((lambda_r + j*lambda_i) * delta)
-    //        = exp(lambda_r * delta) * (cos(lambda_i * delta) + j*sin(lambda_i * delta))
-
-    let mut a_r = vec![0.0f32; state_dim];
-    let mut a_i = vec![0.0f32; state_dim];
-
-    // (A_bar - 1) / lambda per element (complex division)
-    let mut coeff_r = vec![0.0f32; state_dim];
-    let mut coeff_i = vec![0.0f32; state_dim];
-
-    for s in 0..state_dim {
-        let lam_r = -ssm.log_lambda_real[s].exp();
-        let lam_i = ssm.lambda_imag[s];
-        let delta = (scale_factor * ssm.log_delta[s]).exp();
-
-        // A_bar = exp(lambda * delta)
-        let exp_r = lam_r * delta;
-        let exp_i = lam_i * delta;
-        let mag = exp_r.exp();
-        a_r[s] = mag * exp_i.cos();
-        a_i[s] = mag * exp_i.sin();
-
-        // (A_bar - 1) / lambda
-        // numerator: (a_r - 1) + j*a_i
-        // denominator: lam_r + j*lam_i
-        let num_r = a_r[s] - 1.0;
-        let num_i = a_i[s];
-        let denom_sq = lam_r * lam_r + lam_i * lam_i;
-        if denom_sq > 1e-20 {
-            coeff_r[s] = (num_r * lam_r + num_i * lam_i) / denom_sq;
-            coeff_i[s] = (num_i * lam_r - num_r * lam_i) / denom_sq;
-        } else {
-            coeff_r[s] = delta; // limit as lambda → 0
-            coeff_i[s] = 0.0;
-        }
-    }
-
-    // B_bar = coeff * B_tilde (element-wise complex multiply, then multiply B matrix)
-    // B_tilde stored as [state_dim, embed_dim] in GGUF (transposed from Python HF which is also [state_dim, emb])
-    // After load_matrix, candle reverses dims: stored [embed_dim, state_dim] in GGUF → loaded as [state_dim, embed_dim]
-    // Actually GGUF reversal: Python [state_dim, embed_dim] → GGUF stores reversed → candle sees [embed_dim, state_dim]
-    // But we treat it as [state_dim, embed_dim] for the einsum B @ u = u @ B^T
-    // Let's get B_tilde data and compute B_bar manually
-    // The GGUF shape is reversed from Python. Python B_tilde: [512, 512].
-    // So GGUF/candle sees it as [512, 512] (square, so same either way).
-
-    // Get B_tilde data as flat arrays [state_dim * embed_dim]
-    // b_r and b_i are loaded as candle tensors; we need them as flat vecs for the per-state multiply
-    // They were stored with reversed shape but 512×512 is symmetric in shape
-    // We extract row by row: B_tilde[s, :] = row s of the [state_dim, embed_dim] matrix
-
-    // Use flat data from the tensors
-    let b_r_data = get_tensor_data_row_major(&ssm.b_r, state_dim, embed_dim);
-    let b_i_data = get_tensor_data_row_major(&ssm.b_i, state_dim, embed_dim);
-
-    // B_bar[s, e] = coeff_r[s]*b_r[s,e] - coeff_i[s]*b_i[s,e]  (real)
-    //             + coeff_r[s]*b_i[s,e] + coeff_i[s]*b_r[s,e]  (imag)
-    let mut b_bar_r = vec![0.0f32; state_dim * embed_dim];
-    let mut b_bar_i = vec![0.0f32; state_dim * embed_dim];
-
-    for s in 0..state_dim {
-        for e in 0..embed_dim {
-            let br = b_r_data[s * embed_dim + e];
-            let bi = b_i_data[s * embed_dim + e];
-            b_bar_r[s * embed_dim + e] = coeff_r[s] * br - coeff_i[s] * bi;
-            b_bar_i[s * embed_dim + e] = coeff_r[s] * bi + coeff_i[s] * br;
-        }
-    }
-
-    (a_r, a_i, b_bar_r, b_bar_i)
-}
-
-/// Extract tensor data in row-major order as Vec<f32>.
-/// Handles candle's internal memory layout for F32 tensors.
-fn get_tensor_data_row_major(t: &Tensor, rows: usize, cols: usize) -> Vec<f32> {
-    t.to_dtype(DType::F32)
-        .and_then(|t| t.reshape((rows, cols)))
-        .and_then(|t| t.flatten_all())
-        .and_then(|t| t.to_vec1())
-        .unwrap_or_else(|_| vec![0.0f32; rows * cols])
-}
-
-// ---------------------------------------------------------------------------
-// Legendre basis (FlowStateLegendreBasis equivalent)
-// ---------------------------------------------------------------------------
-
-/// Public wrapper for diagnostics.
-pub fn dump_legendre_basis(n_points: usize, degree: usize, range: [f32; 2],
-                           scale: f32, pred_dist: usize) -> Vec<Vec<f32>> {
-    legendre_basis(n_points, degree, range, scale, pred_dist)
-}
-
-/// Compute Legendre polynomial basis matrix.
-/// Returns [n_points][degree+1] scaled by 1/4 (as in get_kernel).
-/// Python formula: t[i] = range[0] + i * scale * (range[1]-range[0]) / pred_dist  (i=1..n_points)
-fn legendre_basis(n_points: usize, degree: usize, range: [f32; 2],
-                  scale: f32, pred_dist: usize) -> Vec<Vec<f32>> {
-    let dt = scale * (range[1] - range[0]) / pred_dist as f32;
-    let t: Vec<f32> = (1..=n_points).map(|i| range[0] + i as f32 * dt).collect();
-
-    // Evaluate Legendre polynomials at each t: [n_points, degree+1]
-    // P_0 = 1, P_1 = x, P_k = ((2k-1)*x*P_{k-1} - (k-1)*P_{k-2}) / k
-    let mut basis = vec![vec![0.0f32; degree + 1]; n_points];
-    for (p, &x) in t.iter().enumerate() {
-        basis[p][0] = 1.0;
-        if degree >= 1 {
-            basis[p][1] = x;
-        }
-        for k in 1..degree {
-            let kf = k as f32;
-            basis[p][k + 1] =
-                ((2.0 * kf + 1.0) * x * basis[p][k] - kf * basis[p][k - 1]) / (kf + 1.0);
-        }
-        // Scale by 1/4 (from get_kernel: return f.float() / 4.0)
-        for d in 0..=degree {
-            basis[p][d] /= 4.0;
-        }
-    }
-
-    basis
-}
-
-// ---------------------------------------------------------------------------
-// Neural network primitives
-// ---------------------------------------------------------------------------
-
-/// y = x @ w^T + b  (w: [out, in])
-fn linear(x: &Tensor, w: &Tensor, b: &[f32], device: &Device) -> anyhow::Result<Tensor> {
-    let out_size = b.len();
-    let y = x.matmul(&w.t()?)?;
-    let b_t = Tensor::from_vec(b.to_vec(), (1, out_size), device)?.broadcast_as(y.shape())?;
-    Ok((y + b_t)?)
-}
-
-/// SELU activation: selu(x) = scale * (x if x>0 else alpha*(exp(x)-1))
-fn selu_tensor(x: &Tensor) -> anyhow::Result<Tensor> {
-    const SCALE: f64 = 1.0507009873554804934193349852946;
-    const ALPHA: f64 = 1.6732632423543772848170429916717;
-    let data: Vec<f32> = x.flatten_all()?.to_vec1()?;
-    let selu_data: Vec<f32> = data.iter().map(|&v| {
-        if v > 0.0 {
-            (SCALE * v as f64) as f32
-        } else {
-            (SCALE * ALPHA * (v.exp() as f64 - 1.0)) as f32
-        }
-    }).collect();
-    Ok(Tensor::from_vec(selu_data, x.shape(), x.device())?)
-}
-
-/// Sigmoid: 1 / (1 + exp(-x))
-fn sigmoid_tensor(x: &Tensor) -> anyhow::Result<Tensor> {
-    let data: Vec<f32> = x.flatten_all()?.to_vec1()?;
-    let sig_data: Vec<f32> = data.iter().map(|&v| 1.0 / (1.0 + (-v).exp())).collect();
-    Ok(Tensor::from_vec(sig_data, x.shape(), x.device())?)
-}
-
-/// LayerNorm: (x - mean) / std * weight + bias (over last dimension)
-fn layer_norm(x: &Tensor, weight: &[f32], bias: &[f32], eps: f32) -> anyhow::Result<Tensor> {
-    let shape = x.shape().clone();
-    let last_dim = *shape.dims().last().context("empty tensor")?;
-    let n = shape.elem_count();
-    let n_rows = n / last_dim;
-
-    let data: Vec<f32> = x.flatten_all()?.to_vec1()?;
-    let mut out = vec![0.0f32; n];
-
-    for r in 0..n_rows {
-        let row = &data[r * last_dim..(r + 1) * last_dim];
-        let mean: f32 = row.iter().sum::<f32>() / last_dim as f32;
-        let var: f32 = row.iter().map(|&v| (v - mean).powi(2)).sum::<f32>() / last_dim as f32;
-        let std_inv = 1.0 / (var + eps).sqrt();
-        for i in 0..last_dim {
-            out[r * last_dim + i] = (row[i] - mean) * std_inv * weight[i] + bias[i];
-        }
-    }
-
-    Ok(Tensor::from_vec(out, shape, x.device())?)
-}

src/main.rs

@@ -1,316 +0,0 @@
-mod config;
-mod convert;
-mod download;
-mod gguf;
-mod infer;
-mod tensor_map;
-mod upload;
-
-use std::collections::BTreeMap;
-use std::path::PathBuf;
-use std::time::{SystemTime, UNIX_EPOCH};
-
-use anyhow::Context;
-use clap::{Parser, Subcommand, ValueEnum};
-use serde::Serialize;
-
-use config::FlowStateConfig;
-use convert::{convert, ConvertOptions};
-use download::download_model;
-use gguf::GGMLType;
-use infer::{FlowStateModel, InferConfig};
-
-#[derive(Parser)]
-#[command(name = "flowstate-r1-rs", about = "Convert and run ibm-granite/granite-timeseries-flowstate-r1")]
-struct Cli {
-    #[command(subcommand)]
-    command: Command,
-}
-
-#[derive(Subcommand)]
-enum Command {
-    /// Upload source + GGUF files to HuggingFace Hub.
-    Upload {
-        /// HuggingFace repo to upload to (owner/name).
-        #[arg(short, long, default_value = "amaye15/flowstate-r1-gguf")]
-        repo: String,
-        /// HuggingFace API token (or set HF_TOKEN env var).
-        #[arg(long, env = "HF_TOKEN")]
-        token: String,
-    },
-    /// Download a FlowState model from HuggingFace and convert it to GGUF.
-    Convert {
-        #[arg(short, long, default_value = "ibm-granite/granite-timeseries-flowstate-r1")]
-        model: String,
-        #[arg(short, long, default_value = "gguf/flowstate-r1-f16.gguf")]
-        output: PathBuf,
-        #[arg(long, default_value = "f16")]
-        dtype: DtypeArg,
-        #[arg(long, default_value = "models")]
-        model_dir: PathBuf,
-        #[arg(long, env = "HF_TOKEN")]
-        token: Option<String>,
-    },
-    /// Print all tensor names in a local safetensors or GGUF file.
-    InspectTensors {
-        path: PathBuf,
-    },
-    /// Dump the Legendre basis matrix to CSV for comparison with Python numpy.
-    DumpBasis {
-        #[arg(long, default_value = "models/config.json")]
-        config: PathBuf,
-        #[arg(short, long, default_value = "24")]
-        prediction_length: usize,
-    },
-    /// Run FlowState univariate forecasting from a GGUF file.
-    ///
-    /// Reads a JSON request from stdin: {"context": [...], "horizon": N}
-    /// Outputs a JSON forecast in OpenAI-compatible format with all quantile levels.
-    /// The `point` field contains the median quantile (q0.5).
-    Infer {
-        /// Path to the GGUF file.
-        #[arg(short, long, default_value = "gguf/flowstate-r1-f16.gguf")]
-        gguf: PathBuf,
-
-        /// Path to config.json from the original HuggingFace model.
-        #[arg(long, default_value = "models/config.json")]
-        config: PathBuf,
-    },
-}
-
-#[derive(Clone, ValueEnum)]
-enum DtypeArg {
-    F32,
-    F16,
-    Q8,
-}
-
-impl From<DtypeArg> for GGMLType {
-    fn from(d: DtypeArg) -> Self {
-        match d {
-            DtypeArg::F32 => GGMLType::F32,
-            DtypeArg::F16 => GGMLType::F16,
-            DtypeArg::Q8  => GGMLType::Q8_0,
-        }
-    }
-}
-
-#[tokio::main]
-async fn main() -> anyhow::Result<()> {
-    let cli = Cli::parse();
-
-    match cli.command {
-        Command::Upload { repo, token } => {
-            let crate_root = std::env::current_dir().context("current dir")?;
-            let model_root = crate_root.parent()
-                .map(|p| p.to_path_buf())
-                .unwrap_or(crate_root);
-            upload::run(&repo, &token, &model_root).await?;
-        }
-
-        Command::Convert { model, output, dtype, model_dir, token } => {
-            println!("Downloading {model} into {} …", model_dir.display());
-            let files = download_model(&model, token.as_deref(), &model_dir)
-                .await
-                .context("download failed")?;
-
-            let config_str = std::fs::read_to_string(&files.config_json)
-                .context("read config.json")?;
-            let config = FlowStateConfig::from_json(&config_str)
-                .context("parse config.json")?;
-
-            println!(
-                "Config: {} encoder layers, embed_dim={}, state_dim={}, decoder_dim={}, {} quantiles",
-                config.encoder_num_layers,
-                config.embedding_feature_dim,
-                config.encoder_state_dim,
-                config.decoder_dim,
-                config.n_quantiles(),
-            );
-
-            let opts = ConvertOptions { output_dtype: dtype.into() };
-            convert(&model, &files, &config, &opts, &output)?;
-            println!("Wrote {}", output.display());
-        }
-
-        Command::InspectTensors { path } => {
-            let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");
-            if ext == "gguf" {
-                let mut file = std::fs::File::open(&path)
-                    .with_context(|| format!("open {}", path.display()))?;
-                let content = candle_core::quantized::gguf_file::Content::read(&mut file)
-                    .context("read GGUF")?;
-                println!("Tensors in {}:", path.display());
-                let mut names: Vec<_> = content.tensor_infos.keys().cloned().collect();
-                names.sort();
-                for name in names {
-                    let ti = &content.tensor_infos[&name];
-                    println!("  {name:80} {:?} {:?}", ti.ggml_dtype, ti.shape);
-                }
-            } else {
-                let bytes = std::fs::read(&path)
-                    .with_context(|| format!("read {}", path.display()))?;
-                let tensors = safetensors::SafeTensors::deserialize(&bytes)
-                    .context("deserialize safetensors")?;
-                println!("Tensors in {}:", path.display());
-                let mut names: Vec<_> = tensors.names().into_iter().collect();
-                names.sort();
-                for name in names {
-                    let t = tensors.tensor(name).unwrap();
-                    println!("  {name:80} {:?} {:?}", t.dtype(), t.shape());
-                }
-            }
-        }
-
-        Command::DumpBasis { config, prediction_length } => {
-            let config_str = std::fs::read_to_string(&config)
-                .with_context(|| format!("read {}", config.display()))?;
-            let fs_config = FlowStateConfig::from_json(&config_str)
-                .context("parse config.json")?;
-            let decoder_dim = fs_config.decoder_dim as usize;
-            let basis_range = fs_config.basis_range();
-            let scale = fs_config.decoder_patch_len as f32 / prediction_length as f32;
-            let basis = infer::dump_legendre_basis(prediction_length, decoder_dim, basis_range, scale, fs_config.decoder_patch_len as usize);
-            let header: Vec<String> = (0..decoder_dim).map(|d| format!("P{d}")).collect();
-            println!("{}", header.join(","));
-            for row in &basis {
-                let vals: Vec<String> = row[..decoder_dim].iter().map(|v| format!("{v:.8}")).collect();
-                println!("{}", vals.join(","));
-            }
-        }
-
-        Command::Infer { gguf, config } => {
-            use std::io::Read;
-            let mut buf = String::new();
-            std::io::stdin().read_to_string(&mut buf).context("read stdin")?;
-            let req: serde_json::Value = serde_json::from_str(&buf).context("parse JSON input")?;
-            let contexts = parse_contexts(req["context"].clone())?;
-            let horizon: usize = req["horizon"].as_u64().context("horizon must be a positive integer")? as usize;
-
-            let config_str = std::fs::read_to_string(&config)
-                .with_context(|| format!("read {}", config.display()))?;
-            let fs_config = FlowStateConfig::from_json(&config_str)
-                .context("parse config.json")?;
-
-            let infer_config = InferConfig {
-                num_layers:        fs_config.encoder_num_layers as usize,
-                embed_dim:         fs_config.embedding_feature_dim as usize,
-                state_dim:         fs_config.encoder_state_dim as usize,
-                n_inputs:          fs_config.n_inputs() as usize,
-                decoder_dim:       fs_config.decoder_dim as usize,
-                decoder_patch_len: fs_config.decoder_patch_len as usize,
-                quantiles:         fs_config.quantiles.clone(),
-                basis_range:       fs_config.basis_range(),
-                context_length:    fs_config.context_length as usize,
-                eps:               1e-5,
-            };
-
-            let quantile_levels = infer_config.quantiles.clone();
-            let median_idx = quantile_levels
-                .iter()
-                .position(|&q| (q - 0.5).abs() < 1e-6)
-                .unwrap_or(quantile_levels.len() / 2);
-
-            eprintln!("Loading model from {} …", gguf.display());
-            let model = FlowStateModel::load(&gguf, infer_config)
-                .context("load model")?;
-
-            let mut fc_choices = Vec::new();
-            for ctx in &contexts {
-                anyhow::ensure!(!ctx.is_empty(), "context series must not be empty");
-                eprintln!("Running forecast ({} context steps → {horizon} future steps) …", ctx.len());
-                let quantile_mat = model.forecast(ctx, horizon).context("forecast")?;
-
-                let point = quantile_mat.get(median_idx).cloned().unwrap_or_default();
-                let mut quantiles = BTreeMap::new();
-                for (i, &level) in quantile_levels.iter().enumerate() {
-                    if let Some(q) = quantile_mat.get(i) {
-                        quantiles.insert(format!("{level:.2}"), q.clone());
-                    }
-                }
-                fc_choices.push((point, quantiles));
-            }
-            let total_ctx: usize = contexts.iter().map(|c| c.len()).sum();
-            println!("{}", forecast_json("flowstate-r1", total_ctx, horizon, fc_choices)?);
-        }
-    }
-
-    Ok(())
-}
-
-fn parse_contexts(val: serde_json::Value) -> anyhow::Result<Vec<Vec<f32>>> {
-    match val {
-        serde_json::Value::Array(arr) if arr.is_empty() => {
-            anyhow::bail!("context must be a non-empty array")
-        }
-        serde_json::Value::Array(arr) => {
-            if arr.first().map(|v| v.is_array()).unwrap_or(false) {
-                arr.into_iter()
-                    .enumerate()
-                    .map(|(i, v)| {
-                        serde_json::from_value::<Vec<f32>>(v)
-                            .with_context(|| format!("context[{i}] must be an array of numbers"))
-                    })
-                    .collect()
-            } else {
-                let ctx = serde_json::from_value::<Vec<f32>>(serde_json::Value::Array(arr))
-                    .context("context must be a JSON array of numbers")?;
-                Ok(vec![ctx])
-            }
-        }
-        _ => anyhow::bail!("context must be a JSON array"),
-    }
-}
-
-fn forecast_json(
-    model_name: &str,
-    context_length: usize,
-    forecast_length: usize,
-    fc_choices: Vec<(Vec<f32>, BTreeMap<String, Vec<f32>>)>,
-) -> anyhow::Result<String> {
-    #[derive(Serialize)]
-    struct ForecastResponse {
-        id: String,
-        object: &'static str,
-        created: u64,
-        model: String,
-        choices: Vec<Choice>,
-        usage: Usage,
-    }
-    #[derive(Serialize)]
-    struct Choice {
-        index: usize,
-        forecast: ForecastOutput,
-        finish_reason: &'static str,
-    }
-    #[derive(Serialize)]
-    struct ForecastOutput {
-        point: Vec<f32>,
-        quantiles: BTreeMap<String, Vec<f32>>,
-    }
-    #[derive(Serialize)]
-    struct Usage {
-        context_length: usize,
-        forecast_length: usize,
-    }
-
-    let created = SystemTime::now()
-        .duration_since(UNIX_EPOCH)
-        .unwrap_or_default()
-        .as_secs();
-
-    let resp = ForecastResponse {
-        id: format!("forecast-{created:016x}"),
-        object: "forecast",
-        created,
-        model: model_name.to_string(),
-        choices: fc_choices.into_iter().enumerate().map(|(i, (point, quantiles))| Choice {
-            index: i,
-            forecast: ForecastOutput { point, quantiles },
-            finish_reason: "stop",
-        }).collect(),
-        usage: Usage { context_length, forecast_length },
-    };
-
-    Ok(serde_json::to_string_pretty(&resp)?)
-}

src/tensor_map.rs

@@ -1,101 +0,0 @@
-/// Map a FlowState HuggingFace tensor name to the GGUF canonical name.
-/// Returns `None` for tensors unused at inference time (caller skips them).
-///
-/// HF prefix in safetensors:  (none — model is saved as FlowStateForPrediction,
-/// which wraps FlowStateModel in `self.model`, but the checkpoint stores
-/// FlowStateModel weights directly without the `.model.` prefix.)
-///
-/// Verified tensor names from model.safetensors header:
-///   embed.embed.weight / embed.embed.bias
-///   encoder.layers.{N}.ssm.{log_Lambda_real, Lambda_imag, B_tilde_r, B_tilde_i,
-///                                C_tilde_r, C_tilde_i, D, log_Delta}
-///   encoder.layers.{N}.out.{weight, bias}
-///   encoder.layers.{N}.norm.{weight, bias}
-///   decoder.lin.{weight, bias}
-pub fn map_tensor_name(hf_name: &str) -> Option<String> {
-    // Embedding
-    match hf_name {
-        "embed.embed.weight" => return Some("embed.weight".into()),
-        "embed.embed.bias"   => return Some("embed.bias".into()),
-        "decoder.lin.weight" => return Some("decoder.weight".into()),
-        "decoder.lin.bias"   => return Some("decoder.bias".into()),
-        _ => {}
-    }
-
-    // encoder.layers.{N}.ssm.*  and  encoder.layers.{N}.{out,norm}.*
-    let rest = hf_name.strip_prefix("encoder.layers.")?;
-    let dot = rest.find('.')?;
-    let n: usize = rest[..dot].parse().ok()?;
-    let suffix = &rest[dot + 1..];
-
-    let name = match suffix {
-        "ssm.log_Lambda_real" => format!("blk.{n}.ssm.log_lambda_real"),
-        "ssm.Lambda_imag"     => format!("blk.{n}.ssm.lambda_imag"),
-        "ssm.B_tilde_r"       => format!("blk.{n}.ssm.b_r"),
-        "ssm.B_tilde_i"       => format!("blk.{n}.ssm.b_i"),
-        "ssm.C_tilde_r"       => format!("blk.{n}.ssm.c_r"),
-        "ssm.C_tilde_i"       => format!("blk.{n}.ssm.c_i"),
-        "ssm.D"               => format!("blk.{n}.ssm.d"),
-        "ssm.log_Delta"       => format!("blk.{n}.ssm.log_delta"),
-        "out.weight"          => format!("blk.{n}.out.weight"),
-        "out.bias"            => format!("blk.{n}.out.bias"),
-        "norm.weight"         => format!("blk.{n}.norm.weight"),
-        "norm.bias"           => format!("blk.{n}.norm.bias"),
-        _ => return None,
-    };
-    Some(name)
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn test_embed() {
-        assert_eq!(map_tensor_name("embed.embed.weight"), Some("embed.weight".into()));
-        assert_eq!(map_tensor_name("embed.embed.bias"),   Some("embed.bias".into()));
-    }
-
-    #[test]
-    fn test_decoder() {
-        assert_eq!(map_tensor_name("decoder.lin.weight"), Some("decoder.weight".into()));
-        assert_eq!(map_tensor_name("decoder.lin.bias"),   Some("decoder.bias".into()));
-    }
-
-    #[test]
-    fn test_ssm_params() {
-        assert_eq!(
-            map_tensor_name("encoder.layers.0.ssm.log_Lambda_real"),
-            Some("blk.0.ssm.log_lambda_real".into())
-        );
-        assert_eq!(
-            map_tensor_name("encoder.layers.3.ssm.B_tilde_r"),
-            Some("blk.3.ssm.b_r".into())
-        );
-        assert_eq!(
-            map_tensor_name("encoder.layers.5.ssm.C_tilde_i"),
-            Some("blk.5.ssm.c_i".into())
-        );
-        assert_eq!(
-            map_tensor_name("encoder.layers.2.ssm.log_Delta"),
-            Some("blk.2.ssm.log_delta".into())
-        );
-    }
-
-    #[test]
-    fn test_layer_mlp_and_norm() {
-        assert_eq!(
-            map_tensor_name("encoder.layers.1.out.weight"),
-            Some("blk.1.out.weight".into())
-        );
-        assert_eq!(
-            map_tensor_name("encoder.layers.4.norm.bias"),
-            Some("blk.4.norm.bias".into())
-        );
-    }
-
-    #[test]
-    fn test_unknown_returns_none() {
-        assert_eq!(map_tensor_name("something.unknown"), None);
-    }
-}

src/upload.rs

@@ -1,625 +0,0 @@
-use std::io::Read;
-use std::path::{Path, PathBuf};
-
-use anyhow::Context;
-use base64::Engine as _;
-use futures_util::StreamExt;
-use indicatif::{ProgressBar, ProgressStyle};
-use sha2::{Digest, Sha256};
-use tokio_util::io::ReaderStream;
-
-const HF_BASE: &str = "https://huggingface.co";
-const PREUPLOAD_SAMPLE: usize = 512;
-
-const SKIP_DIRS: &[&str] = &["target", ".git", "__pycache__", ".venv", "models"];
-
-// HF manages these files automatically — never delete them
-const KEEP_ALWAYS: &[&str] = &[".gitattributes", ".gitignore"];
-
-struct RegularFile {
-    remote: String,
-    content_b64: String,
-}
-
-struct LfsFile {
-    local: PathBuf,
-    remote: String,
-    size: u64,
-    oid: String, // sha256 hex
-}
-
-pub async fn run(repo_id: &str, token: &str, root: &Path) -> anyhow::Result<()> {
-    let client = build_client(token)?;
-
-    ensure_repo(&client, repo_id).await?;
-
-    // List existing repo files so we can delete stale ones in the same commit
-    let existing_files = list_repo_files(&client, repo_id).await?;
-    if !existing_files.is_empty() {
-        println!("  {} file(s) currently in repo", existing_files.len());
-    }
-
-    let files = gather_files(root)?;
-    println!("Gathered {} file(s) to upload …", files.len());
-
-    let (regular, lfs) = classify_files(&client, repo_id, &files).await?;
-    println!("  {} regular, {} LFS", regular.len(), lfs.len());
-
-    if !lfs.is_empty() {
-        upload_lfs(&client, repo_id, &lfs).await?;
-    }
-
-    // Delete anything currently in the repo that we're not re-uploading
-    let upload_paths: std::collections::HashSet<&str> = regular
-        .iter()
-        .map(|f| f.remote.as_str())
-        .chain(lfs.iter().map(|f| f.remote.as_str()))
-        .collect();
-    let to_delete: Vec<String> = existing_files
-        .into_iter()
-        .filter(|p| !upload_paths.contains(p.as_str()) && !KEEP_ALWAYS.contains(&p.as_str()))
-        .collect();
-    if !to_delete.is_empty() {
-        println!("  Deleting {} stale file(s)", to_delete.len());
-    }
-
-    make_commit(&client, repo_id, &regular, &lfs, &to_delete).await?;
-
-    println!("Uploaded → https://huggingface.co/{repo_id}");
-    Ok(())
-}
-
-fn build_client(token: &str) -> anyhow::Result<reqwest::Client> {
-    let mut headers = reqwest::header::HeaderMap::new();
-    headers.insert(
-        reqwest::header::AUTHORIZATION,
-        format!("Bearer {token}").parse().context("invalid HF token")?,
-    );
-    headers.insert(
-        reqwest::header::USER_AGENT,
-        "flowstate-r1-rs/0.1".parse().unwrap(),
-    );
-    Ok(reqwest::Client::builder().default_headers(headers).build()?)
-}
-
-async fn ensure_repo(client: &reqwest::Client, repo_id: &str) -> anyhow::Result<()> {
-    let name = repo_id.split('/').nth(1).context("repo_id must be owner/name")?;
-    let resp: reqwest::Response = client
-        .post(format!("{HF_BASE}/api/repos/create"))
-        .json(&serde_json::json!({ "name": name, "type": "model", "private": false }))
-        .send()
-        .await
-        .context("create repo")?;
-
-    let status = resp.status();
-    if status.is_success() {
-        println!("Created repo {repo_id}");
-    } else if status.as_u16() == 409 {
-        // already exists — fine
-    } else {
-        let body: String = resp.text().await.unwrap_or_default();
-        anyhow::bail!("create repo: HTTP {status}: {body}");
-    }
-    Ok(())
-}
-
-/// List all file (blob) paths currently in the repo, following Link header pagination.
-async fn list_repo_files(client: &reqwest::Client, repo_id: &str) -> anyhow::Result<Vec<String>> {
-    let mut files = Vec::new();
-    let mut url = format!("{HF_BASE}/api/models/{repo_id}/tree/main?recursive=true&limit=1000");
-
-    loop {
-        let resp: reqwest::Response = client.get(&url).send().await.context("list repo tree")?;
-        let status = resp.status();
-        if status.as_u16() == 404 {
-            return Ok(files); // repo is new / empty
-        }
-        if !status.is_success() {
-            let body = resp.text().await.unwrap_or_default();
-            anyhow::bail!("list repo tree: HTTP {status}: {body}");
-        }
-
-        // Extract next-page URL from Link header before consuming body
-        let next_url = resp
-            .headers()
-            .get("link")
-            .and_then(|v| v.to_str().ok())
-            .and_then(parse_next_link);
-
-        #[derive(serde::Deserialize)]
-        struct TreeEntry { r#type: String, path: String }
-        let entries: Vec<TreeEntry> = resp.json().await.context("parse tree response")?;
-        for e in entries {
-            if e.r#type == "file" {
-                files.push(e.path);
-            }
-        }
-
-        match next_url {
-            Some(next) => url = next,
-            None => break,
-        }
-    }
-
-    Ok(files)
-}
-
-/// Parse `<url>; rel="next"` from a Link header value.
-fn parse_next_link(header: &str) -> Option<String> {
-    for part in header.split(',') {
-        let part = part.trim();
-        if part.contains(r#"rel="next""#) {
-            if let Some(url_part) = part.split(';').next() {
-                let url = url_part.trim().trim_start_matches('<').trim_end_matches('>');
-                return Some(url.to_string());
-            }
-        }
-    }
-    None
-}
-
-/// Ask HF which files need LFS vs regular inline upload, then prepare both lists.
-async fn classify_files(
-    client: &reqwest::Client,
-    repo_id: &str,
-    files: &[(PathBuf, String)],
-) -> anyhow::Result<(Vec<RegularFile>, Vec<LfsFile>)> {
-    let mut preupload_entries: Vec<serde_json::Value> = Vec::new();
-    for (local, remote) in files {
-        let size = std::fs::metadata(local)
-            .with_context(|| format!("stat {}", local.display()))?
-            .len();
-        let sample = {
-            let bytes = std::fs::read(local)
-                .with_context(|| format!("read {}", local.display()))?;
-            let n = bytes.len().min(PREUPLOAD_SAMPLE);
-            base64::engine::general_purpose::STANDARD.encode(&bytes[..n])
-        };
-        preupload_entries.push(serde_json::json!({
-            "path": remote,
-            "size": size,
-            "sample": sample,
-        }));
-    }
-
-    let url = format!("{HF_BASE}/api/models/{repo_id}/preupload/main");
-    let resp: reqwest::Response = client
-        .post(&url)
-        .json(&serde_json::json!({ "files": preupload_entries }))
-        .send()
-        .await
-        .context("preupload request")?;
-
-    let status = resp.status();
-    if !status.is_success() {
-        let body: String = resp.text().await.unwrap_or_default();
-        anyhow::bail!("preupload: HTTP {status}: {body}");
-    }
-
-    #[derive(serde::Deserialize)]
-    struct PreuploadFile {
-        path: String,
-        #[serde(rename = "uploadMode")]
-        upload_mode: String,
-        #[serde(rename = "shouldIgnore", default)]
-        _should_ignore: bool,
-    }
-    #[derive(serde::Deserialize)]
-    struct PreuploadResp { files: Vec<PreuploadFile> }
-
-    let preupload: PreuploadResp = resp.json().await.context("parse preupload response")?;
-    let modes: std::collections::HashMap<String, String> = preupload.files
-        .into_iter()
-        .map(|f| (f.path, f.upload_mode))
-        .collect();
-
-    let mut regular: Vec<RegularFile> = Vec::new();
-    let mut lfs: Vec<LfsFile> = Vec::new();
-
-    for (local, remote) in files {
-        let mode = modes.get(remote).map(|s| s.as_str()).unwrap_or("regular");
-        if mode == "lfs" {
-            let size = std::fs::metadata(local)?.len();
-            let oid = sha256_file(local)?;
-            lfs.push(LfsFile { local: local.clone(), remote: remote.clone(), size, oid });
-        } else {
-            let bytes = std::fs::read(local)
-                .with_context(|| format!("read {}", local.display()))?;
-            let content_b64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
-            regular.push(RegularFile { remote: remote.clone(), content_b64 });
-        }
-    }
-
-    Ok((regular, lfs))
-}
-
-fn gather_files(root: &Path) -> anyhow::Result<Vec<(PathBuf, String)>> {
-    let mut out = Vec::new();
-    walk_dir(root, "", &mut out)?;
-
-    // Promote the README.md from the first-level crate subdirectory to the repo root
-    // so it renders on the HF repo page (e.g. "toto-rs/README.md" → "README.md").
-    let has_root_readme = out.iter().any(|(_, r)| r == "README.md");
-    if !has_root_readme {
-        for (_, remote) in &mut out {
-            let parts: Vec<&str> = remote.splitn(3, '/').collect();
-            if parts.len() == 2 && parts[1] == "README.md" {
-                *remote = "README.md".to_string();
-                break;
-            }
-        }
-    }
-
-    Ok(out)
-}
-
-fn walk_dir(dir: &Path, prefix: &str, out: &mut Vec<(PathBuf, String)>) -> anyhow::Result<()> {
-    let mut entries: Vec<_> = std::fs::read_dir(dir)
-        .with_context(|| format!("read_dir {}", dir.display()))?
-        .collect::<Result<_, _>>()?;
-    entries.sort_by_key(|e| e.file_name());
-    for entry in entries {
-        let path = entry.path();
-        let name = entry
-            .file_name()
-            .into_string()
-            .map_err(|_| anyhow::anyhow!("non-UTF-8 filename"))?;
-        if SKIP_DIRS.contains(&name.as_str()) {
-            continue;
-        }
-        let remote = if prefix.is_empty() { name.clone() } else { format!("{prefix}/{name}") };
-        if path.is_file() {
-            out.push((path, remote));
-        } else if path.is_dir() {
-            walk_dir(&path, &remote, out)?;
-        }
-    }
-    Ok(())
-}
-
-fn sha256_file(path: &Path) -> anyhow::Result<String> {
-    let mut f =
-        std::fs::File::open(path).with_context(|| format!("open {}", path.display()))?;
-    let mut hasher = Sha256::new();
-    let mut buf = vec![0u8; 8 * 1024 * 1024];
-    loop {
-        let n = f.read(&mut buf)?;
-        if n == 0 { break; }
-        hasher.update(&buf[..n]);
-    }
-    Ok(format!("{:x}", hasher.finalize()))
-}
-
-async fn upload_lfs(
-    client: &reqwest::Client,
-    repo_id: &str,
-    files: &[LfsFile],
-) -> anyhow::Result<()> {
-    let objects: Vec<_> = files
-        .iter()
-        .map(|f| serde_json::json!({ "oid": f.oid, "size": f.size }))
-        .collect();
-
-    // Request both multipart (for files >5 GB) and basic transfers
-    let url = format!("{HF_BASE}/{repo_id}.git/info/lfs/objects/batch");
-    let lfs_body = serde_json::to_string(&serde_json::json!({
-        "operation": "upload",
-        "transfers": ["multipart", "basic"],
-        "objects": objects,
-    }))?;
-    let resp: reqwest::Response = client
-        .post(&url)
-        .header("Content-Type", "application/vnd.git-lfs+json")
-        .header("Accept", "application/vnd.git-lfs+json")
-        .body(lfs_body)
-        .send()
-        .await
-        .context("LFS batch request")?;
-
-    let status = resp.status();
-    if !status.is_success() {
-        let body: String = resp.text().await.unwrap_or_default();
-        anyhow::bail!("LFS batch: HTTP {status}: {body}");
-    }
-
-    #[derive(serde::Deserialize)]
-    struct BatchResp { objects: Vec<serde_json::Value> }
-    let batch: BatchResp = resp.json().await.context("parse LFS batch response")?;
-
-    for (file, obj) in files.iter().zip(batch.objects.iter()) {
-        let Some(upload_href) =
-            obj.pointer("/actions/upload/href").and_then(|v: &serde_json::Value| v.as_str())
-        else {
-            println!("  (already on LFS) {}", file.remote);
-            continue;
-        };
-
-        // Multipart if the server provided chunk_size in the header
-        let is_multipart = obj.pointer("/actions/upload/header/chunk_size").is_some();
-
-        if is_multipart {
-            upload_lfs_object_multipart(file, upload_href, obj).await
-                .with_context(|| format!("upload (multipart) {}", file.remote))?;
-        } else {
-            upload_lfs_object(file, upload_href, obj).await
-                .with_context(|| format!("upload {}", file.remote))?;
-        }
-
-        // Verify step (optional but recommended by Git LFS spec)
-        if let Some(verify_href) =
-            obj.pointer("/actions/verify/href").and_then(|v: &serde_json::Value| v.as_str())
-        {
-            let verify_headers = obj
-                .pointer("/actions/verify/header")
-                .and_then(|v: &serde_json::Value| v.as_object())
-                .cloned()
-                .unwrap_or_default();
-
-            let mut vreq: reqwest::RequestBuilder = reqwest::Client::new()
-                .post(verify_href)
-                .header("Content-Type", "application/vnd.git-lfs+json")
-                .json(&serde_json::json!({ "oid": file.oid, "size": file.size }));
-            for (k, v) in &verify_headers {
-                if let Some(val) = v.as_str() {
-                    vreq = vreq.header(k.as_str(), val);
-                }
-            }
-            let vresp: reqwest::Response = vreq.send().await.context("LFS verify")?;
-            if !vresp.status().is_success() {
-                eprintln!("  Warning: LFS verify returned {}", vresp.status());
-            }
-        }
-    }
-
-    Ok(())
-}
-
-async fn upload_lfs_object(
-    file: &LfsFile,
-    href: &str,
-    obj: &serde_json::Value,
-) -> anyhow::Result<()> {
-    let pb = ProgressBar::new(file.size);
-    pb.set_style(
-        ProgressStyle::with_template(
-            "  {msg} [{bar:40}] {bytes}/{total_bytes} ({bytes_per_sec}, eta {eta})",
-        )
-        .unwrap()
-        .progress_chars("=>-"),
-    );
-    pb.set_message(file.remote.clone());
-
-    let f = tokio::fs::File::open(&file.local)
-        .await
-        .with_context(|| format!("open {}", file.local.display()))?;
-
-    let pb2 = pb.clone();
-    let stream = ReaderStream::new(f).map(move |chunk| {
-        if let Ok(ref b) = chunk {
-            pb2.inc(b.len() as u64);
-        }
-        chunk
-    });
-
-    // LFS upload goes to S3 / Azure — use a plain client (no HF auth header)
-    let mut req = reqwest::Client::new()
-        .put(href)
-        .header("Content-Length", file.size.to_string());
-
-    if let Some(extra) = obj.pointer("/actions/upload/header").and_then(|v| v.as_object()) {
-        for (k, v) in extra {
-            if let Some(val) = v.as_str() {
-                req = req.header(k.as_str(), val);
-            }
-        }
-    }
-
-    let resp = req
-        .body(reqwest::Body::wrap_stream(stream))
-        .send()
-        .await
-        .context("PUT LFS object")?;
-
-    pb.finish_and_clear();
-
-    let status = resp.status();
-    if !status.is_success() {
-        let body = resp.text().await.unwrap_or_default();
-        anyhow::bail!("LFS PUT HTTP {status}: {body}");
-    }
-
-    Ok(())
-}
-
-async fn upload_lfs_object_multipart(
-    file: &LfsFile,
-    complete_href: &str,
-    obj: &serde_json::Value,
-) -> anyhow::Result<()> {
-    use tokio::io::AsyncReadExt;
-
-    let header = obj
-        .pointer("/actions/upload/header")
-        .and_then(|v| v.as_object())
-        .ok_or_else(|| anyhow::anyhow!("no header in multipart LFS response"))?;
-
-    let chunk_size: usize = header
-        .get("chunk_size")
-        .and_then(|v| v.as_str())
-        .and_then(|s| s.parse().ok())
-        .ok_or_else(|| anyhow::anyhow!("missing chunk_size in multipart header"))?;
-
-    // Collect part URLs sorted numerically by key ("00001", "00002", …)
-    let mut parts: Vec<(u32, String)> = header
-        .iter()
-        .filter_map(|(k, v)| {
-            let n: u32 = k.parse().ok()?;
-            Some((n, v.as_str()?.to_string()))
-        })
-        .collect();
-    parts.sort_by_key(|(n, _)| *n);
-
-    let pb = ProgressBar::new(file.size);
-    pb.set_style(
-        ProgressStyle::with_template(
-            "  {msg} [{bar:40}] {bytes}/{total_bytes} ({bytes_per_sec}, eta {eta})",
-        )
-        .unwrap()
-        .progress_chars("=>-"),
-    );
-    pb.set_message(file.remote.clone());
-
-    let mut f = tokio::fs::File::open(&file.local)
-        .await
-        .with_context(|| format!("open {}", file.local.display()))?;
-
-    let s3 = reqwest::Client::new(); // plain client — S3 parts use pre-signed URLs
-    let mut etags: Vec<(u32, String)> = Vec::with_capacity(parts.len());
-
-    for (part_num, url) in &parts {
-        // Read up to chunk_size bytes for this part
-        let mut buf = vec![0u8; chunk_size];
-        let mut pos = 0;
-        while pos < chunk_size {
-            let n = f.read(&mut buf[pos..]).await?;
-            if n == 0 { break; }
-            pos += n;
-        }
-        if pos == 0 { break; }
-        buf.truncate(pos);
-        let len = buf.len();
-
-        // Retry up to 3 times on transient connection errors
-        const MAX_RETRIES: usize = 3;
-        let mut last_err: Option<anyhow::Error> = None;
-        let mut etag_opt: Option<String> = None;
-        for attempt in 0..MAX_RETRIES {
-            if attempt > 0 {
-                tokio::time::sleep(std::time::Duration::from_secs(2)).await;
-                eprintln!("  Retrying part {part_num} (attempt {})…", attempt + 1);
-            }
-            match s3
-                .put(url.as_str())
-                .header("Content-Length", len.to_string())
-                .body(buf.clone())
-                .send()
-                .await
-            {
-                Err(e) => {
-                    last_err = Some(anyhow::anyhow!("PUT part {part_num}: {e}"));
-                }
-                Ok(resp) => {
-                    let status = resp.status();
-                    let tag = resp.headers().get("etag")
-                        .and_then(|v| v.to_str().ok())
-                        .map(|s| s.to_string());
-                    if status.is_success() {
-                        if let Some(t) = tag {
-                            etag_opt = Some(t);
-                            last_err = None;
-                            break;
-                        } else {
-                            last_err = Some(anyhow::anyhow!("PUT part {part_num}: no ETag in response"));
-                        }
-                    } else {
-                        let body = resp.text().await.unwrap_or_default();
-                        last_err = Some(anyhow::anyhow!("PUT part {part_num}: HTTP {status}: {body}"));
-                    }
-                }
-            }
-        }
-        let etag = etag_opt.ok_or_else(|| {
-            last_err.unwrap_or_else(|| anyhow::anyhow!("PUT part {part_num}: exhausted retries"))
-        })?;
-
-        pb.inc(len as u64);
-        etags.push((*part_num, etag));
-    }
-
-    pb.finish_and_clear();
-
-    // Tell HF to assemble the parts on S3
-    let parts_json: Vec<serde_json::Value> = etags
-        .iter()
-        .map(|(n, e)| serde_json::json!({ "partNumber": n, "etag": e }))
-        .collect();
-
-    let resp = s3
-        .post(complete_href)
-        .json(&serde_json::json!({ "oid": file.oid, "parts": parts_json }))
-        .send()
-        .await
-        .context("complete multipart")?;
-
-    let status = resp.status();
-    if !status.is_success() {
-        let body = resp.text().await.unwrap_or_default();
-        anyhow::bail!("complete multipart: HTTP {status}: {body}");
-    }
-
-    Ok(())
-}
-
-async fn make_commit(
-    client: &reqwest::Client,
-    repo_id: &str,
-    regular: &[RegularFile],
-    lfs: &[LfsFile],
-    to_delete: &[String],
-) -> anyhow::Result<()> {
-    let mut lines: Vec<String> = Vec::new();
-
-    lines.push(serde_json::to_string(&serde_json::json!({
-        "key": "header",
-        "value": { "summary": "Upload model files", "description": "" },
-    }))?);
-
-    for rf in regular {
-        lines.push(serde_json::to_string(&serde_json::json!({
-            "key": "file",
-            "value": {
-                "path": rf.remote,
-                "encoding": "base64",
-                "content": rf.content_b64,
-            },
-        }))?);
-    }
-
-    for lf in lfs {
-        lines.push(serde_json::to_string(&serde_json::json!({
-            "key": "lfsFile",
-            "value": {
-                "path": lf.remote,
-                "algo": "sha256",
-                "oid": lf.oid,
-                "size": lf.size,
-            },
-        }))?);
-    }
-
-    for path in to_delete {
-        lines.push(serde_json::to_string(&serde_json::json!({
-            "key": "deletedEntry",
-            "value": { "path": path },
-        }))?);
-    }
-
-    let body = lines.join("\n");
-
-    let url = format!("{HF_BASE}/api/models/{repo_id}/commit/main");
-    let resp = client
-        .post(&url)
-        .header("Content-Type", "application/x-ndjson")
-        .body(body)
-        .send()
-        .await
-        .context("POST commit")?;
-
-    let status = resp.status();
-    if !status.is_success() {
-        let body = resp.text().await.unwrap_or_default();
-        anyhow::bail!("commit: HTTP {status}: {body}");
-    }
-
-    Ok(())
-}