Skip to main content

oxibonsai_model/convert/mlx_image/
error.rs

1//! Error types for the MLX (FLUX.2 DiT) → OxiBonsai GGUF conversion pipeline.
2
3use std::path::PathBuf;
4
5use thiserror::Error;
6
7/// Errors raised while packing a single MLX-quantized linear module into
8/// `BlockTQ2_0_g128` blocks.
9///
10/// These are the *parity guards* described in the converter design: any
11/// violation means the MLX tensor does not match the validated ternary
12/// assumptions, so we refuse to produce a silently-wrong GGUF file.
13#[derive(Debug, Error)]
14pub enum PackError {
15    /// The packed-weight column count does not equal `in_features / 16`.
16    #[error(
17        "module '{module}': weight has {got} columns, expected in/16 = {expected} \
18         (in_features = {in_features})"
19    )]
20    WeightColumnsMismatch {
21        /// Diffusers module name.
22        module: String,
23        /// Observed `weight` column count.
24        got: usize,
25        /// Expected column count (`in_features / 16`).
26        expected: usize,
27        /// Logical input-feature dimension.
28        in_features: usize,
29    },
30
31    /// The scales/biases column count does not equal `in_features / 128`.
32    #[error(
33        "module '{module}': {which} has {got} columns, expected in/128 = {expected} \
34         (in_features = {in_features})"
35    )]
36    GroupColumnsMismatch {
37        /// Diffusers module name.
38        module: String,
39        /// Which sub-tensor (`"scales"` or `"biases"`).
40        which: &'static str,
41        /// Observed column count.
42        got: usize,
43        /// Expected column count (`in_features / 128`).
44        expected: usize,
45        /// Logical input-feature dimension.
46        in_features: usize,
47    },
48
49    /// A sub-tensor buffer had an unexpected element count for the stated shape.
50    #[error("module '{module}': {which} has {got} elements, expected {expected}")]
51    BufferLengthMismatch {
52        /// Diffusers module name.
53        module: String,
54        /// Which sub-tensor (`"weight"`, `"scales"`, `"biases"`).
55        which: &'static str,
56        /// Observed element count.
57        got: usize,
58        /// Expected element count from `out × cols`.
59        expected: usize,
60    },
61
62    /// `in_features` is not a positive multiple of 128 (the TQ2_0_g128 group size).
63    #[error("module '{module}': in_features = {in_features} is not a positive multiple of 128")]
64    InFeaturesNotAligned {
65        /// Diffusers module name.
66        module: String,
67        /// Logical input-feature dimension.
68        in_features: usize,
69    },
70
71    /// A 2-bit MLX code exceeded 2 (i.e. a reserved `q=3` was found), which is
72    /// inconsistent with the validated ternary assumption (`q ∈ {0, 1, 2}`).
73    #[error(
74        "module '{module}' [row {row}, group {group}]: 2-bit code value {value} > 2 \
75         (reserved q=3 found; tensor is not ternary)"
76    )]
77    CodeOutOfRange {
78        /// Diffusers module name.
79        module: String,
80        /// Output-feature row index.
81        row: usize,
82        /// 128-element group index along the input dimension.
83        group: usize,
84        /// Observed code value (always > 2 when this error is raised).
85        value: u8,
86    },
87
88    /// The affine bias was not exactly `-scale`, breaking the symmetric-ternary
89    /// assumption (`w = scale·(q-1)`).
90    #[error(
91        "module '{module}' [row {row}, group {group}]: bias ({bias}) != -scale (-{scale}); \
92         affine quantization is not symmetric ternary"
93    )]
94    AsymmetricBias {
95        /// Diffusers module name.
96        module: String,
97        /// Output-feature row index.
98        row: usize,
99        /// 128-element group index along the input dimension.
100        group: usize,
101        /// Decoded bias value (f32).
102        bias: f32,
103        /// Decoded scale value (f32).
104        scale: f32,
105    },
106}
107
108/// Top-level error for the MLX FLUX.2 DiT importer.
109#[derive(Debug, Error)]
110pub enum MlxImageImportError {
111    /// I/O failure while opening or memory-mapping the safetensors input.
112    #[error("I/O error for {path:?}: {source}")]
113    Io {
114        /// Path that was being accessed when the error occurred.
115        path: PathBuf,
116        /// Underlying I/O error.
117        #[source]
118        source: std::io::Error,
119    },
120
121    /// The safetensors container could not be parsed.
122    #[error("failed to parse safetensors file {path:?}: {msg}")]
123    Parse {
124        /// Path to the safetensors file.
125        path: PathBuf,
126        /// Human-readable parser message.
127        msg: String,
128    },
129
130    /// A sub-tensor of a quantized module was missing or had the wrong dtype.
131    #[error("module '{module}': {reason}")]
132    BadModule {
133        /// Diffusers module name.
134        module: String,
135        /// Human-readable explanation.
136        reason: String,
137    },
138
139    /// Packing a quantized module into ternary blocks failed a parity guard.
140    #[error("packing failed: {0}")]
141    Pack(#[from] PackError),
142
143    /// The underlying GGUF writer failed.
144    #[error("GGUF writer error: {0}")]
145    GgufWrite(String),
146
147    /// An unsupported quantisation format string was requested.
148    #[error("unsupported quantisation format '{0}'; only 'tq2_0_g128' is supported")]
149    UnsupportedQuant(String),
150}