Skip to content

Commit

Permalink
feat: use URL hash as filename for fetched models
Browse files Browse the repository at this point in the history
  • Loading branch information
decahedron1 committed Jul 31, 2024
1 parent 654a609 commit ecf76f9
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 4 deletions.
3 changes: 2 additions & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ training = [ "ort-sys/training" ]

operator-libraries = [ "libc", "winapi" ]

fetch-models = [ "ureq" ]
fetch-models = [ "ureq", "sha2" ]
download-binaries = [ "ort-sys/download-binaries" ]
load-dynamic = [ "libloading", "ort-sys/load-dynamic" ]
copy-dylibs = [ "ort-sys/copy-dylibs" ]
Expand Down Expand Up @@ -89,6 +89,7 @@ ort-sys = { version = "2.0.0-rc.4", path = "ort-sys" }
libloading = { version = "0.8", optional = true }

ureq = { version = "2.1", optional = true, default-features = false, features = [ "tls" ] }
sha2 = { version = "0.10", optional = true }
tracing = "0.1"
half = { version = "2.1", optional = true }

Expand Down
9 changes: 6 additions & 3 deletions src/session/builder.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
#[cfg(any(feature = "operator-libraries", not(windows)))]
use std::ffi::CString;
#[cfg(feature = "fetch-models")]
use std::fmt::Write;
#[cfg(not(target_arch = "wasm32"))]
use std::path::Path;
#[cfg(feature = "fetch-models")]
use std::path::PathBuf;
use std::{
any::Any,
marker::PhantomData,
Expand Down Expand Up @@ -259,7 +259,10 @@ impl SessionBuilder {
}

let url = model_url.as_ref();
let model_filename = PathBuf::from(url.split('/').last().expect("Missing filename in model URL"));
let model_filename = <sha2::Sha256 as sha2::Digest>::digest(url).into_iter().fold(String::new(), |mut s, b| {
let _ = write!(&mut s, "{:02x}", b);
s
});
let model_filepath = download_dir.join(model_filename);
let downloaded_path = if model_filepath.exists() {
tracing::info!(model_filepath = format!("{}", model_filepath.display()).as_str(), "Model already exists, skipping download");
Expand Down

0 comments on commit ecf76f9

Please sign in to comment.