aleleba/lightcraft
1
0
Fork 0
mirror of https://github.com/storytold/lightcraft.git synced 2026-10-29 12:37:51 -06:00
Code Issues Packages Projects Releases Wiki Activity

M12: lightcraft-segment: Apache-2.0, checked weights, pure-Rust model download

Browse Source 
- The crate ports Apache-2.0-only Hugging Face transformers code: it is
  now `license = "Apache-2.0"`, and every ported file carries the
  modification notice Apache-2.0 §4(b) asks for.
- Weights::open checks every tensor's byte range (and the header)
  against the file's length: a truncated or damaged download is an
  error at load, never a short read later; regression test with
  hostile offsets.
- fetch/: the model download, only when asked for. An ordered mirror
  list (env LIGHTCRAFT_SAM3_MIRRORS, a mirrors file, DEFAULT_MIRRORS —
  empty until LightCraft's CDN exists, TODO for the maintainer),
  fallback to the next mirror, resume via Range, redirects, connect and
  stall timeouts, cancel, size caps, exact size + SHA-256 (pinned from
  tools/install-sam3.sh) before an atomic rename; damaged files are
  deleted, HTTP error bodies never written.
- HTTP/1.1 client over std sockets + rustls with the RustCrypto provider
  (rustls-rustcrypto) and webpki-roots: no ring / aws-lc-rs, nothing
  compiled from C (no new `cc` users in the app's build graph).
- Tests against a local TcpListener server (no internet): fallback,
  redirects, chunked bodies, resume, 416, hash mismatch, too much data,
  stalls, cancel.
- *.safetensors are gitignored.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Brandon Thomas
2026-10-07 09:56:49 -04:00
co-authored by Claude Opus 5.5
parent fefd61c967
commit b52f43d1d5
16 changed files with 2119 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+3
View File
@@ -8,3 +8,6 @@ plan/
.DS_Store
.claude/
/craft-fonts
# model weights are never committed (SAM 3: docs/ai-masks.md)
*.safetensors
*.safetensors.part
Cargo.lock
Generated
+743 -12
View File
File diff suppressed because it is too large Load Diff
crates/segment/Cargo.toml
+5 -1
View File
@@ -2,7 +2,7 @@
name = "lightcraft-segment"
version.workspace = true
edition.workspace = true
license.workspace = true
license = "Apache-2.0"
description = "LightCraft segmentation: SAM 3 (Segment Anything with Concepts) inference in pure Rust (L3)"
[lints]
@@ -18,6 +18,10 @@ log = { workspace = true }
[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
candle-core = { workspace = true }
candle-nn = { workspace = true }
rustls = { version = "0.23", default-features = false, features = ["std", "tls12"] }
rustls-rustcrypto = "0.0.2-alpha"
webpki-roots = "1"
sha2 = { version = "0.10", default-features = false, features = ["std"] }
[target.'cfg(target_os = "macos")'.dependencies]
candle-core = { workspace = true, features = ["metal"] }
crates/segment/src/clip.rs
+4
View File
@@ -1,5 +1,9 @@
//! The CLIP text encoder (24 layers, 1024 wide, causal) that turns a prompt into the token
//! features SAM 3's detector attends to.
//!
//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from
//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`, CLIP text model), Copyright The HuggingFace
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
use candle_core::{Device, Module, Tensor};
use candle_nn::{Embedding, LayerNorm, VarBuilder};
crates/segment/src/detector.rs
+4
View File
@@ -1,6 +1,10 @@
//! Text prompts (SAM 3's concept detector): the prompt's CLIP features fused into the image
//! features by a DETR encoder, 200 object queries refined by a DETR decoder (with box
//! relative-position bias and a presence token), per-query masks from a pixel decoder.
//!
//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from
//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`), Copyright The HuggingFace
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
use candle_core::{D, Device, IndexOp, Module, Tensor};
use candle_nn::{Conv2d, Conv2dConfig, GroupNorm, LayerNorm, Linear, VarBuilder};
crates/segment/src/fetch/http.rs
+541
View File
@@ -0,0 +1,541 @@
//! A small HTTP/1.1 GET client for model downloads, in pure Rust: `std::net` sockets and, for
//! `https://`, rustls with the RustCrypto provider (`rustls-rustcrypto`; no `ring` or
//! `aws-lc-rs`, so nothing is compiled from C or assembly) and the Mozilla root certificates
//! (`webpki-roots`).
//!
//! Only what a large-file download needs: one request per connection (`Connection: close`),
//! `Range` requests for resuming, redirects (followed by the caller), `Content-Length`,
//! chunked and read-to-close bodies. Every read waits at most `stall` for data and checks the
//! cancel flag at least twice a second, so a dead server or a cancelled download never hangs
//! the downloading thread. No proxies (see docs/ai-masks.md).
use std::io::{ErrorKind, Read, Write};
use std::net::{TcpStream, ToSocketAddrs};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::{Duration, Instant};
/// Why a request failed.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum HttpError {
/// The URL can't be used (scheme, host, characters).
BadUrl(String),
/// Name resolution or connecting failed (or took longer than the connect timeout).
Connect(String),
/// No data arrived for the stall timeout.
Stalled,
/// The cancel flag was raised.
Cancelled,
/// TLS failed (certificate, handshake).
Tls(String),
/// The server's response could not be understood.
Protocol(String),
/// Any other I/O error.
Io(String),
}
impl std::fmt::Display for HttpError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
HttpError::BadUrl(e) => write!(f, "unusable URL: {e}"),
HttpError::Connect(e) => write!(f, "could not connect: {e}"),
HttpError::Stalled => write!(f, "the server stopped sending data"),
HttpError::Cancelled => write!(f, "cancelled"),
HttpError::Tls(e) => write!(f, "secure connection failed: {e}"),
HttpError::Protocol(e) => write!(f, "unexpected response: {e}"),
HttpError::Io(e) => write!(f, "{e}"),
}
}
}
/// A parsed `http://` or `https://` URL.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Url {
pub tls: bool,
pub host: String,
pub port: u16,
/// Path and query, starting with `/`.
pub path: String,
}
impl Url {
pub fn parse(s: &str) -> Result<Url, HttpError> {
let bad = |why: &str| HttpError::BadUrl(format!("{s}: {why}"));
let s = s.trim();
if s.chars().any(|c| c.is_control() || c == ' ') {
return Err(bad("contains spaces or control characters"));
}
let (tls, rest) = if let Some(r) = s.strip_prefix("https://") {
(true, r)
} else if let Some(r) = s.strip_prefix("http://") {
(false, r)
} else {
return Err(bad("only http:// and https:// are supported"));
};
let rest = rest.split('#').next().unwrap_or_default();
let (authority, path) = match rest.find(['/', '?']) {
Some(i) => (rest.get(..i).unwrap_or_default(), rest.get(i..).unwrap_or_default()),
None => (rest, ""),
};
let path = if path.starts_with('/') { path.to_string() } else { format!("/{path}") };
if authority.contains('@') {
return Err(bad("user names and passwords in URLs are not supported"));
}
let (host, port) = if let Some(v6) = authority.strip_prefix('[') {
let end = v6.find(']').ok_or_else(|| bad("unterminated IPv6 address"))?;
let host = v6.get(..end).unwrap_or_default();
let after = v6.get(end + 1..).unwrap_or_default();
(host.to_string(), after.strip_prefix(':'))
} else {
match authority.rsplit_once(':') {
Some((h, p)) => (h.to_string(), Some(p)),
None => (authority.to_string(), None),
}
};
let port = match port {
Some(p) => p.parse::<u16>().ok().filter(|p| *p != 0).ok_or_else(|| bad("bad port"))?,
None if tls => 443,
None => 80,
};
if host.is_empty() || !host.chars().all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '.' | ':' | '_')) {
return Err(bad("bad host name"));
}
Ok(Url { tls, host: host.to_ascii_lowercase(), port, path })
}
/// `Location` of a redirect, resolved against this URL.
pub fn join(&self, location: &str) -> Result<Url, HttpError> {
let l = location.trim();
if l.starts_with("http://") || l.starts_with("https://") {
return Url::parse(l);
}
// another scheme (`ftp:`, `file:`, `data:`…)
if let Some((scheme, _)) = l.split_once(':')
&& !scheme.is_empty()
&& scheme.chars().all(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '-' | '.'))
{
return Err(HttpError::BadUrl(format!("redirect to an unsupported location `{scheme}:`")));
}
let scheme = if self.tls { "https" } else { "http" };
let host = if self.host.contains(':') { format!("[{}]", self.host) } else { self.host.clone() };
let base = format!("{scheme}://{host}:{}", self.port);
if let Some(rest) = l.strip_prefix("//") {
return Url::parse(&format!("{scheme}://{rest}"));
}
if l.starts_with('/') {
return Url::parse(&format!("{base}{l}"));
}
// relative to the current directory
let dir = self.path.split('?').next().unwrap_or("/");
let dir = dir.rsplit_once('/').map_or("", |(d, _)| d);
Url::parse(&format!("{base}{dir}/{l}"))
}
fn host_header(&self) -> String {
let host = if self.host.contains(':') { format!("[{}]", self.host) } else { self.host.clone() };
let default = if self.tls { 443 } else { 80 };
if self.port == default { host } else { format!("{host}:{}", self.port) }
}
/// The URL without its query (for messages: a query may hold a signed token).
pub fn display(&self) -> String {
let scheme = if self.tls { "https" } else { "http" };
format!("{scheme}://{}{}", self.host_header(), self.path.split('?').next().unwrap_or("/"))
}
}
/// Timeouts and the cancel flag for one request.
pub struct Limits<'a> {
pub connect: Duration,
pub stall: Duration,
pub cancel: &'a AtomicBool,
}
/// How often a blocked read wakes up to check the cancel flag.
const TICK: Duration = Duration::from_millis(250);
/// Longest status line + headers accepted.
const MAX_HEAD: usize = 64 * 1024;
enum Stream {
Plain(TcpStream),
Tls(Box<rustls::ClientConnection>, TcpStream),
}
fn is_timeout(e: &std::io::Error) -> bool {
matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut | ErrorKind::Interrupted)
}
fn tls_config() -> Result<Arc<rustls::ClientConfig>, HttpError> {
let roots = rustls::RootCertStore { roots: webpki_roots::TLS_SERVER_ROOTS.to_vec() };
let config = rustls::ClientConfig::builder_with_provider(Arc::new(rustls_rustcrypto::provider()))
.with_safe_default_protocol_versions()
.map_err(|e| HttpError::Tls(e.to_string()))?
.with_root_certificates(roots)
.with_no_client_auth();
Ok(Arc::new(config))
}
impl Stream {
fn connect(url: &Url, limits: &Limits) -> Result<Stream, HttpError> {
let addrs: Vec<_> = (url.host.as_str(), url.port).to_socket_addrs().map_err(|e| HttpError::Connect(format!("{}: {e}", url.host)))?.collect();
if addrs.is_empty() {
return Err(HttpError::Connect(format!("{}: no address", url.host)));
}
let mut last = String::new();
let mut sock = None;
for a in addrs {
if limits.cancel.load(Ordering::Relaxed) {
return Err(HttpError::Cancelled);
}
match TcpStream::connect_timeout(&a, limits.connect) {
Ok(s) => {
sock = Some(s);
break;
}
Err(e) => last = format!("{}: {e}", url.host),
}
}
let sock = sock.ok_or(HttpError::Connect(last))?;
sock.set_read_timeout(Some(TICK)).map_err(|e| HttpError::Io(e.to_string()))?;
sock.set_write_timeout(Some(TICK)).map_err(|e| HttpError::Io(e.to_string()))?;
let _ = sock.set_nodelay(true);
if !url.tls {
return Ok(Stream::Plain(sock));
}
let name = rustls::pki_types::ServerName::try_from(url.host.clone()).map_err(|e| HttpError::Tls(format!("{}: {e}", url.host)))?;
let conn = rustls::ClientConnection::new(tls_config()?, name).map_err(|e| HttpError::Tls(e.to_string()))?;
let mut s = Stream::Tls(Box::new(conn), sock);
s.handshake(limits)?;
Ok(s)
}
/// Wait for `f` to make progress, retrying timeouts until `stall` passes without any.
fn retry<T>(limits: &Limits, mut f: impl FnMut() -> std::io::Result<T>) -> Result<T, HttpError> {
let start = Instant::now();
loop {
if limits.cancel.load(Ordering::Relaxed) {
return Err(HttpError::Cancelled);
}
match f() {
Ok(v) => return Ok(v),
Err(e) if is_timeout(&e) => {
if start.elapsed() >= limits.stall {
return Err(HttpError::Stalled);
}
}
Err(e) if e.kind() == ErrorKind::InvalidData => return Err(HttpError::Tls(e.to_string())),
Err(e) => return Err(HttpError::Io(e.to_string())),
}
}
}
fn handshake(&mut self, limits: &Limits) -> Result<(), HttpError> {
let Stream::Tls(conn, sock) = self else { return Ok(()) };
while conn.is_handshaking() {
Self::retry(limits, || conn.complete_io(sock))?;
}
Ok(())
}
fn write_all(&mut self, data: &[u8], limits: &Limits) -> Result<(), HttpError> {
match self {
Stream::Plain(sock) => {
let mut rest = data;
while !rest.is_empty() {
let n = Self::retry(limits, || sock.write(rest))?;
if n == 0 {
return Err(HttpError::Io("connection closed while sending the request".into()));
}
rest = rest.get(n..).unwrap_or_default();
}
Ok(())
}
Stream::Tls(conn, sock) => {
// buffered by rustls, then sent as TLS records
conn.writer().write_all(data).map_err(|e| HttpError::Tls(e.to_string()))?;
while conn.wants_write() {
Self::retry(limits, || conn.write_tls(sock))?;
}
Ok(())
}
}
}
/// Read some bytes (0 at the end of the stream), waiting at most `stall` for them.
fn read(&mut self, buf: &mut [u8], limits: &Limits) -> Result<usize, HttpError> {
match self {
Stream::Plain(sock) => Self::retry(limits, || sock.read(buf)),
Stream::Tls(conn, sock) => {
let start = Instant::now();
loop {
if limits.cancel.load(Ordering::Relaxed) {
return Err(HttpError::Cancelled);
}
match conn.reader().read(buf) {
Ok(n) => return Ok(n),
// closed without close_notify: the caller checks the length it got
Err(e) if e.kind() == ErrorKind::UnexpectedEof => return Ok(0),
Err(e) if e.kind() == ErrorKind::WouldBlock => {}
Err(e) => return Err(HttpError::Tls(e.to_string())),
}
// no plaintext buffered: more TLS records from the socket
match conn.read_tls(sock) {
Ok(0) => return Ok(0),
Ok(_) => {
conn.process_new_packets().map_err(|e| HttpError::Tls(e.to_string()))?;
while conn.wants_write() {
Self::retry(limits, || conn.write_tls(sock))?;
}
}
Err(e) if is_timeout(&e) => {
if start.elapsed() >= limits.stall {
return Err(HttpError::Stalled);
}
}
Err(e) => return Err(HttpError::Io(e.to_string())),
}
}
}
}
}
}
/// The connection with a read buffer (headers and chunk sizes are read line by line).
struct Conn {
stream: Stream,
buf: Vec<u8>,
pos: usize,
}
impl Conn {
fn fill(&mut self, limits: &Limits) -> Result<usize, HttpError> {
if self.pos >= self.buf.len() {
self.buf.clear();
self.pos = 0;
}
let mut tmp = [0u8; 64 * 1024];
let n = self.stream.read(&mut tmp, limits)?;
self.buf.extend_from_slice(tmp.get(..n).unwrap_or_default());
Ok(n)
}
/// One line without its CRLF; at most `max` bytes.
fn line(&mut self, max: usize, limits: &Limits) -> Result<String, HttpError> {
loop {
let avail = self.buf.get(self.pos..).unwrap_or_default();
if let Some(i) = avail.iter().position(|b| *b == b'\n') {
let line = avail.get(..i).unwrap_or_default();
let line = line.strip_suffix(b"\r").unwrap_or(line);
let s = String::from_utf8_lossy(line).into_owned();
self.pos = self.pos.saturating_add(i + 1);
return Ok(s);
}
if avail.len() > max {
return Err(HttpError::Protocol("header line too long".into()));
}
// keep the partial line, read more after it
if self.pos > 0 {
self.buf.drain(..self.pos.min(self.buf.len()));
self.pos = 0;
}
let mut tmp = [0u8; 16 * 1024];
let n = self.stream.read(&mut tmp, limits)?;
if n == 0 {
return Err(HttpError::Protocol("connection closed in the middle of the response".into()));
}
self.buf.extend_from_slice(tmp.get(..n).unwrap_or_default());
}
}
/// Up to `out.len()` body bytes (0 at the end of the stream).
fn read(&mut self, out: &mut [u8], limits: &Limits) -> Result<usize, HttpError> {
if self.pos >= self.buf.len() && self.fill(limits)? == 0 {
return Ok(0);
}
let avail = self.buf.get(self.pos..).unwrap_or_default();
let n = avail.len().min(out.len());
if let (Some(dst), Some(src)) = (out.get_mut(..n), avail.get(..n)) {
dst.copy_from_slice(src);
}
self.pos = self.pos.saturating_add(n);
Ok(n)
}
}
/// A response's status, headers and body.
pub struct Response {
pub status: u16,
headers: Vec<(String, String)>,
conn: Conn,
body: BodyMode,
}
enum BodyMode {
Length(u64),
Chunked { left: u64, done: bool },
Close,
Done,
}
impl Response {
/// The first header named `name` (case-insensitive).
pub fn header(&self, name: &str) -> Option<&str> {
self.headers.iter().find(|(k, _)| k.eq_ignore_ascii_case(name)).map(|(_, v)| v.as_str())
}
/// The body's length when the server announced it.
pub fn content_length(&self) -> Option<u64> {
match self.body {
BodyMode::Length(n) => Some(n),
_ => None,
}
}
/// Read body bytes into `out` (0 at the end of the body).
pub fn read(&mut self, out: &mut [u8], limits: &Limits) -> Result<usize, HttpError> {
match &mut self.body {
BodyMode::Done => Ok(0),
BodyMode::Length(left) => {
if *left == 0 {
self.body = BodyMode::Done;
return Ok(0);
}
let want = usize::try_from(*left).unwrap_or(usize::MAX).min(out.len());
let n = self.conn.read(out.get_mut(..want).unwrap_or_default(), limits)?;
if n == 0 {
return Err(HttpError::Protocol("the connection closed before the end of the file".into()));
}
*left = left.saturating_sub(n as u64);
Ok(n)
}
BodyMode::Close => {
let n = self.conn.read(out, limits)?;
if n == 0 {
self.body = BodyMode::Done;
}
Ok(n)
}
BodyMode::Chunked { left, done } => {
if *done {
return Ok(0);
}
if *left == 0 {
let line = self.conn.line(1024, limits)?;
let hex = line.split(';').next().unwrap_or_default().trim();
let size = u64::from_str_radix(hex, 16).map_err(|_| HttpError::Protocol(format!("bad chunk size `{hex}`")))?;
if size == 0 {
// trailers, up to the empty line
let mut n = 0;
while !self.conn.line(MAX_HEAD, limits)?.is_empty() {
n += 1;
if n > 100 {
return Err(HttpError::Protocol("too many trailers".into()));
}
}
*done = true;
return Ok(0);
}
*left = size;
}
let want = usize::try_from(*left).unwrap_or(usize::MAX).min(out.len());
let n = self.conn.read(out.get_mut(..want).unwrap_or_default(), limits)?;
if n == 0 {
return Err(HttpError::Protocol("the connection closed in the middle of a chunk".into()));
}
*left = left.saturating_sub(n as u64);
if *left == 0 && !self.conn.line(16, limits)?.is_empty() {
return Err(HttpError::Protocol("missing CRLF after a chunk".into()));
}
Ok(n)
}
}
}
}
/// Send `GET url` with extra `headers` and read the status line and headers.
pub fn get(url: &Url, headers: &[(&str, String)], limits: &Limits) -> Result<Response, HttpError> {
let mut req = format!(
"GET {} HTTP/1.1\r\nHost: {}\r\nUser-Agent: LightCraft/{}\r\nAccept: */*\r\nAccept-Encoding: identity\r\nConnection: close\r\n",
url.path,
url.host_header(),
env!("CARGO_PKG_VERSION")
);
for (k, v) in headers {
if v.chars().any(|c| c == '\r' || c == '\n') || k.chars().any(|c| c == '\r' || c == '\n' || c == ':') {
return Err(HttpError::BadUrl("header with a line break".into()));
}
req.push_str(&format!("{k}: {v}\r\n"));
}
req.push_str("\r\n");
let mut stream = Stream::connect(url, limits)?;
stream.write_all(req.as_bytes(), limits)?;
let mut conn = Conn { stream, buf: Vec::new(), pos: 0 };
let status_line = conn.line(MAX_HEAD, limits)?;
let mut parts = status_line.split_whitespace();
let version = parts.next().unwrap_or_default();
if !version.starts_with("HTTP/1.") {
return Err(HttpError::Protocol(format!("not an HTTP/1.x response: {}", status_line.chars().take(60).collect::<String>())));
}
let status = parts.next().and_then(|s| s.parse::<u16>().ok()).ok_or_else(|| HttpError::Protocol("bad status line".into()))?;
let mut headers = Vec::new();
let mut total = status_line.len();
loop {
let line = conn.line(MAX_HEAD, limits)?;
if line.is_empty() {
break;
}
total = total.saturating_add(line.len());
if total > MAX_HEAD || headers.len() > 200 {
return Err(HttpError::Protocol("response headers too large".into()));
}
if let Some((k, v)) = line.split_once(':') {
headers.push((k.trim().to_string(), v.trim().to_string()));
}
}
let find = |name: &str| headers.iter().find(|(k, _)| k.eq_ignore_ascii_case(name)).map(|(_, v)| v.as_str());
let body = if status == 204 || status == 304 || (100..200).contains(&status) {
BodyMode::Done
} else if find("transfer-encoding").is_some_and(|v| v.to_ascii_lowercase().contains("chunked")) {
BodyMode::Chunked { left: 0, done: false }
} else if let Some(n) = find("content-length") {
BodyMode::Length(n.trim().parse::<u64>().map_err(|_| HttpError::Protocol(format!("bad Content-Length `{n}`")))?)
} else {
BodyMode::Close
};
Ok(Response { status, headers, conn, body })
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parses_urls() {
let u = Url::parse("https://cdn.example.com/models/sam3/model.safetensors?sig=1#x").unwrap();
assert_eq!((u.tls, u.host.as_str(), u.port, u.path.as_str()), (true, "cdn.example.com", 443, "/models/sam3/model.safetensors?sig=1"));
assert_eq!(u.display(), "https://cdn.example.com/models/sam3/model.safetensors");
let u = Url::parse("http://127.0.0.1:8080").unwrap();
assert_eq!((u.tls, u.port, u.path.as_str()), (false, 8080, "/"));
let u = Url::parse("http://[::1]:9/a").unwrap();
assert_eq!((u.host.as_str(), u.port), ("::1", 9));
for bad in ["ftp://x/y", "https://", "https://a b/", "https://u:p@h/", "http://h:0/", "http://h:99999/", "https://h/\r\nX: y", "file:///etc"]
{
assert!(Url::parse(bad).is_err(), "{bad}");
}
}
#[test]
fn resolves_redirects() {
let u = Url::parse("https://a.example/x/y/file?q").unwrap();
assert_eq!(u.join("https://b.example/z").unwrap().host, "b.example");
assert_eq!(u.join("/root").unwrap().path, "/root");
assert_eq!(u.join("other").unwrap().path, "/x/y/other");
assert_eq!(u.join("//c.example/p").unwrap().host, "c.example");
assert!(u.join("ftp://nope").is_err());
}
#[test]
fn the_tls_config_builds_with_the_pure_rust_provider() {
assert!(tls_config().is_ok());
}
}
crates/segment/src/fetch/mod.rs
+424
View File
@@ -0,0 +1,424 @@
//! Downloading the SAM 3 model, only when the user asked for it: from an ordered list of mirrors
//! (the next one when a mirror fails), resuming partial files, with connect and stall timeouts,
//! a size cap, progress and cancellation, and checked (exact size, SHA-256 where pinned) before
//! each file is moved into the model folder. Runs on whatever thread calls [`download`] (the
//! engine uses a background thread; the UI thread never waits for the network).
//!
//! Files are written to `<name>.part` next to their final place and renamed when verified, so a
//! model folder never holds a half-written file under a real name, and an interrupted download
//! resumes from where it stopped. A file that fails its check is deleted, never kept.
//!
//! The weights are never part of LightCraft (SAM License, see docs/ai-masks.md); where they are
//! downloaded from is configured by [`DEFAULT_MIRRORS`] and the user's own list.
pub mod http;
use std::fs::{File, OpenOptions};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use sha2::{Digest, Sha256};
use http::{HttpError, Limits, Url};
/// One file of the model.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct FileSpec {
pub name: &'static str,
/// Exact size in bytes, when pinned (then the download resumes and is checked against it).
pub size: Option<u64>,
/// SHA-256 (lower-case hex), when pinned.
pub sha256: Option<&'static str>,
/// Largest size accepted when `size` is not pinned.
pub max: u64,
}
/// The `facebook/sam3` checkpoint files [`crate::Sam3::load`] needs.
///
/// `model.safetensors` is pinned to the official checkpoint (Hugging Face LFS SHA-256). The two
/// tokenizer files are small text files without a pinned hash yet: they are only parsed by the
/// tokenizer (never executed) and capped in size.
// TODO(maintainer): pin `vocab.json` and `merges.txt` (size + SHA-256) when the files are
// uploaded to LightCraft's CDN.
pub const SAM3_FILES: &[FileSpec] = &[
FileSpec { name: "vocab.json", size: None, sha256: None, max: 16 << 20 },
FileSpec { name: "merges.txt", size: None, sha256: None, max: 16 << 20 },
FileSpec {
name: crate::WEIGHTS_FILE,
size: Some(SAM3_WEIGHTS_SIZE),
sha256: Some("6d06f0a5f84e435071fe6603e61d0b4cc7b40e0d39d487cfd4d67d8cc11cc14a"),
max: SAM3_WEIGHTS_SIZE,
},
];
/// Size of the official `model.safetensors` (3.44 GB).
pub const SAM3_WEIGHTS_SIZE: u64 = 3_439_938_512;
/// Where LightCraft downloads the model from, in order: each is a base URL, and file `f` is at
/// `<base>/<f>`. **Empty for now**: LightCraft has no download location of its own yet, so the
/// in-app download needs the user's list (`LIGHTCRAFT_SAM3_MIRRORS` or the mirrors file, see
/// [`mirrors`] and docs/ai-masks.md) until this is filled in.
///
/// Hugging Face's `facebook/sam3` can't be a default: it is gated (each user must accept the
/// SAM License there and download with their own token).
// TODO(maintainer): add LightCraft's CDN locations (https, primary first), e.g.
// "https://<cdn-host>/models/sam3/<revision>" — the files there must match `SAM3_FILES`.
pub const DEFAULT_MIRRORS: &[&str] = &[];
/// Environment variable with extra mirrors (base URLs separated by commas, spaces or newlines),
/// tried before the defaults.
pub const MIRRORS_ENV: &str = "LIGHTCRAFT_SAM3_MIRRORS";
/// Most mirrors used (a hostile list can't make a download loop forever).
const MAX_MIRRORS: usize = 16;
/// Redirects followed per request.
const MAX_REDIRECTS: usize = 8;
/// The mirrors to try, in order: the environment variable's, then the mirrors file's (one base
/// URL per line, `#` comments), then [`DEFAULT_MIRRORS`]; duplicates and unusable URLs dropped.
pub fn mirrors(env: Option<&str>, file: Option<&Path>) -> Vec<String> {
let mut out: Vec<String> = Vec::new();
let mut add = |s: &str| {
let s = s.trim().trim_end_matches('/');
if !s.is_empty() && Url::parse(s).is_ok() && !out.iter().any(|m| m == s) && out.len() < MAX_MIRRORS {
out.push(s.to_string());
}
};
if let Some(env) = env {
env.split([',', ' ', '\n', '\t']).for_each(&mut add);
}
// a small text file; anything larger isn't a mirror list
if let Some(text) = file.and_then(|f| std::fs::metadata(f).ok().filter(|m| m.len() < 64 * 1024).and_then(|_| std::fs::read_to_string(f).ok())) {
text.lines().map(|l| l.split('#').next().unwrap_or_default()).for_each(&mut add);
}
DEFAULT_MIRRORS.iter().for_each(|m| add(m));
out
}
/// Download settings.
#[derive(Clone, Debug)]
pub struct Options {
/// Longest wait to connect to a mirror.
pub connect_timeout: Duration,
/// Longest wait for the next bytes once connected.
pub stall_timeout: Duration,
/// Tries per mirror and file (each resumes where the last one stopped).
pub attempts: u32,
}
impl Default for Options {
fn default() -> Self {
Options { connect_timeout: Duration::from_secs(15), stall_timeout: Duration::from_secs(30), attempts: 3 }
}
}
/// Where a download is.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Progress {
pub file: String,
/// Bytes of all files on disk so far, and the total expected (pinned sizes + announced ones).
pub done: u64,
pub total: u64,
/// The mirror in use (without query).
pub mirror: String,
}
/// Why a download failed.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DownloadError {
/// No mirror is configured.
NoMirrors,
Cancelled,
/// A local file problem (folder not writable, disk full…).
Disk(String),
/// Every mirror failed for `file`; the reason for each.
AllMirrorsFailed {
file: String,
errors: Vec<String>,
},
}
impl std::fmt::Display for DownloadError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
DownloadError::NoMirrors => write!(
f,
"no download location is configured for the SAM 3 model in this build (set {MIRRORS_ENV}, or put the files in the model folder yourself; see docs/ai-masks.md)"
),
DownloadError::Cancelled => write!(f, "the download was cancelled"),
DownloadError::Disk(e) => write!(f, "{e}"),
DownloadError::AllMirrorsFailed { file, errors } => write!(f, "could not download {file}: {}", errors.join("; ")),
}
}
}
/// One mirror's failure for one file; `Retry` resumes on the same mirror.
enum Fail {
Retry(String),
NextMirror(String),
Cancelled,
Disk(String),
}
/// Download `files` into `dir` from `mirrors` (in order; see [`mirrors`]). Files already there
/// and verified are kept. `progress` is called as bytes arrive (often: throttle in it).
pub fn download(
files: &[FileSpec],
mirrors: &[String],
dir: &Path,
opts: &Options,
cancel: &AtomicBool,
progress: &mut dyn FnMut(&Progress),
) -> Result<(), DownloadError> {
std::fs::create_dir_all(dir).map_err(|e| DownloadError::Disk(format!("can't create {}: {e}", dir.display())))?;
// what's left to fetch
let mut todo = Vec::new();
for f in files {
if !verified(f, &dir.join(f.name))? {
todo.push(f);
}
}
if todo.is_empty() {
return Ok(());
}
if mirrors.is_empty() {
return Err(DownloadError::NoMirrors);
}
let total: u64 = files.iter().map(|f| f.size.unwrap_or(0)).sum();
let mut before: u64 = files.iter().filter(|f| !todo.contains(f)).filter_map(|f| f.size).sum();
for f in todo {
let mut errors = Vec::new();
let mut ok = false;
'mirrors: for m in mirrors {
let url = format!("{m}/{}", f.name);
for _ in 0..opts.attempts.max(1) {
if cancel.load(Ordering::Relaxed) {
return Err(DownloadError::Cancelled);
}
let mut report = |n: u64, announced: u64| {
let mirror = Url::parse(m).map(|u| u.display()).unwrap_or_default();
progress(&Progress {
file: f.name.to_string(),
done: before.saturating_add(n),
total: total.max(before.saturating_add(announced)),
mirror,
})
};
match fetch_file(f, &url, dir, opts, cancel, &mut report) {
Ok(()) => {
ok = true;
break 'mirrors;
}
Err(Fail::Cancelled) => return Err(DownloadError::Cancelled),
Err(Fail::Disk(e)) => return Err(DownloadError::Disk(e)),
Err(Fail::Retry(e)) => {
log::warn!("SAM 3 download of {url}: {e} (retrying)");
errors.push(format!("{}: {e}", short(m)));
}
Err(Fail::NextMirror(e)) => {
log::warn!("SAM 3 download of {url}: {e}");
errors.push(format!("{}: {e}", short(m)));
continue 'mirrors;
}
}
}
}
if !ok {
return Err(DownloadError::AllMirrorsFailed { file: f.name.to_string(), errors });
}
before = before.saturating_add(std::fs::metadata(dir.join(f.name)).map(|m| m.len()).unwrap_or(0));
}
Ok(())
}
/// A mirror for messages (no query string: it may hold a token).
fn short(m: &str) -> String {
Url::parse(m).map(|u| u.display()).unwrap_or_else(|_| "mirror".into())
}
/// Whether `path` is already the right file (exact size and hash when pinned).
fn verified(f: &FileSpec, path: &Path) -> Result<bool, DownloadError> {
let Ok(meta) = std::fs::metadata(path) else { return Ok(false) };
if !meta.is_file() || meta.len() == 0 {
return Ok(false);
}
if f.size.is_some_and(|s| s != meta.len()) || meta.len() > f.max {
return Ok(false);
}
match f.sha256 {
Some(want) => Ok(sha256_file(path).map_err(DownloadError::Disk)? == want),
None => Ok(true),
}
}
fn sha256_file(path: &Path) -> Result<String, String> {
let mut file = File::open(path).map_err(|e| format!("{}: {e}", path.display()))?;
let mut h = Sha256::new();
let mut buf = vec![0u8; 1 << 20];
loop {
let n = file.read(&mut buf).map_err(|e| format!("{}: {e}", path.display()))?;
if n == 0 {
break;
}
h.update(buf.get(..n).unwrap_or_default());
}
Ok(h.finalize().iter().map(|b| format!("{b:02x}")).collect())
}
fn part_path(dir: &Path, name: &str) -> PathBuf {
dir.join(format!("{name}.part"))
}
/// Fetch one file from one mirror into `<name>.part` (resuming it when its size is pinned),
/// check it and move it into place.
fn fetch_file(f: &FileSpec, url: &str, dir: &Path, opts: &Options, cancel: &AtomicBool, report: &mut dyn FnMut(u64, u64)) -> Result<(), Fail> {
let part = part_path(dir, f.name);
let limit = f.size.unwrap_or(f.max);
let disk = |e: std::io::Error| Fail::Disk(format!("{}: {e}", part.display()));
// resume only a pinned file (an unpinned one can't be checked after resuming)
let mut have = match std::fs::metadata(&part) {
Ok(m) if f.size.is_some() && m.len() <= limit => m.len(),
Ok(_) => {
std::fs::remove_file(&part).map_err(disk)?;
0
}
Err(_) => 0,
};
let fetched = f.size != Some(have) || have == 0;
if fetched {
let limits = Limits { connect: opts.connect_timeout, stall: opts.stall_timeout, cancel };
let mut url = Url::parse(url).map_err(|e| Fail::NextMirror(e.to_string()))?;
let origin = url.host.clone();
let mut redirects = 0;
let mut resp = loop {
let mut headers = vec![];
if have > 0 {
headers.push(("Range", format!("bytes={have}-")));
}
// a token for the configured host only (never sent on to a redirect's host)
if let Ok(t) = std::env::var("LIGHTCRAFT_SAM3_TOKEN")
&& !t.trim().is_empty()
&& url.tls
&& url.host == origin
{
headers.push(("Authorization", format!("Bearer {}", t.trim())));
}
let r = http::get(&url, &headers, &limits).map_err(http_fail)?;
match r.status {
301 | 302 | 303 | 307 | 308 => {
redirects += 1;
if redirects > MAX_REDIRECTS {
return Err(Fail::NextMirror("too many redirects".into()));
}
let loc = r.header("location").ok_or_else(|| Fail::NextMirror("redirect without a location".into()))?;
url = url.join(loc).map_err(|e| Fail::NextMirror(e.to_string()))?;
}
_ => break r,
}
};
let mut file = match resp.status {
200 => {
// the whole file (from the start, also when a range was asked for)
if resp.content_length().is_some_and(|n| n > limit || f.size.is_some_and(|s| s != n)) {
return Err(Fail::NextMirror(format!("the server's file has the wrong size ({:?} bytes)", resp.content_length())));
}
have = 0;
File::create(&part).map_err(disk)?
}
206 if have > 0 => {
let range = resp.header("content-range").unwrap_or_default();
let (start, total) = parse_content_range(range).ok_or_else(|| Fail::NextMirror(format!("bad Content-Range `{range}`")))?;
if start != have || total.is_some_and(|t| Some(t) != f.size) {
return Err(Fail::NextMirror(format!("the server resumed at the wrong place ({range})")));
}
OpenOptions::new().append(true).open(&part).map_err(disk)?
}
416 => {
// the partial file doesn't fit what the server has: start over
std::fs::remove_file(&part).map_err(disk)?;
return Err(Fail::Retry("the partial download didn't match; starting over".into()));
}
401 | 403 => return Err(Fail::NextMirror(format!("access denied (HTTP {})", resp.status))),
404 | 410 => return Err(Fail::NextMirror(format!("not found (HTTP {})", resp.status))),
s if s >= 500 => return Err(Fail::Retry(format!("server error (HTTP {s})"))),
s => return Err(Fail::NextMirror(format!("HTTP {s}"))),
};
let announced = f.size.unwrap_or_else(|| resp.content_length().unwrap_or(0));
report(have, announced);
let mut buf = vec![0u8; 256 * 1024];
loop {
let n = match resp.read(&mut buf, &limits) {
Ok(n) => n,
Err(e) => {
let _ = file.flush();
if f.size.is_none() {
drop(file);
let _ = std::fs::remove_file(&part);
}
return Err(http_fail(e));
}
};
if n == 0 {
break;
}
if have.saturating_add(n as u64) > limit {
drop(file);
let _ = std::fs::remove_file(&part);
return Err(Fail::NextMirror(format!("the server sent more than the expected {limit} bytes")));
}
file.write_all(buf.get(..n).unwrap_or_default()).map_err(disk)?;
have = have.saturating_add(n as u64);
report(have, announced);
}
file.sync_all().map_err(disk)?;
}
// checks: exact size, hash
let len = std::fs::metadata(&part).map_err(disk)?.len();
if f.size.is_some_and(|s| s != len) {
return Err(Fail::Retry(format!("incomplete ({len} of {} bytes)", f.size.unwrap_or(0))));
}
if len == 0 {
let _ = std::fs::remove_file(&part);
return Err(Fail::NextMirror("the server sent an empty file".into()));
}
if let Some(want) = f.sha256 {
let got = sha256_file(&part).map_err(Fail::Disk)?;
if got != want {
std::fs::remove_file(&part).map_err(disk)?;
let why = format!("the file is damaged or not the official one (SHA-256 {got})");
// a damaged partial file from an earlier run isn't this mirror's fault
return Err(if fetched { Fail::NextMirror(why) } else { Fail::Retry(why) });
}
}
std::fs::rename(&part, dir.join(f.name)).map_err(disk)?;
Ok(())
}
fn http_fail(e: HttpError) -> Fail {
match e {
HttpError::Cancelled => Fail::Cancelled,
HttpError::Stalled | HttpError::Io(_) | HttpError::Protocol(_) => Fail::Retry(e.to_string()),
HttpError::BadUrl(_) | HttpError::Connect(_) | HttpError::Tls(_) => Fail::NextMirror(e.to_string()),
}
}
/// `bytes a-b/total` → (a, total) (`*` total → None).
fn parse_content_range(v: &str) -> Option<(u64, Option<u64>)> {
let v = v.trim().strip_prefix("bytes")?.trim();
let (range, total) = v.split_once('/')?;
let (a, b) = range.split_once('-')?;
let (a, b) = (a.trim().parse::<u64>().ok()?, b.trim().parse::<u64>().ok()?);
if b < a {
return None;
}
let total = match total.trim() {
"*" => None,
t => Some(t.parse::<u64>().ok()?),
};
Some((a, total))
}
#[cfg(test)]
mod tests;
crates/segment/src/fetch/tests.rs
+314
View File
@@ -0,0 +1,314 @@
//! The downloader against a local HTTP server (std `TcpListener`; never the internet): mirror
//! fallback, redirects, resuming, chunked bodies, stalls, hash mismatches, size caps, cancel.
use std::io::{BufRead, BufReader, Write};
use std::net::{TcpListener, TcpStream};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use super::*;
/// The served file: 300 kB of a deterministic pattern.
fn payload() -> Vec<u8> {
(0..300_000u32).map(|i| (i.wrapping_mul(2_654_435_761) >> 13) as u8).collect()
}
fn sha_hex(b: &[u8]) -> String {
Sha256::digest(b).iter().map(|b| format!("{b:02x}")).collect()
}
fn leak(s: String) -> &'static str {
Box::leak(s.into_boxed_str())
}
fn pinned() -> FileSpec {
let p = payload();
FileSpec { name: "model.bin", size: Some(p.len() as u64), sha256: Some(leak(sha_hex(&p))), max: p.len() as u64 }
}
struct Server {
base: String,
hits: Arc<AtomicUsize>,
}
/// Serve on 127.0.0.1; the first path segment picks the behaviour.
fn server() -> Server {
let l = TcpListener::bind("127.0.0.1:0").unwrap();
let base = format!("http://{}", l.local_addr().unwrap());
let hits = Arc::new(AtomicUsize::new(0));
let h = hits.clone();
std::thread::spawn(move || {
for s in l.incoming().flatten() {
let h = h.clone();
std::thread::spawn(move || handle(s, &h));
}
});
Server { base, hits }
}
fn handle(mut s: TcpStream, hits: &AtomicUsize) {
let n = hits.fetch_add(1, Ordering::SeqCst);
let mut r = BufReader::new(s.try_clone().unwrap());
let mut first = String::new();
if r.read_line(&mut first).is_err() {
return;
}
let mut range: Option<u64> = None;
loop {
let mut l = String::new();
if r.read_line(&mut l).unwrap_or(0) == 0 || l.trim().is_empty() {
break;
}
if let Some(v) = l.to_ascii_lowercase().strip_prefix("range: bytes=") {
range = v.trim().trim_end_matches('-').parse().ok();
}
}
let path = first.split_whitespace().nth(1).unwrap_or("/").to_string();
let mode = path.trim_start_matches('/').split('/').next().unwrap_or("").to_string();
let body = payload();
let send = |s: &mut TcpStream, status: &str, headers: &str, data: &[u8]| {
let _ = s.write_all(format!("HTTP/1.1 {status}\r\n{headers}Connection: close\r\n\r\n").as_bytes());
let _ = s.write_all(data);
};
match mode.as_str() {
"missing" => send(&mut s, "404 Not Found", "Content-Length: 9\r\n", b"not found"),
"error" => send(&mut s, "500 Oops", "Content-Length: 25\r\n", b"<html>server error</html>"),
"redirect" => send(&mut s, "302 Found", &format!("Location: {}\r\nContent-Length: 0\r\n", path.replacen("/redirect", "/good", 1)), b""),
"wrong" => {
let mut b = body.clone();
b[1000] ^= 0xff;
send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", b.len()), &b)
}
"toolong" => {
// no length announced, more bytes than the file can have
let mut b = body.clone();
b.extend_from_slice(&[7u8; 5000]);
send(&mut s, "200 OK", "", &b)
}
"chunked" => {
let _ = s.write_all(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n");
for c in body.chunks(70_001) {
let _ = s.write_all(format!("{:x};ext=1\r\n", c.len()).as_bytes());
let _ = s.write_all(c);
let _ = s.write_all(b"\r\n");
}
let _ = s.write_all(b"0\r\nX-Trailer: 1\r\n\r\n");
}
"stall" => {
// half the file, then nothing
let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes());
let _ = s.write_all(&body[..body.len() / 2]);
std::thread::sleep(Duration::from_secs(5));
}
"slow" => {
let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes());
for c in body.chunks(1000) {
if s.write_all(c).is_err() {
return;
}
std::thread::sleep(Duration::from_millis(20));
}
}
"flaky" if n.is_multiple_of(2) && range.is_none() => {
// the first try breaks off after a third
let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes());
let _ = s.write_all(&body[..body.len() / 3]);
}
"range416" if range.is_some() => send(&mut s, "416 Range Not Satisfiable", "Content-Length: 0\r\n", b""),
"norange" | "range416" => send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", body.len()), &body),
_ => match range {
Some(from) if from as usize >= body.len() => send(&mut s, "416 Range Not Satisfiable", "Content-Length: 0\r\n", b""),
Some(from) => send(
&mut s,
"206 Partial Content",
&format!("Content-Range: bytes {from}-{}/{}\r\nContent-Length: {}\r\n", body.len() - 1, body.len(), body.len() - from as usize),
&body[from as usize..],
),
None => send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", body.len()), &body),
},
}
}
fn tmp(name: &str) -> PathBuf {
let d = std::env::temp_dir().join(format!("lc-sam3-fetch-{name}-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&d);
d
}
fn quick() -> Options {
Options { connect_timeout: Duration::from_secs(2), stall_timeout: Duration::from_millis(600), attempts: 2 }
}
fn run(files: &[FileSpec], mirrors: &[String], dir: &Path, cancel: &AtomicBool) -> (Result<(), DownloadError>, Vec<Progress>) {
let mut seen = Vec::new();
let r = download(files, mirrors, dir, &quick(), cancel, &mut |p| seen.push(p.clone()));
(r, seen)
}
#[test]
fn falls_back_to_the_next_mirror_and_verifies() {
let srv = server();
let dir = tmp("fallback");
let mirrors = vec![
format!("{}/missing", srv.base),
format!("{}/error", srv.base),
"http://127.0.0.1:1/unreachable".to_string(),
format!("{}/redirect", srv.base),
];
let (r, seen) = run(&[pinned()], &mirrors, &dir, &AtomicBool::new(false));
r.unwrap();
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
assert!(!part_path(&dir, "model.bin").exists());
let last = seen.last().unwrap();
assert_eq!((last.done, last.total), (300_000, 300_000));
// an error page never ends up in the file: nothing left over from the failed mirrors
assert_eq!(std::fs::read_dir(&dir).unwrap().count(), 1);
// installed: nothing is fetched again
let hits = srv.hits.load(Ordering::SeqCst);
run(&[pinned()], &mirrors, &dir, &AtomicBool::new(false)).0.unwrap();
assert_eq!(srv.hits.load(Ordering::SeqCst), hits);
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn a_hash_mismatch_is_deleted_and_the_next_mirror_used() {
let srv = server();
let dir = tmp("hash");
let (r, _) = run(&[pinned()], &[format!("{}/wrong", srv.base)], &dir, &AtomicBool::new(false));
let e = r.unwrap_err().to_string();
assert!(e.contains("SHA-256"), "{e}");
assert!(!dir.join("model.bin").exists() && !part_path(&dir, "model.bin").exists(), "a damaged file is never kept");
run(&[pinned()], &[format!("{}/wrong", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn resumes_partial_files_and_restarts_bad_ones() {
let srv = server();
let dir = tmp("resume");
std::fs::create_dir_all(&dir).unwrap();
// a third already there
std::fs::write(part_path(&dir, "model.bin"), &payload()[..100_000]).unwrap();
run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
// a partial file with the wrong bytes: the hash fails, it's deleted, the next try starts over
std::fs::remove_file(dir.join("model.bin")).unwrap();
std::fs::write(part_path(&dir, "model.bin"), vec![0u8; 100_000]).unwrap();
run(&[pinned()], &[format!("{}/good", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
// a server ignoring Range: the whole file again, not appended
std::fs::remove_file(dir.join("model.bin")).unwrap();
std::fs::write(part_path(&dir, "model.bin"), &payload()[..5]).unwrap();
run(&[pinned()], &[format!("{}/norange", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
// a connection breaking off is resumed on the same mirror
std::fs::remove_file(dir.join("model.bin")).unwrap();
run(&[pinned()], &[format!("{}/flaky", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn a_416_or_a_bad_complete_part_starts_over() {
let srv = server();
let dir = tmp("416");
std::fs::create_dir_all(&dir).unwrap();
// the server refuses the range: the part is dropped and the file fetched whole (no loop)
std::fs::write(part_path(&dir, "model.bin"), &payload()[..100]).unwrap();
run(&[pinned()], &[format!("{}/range416", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
// as long as the file but wrong: hash fails → deleted → downloaded again from the same mirror
std::fs::remove_file(dir.join("model.bin")).unwrap();
std::fs::write(part_path(&dir, "model.bin"), vec![1u8; 300_000]).unwrap();
run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn chunked_bodies_and_unpinned_files() {
let srv = server();
let dir = tmp("chunked");
let f = FileSpec { name: "vocab.json", size: None, sha256: None, max: 1 << 20 };
run(std::slice::from_ref(&f), &[format!("{}/chunked", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
assert_eq!(std::fs::read(dir.join("vocab.json")).unwrap(), payload());
// more than the cap: refused, nothing kept
let small = FileSpec { max: 1000, ..f };
let dir2 = tmp("cap");
let e = run(&[small], &[format!("{}/good", srv.base)], &dir2, &AtomicBool::new(false)).0.unwrap_err().to_string();
assert!(e.contains("wrong size") || e.contains("more than"), "{e}");
assert!(!dir2.join("vocab.json").exists() && !part_path(&dir2, "vocab.json").exists());
let _ = std::fs::remove_dir_all(&dir);
let _ = std::fs::remove_dir_all(&dir2);
}
#[test]
fn a_server_sending_too_much_is_cut_off() {
let srv = server();
let dir = tmp("toolong");
let e = run(&[pinned()], &[format!("{}/toolong", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap_err().to_string();
assert!(e.contains("more than"), "{e}");
assert!(!dir.join("model.bin").exists() && !part_path(&dir, "model.bin").exists());
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn a_stalled_server_times_out_and_the_next_mirror_finishes() {
let srv = server();
let dir = tmp("stall");
let t = Instant::now();
let (r, _) = run(&[pinned()], &[format!("{}/stall", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false));
r.unwrap();
assert!(t.elapsed() < Duration::from_secs(4), "{:?}", t.elapsed());
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn cancel_stops_quickly_and_keeps_the_part_for_later() {
let srv = server();
let dir = tmp("cancel");
let cancel = Arc::new(AtomicBool::new(false));
let c = cancel.clone();
std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(400));
c.store(true, Ordering::SeqCst);
});
let t = Instant::now();
let (r, _) = run(&[pinned()], &[format!("{}/slow", srv.base)], &dir, &cancel);
assert_eq!(r, Err(DownloadError::Cancelled));
assert!(t.elapsed() < Duration::from_secs(3), "{:?}", t.elapsed());
assert!(!dir.join("model.bin").exists());
let kept = std::fs::metadata(part_path(&dir, "model.bin")).map(|m| m.len()).unwrap_or(0);
assert!(kept > 0 && kept < 300_000, "{kept}");
// and the next run resumes it
run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap();
assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload());
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn no_mirrors_is_a_clear_error_and_mirror_lists_are_parsed() {
let dir = tmp("nomirror");
let e = run(&[pinned()], &[], &dir, &AtomicBool::new(false)).0.unwrap_err();
assert_eq!(e, DownloadError::NoMirrors);
assert!(e.to_string().contains(MIRRORS_ENV));
let file = std::env::temp_dir().join(format!("lc-sam3-mirrors-{}.txt", std::process::id()));
std::fs::write(&file, "# mine\nhttps://b.example/m/ # second\n\nnot a url\nhttps://a.example/x\n").unwrap();
let m = mirrors(Some("https://a.example/x/, ftp://no"), Some(&file));
assert_eq!(m[..2], ["https://a.example/x".to_string(), "https://b.example/m".to_string()]);
assert_eq!(m.len(), 2 + DEFAULT_MIRRORS.len());
assert_eq!(mirrors(None, Some(Path::new("/nonexistent/mirrors.txt"))).len(), DEFAULT_MIRRORS.len());
let _ = std::fs::remove_file(&file);
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn content_ranges() {
assert_eq!(parse_content_range("bytes 100-299/300"), Some((100, Some(300))));
assert_eq!(parse_content_range("bytes 0-0/*"), Some((0, None)));
assert_eq!(parse_content_range("bytes 9-1/300"), None);
assert_eq!(parse_content_range("items 1-2/3"), None);
}
crates/segment/src/lib.rs
+12 -1
View File
@@ -13,7 +13,16 @@
//! read from a directory the user downloads them to ([`Sam3::load`]).
//!
//! Ported from the Hugging Face `transformers` implementation (Apache-2.0); see `NOTICE`.
//! No UI dependencies (L3).
//! This crate is therefore licensed under the Apache License 2.0 only (not the MIT option of
//! the rest of LightCraft). No UI dependencies (L3).
//!
//! [`fetch`] downloads the model, only when the user asks for it (from configurable mirrors,
//! verified before use). Nothing in LightCraft requires the model: without it, AI masks report
//! that it isn't installed and everything else works.
//!
//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from
//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`src/transformers/models/sam3` and `sam3_tracker`), Copyright The HuggingFace
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
#![forbid(unsafe_code)]
#![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)]
@@ -21,6 +30,8 @@
mod clip;
#[cfg(not(target_arch = "wasm32"))]
mod detector;
#[cfg(not(target_arch = "wasm32"))]
pub mod fetch;
pub mod mask;
#[cfg(not(target_arch = "wasm32"))]
mod neck;
crates/segment/src/mask.rs
+4
View File
@@ -1,4 +1,8 @@
//! Image preprocessing and mask post-processing (CPU, no model needed).
//!
//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from
//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/image_processing_sam3*.py`), Copyright The HuggingFace
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
use crate::{Error, Result};
crates/segment/src/neck.rs
+4
View File
@@ -1,5 +1,9 @@
//! The feature pyramid (`Sam3VisionNeck`): the backbone's 72 × 72 map at four scales
//! (288, 144, 72 and 36 per side), each projected to 256 channels.
//!
//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from
//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`), Copyright The HuggingFace
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
use candle_core::{Module, Tensor};
use candle_nn::{Conv2d, Conv2dConfig, ConvTranspose2d, ConvTranspose2dConfig, VarBuilder};
crates/segment/src/nn.rs
+4
View File
@@ -1,4 +1,8 @@
//! Shared building blocks: attention, MLPs, layer norms, sine position encodings.
//!
//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from
//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`, `models/sam3_tracker`), Copyright The HuggingFace
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
use candle_core::{D, DType, Device, Module, Tensor};
use candle_nn::{LayerNorm, Linear, VarBuilder};
crates/segment/src/tokenizer.rs
+4
View File
@@ -1,5 +1,9 @@
//! The CLIP byte-pair-encoding tokenizer (lower-cased, `</w>` word ends), padded to the
//! 32-token context SAM 3's text encoder takes.
//!
//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from
//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (the CLIP tokenizer used by `models/sam3`), Copyright The HuggingFace
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
use std::collections::HashMap;
use std::path::Path;
crates/segment/src/tracker.rs
+4
View File
@@ -1,5 +1,9 @@
//! Point prompts (SAM 3's interactive "tracker" head, SAM 2 style): a prompt encoder for
//! positive/negative clicks and a two-way-transformer mask decoder over the image features.
//!
//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from
//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3_tracker/modeling_sam3_tracker.py`), Copyright The HuggingFace
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
use candle_core::{DType, Device, IndexOp, Module, Tensor};
use candle_nn::{Conv2d, Conv2dConfig, ConvTranspose2d, ConvTranspose2dConfig, LayerNorm, VarBuilder};
crates/segment/src/vit.rs
+4
View File
@@ -1,5 +1,9 @@
//! The SAM 3 vision backbone: a 32-layer ViT (Perception Encoder) with 2-D axial rotary
//! position encoding, windowed attention (24 × 24) in most layers and global attention in four.
//!
//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from
//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`), Copyright The HuggingFace
//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE.
use candle_core::{D, Device, Module, Tensor};
use candle_nn::{Conv2d, Conv2dConfig, LayerNorm, VarBuilder};
crates/segment/src/weights.rs
+45 -2
View File
@@ -31,10 +31,11 @@ pub struct Weights {
impl Weights {
pub fn open(path: &Path) -> Result<Arc<Self>> {
let mut file = File::open(path).map_err(|e| Error::Model(format!("{}: {e}", path.display())))?;
let file_len = file.metadata().map_err(|e| Error::Model(format!("{}: {e}", path.display())))?.len();
let mut len = [0u8; 8];
file.read_exact(&mut len).map_err(|e| Error::Model(format!("{}: {e}", path.display())))?;
let n = u64::from_le_bytes(len);
if n > 100 << 20 {
if n > 100 << 20 || n.saturating_add(8) > file_len {
return Err(Error::Model(format!("{}: not a safetensors file (header of {n} bytes)", path.display())));
}
let mut header = vec![0u8; n as usize];
@@ -62,6 +63,10 @@ impl Weights {
if end.checked_sub(start) != want.map(|w| w as u64) {
return Err(Error::Model(format!("{name}: size does not match its shape")));
}
// the bytes must be in the file (a damaged or truncated download fails here, not later)
if (8 + n).checked_add(end).is_none_or(|e| e > file_len) {
return Err(Error::Model(format!("{}: tensor {name} lies outside the file (incomplete or damaged download?)", path.display())));
}
entries.insert(name, Entry { shape, start, end });
}
Ok(Arc::new(Self { file: Mutex::new(file), entries, data_start: 8 + n }))
@@ -74,7 +79,9 @@ impl Weights {
/// Tensor `name` as stored, on the CPU.
pub fn read(&self, name: &str) -> Result<Tensor> {
let e = self.entries.get(name).ok_or_else(|| Error::Model(format!("missing tensor `{name}`")))?;
let mut bytes = vec![0u8; e.end.saturating_sub(e.start) as usize];
// checked against the file's length in `open`
let len = usize::try_from(e.end.saturating_sub(e.start)).map_err(|_| Error::Model(format!("{name}: too large")))?;
let mut bytes = vec![0u8; len];
{
let mut f = self.file.lock().unwrap_or_else(std::sync::PoisonError::into_inner);
f.seek(SeekFrom::Start(self.data_start + e.start)).map_err(|err| Error::Model(format!("{name}: {err}")))?;
@@ -106,3 +113,39 @@ impl SimpleBackend for Backend {
self.0.contains(name)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn write(name: &str, header: &str, data_len: usize) -> std::path::PathBuf {
let p = std::env::temp_dir().join(format!("lc-weights-{name}-{}.safetensors", std::process::id()));
let mut b = (header.len() as u64).to_le_bytes().to_vec();
b.extend_from_slice(header.as_bytes());
b.extend(std::iter::repeat_n(0u8, data_len));
std::fs::write(&p, b).unwrap();
p
}
#[test]
fn tensors_must_lie_inside_the_file() {
let ok = write("ok", r#"{"a":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]}}"#, 16);
let w = Weights::open(&ok).unwrap();
assert_eq!(w.read("a").unwrap().dims(), &[2, 2]);
// a truncated download: the header promises bytes that aren't there
let short = write("short", r#"{"a":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]}}"#, 8);
assert!(Weights::open(&short).err().unwrap().to_string().contains("outside the file"));
// hostile offsets: huge, or beyond u64
let huge = write("huge", r#"{"a":{"dtype":"F32","shape":[1073741824,1073741824],"data_offsets":[0,4611686018427387904]}}"#, 16);
assert!(Weights::open(&huge).is_err());
let wrap = write("wrap", r#"{"a":{"dtype":"F32","shape":[1],"data_offsets":[18446744073709551612,18446744073709551616]}}"#, 16);
assert!(Weights::open(&wrap).is_err());
// a header longer than the file
let p = std::env::temp_dir().join(format!("lc-weights-hdr-{}.safetensors", std::process::id()));
std::fs::write(&p, 1000u64.to_le_bytes()).unwrap();
assert!(Weights::open(&p).is_err());
for f in [ok, short, huge, wrap, p] {
let _ = std::fs::remove_file(f);
}
}
}