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}