Merge pull request #209 from storytold/feature/weights_seed_v2

Seeding model weights

Cargo.lock

@@ -1751,6 +1751,7 @@ dependencies = [
  "mimetypes",
  "my-workspace-hack",
  "mysql_queries",
+ "rand 0.8.5",
  "sqlx",
  "storyteller_root",
  "tokens",

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"] }

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(())

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;
+pub mod zero_shot_tts;
+pub mod seed_model_weights;

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<MySql>,
+    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<String>,
+    maybe_bucket_extension: Option<String>,
+}
+
+async fn seed_model(
+    mysql_pool: &Pool<MySql>,
+    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<BucketDetails> {
+    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()),
+    })
+}