From 93f8a952985d98c42d7c47e1f7dd8037838634e0 Mon Sep 17 00:00:00 2001 From: saltacc Date: Mon, 22 Jan 2024 12:01:45 -0800 Subject: [PATCH] make seed tool for media uploads --- crates/cli/dev-database-seed/src/main.rs | 2 + .../src/seeding/media_files.rs | 56 +++++++++++++++++++ .../cli/dev-database-seed/src/seeding/mod.rs | 3 +- .../remote_file_manager/media_descriptor.rs | 18 ++++++ .../remote_cloud_bucket_details.rs | 15 +++++ 5 files changed, 93 insertions(+), 1 deletion(-) create mode 100644 crates/cli/dev-database-seed/src/seeding/media_files.rs diff --git a/crates/cli/dev-database-seed/src/main.rs b/crates/cli/dev-database-seed/src/main.rs index e123dc3c24..6ab4c2710c 100644 --- a/crates/cli/dev-database-seed/src/main.rs +++ b/crates/cli/dev-database-seed/src/main.rs @@ -8,6 +8,7 @@ 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::media_files::seed_media_seedtool; use crate::seeding::seed_media_files::seed_media_files; use crate::seeding::tts_tacotron2::seed_tts_tacotron2; use crate::seeding::users::seed_user_accounts; @@ -56,6 +57,7 @@ pub async fn main() -> AnyhowResult<()> { //seed_zero_shot_tts(&pool, maybe_bucket_clients.as_ref()).await?; seed_voice_conversion(&pool).await?; seed_weights(&pool).await?; + seed_media_seedtool(&pool).await?; seed_tts_tacotron2(&pool, maybe_bucket_clients.as_ref()).await?; // should seed the weights with a few files for hanashi diff --git a/crates/cli/dev-database-seed/src/seeding/media_files.rs b/crates/cli/dev-database-seed/src/seeding/media_files.rs new file mode 100644 index 0000000000..d9df7b7c3a --- /dev/null +++ b/crates/cli/dev-database-seed/src/seeding/media_files.rs @@ -0,0 +1,56 @@ +use log::info; +use sqlx::{MySql, Pool}; +use cloud_storage::remote_file_manager::media_descriptor::MediaVideoMp4Descriptor; +use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; +use enums::by_table::media_files::media_file_type::MediaFileType; +use enums::common::visibility::Visibility; +use mysql_queries::queries::media_files::create::insert_media_file_from_cli_tool::{insert_media_file_from_cli_tool, InsertArgs}; +use errors::{anyhow, AnyhowResult}; +use mysql_queries::queries::users::user::get_user_token_by_username::get_user_token_by_username; +use storyteller_root::get_seed_tool_data_root; +use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::users::UserToken; +use crate::seeding::users::HANASHI_USERNAME; + + + +pub async fn seed_test_videos(mysql_pool: &Pool, user_token: UserToken) -> AnyhowResult<()> { + let model_weight_token1 = MediaFileToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let seed_tool_data_root = get_seed_tool_data_root(); + let mut media_file_path = seed_tool_data_root.clone(); + media_file_path.push("media/video/video-style-transfer/miku.mp4"); + let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?; + let video_media_descriptor = Box::new(MediaVideoMp4Descriptor {}); + let metadata1 = remote_cloud_file_client.upload_file(video_media_descriptor, media_file_path.as_path().to_str().unwrap()).await?; + let bucket_details = metadata1.bucket_details.as_ref().unwrap().clone(); + let media1 = InsertArgs { + pool: &mysql_pool, + maybe_use_apriori_media_token: Some(&model_weight_token1), + media_file_type: MediaFileType::Video, + maybe_mime_type: Some(&*metadata1.mimetype), + file_size_bytes: metadata1.file_size_bytes, + sha256_checksum: &metadata1.sha256_checksum, + maybe_origin_filename: Some("miku.mp4"), + maybe_creator_user_token: Some(&user_token), + creator_set_visibility: Visibility::Public, + public_bucket_directory_hash: &bucket_details.object_hash, + maybe_public_bucket_prefix: Some(&bucket_details.prefix), + maybe_public_bucket_extension: Some(&bucket_details.suffix), + }; + let (_media_file_token1, _) = insert_media_file_from_cli_tool(media1).await.unwrap(); + + Ok(()) +} + +pub async fn seed_media_seedtool(mysql_pool: &Pool) -> AnyhowResult<()> { + info!("Seeding weights..."); + + let user_token = match get_user_token_by_username(HANASHI_USERNAME, mysql_pool).await? { + None => { + return Err(anyhow!("could not find user hanashi")); + } + Some(token) => token, + }; + seed_test_videos(mysql_pool, user_token.clone()).await?; + Ok(()) +} diff --git a/crates/cli/dev-database-seed/src/seeding/mod.rs b/crates/cli/dev-database-seed/src/seeding/mod.rs index 7bc49a4b9a..8ecea17e25 100644 --- a/crates/cli/dev-database-seed/src/seeding/mod.rs +++ b/crates/cli/dev-database-seed/src/seeding/mod.rs @@ -4,4 +4,5 @@ pub mod tts_tacotron2; pub mod users; pub mod voice_conversion; pub mod zero_shot_tts; -pub mod seed_model_weights; \ No newline at end of file +pub mod seed_model_weights; +pub(crate) mod media_files; \ No newline at end of file diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/media_descriptor.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/media_descriptor.rs index d969193271..f4edf45da2 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/media_descriptor.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/media_descriptor.rs @@ -20,3 +20,21 @@ impl FileDescriptor for MediaImagePngDescriptor { true } } + +pub struct MediaVideoMp4Descriptor; +impl FileDescriptor for MediaVideoMp4Descriptor { + fn remote_directory_path(&self) -> &str { + MEDIA_FILE_DIRECTORY + } + fn get_suffix(&self)->String { + "mp4".to_string() + } + + fn get_prefix(&self)->String { + "video".to_string() + } + + fn is_public(&self) -> bool { + true + } +} diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs index 9ea88495a7..077b0265fc 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs @@ -1,3 +1,4 @@ +use crate::remote_file_manager::media_descriptor; use crate::remote_file_manager::weights_descriptor::{self}; use super::file_descriptor::FileDescriptor; @@ -28,6 +29,7 @@ impl RemoteCloudBucketDetails { pub fn file_descriptor_from_bucket_details(&self) -> Box { match self.prefix.as_str() { + // Weights "loRA" => Box::new(weights_descriptor::WeightsLoRADescriptor {}), "sd15" => Box::new(weights_descriptor::WeightsSD15Descriptor {}), "sdxl" => Box::new(weights_descriptor::WeightsSDXLDescriptor {}), @@ -41,6 +43,19 @@ impl RemoteCloudBucketDetails { }, "svc" => Box::new(weights_descriptor::WeightsSVCDescriptor {}), "workflow" => Box::new(weights_descriptor::WeightsWorkflowDescriptor {}), + // Media + "image" => { + match self.suffix.as_str() { + "png" => Box::new(media_descriptor::MediaImagePngDescriptor {}), + _ => panic!("Unknown suffix: {}",self.suffix) + } + }, + "video" => { + match self.suffix.as_str() { + "mp4" => Box::new(media_descriptor::MediaVideoMp4Descriptor {}), + _ => panic!("Unknown suffix: {}",self.suffix) + } + }, _ => panic!("Unknown prefix: {}", self.prefix) } }