switch to macro-based definition for now

This commit is contained in:
Brandon Thomas committed 2025-03-13 22:47:33 -04:00
1 parent 7d0ee6194e
commit fa244a4a70
3 files changed
+21 -156

No files matched your search

@@ -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,
@@ -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<WeightFiletype> {
const JSON : &[u8] = ".json".as_bytes();
@@ -35,132 +27,15 @@ pub const fn maybe_filetype(url_or_filename: &'static str) -> Option<WeightFilet
macro_rules! weight {
($name: literal, $file_name:literal, $weight_function:expr) => {
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<WeightFiletype>,
function: Option<WeightFunction>,
}
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) }
}
}
@@ -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!(