diff --git a/crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor.rs b/crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor.rs index 78bb6326ba..70387cc855 100644 --- a/crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor.rs +++ b/crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor.rs @@ -3,8 +3,6 @@ // http::header::name::HeaderName: https://github.com/hyperium/http/blob/master/src/header/name.rs // as a pattern that will support both compile-time constants and dynamically allocated String-based variants. -use crate::ml::weights_registry::weight_descriptor_builder::WeightBuilder; - #[derive(Clone, Copy)] pub enum WeightFiletype { Json, @@ -19,15 +17,6 @@ pub enum WeightFunction { Vae, } -pub struct WeightDescriptor2 { - /// Name of the file on the filesystem, eg. lykon_dreamshaper_7_text_encoder.fp16.safetensors - pub filename: &'static str, - - /// URL of the file in the R2 bucket, - /// eg. https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/lykon_dreamshaper_7_text_encoder.fp16.safetensors - pub r2_download_url: &'static str, -} - pub struct WeightDescriptor { /// Name of the weight, eg. "Lykon Dreamshaper 7" pub name: &'static str, diff --git a/crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor_builder.rs b/crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor_builder.rs index f3bcf4d31d..4d70116bf9 100644 --- a/crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor_builder.rs +++ b/crates/desktop/src-tauri/src/ml/weights_registry/weight_descriptor_builder.rs @@ -1,15 +1,7 @@ -use crate::ml::weights_registry::weight_descriptor::{WeightDescriptor, WeightFiletype, WeightFunction, WeightDescriptor2}; -use const_format::concatcp; -const BUILDER_ERROR_MESSAGE : &str = "Program will not compile if builder isn't used correctly. Check that all arguments are passed to the builder."; +use crate::ml::weights_registry::weight_descriptor::WeightFiletype; +pub (super) const BUILDER_ERROR_MESSAGE : &str = "Program will not compile if builder isn't used correctly. Check that all arguments are passed to the builder."; -pub const R2_BUCKET_URL : &str = "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/"; - -/// r2_url!("foo.txt") --> https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/foo.txt (&'static str) -//macro_rules! r2_url { -// ($ele: expr) => { -// concatcp!(R2_BUCKET_URL, $ele) -// }; -//} +pub (super) const R2_BUCKET_URL : &str = "https://pub-bc5e2bc0cdee4bb5ae8fca9d641ca0d6.r2.dev/"; pub const fn maybe_filetype(url_or_filename: &'static str) -> Option { const JSON : &[u8] = ".json".as_bytes(); @@ -35,132 +27,15 @@ pub const fn maybe_filetype(url_or_filename: &'static str) -> Option { - WeightDescriptor { + $crate::ml::weights_registry::weight_descriptor::WeightDescriptor { name: $name, filename: $file_name, r2_download_url: const_format::concatcp!(crate::ml::weights_registry::weight_descriptor_builder::R2_BUCKET_URL, $file_name), - filetype: crate::ml::weights_registry::weight_descriptor_builder::maybe_filetype($file_name).expect("SHOULD WORK"), + filetype: $crate::ml::weights_registry::weight_descriptor_builder::maybe_filetype($file_name) + .expect($crate::ml::weights_registry::weight_descriptor_builder::BUILDER_ERROR_MESSAGE), function: $weight_function, } }; } pub(crate) use weight; - -#[derive(Default)] -pub(super) struct WeightBuilder { - name: Option<&'static str>, - filename: Option<&'static str>, - r2_download_url: Option<&'static str>, - filetype: Option, - function: Option, -} - -impl WeightBuilder { - pub (super) const fn name(name: &'static str) -> Self { - Self { - name: Some(name), - filename: None, - r2_download_url: None, - filetype: None, - function: None, - } - } - pub (super) const fn filetype(self, filetype: WeightFiletype) -> Self { - Self { - filetype: Some(filetype), - .. self - } - } - pub (super) const fn function(self, weight_classification: WeightFunction) -> Self { - Self { - function: Some(weight_classification), - .. self - } - } - pub (super) const fn filename(self, filename: &'static str) -> Self { - Self { - filename: Some(filename), - r2_download_url: Some("todo"), - filetype: Some(WeightFiletype::Json), - .. self - } - } - //pub (super) const fn filename(self, filename: &'static str) -> Self { - // let s = ConstStr::new().append_str(R2_BUCKET_URL).append_str(filename); - // let url = s.as_str(); - // Self { - // r2_download_url: Some(url), - // .. self - // } - //} - pub (super) const fn r2_download_url(self, r2_download_url: &'static str) -> Self { - let maybe_filetype = maybe_filetype(r2_download_url); - match maybe_filetype { - Some(filetype) => { - Self { - r2_download_url: Some(r2_download_url), - filetype: Some(filetype), - .. self - } - }, - None => { - Self { - r2_download_url: Some(r2_download_url), - .. self - } - } - } - } - pub (super) const fn build(self) -> WeightDescriptor { - WeightDescriptor { - name: self.name.expect(BUILDER_ERROR_MESSAGE), - filename: self.filename.expect(BUILDER_ERROR_MESSAGE), - r2_download_url: self.r2_download_url.expect(BUILDER_ERROR_MESSAGE), - filetype: self.filetype.expect(BUILDER_ERROR_MESSAGE), - function: self.function.expect(BUILDER_ERROR_MESSAGE), - } - } -} - - -const BUFFER_SIZE : usize = 8192; - -/// This is a hack since we can't use macros to append &'static str variables, nor have const generics landed. -struct ConstStr { - data: [u8; BUFFER_SIZE], - len: usize, -} - -impl ConstStr { - pub const fn new() -> Self { - Self { - data: [0u8; BUFFER_SIZE], - len: 0, - } - } - - pub const fn append_str(mut self, s: &str) -> Self { - let b = s.as_bytes(); - let mut index = 0; - while index < b.len() { - self.data[self.len] = b[index]; - self.len += 1; - index += 1; - } - self - } - - pub const fn as_str<'a>(&'a self) -> &'a str { - let mut data: &[u8] = &self.data; - let mut n = data.len() - self.len; - while n > 0 { - n -= 1; - match data.split_last() { - Some((_, rest)) => data = rest, - None => panic!(), - } - } - unsafe { std::str::from_utf8_unchecked(data) } - } -} \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/ml/weights_registry/weights.rs b/crates/desktop/src-tauri/src/ml/weights_registry/weights.rs index 3b6721ae5b..9115ba90f6 100644 --- a/crates/desktop/src-tauri/src/ml/weights_registry/weights.rs +++ b/crates/desktop/src-tauri/src/ml/weights_registry/weights.rs @@ -1,24 +1,25 @@ use crate::ml::weights_registry::weight_descriptor::{WeightDescriptor, WeightFunction}; -use crate::ml::weights_registry::weight_descriptor_builder::WeightBuilder; - use crate::ml::weights_registry::weight_descriptor_builder::weight; -pub const CLIP_JSON : WeightDescriptor = WeightBuilder::name("CLIP") - .filename("clip_vit_bigg_14_lion2b_39b_b160k.tokenizer.json") - .function(WeightFunction::TextTokenizer) - .build(); +pub const CLIP_JSON : WeightDescriptor = weight!( + "CLIP", + "clip_vit_bigg_14_lion2b_39b_b160k.tokenizer.json", + WeightFunction::TextTokenizer +); #[deprecated] -pub const SDXL_TURBO_CLIP_TEXT_ENCODER : WeightDescriptor = WeightBuilder::name("SDXL Turbo CLIP") - .filename("clip_text_encoder.safetensors") - .function(WeightFunction::TextEncoder) - .build(); +pub const SDXL_TURBO_CLIP_TEXT_ENCODER : WeightDescriptor = weight!( + "SDXL Turbo CLIP", + "clip_text_encoder.safetensors", + WeightFunction::TextEncoder +); /// For LCM-based image generation -pub const LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16 : WeightDescriptor = WeightBuilder::name("Lykon Dreamshaper 7") - .filename("lykon_dreamshaper_7_text_encoder.fp16.safetensors") - .function(WeightFunction::TextEncoder) - .build(); +pub const LYKON_DEAMSHAPER_7_TEXT_ENCODER_FP16 : WeightDescriptor = weight!( + "Lykon Dreamshaper 7", + "lykon_dreamshaper_7_text_encoder.fp16.safetensors", + WeightFunction::TextEncoder +); /// For LCM-based image generation pub const LYKON_DEAMSHAPER_7_VAE : WeightDescriptor = weight!(