andodeki/makepad
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
makepad/libs/ai/models/paint/src/safetensors.rs
Admin 7f59912916 libs/ai: one AI stack, replacing libs/ggml, llama, mlx, cuda, tts, voice2 and pbr_paint 
The model code was spread across eight crates that had grown into each other:
ggml and cuda and mlx each owned part of a tensor runtime, llama and tts and
voice2 each owned part of a model, and libs/diffusion owned everything else.
They are now one tree with an explicit shape:

  libs/ai/cuda     — kernels and launch surface
  libs/ai/metal    — Metal shaders and the shim
  libs/ai/llm      — the language-model runtime (sessions, lanes, contexts,
                     the CUDA and Metal executors, the compiled Metal path)
  libs/ai/models/  — common, flux, h3, music, paint, speech, stems, vision

libs/diffusion is not deleted but demoted: what remains is the VALIDATOR
crate — several dozen `*_validate.rs` oracles that check a native
implementation against a reference, which is where they belong now that the
implementations live next door.

The functional work inside the move is mostly in the LLM runtime: N lanes that
draft while one verify batch serves all of them, per-slot prefill over a shared
folded attention arena, speculation that survives batching, and a scheduler
that reports rather than publishes. And in the CUDA build: a machine without
usable CUDA must still LINK (and say so), the default kernel arch is the
building machine's GPU, `NO_CUDA` forces the stub even where the toolkit
exists, and kernels compile in parallel with progress.

libs/video_flow is new here: classical optical flow estimation and the `mkfl`
motion-field payload — a flow field measured from a clip without a model,
which is what drives free-rate bounce-looping playback and the uprez/tween
enhance pipe.
2026-08-23 01:34:35 +02:00

659 lines
24 KiB
Rust
Raw Permalink Blame History

