diff --git a/Cargo.lock b/Cargo.lock index 9f21b922eb..19da97d557 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1751,6 +1751,7 @@ dependencies = [ "mimetypes", "my-workspace-hack", "mysql_queries", + "rand 0.8.5", "sqlx", "storyteller_root", "tokens", diff --git a/crates/cli/dev_database_seed/Cargo.toml b/crates/cli/dev_database_seed/Cargo.toml index a771241cf5..a6e1bfc853 100644 --- a/crates/cli/dev_database_seed/Cargo.toml +++ b/crates/cli/dev_database_seed/Cargo.toml @@ -34,7 +34,7 @@ my-workspace-hack = { version = "0.1", path = "../../my-workspace-hack" } # External log = "0.4.14" clap = "2.34.0" # TODO: This is really out of date, but has conflicts when updated - +rand = "0.8.5" # Runtime # NB: We switched to Tokio from `async-std` for broader ecosystem support. tokio = { version = "1.34.0", features = ["macros"] } diff --git a/crates/cli/dev_database_seed/src/main.rs b/crates/cli/dev_database_seed/src/main.rs index 558478772e..6a2290bea5 100644 --- a/crates/cli/dev_database_seed/src/main.rs +++ b/crates/cli/dev_database_seed/src/main.rs @@ -7,7 +7,8 @@ use errors::AnyhowResult; use crate::bucket_clients::get_bucket_clients; use crate::cli_args::parse_cli_args; -use crate::seeding::model_weights::seed_weights; +use crate::seeding::model_weights::seed_weights; // V1 +use crate::seeding::seed_model_weights::seed_weights_files; // V2 v3 will come later on when everything is more conslidated and clear. use crate::seeding::seed_media_files::seed_media_files; use crate::seeding::tts_tacotron2::seed_tts_tacotron2; use crate::seeding::users::seed_user_accounts; @@ -60,6 +61,9 @@ pub async fn main() -> AnyhowResult<()> { seed_weights(&pool).await?; seed_tts_tacotron2(&pool, maybe_bucket_clients.as_ref()).await?; + + // should seed the weights with a few files for hanashi + seed_weights_files(&pool, maybe_bucket_clients.as_ref()).await?; info!("Done!"); Ok(()) diff --git a/crates/cli/dev_database_seed/src/seeding/mod.rs b/crates/cli/dev_database_seed/src/seeding/mod.rs index c1100805b6..7bc49a4b9a 100644 --- a/crates/cli/dev_database_seed/src/seeding/mod.rs +++ b/crates/cli/dev_database_seed/src/seeding/mod.rs @@ -3,4 +3,5 @@ pub mod seed_media_files; pub mod tts_tacotron2; pub mod users; pub mod voice_conversion; -pub mod zero_shot_tts; \ No newline at end of file +pub mod zero_shot_tts; +pub mod seed_model_weights; \ No newline at end of file diff --git a/crates/cli/dev_database_seed/src/seeding/seed_model_weights.rs b/crates/cli/dev_database_seed/src/seeding/seed_model_weights.rs new file mode 100644 index 0000000000..72ab9badd3 --- /dev/null +++ b/crates/cli/dev_database_seed/src/seeding/seed_model_weights.rs @@ -0,0 +1,241 @@ +use std::path::Path; + +use log::{ info, warn }; +use sqlx::{ MySql, Pool }; + +use buckets::{ public::{ media_files::bucket_file_path::MediaFileBucketPath, self }, private }; + +use filesys::file_read_bytes::file_read_bytes; +use filesys::file_size::file_size; +use filesys::path_to_string::path_to_string; +use hashing::sha256::sha256_hash_file::sha256_hash_file; +use mimetypes::mimetype_for_bytes::get_mimetype_for_bytes; +use mimetypes::mimetype_for_file::get_mimetype_for_file; + +use storyteller_root::get_seed_tool_data_root; +use tokens::tokens::users::UserToken; +use crate::bucket_clients::BucketClients; +use crate::seeding::users::HANASHI_USER_TOKEN; + +use enums::{ + by_table::model_weights::{ weights_category::WeightsCategory, weights_types::WeightsType }, + common::visibility::Visibility, +}; + +use rand::Rng; +use errors::{ anyhow, AnyhowResult }; +use mysql_queries::queries::model_weights::create_weight::{ + create_weight, + CreateModelWeightsArgs, + self, +}; +use tokens::tokens::model_weights::ModelWeightToken; + +pub async fn seed_weights_files( + mysql_pool: &Pool, + maybe_bucket_clients: Option<&BucketClients> +) -> AnyhowResult<()> { + info!("Seeding weights files..."); + + let seeded_weights = [ + ( + HANASHI_USER_TOKEN, + "m_21jxt406aattq2cna1xtgf5m4mjyy1", + WeightsType::LoRA, + WeightsCategory::ImageGeneration, + "Niji Mecha", + "models/imagegen/loRA/nijiMecha.safetensors", + ), + ( + HANASHI_USER_TOKEN, + "m_cxzabpxf88h10rshqraksxbqmhs412", + WeightsType::StableDiffusion15, + WeightsCategory::ImageGeneration, + "Niji Diffused Mix", + "models/imagegen/sd15/nijidiffusedmix_v40.safetensors", + ), + ( + HANASHI_USER_TOKEN, + "m_21jxt406aattq2cna1xtgf5m4mjyy9", + WeightsType::LoRA, + WeightsCategory::ImageGeneration, + "Niji Mecha", + "models/imagegen/loRA/nijiMecha.safetensors", + ), + ( + HANASHI_USER_TOKEN, + "m_cxzabpxf88h10rshqraksxbqmhs413", + WeightsType::StableDiffusion15, + WeightsCategory::ImageGeneration, + "Niji Diffused Mix", + "models/imagegen/sd15/nijidiffusedmix_v40.safetensors", + ), + ]; + + let seed_tool_data_root = get_seed_tool_data_root(); + + for ( + user_token, + weights_file_token, + weights_type, + weights_category, + model_name, + subdirectory_path, + ) in seeded_weights { + let user_token = UserToken::new_from_str(user_token); + // prepare the inputs for a model + let mut rng = rand::thread_rng(); + let n: u32 = rng.gen(); + + let model_weight_token = ModelWeightToken(format!("{}", n)); // random token + let model_weight_token = model_weight_token; + + let title = model_name; + let description = "This is a description"; + let description_rendered_html = "This is a description rendered html"; + let original_filename = "model.safetensors"; + let original_download_url = "https://example.com/model.safetensors"; + let thumbnail_token = format!("{}", n); + + let public_bucket_hash: String = weights_file_token.to_string(); + let public_bucket_prefix = weights_type.to_str(); + let public_bucket_extension = weights_category.to_str(); + + let cached_user_ratings_total_count = n.clone(); + let cached_user_ratings_positive_count = n.clone(); + let cached_user_ratings_negative_count = n.clone(); + let cached_user_ratings_ratio = 100.0 as f32; + let version = 1; + + // get the abs path of the weights file to upload. + let mut weight_file_path = seed_tool_data_root.clone(); + weight_file_path.push(subdirectory_path); + + let file_size_bytes = file_size(weight_file_path.clone())?; + let sha256_checksum = sha256_hash_file(weight_file_path.clone())?; + + let result = seed_model( + &mysql_pool, + &model_weight_token, + &user_token, + weights_type, + &weight_file_path, + maybe_bucket_clients + ).await; + + let create_model_weights_args = CreateModelWeightsArgs { + token: &model_weight_token, + weights_type: weights_type, // Assuming weights_type is defined elsewhere + weights_category: weights_category, // Assuming weights_category is defined elsewhere + title: title.to_string(), + maybe_thumbnail_token: Some(thumbnail_token), + description_markdown: description.to_string(), + description_rendered_html: description_rendered_html.to_string(), + creator_user_token: Some(&user_token), + creator_ip_address: "127.0.0.1", // Assuming the IP address is localhost + creator_set_visibility: Visibility::Public, // Assuming the visibility is public + maybe_last_update_user_token: None, // Assuming there is no last update user token + original_download_url: Some(original_download_url.to_string()), + original_filename: Some(original_filename.to_string()), + file_size_bytes: file_size_bytes as i32, // Assuming the file size is 0 + file_checksum_sha2: sha256_checksum.to_string(), // Assuming the file checksum is an empty string + private_bucket_hash: public_bucket_hash, + maybe_private_bucket_prefix: Some(public_bucket_prefix.to_string()), + maybe_private_bucket_extension: Some(public_bucket_extension.to_string()), + cached_user_ratings_total_count: cached_user_ratings_total_count, + cached_user_ratings_positive_count: cached_user_ratings_positive_count, + cached_user_ratings_negative_count: cached_user_ratings_negative_count, + maybe_cached_user_ratings_ratio: Some(cached_user_ratings_ratio), + version: version, + mysql_pool: &mysql_pool, // Assuming mysql_pool is defined elsewhere + }; + + create_weight(create_model_weights_args).await?; + + match result { + Ok(_) => info!("Seeded {:?}", weight_file_path), + Err(err) => + warn!( + r#" + Could not seed media file {} , {:?} : {:?} + (No worries: if there's a duplicate key error, we probably already + seeded the media file on a previous invocation!) + "#, + model_weight_token, + subdirectory_path, + err + ), + } + } + + Ok(()) +} +struct BucketDetails { + bucket_hash: String, + maybe_bucket_prefix: Option, + maybe_bucket_extension: Option, +} + +async fn seed_model( + mysql_pool: &Pool, + model_weight_token: &ModelWeightToken, + user_token: &UserToken, + weight_type: WeightsType, + weight_file_path: &Path, + maybe_bucket_clients: Option<&BucketClients> +) -> AnyhowResult<()> { + info!("Seeding Media file {:?} ...", weight_file_path); + // mutable variables + let mut bucket_hash = "NOT_UPLOADED_BY_SEED_TOOL".to_string(); + let mut maybe_bucket_prefix = ""; + let mut maybe_bucket_extension = ""; + + if let Some(bucket_clients) = maybe_bucket_clients { + let bucket_details = seed_file_to_bucket(weight_file_path, bucket_clients).await?; + bucket_hash = bucket_details.bucket_hash; + maybe_bucket_prefix = bucket_details.maybe_bucket_prefix.as_deref().unwrap(); + maybe_bucket_extension = bucket_details.maybe_bucket_extension.as_deref().unwrap(); + } + + Ok(()) +} + +async fn seed_file_to_bucket( + weight_file_path: &Path, + bucket_clients: &BucketClients +) -> AnyhowResult { + info!("Uploading weights file {:?} ...", weight_file_path); + + let maybe_bucket_prefix = Some("fakeyou_"); + + let maybe_bucket_extension = weight_file_path + .extension() + .map(|extension| extension.to_str()) + .flatten(); + // we should just have this be file bucket path and then we can use a file descriptor to turn it into a specific type of path + let bucket_location = MediaFileBucketPath::generate_new( + maybe_bucket_prefix, + maybe_bucket_extension + ); + + let bucket_path = path_to_string(bucket_location.to_full_object_pathbuf()); + + info!("Reading weights file: {:?}", weight_file_path); + + let bytes = file_read_bytes(weight_file_path)?; + let mimetype = get_mimetype_for_bytes(&bytes).unwrap_or("application/octet-stream"); + + info!("Uploading weights file to bucket path: {:?}", bucket_path); + + let _r = bucket_clients.public.upload_file_with_content_type( + &bucket_path, + bytes.as_ref(), + mimetype + ).await?; + + Ok(BucketDetails { + bucket_hash: bucket_location.get_object_hash().to_string(), + maybe_bucket_prefix: maybe_bucket_prefix.map(|s| s.to_string()), + maybe_bucket_extension: maybe_bucket_extension.map(|s| s.to_string()), + }) +}