//! Bounded, dependency-free safetensors indexer for the pinned DINOv2-Giant
//! checkpoint. The 4.55 GB payload is never copied wholesale: only the JSON
//! header is parsed, then tensors are sought and read one at a time in archive
//! offset order.
//!
//! The parser intentionally accepts only the scalar types the native Paint
//! graph consumes. Unsupported dtypes, malformed shapes, overlapping ranges,
//! duplicate names, and out-of-file offsets fail closed before any upload.
use crate::torch_bin::TorchDtype;
use std::collections::BTreeSet;
use std::io::{Read, Seek, SeekFrom};
const MAX_HEADER_BYTES: usize = 64 * 1024 * 1024;
const MAX_TENSORS: usize = 16_384;
const MAX_TENSOR_RANK: usize = 16;
const MAX_TENSOR_NAME_BYTES: usize = 4_096;
const MAX_JSON_STRING_BYTES: usize = 1024 * 1024;
const MAX_JSON_DEPTH: usize = 64;
const MAX_JSON_VALUES: usize = 1_000_000;
const MAX_TENSOR_BYTES: usize = 8 * 1024 * 1024 * 1024;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SafeTensorError {
Io(String),
Header(String),
Unsupported(String),
Tensor(String),
}
impl std::fmt::Display for SafeTensorError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
SafeTensorError::Io(message) => write!(f, "io: {message}"),
SafeTensorError::Header(message) => write!(f, "header: {message}"),
SafeTensorError::Unsupported(message) => write!(f, "unsupported: {message}"),
SafeTensorError::Tensor(message) => write!(f, "tensor: {message}"),
}
}
}
impl std::error::Error for SafeTensorError {}
impl From<std::io::Error> for SafeTensorError {
fn from(value: std::io::Error) -> Self {
Self::Io(value.to_string())
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SafeTensorRecord {
pub name: String,
pub dtype: TorchDtype,
pub shape: Vec<usize>,
pub numel: usize,
/// Relative to the safetensors data section (immediately after header).
pub data_start: usize,
pub data_end: usize,
}
#[derive(Debug)]
pub struct SafeTensorIndex {
pub tensors: Vec<SafeTensorRecord>,
data_offset: usize,
file_len: usize,
}
impl SafeTensorIndex {
pub fn find(&self, name: &str) -> Option<&SafeTensorRecord> {
self.tensors.iter().find(|record| record.name == name)
}
pub fn archive_offset(&self, record: &SafeTensorRecord) -> Result<usize, SafeTensorError> {
self.checked_record_range(record).map(|(start, _)| start)
}
fn checked_record_range(
&self,
record: &SafeTensorRecord,
) -> Result<(usize, usize), SafeTensorError> {
if record.name.len() > MAX_TENSOR_NAME_BYTES || record.shape.len() > MAX_TENSOR_RANK {
return Err(SafeTensorError::Tensor(format!(
"{} name or rank exceeds bounded limits",
record.name
)));
}
if record.data_end < record.data_start {
return Err(SafeTensorError::Tensor(format!(
"{} has reversed data range",
record.name
)));
}
let numel = record.shape.iter().try_fold(1usize, |product, dimension| {
product.checked_mul(*dimension).ok_or_else(|| {
SafeTensorError::Tensor(format!("{} shape product overflow", record.name))
})
})?;
if numel != record.numel {
return Err(SafeTensorError::Tensor(format!(
"{} numel does not match shape",
record.name
)));
}
let expected = numel.checked_mul(record.dtype.elem_size()).ok_or_else(|| {
SafeTensorError::Tensor(format!("{} byte length overflow", record.name))
})?;
let len = record
.data_end
.checked_sub(record.data_start)
.ok_or_else(|| SafeTensorError::Tensor(format!("{} range underflow", record.name)))?;
if len != expected {
return Err(SafeTensorError::Tensor(format!(
"{} range length does not match shape/dtype",
record.name
)));
}
if len > MAX_TENSOR_BYTES {
return Err(SafeTensorError::Unsupported(format!(
"{} tensor is {len} bytes (limit {MAX_TENSOR_BYTES})",
record.name
)));
}
let start = self
.data_offset
.checked_add(record.data_start)
.ok_or_else(|| SafeTensorError::Tensor(format!("{} offset overflow", record.name)))?;
if start.checked_add(len).is_none_or(|end| end > self.file_len) {
return Err(SafeTensorError::Tensor(format!(
"{} range outside archive",
record.name
)));
}
Ok((start, len))
}
pub fn read_tensor_into<R: Read + Seek>(
&self,
archive: &mut R,
record: &SafeTensorRecord,
out: &mut Vec<u8>,
) -> Result<(), SafeTensorError> {
let (start, len) = self.checked_record_range(record)?;
archive.seek(SeekFrom::Start(start as u64))?;
out.clear();
out.try_reserve_exact(len)
.map_err(|_| SafeTensorError::Tensor(format!("{} allocation failed", record.name)))?;
out.resize(len, 0);
archive.read_exact(out)?;
Ok(())
}
}
struct Parser<'a> {
bytes: &'a [u8],
at: usize,
values_seen: usize,
}
impl<'a> Parser<'a> {
fn new(bytes: &'a [u8]) -> Self {
Self {
bytes,
at: 0,
values_seen: 0,
}
}
fn error(&self, message: impl AsRef<str>) -> SafeTensorError {
SafeTensorError::Header(format!("{} at byte {}", message.as_ref(), self.at))
}
fn whitespace(&mut self) {
while self
.bytes
.get(self.at)
.is_some_and(|byte| byte.is_ascii_whitespace())
{
self.at += 1;
}
}
fn take(&mut self, expected: u8) -> Result<(), SafeTensorError> {
self.whitespace();
if self.bytes.get(self.at) != Some(&expected) {
return Err(self.error(format!("expected '{}'", expected as char)));
}
self.at += 1;
Ok(())
}
fn consume(&mut self, byte: u8) -> bool {
self.whitespace();
if self.bytes.get(self.at) == Some(&byte) {
self.at += 1;
true
} else {
false
}
}
fn string(&mut self) -> Result<String, SafeTensorError> {
self.take(b'"')?;
let mut out = Vec::new();
loop {
let byte = *self
.bytes
.get(self.at)
.ok_or_else(|| self.error("unterminated string"))?;
self.at += 1;
match byte {
b'"' => {
return String::from_utf8(out)
.map_err(|_| self.error("string is not valid UTF-8"));
}
b'\\' => {
let escaped = *self
.bytes
.get(self.at)
.ok_or_else(|| self.error("unterminated escape"))?;
self.at += 1;
match escaped {
b'"' | b'\\' | b'/' => out.push(escaped),
b'b' => out.push(8),
b'f' => out.push(12),
b'n' => out.push(b'\n'),
b'r' => out.push(b'\r'),
b't' => out.push(b'\t'),
b'u' => {
let end = self
.at
.checked_add(4)
.ok_or_else(|| self.error("unicode escape overflow"))?;
let digits = self
.bytes
.get(self.at..end)
.ok_or_else(|| self.error("short unicode escape"))?;
let text = std::str::from_utf8(digits)
.map_err(|_| self.error("bad unicode escape"))?;
let value = u16::from_str_radix(text, 16)
.map_err(|_| self.error("bad unicode escape"))?;
let ch = char::from_u32(value as u32)
.ok_or_else(|| self.error("surrogate escape is unsupported"))?;
let mut encoded = [0u8; 4];
out.extend_from_slice(ch.encode_utf8(&mut encoded).as_bytes());
self.at = end;
}
_ => return Err(self.error("unknown string escape")),
}
}
0..=31 => return Err(self.error("control byte in string")),
_ => out.push(byte),
}
if out.len() > MAX_JSON_STRING_BYTES {
return Err(self.error("JSON string exceeds bounded length"));
}
}
}
fn usize(&mut self) -> Result<usize, SafeTensorError> {
self.whitespace();
let start = self.at;
while self
.bytes
.get(self.at)
.is_some_and(|byte| byte.is_ascii_digit())
{
self.at += 1;
}
if start == self.at {
return Err(self.error("expected unsigned integer"));
}
let text = std::str::from_utf8(&self.bytes[start..self.at])
.map_err(|_| self.error("invalid integer"))?;
text.parse::<usize>()
.map_err(|_| self.error("integer exceeds address space"))
}
fn usize_array(&mut self) -> Result<Vec<usize>, SafeTensorError> {
self.take(b'[')?;
let mut out = Vec::new();
if self.consume(b']') {
return Ok(out);
}
loop {
if out.len() >= MAX_TENSOR_RANK {
return Err(self.error("integer array exceeds rank/value limit"));
}
out.push(self.usize()?);
if self.consume(b']') {
return Ok(out);
}
self.take(b',')?;
}
}
fn literal(&mut self, literal: &[u8]) -> Result<(), SafeTensorError> {
self.whitespace();
let end = self
.at
.checked_add(literal.len())
.ok_or_else(|| self.error("JSON literal offset overflow"))?;
if self.bytes.get(self.at..end) != Some(literal) {
return Err(self.error("bad JSON literal"));
}
self.at = end;
Ok(())
}
fn skip_number(&mut self) -> Result<(), SafeTensorError> {
self.whitespace();
let start = self.at;
while self.bytes.get(self.at).is_some_and(|byte| {
byte.is_ascii_digit() || matches!(*byte, b'-' | b'+' | b'.' | b'e' | b'E')
}) {
self.at += 1;
}
if start == self.at {
return Err(self.error("expected JSON number"));
}
Ok(())
}
fn skip_value(&mut self, depth: usize) -> Result<(), SafeTensorError> {
if depth > MAX_JSON_DEPTH {
return Err(self.error("JSON nesting exceeds bounded depth"));
}
self.values_seen = self
.values_seen
.checked_add(1)
.ok_or_else(|| self.error("JSON value count overflow"))?;
if self.values_seen > MAX_JSON_VALUES {
return Err(self.error("JSON value count exceeds limit"));
}
self.whitespace();
match self.bytes.get(self.at).copied() {
Some(b'"') => {
self.string()?;
Ok(())
}
Some(b'{') => {
self.take(b'{')?;
if self.consume(b'}') {
return Ok(());
}
loop {
self.string()?;
self.take(b':')?;
self.skip_value(depth + 1)?;
if self.consume(b'}') {
return Ok(());
}
self.take(b',')?;
}
}
Some(b'[') => {
self.take(b'[')?;
if self.consume(b']') {
return Ok(());
}
loop {
self.skip_value(depth + 1)?;
if self.consume(b']') {
return Ok(());
}
self.take(b',')?;
}
}
Some(b't') => self.literal(b"true"),
Some(b'f') => self.literal(b"false"),
Some(b'n') => self.literal(b"null"),
Some(_) => self.skip_number(),
None => Err(self.error("expected JSON value")),
}
}
fn tensor(&mut self, name: String) -> Result<SafeTensorRecord, SafeTensorError> {
if name.len() > MAX_TENSOR_NAME_BYTES {
return Err(self.error("tensor name exceeds bounded length"));
}
self.take(b'{')?;
let mut dtype = None;
let mut shape = None;
let mut offsets = None;
if !self.consume(b'}') {
loop {
let field = self.string()?;
self.take(b':')?;
match field.as_str() {
"dtype" => {
if dtype.is_some() {
return Err(self.error(format!("{name}: duplicate dtype")));
}
let text = self.string()?;
dtype = Some(match text.as_str() {
"F32" => TorchDtype::F32,
"F16" => TorchDtype::F16,
"BF16" => TorchDtype::BF16,
"I64" => TorchDtype::I64,
_ => {
return Err(SafeTensorError::Unsupported(format!(
"{name}: safetensors dtype {text}"
)))
}
});
}
"shape" => {
if shape.is_some() {
return Err(self.error(format!("{name}: duplicate shape")));
}
shape = Some(self.usize_array()?);
}
"data_offsets" => {
if offsets.is_some() {
return Err(self.error(format!("{name}: duplicate data_offsets")));
}
let values = self.usize_array()?;
if values.len() != 2 {
return Err(self.error(format!(
"{name}: data_offsets needs two integers"
)));
}
offsets = Some((values[0], values[1]));
}
_ => self.skip_value(0)?,
}
if self.consume(b'}') {
break;
}
self.take(b',')?;
}
}
let dtype = dtype.ok_or_else(|| self.error(format!("{name}: missing dtype")))?;
let shape = shape.ok_or_else(|| self.error(format!("{name}: missing shape")))?;
let (data_start, data_end) =
offsets.ok_or_else(|| self.error(format!("{name}: missing data_offsets")))?;
if data_end < data_start {
return Err(self.error(format!("{name}: reversed data_offsets")));
}
let numel = shape.iter().try_fold(1usize, |product, dim| {
product.checked_mul(*dim).ok_or_else(|| {
SafeTensorError::Tensor(format!("{name}: shape product overflow"))
})
})?;
let expected = numel
.checked_mul(dtype.elem_size())
.ok_or_else(|| SafeTensorError::Tensor(format!("{name}: byte length overflow")))?;
if data_end - data_start != expected {
return Err(SafeTensorError::Tensor(format!(
"{name}: range is {} bytes, shape/dtype require {expected}",
data_end - data_start
)));
}
Ok(SafeTensorRecord {
name,
dtype,
shape,
numel,
data_start,
data_end,
})
}
fn header(mut self) -> Result<Vec<SafeTensorRecord>, SafeTensorError> {
self.take(b'{')?;
let mut tensors = Vec::new();
let mut names = BTreeSet::new();
if !self.consume(b'}') {
loop {
let name = self.string()?;
self.take(b':')?;
if name == "__metadata__" {
self.skip_value(0)?;
} else {
if !names.insert(name.clone()) {
return Err(self.error(format!("duplicate tensor {name}")));
}
if tensors.len() >= MAX_TENSORS {
return Err(self.error("tensor count exceeds bounded limit"));
}
tensors.push(self.tensor(name)?);
}
if self.consume(b'}') {
break;
}
self.take(b',')?;
}
}
self.whitespace();
if self.at != self.bytes.len() {
return Err(self.error("trailing header data"));
}
Ok(tensors)
}
}
pub fn read_index_from<R: Read + Seek>(archive: &mut R) -> Result<SafeTensorIndex, SafeTensorError> {
let file_len_u64 = archive.seek(SeekFrom::End(0))?;
let file_len = usize::try_from(file_len_u64)
.map_err(|_| SafeTensorError::Unsupported("archive exceeds address space".to_string()))?;
if file_len < 8 {
return Err(SafeTensorError::Header("file shorter than length prefix".to_string()));
}
archive.seek(SeekFrom::Start(0))?;
let mut prefix = [0u8; 8];
archive.read_exact(&mut prefix)?;
let header_len_u64 = u64::from_le_bytes(prefix);
let header_len = usize::try_from(header_len_u64)
.map_err(|_| SafeTensorError::Unsupported("header exceeds address space".to_string()))?;
if header_len == 0 || header_len > MAX_HEADER_BYTES {
return Err(SafeTensorError::Unsupported(format!(
"safetensors header is {header_len} bytes (limit {MAX_HEADER_BYTES})"
)));
}
let data_offset = 8usize
.checked_add(header_len)
.ok_or_else(|| SafeTensorError::Header("header offset overflow".to_string()))?;
if data_offset > file_len {
return Err(SafeTensorError::Header("header extends beyond file".to_string()));
}
let mut header = Vec::new();
header
.try_reserve_exact(header_len)
.map_err(|_| SafeTensorError::Header("header allocation failed".to_string()))?;
header.resize(header_len, 0);
archive.read_exact(&mut header)?;
let tensors = Parser::new(&header).header()?;
let data_len = file_len - data_offset;
let mut ranges: Vec<(&str, usize, usize)> = tensors
.iter()
.map(|record| (record.name.as_str(), record.data_start, record.data_end))
.collect();
ranges.sort_by_key(|(_, start, end)| (*start, *end));
let mut prior_end = 0usize;
for (name, start, end) in ranges {
if end > data_len {
return Err(SafeTensorError::Tensor(format!(
"{name}: data range {start}..{end} exceeds payload {data_len}"
)));
}
if start < prior_end && start != end {
return Err(SafeTensorError::Tensor(format!(
"{name}: data range overlaps a previous tensor"
)));
}
prior_end = prior_end.max(end);
}
Ok(SafeTensorIndex {
tensors,
data_offset,
file_len,
})
}
#[cfg(test)]
mod tests {
use super::*;
fn archive(header: &str, data: &[u8]) -> Vec<u8> {
let mut header = header.as_bytes().to_vec();
while header.len() % 8 != 0 {
header.push(b' ');
}
let mut out = Vec::new();
out.extend_from_slice(&(header.len() as u64).to_le_bytes());
out.extend_from_slice(&header);
out.extend_from_slice(data);
out
}
#[test]
fn indexes_and_streams_in_data_offset_order() {
let bytes = archive(
r#"{"z.weight":{"dtype":"F32","shape":[2],"data_offsets":[4,12]},"__metadata__":{"format":"pt"},"a.bias":{"dtype":"F16","shape":[2],"data_offsets":[0,4]}}"#,
&[0x00, 0x3c, 0x00, 0xc0, 0, 0, 0x80, 0x3f, 0, 0, 0, 0xc0],
);
let mut cursor = std::io::Cursor::new(bytes);
let index = read_index_from(&mut cursor).unwrap();
assert_eq!(index.tensors.len(), 2);
assert_eq!(index.find("z.weight").unwrap().shape, vec![2]);
assert_eq!(index.find("a.bias").unwrap().dtype, TorchDtype::F16);
let mut data = Vec::new();
index
.read_tensor_into(&mut cursor, index.find("a.bias").unwrap(), &mut data)
.unwrap();
assert_eq!(data, vec![0x00, 0x3c, 0x00, 0xc0]);
assert!(
index.archive_offset(index.find("a.bias").unwrap()).unwrap()
< index.archive_offset(index.find("z.weight").unwrap()).unwrap()
);
}
#[test]
fn rejects_shape_range_mismatch() {
let bytes = archive(
r#"{"bad":{"dtype":"F32","shape":[2],"data_offsets":[0,4]}}"#,
&[0; 4],
);
let error = read_index_from(&mut std::io::Cursor::new(bytes)).unwrap_err();
assert!(matches!(error, SafeTensorError::Tensor(message) if message.contains("require 8")));
}
#[test]
fn rejects_overlapping_ranges() {
let bytes = archive(
r#"{"a":{"dtype":"F32","shape":[2],"data_offsets":[0,8]},"b":{"dtype":"F32","shape":[2],"data_offsets":[4,12]}}"#,
&[0; 12],
);
let error = read_index_from(&mut std::io::Cursor::new(bytes)).unwrap_err();
assert!(matches!(error, SafeTensorError::Tensor(message) if message.contains("overlaps")));
}
#[test]
fn rejects_unsupported_dtype() {
let bytes = archive(
r#"{"bad":{"dtype":"F64","shape":[1],"data_offsets":[0,8]}}"#,
&[0; 8],
);
let error = read_index_from(&mut std::io::Cursor::new(bytes)).unwrap_err();
assert!(matches!(error, SafeTensorError::Unsupported(message) if message.contains("F64")));
}
#[test]
fn rejects_deep_metadata_and_excess_tensor_rank() {
let nested = format!(
"{{\"__metadata__\":{}0{}}}",
"[".repeat(MAX_JSON_DEPTH + 2),
"]".repeat(MAX_JSON_DEPTH + 2)
);
let error = read_index_from(&mut std::io::Cursor::new(archive(&nested, &[]))).unwrap_err();
assert!(matches!(error, SafeTensorError::Header(message) if message.contains("depth")));
let dimensions = std::iter::repeat_n("1", MAX_TENSOR_RANK + 1)
.collect::<Vec<_>>()
.join(",");
let header = format!(
"{{\"bad\":{{\"dtype\":\"F16\",\"shape\":[{dimensions}],\"data_offsets\":[0,2]}}}}"
);
let error =
read_index_from(&mut std::io::Cursor::new(archive(&header, &[0; 2]))).unwrap_err();
assert!(matches!(error, SafeTensorError::Header(message) if message.contains("limit")));
}
#[test]
fn public_record_ranges_are_revalidated_before_reading() {
let bytes = archive(
r#"{"a":{"dtype":"F16","shape":[2],"data_offsets":[0,4]}}"#,
&[0; 4],
);
let mut cursor = std::io::Cursor::new(bytes);
let index = read_index_from(&mut cursor).unwrap();
let mut hostile = index.find("a").unwrap().clone();
hostile.data_start = 4;
hostile.data_end = 1;
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
index.read_tensor_into(&mut cursor, &hostile, &mut Vec::new())
}));
assert!(matches!(result, Ok(Err(SafeTensorError::Tensor(_)))));
}
}
Reference in a new issue View git blame Copy permalink