diff --git a/.env b/.env index a39f1b1d8d..28ae0591ed 100644 --- a/.env +++ b/.env @@ -64,6 +64,8 @@ HIFIGAN_ROOT_CODE_DIRECTORY=/home/bt/dev/storyteller/storyteller-ml/tts HIFIGAN_VIRTUAL_ENV_ACTIVATION_COMMAND="source python/bin/activate" HIFIGAN_MODEL_CHECK_SCRIPT_NAME="vocodes_model_check_hifigan.py" +DOWNLOAD_SCRIPT="/home/tensor/code/storyteller/web-downloader/.venv/bin/python3.11 /home/tensor/code/storyteller/web-downloader/download_internet_file.py" + TTS_INFERENCE_SIDECAR_HOSTNAME=localhost:8000 #TTS_WAVEGLOW_VOCODER_MODEL_FILENAME=waveglow_256channels_universal_v5.pt # THIS SOUNDS HORRID @@ -76,4 +78,4 @@ TTS_HIFIGAN_SUPERRES_VOCODER_MODEL_FILENAME=hifigan_superres_twilight_33000.pt # Development key. Permissible to check in. NEWRELIC_API_KEY=c402be1c4c4216ec83d0cfc0851520cfFFFFNRAL -#STORYTELLER_ROOT="/home/tensor/code/storyteller/" # Put this in .env-secrets so that not everyone gets this setting. +STORYTELLER_ROOT="/home/tensor/code/storyteller/" # Put this in .env-secrets so that not everyone gets this setting. diff --git a/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs b/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs index c91f161893..91817ee1d9 100644 --- a/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs +++ b/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs @@ -44,6 +44,9 @@ pub enum MediaFileOriginModelType { #[serde(rename = "styletts2")] StyleTTS2, + + #[serde(rename = "stable_diffusion_1_5")] + StableDiffusion15, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -64,6 +67,7 @@ impl MediaFileOriginModelType { Self::MocapNet => "mocap_net", Self::ComfyUi => "comfy_ui", Self::StyleTTS2 => "styletts2", + Self::StableDiffusion15 => "stable_diffusion_1_5" } } @@ -78,6 +82,7 @@ impl MediaFileOriginModelType { "mocap_net" => Ok(Self::MocapNet), "comfy_ui" => Ok(Self::ComfyUi), "styletts2" => Ok(Self::StyleTTS2), + "stable_diffusion_1_5" => Ok(Self::StableDiffusion15), _ => Err(format!("invalid value: {:?}", value)), } } @@ -95,6 +100,7 @@ impl MediaFileOriginModelType { Self::MocapNet, Self::ComfyUi, Self::StyleTTS2, + Self::StableDiffusion15 ]) } } @@ -123,6 +129,7 @@ mod tests { assert_eq!(MediaFileOriginModelType::MocapNet.to_str(), "mocap_net"); assert_eq!(MediaFileOriginModelType::ComfyUi.to_str(), "comfy_ui"); assert_eq!(MediaFileOriginModelType::StyleTTS2.to_str(), "styletts2"); + assert_eq!(MediaFileOriginModelType::StableDiffusion15.to_str(), "stable_diffusion_1_5"); } #[test] @@ -136,6 +143,7 @@ mod tests { assert_eq!(MediaFileOriginModelType::from_str("mocap_net").unwrap(), MediaFileOriginModelType::MocapNet); assert_eq!(MediaFileOriginModelType::from_str("comfy_ui").unwrap(), MediaFileOriginModelType::ComfyUi); assert_eq!(MediaFileOriginModelType::from_str("styletts2").unwrap(), MediaFileOriginModelType::StyleTTS2); + assert_eq!(MediaFileOriginModelType::from_str("stable_diffusion_1_5").unwrap(), MediaFileOriginModelType::StyleTTS2); assert!(MediaFileOriginModelType::from_str("foo").is_err()); } @@ -152,6 +160,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::MocapNet)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::ComfyUi)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::StyleTTS2)); + assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::StableDiffusion15)); assert_eq!(variants.pop_first(), None); } } diff --git a/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs b/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs index 0efb1294c4..665b412ba5 100644 --- a/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs +++ b/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs @@ -47,6 +47,8 @@ pub enum MediaFileOriginProductCategory { #[serde(rename = "mocap")] Mocap, + #[serde(rename = "image_gen")] + ImageGeneration, // Workflow #[serde(rename = "workflow")] Workflow, @@ -68,6 +70,7 @@ impl MediaFileOriginProductCategory { Self::ZeroShotVoice => "zs_voice", Self::VideoFilter => "video_filter", Self::Mocap => "mocap", + Self::ImageGeneration => "image_gen", Self::Workflow => "workflow", } } @@ -81,6 +84,7 @@ impl MediaFileOriginProductCategory { "zs_voice" => Ok(Self::ZeroShotVoice), "video_filter" => Ok(Self::VideoFilter), "mocap" => Ok(Self::Mocap), + "image_gen" => Ok(Self::ImageGeneration), "workflow" => Ok(Self::Workflow), _ => Err(format!("invalid value: {:?}", value)), } @@ -97,6 +101,7 @@ impl MediaFileOriginProductCategory { Self::VoiceConversion, Self::ZeroShotVoice, Self::Mocap, + Self::ImageGeneration, Self::Workflow, ]) } @@ -119,6 +124,7 @@ mod tests { assert_serialization(MediaFileOriginProductCategory::ZeroShotVoice, "zs_voice"); assert_serialization(MediaFileOriginProductCategory::VideoFilter, "video_filter"); assert_serialization(MediaFileOriginProductCategory::Mocap, "mocap"); + assert_serialization(MediaFileOriginProductCategory::Mocap, "image_gen"); assert_serialization(MediaFileOriginProductCategory::Workflow, "workflow"); } @@ -131,6 +137,7 @@ mod tests { assert_eq!(MediaFileOriginProductCategory::ZeroShotVoice.to_str(), "zs_voice"); assert_eq!(MediaFileOriginProductCategory::VideoFilter.to_str(), "video_filter"); assert_eq!(MediaFileOriginProductCategory::Mocap.to_str(), "mocap"); + assert_eq!(MediaFileOriginProductCategory::ImageGeneration.to_str(), "image_gen"); assert_eq!(MediaFileOriginProductCategory::Workflow.to_str(), "workflow"); } @@ -143,13 +150,14 @@ mod tests { assert_eq!(MediaFileOriginProductCategory::from_str("zs_voice").unwrap(), MediaFileOriginProductCategory::ZeroShotVoice); assert_eq!(MediaFileOriginProductCategory::from_str("video_filter").unwrap(), MediaFileOriginProductCategory::VideoFilter); assert_eq!(MediaFileOriginProductCategory::from_str("mocap").unwrap(), MediaFileOriginProductCategory::Mocap); + assert_eq!(MediaFileOriginProductCategory::from_str("image_gen").unwrap(),MediaFileOriginProductCategory::ImageGeneration); assert_eq!(MediaFileOriginProductCategory::from_str("workflow").unwrap(), MediaFileOriginProductCategory::Workflow); } #[test] fn all_variants() { let mut variants = MediaFileOriginProductCategory::all_variants(); - assert_eq!(variants.len(), 8); + assert_eq!(variants.len(), 9); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Unknown)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::FaceAnimator)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::TextToSpeech)); @@ -157,6 +165,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::ZeroShotVoice)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::VideoFilter)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Mocap)); + assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::ImageGeneration)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Workflow)); assert_eq!(variants.pop_first(), None); } diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_file_manager.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_file_manager.rs index e7d4dd0eb7..a30bac13fb 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_file_manager.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_file_manager.rs @@ -38,7 +38,7 @@ impl RemoteCloudFileClient { } } - + // Where the to_system_file_path is let sd_checkpoint_path = work_temp_dir.path().join("sd_checkpoint.safetensors"); path and file name and extension pub async fn download_file(&self, remote_cloud_bucket_details:RemoteCloudBucketDetails, to_system_file_path:String) -> AnyhowResult<()> { let file_descriptor = remote_cloud_bucket_details.file_descriptor_from_bucket_details(); let file_bucket_directory = FileBucketDirectory::from_existing_bucket_details(remote_cloud_bucket_details); @@ -85,7 +85,7 @@ impl RemoteCloudFileClient { } // Retrieve the metadata from the file - fn get_file_meta_data(system_file_path:&str) -> AnyhowResult { + pub fn get_file_meta_data(system_file_path:&str) -> AnyhowResult { let file_size_bytes = file_size(system_file_path.clone())?; let sha256_checksum = sha256_hash_file(system_file_path.clone())?; diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs index b23138e940..6d99c36b14 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs @@ -230,6 +230,8 @@ mod tests { let lora_upload_path = "lora_upload_path".to_string(); let checkpoint = "checkpoint".to_string(); let type_of_inference = "type_of_inference".to_string(); + let version: u32 = 0; + let args = GenericInferenceArgs { inference_category: Some(InferenceCategoryAbbreviated::ImageGeneration), args: Some(PolymorphicInferenceArgs::Ig(StableDiffusionArgs { @@ -247,8 +249,9 @@ mod tests { maybe_upload_path: Some(upload_path), maybe_lora_upload_path: Some(lora_upload_path), type_of_inference, - description: Some("Option".to_string()), - name: Some("Model Name".to_string()) + maybe_description: Some("Option".to_string()), + maybe_name: Some("Model Name".to_string()), + maybe_version: Some(version), })), }; diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs index cdae13bdd2..06d1b0761b 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs @@ -31,11 +31,11 @@ pub struct StableDiffusionArgs { #[serde(rename = "w")] #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_width: Option, + pub maybe_width: Option, #[serde(rename = "h")] #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_height: Option, + pub maybe_height: Option, #[serde(rename = "s")] #[serde(skip_serializing_if = "Option::is_none")] @@ -59,7 +59,7 @@ pub struct StableDiffusionArgs { #[serde(rename = "cf")] #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_cfg_scale: Option, + pub maybe_cfg_scale: Option, #[serde(rename = "lu")] #[serde(skip_serializing_if = "Option::is_none")] @@ -67,10 +67,10 @@ pub struct StableDiffusionArgs { #[serde(rename = "sa")] #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_number_of_samples: Option, + pub maybe_number_of_samples: Option, #[serde(rename = "bc")] - pub maybe_batch_count: Option, + pub maybe_batch_count: Option, #[serde(rename = "t")] pub type_of_inference: String, @@ -78,11 +78,15 @@ pub struct StableDiffusionArgs { // loRA / checkpoint description #[serde(rename = "de")] #[serde(skip_serializing_if = "Option::is_none")] - pub description: Option, + pub maybe_description: Option, // loRA / checkpoint name and description #[serde(rename = "na")] #[serde(skip_serializing_if = "Option::is_none")] - pub name: Option + pub maybe_name: Option, + + #[serde(rename = "ve")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_version: Option } diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs index ec51b21e76..c740d8b0c5 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs @@ -1,172 +1,198 @@ use anyhow::anyhow; +use enums::common::visibility::Visibility; use sqlx; use sqlx::MySqlPool; use enums::by_table::generic_synthetic_ids::id_category::IdCategory; -use enums::by_table::media_files::media_file_origin_category::MediaFileOriginCategory; +use enums::by_table::media_files::media_file_origin_category::{MediaFileOriginCategory, self}; use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType; use enums::by_table::media_files::media_file_origin_product_category::MediaFileOriginProductCategory; use enums::by_table::media_files::media_file_type::MediaFileType; + use errors::AnyhowResult; +use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken; use tokens::tokens::media_files::MediaFileToken; -use tokens::tokens::weight_model_weight; +use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::users::UserToken; use crate::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; use crate::queries::generic_synthetic_ids::transactional_increment_generic_synthetic_id::transactional_increment_generic_synthetic_id; +use crate::queries::media_files::create::insert_media_file_from_file_upload::UploadType; // thought about this it seems like this can be a bit more geneneric instead of having this we can ... pub struct InsertArgs<'a> { pub pool: &'a MySqlPool, pub job: &'a AvailableInferenceJob, - pub maybe_mime_type: Option<&'a str>, - pub file_size_bytes: u64, - pub sha256_checksum: &'a str, + pub media_type: MediaFileType, + // Origin and categorization + pub origin_category: MediaFileOriginCategory, + pub origin_product_category: MediaFileOriginProductCategory, + pub maybe_origin_model_type: Option, + pub maybe_origin_model_token: Option, + pub maybe_origin_filename: Option, + + // Generation flags and media details + pub is_batch_generated: bool, // Assuming you set this in your application logic + + pub maybe_mime_type: Option<&'a str>, + + pub file_size_bytes: u64, + pub maybe_duration_millis: Option, + pub maybe_audio_encoding: Option<&'a str>, + pub maybe_video_encoding: Option<&'a str>, + pub maybe_frame_width: Option, + pub maybe_frame_height: Option, + pub checksum_sha2: &'a str, + + // Storage details pub public_bucket_directory_hash: &'a str, pub maybe_public_bucket_prefix: Option<&'a str>, pub maybe_public_bucket_extension: Option<&'a str>, + pub extra_file_modification_info: Option<&'a str>, // Assuming TEXT type can be represented by &str - pub is_on_prem: bool, - - pub worker_hostname: &'a str, - pub worker_cluster: &'a str, - + // Creator information pub maybe_creator_user_token: Option<&'a UserToken>, pub maybe_creator_anonymous_visitor_token: Option<&'a AnonymousVisitorTrackingToken>, + pub creator_ip_address: &'a str, pub creator_set_visibility: Visibility, - pub upload_type: UploadType, - pub media_file_type: MediaFileType, - pub duration_millis: u64, - pub media_file_origin: MediaFileOriginProductCategory, - pub media_file_origin_product_category: MediaFileOriginProductCategory, - pub media_file_origin_model_type: Mediafile, - pub maybe_origin_filename: String, - pub maybe_origin_model_token:WeightsModelToken + pub maybe_creator_file_synthetic_id_category: IdCategory, + pub maybe_creator_category_synthetic_id_category: IdCategory, + + // Modification and generation info + pub maybe_mod_user_token: Option<&'a UserToken>, + pub is_generated_on_prem: bool, + pub generated_by_worker: Option<&'a str>, + pub generated_by_cluster: Option<&'a str>, } + pub async fn insert_media_file_generic( args: InsertArgs<'_> ) -> AnyhowResult<(MediaFileToken, u64)> { - let result_token = MediaFileToken::generate(); let mut maybe_creator_file_synthetic_id : Option = None; let mut maybe_creator_category_synthetic_id : Option = None; let mut transaction = args.pool.begin().await?; - + // create a user token if let Some(creator_user_token) = args.job.maybe_creator_user_token.as_deref() { let user_token = UserToken::new_from_str(creator_user_token); - + let next_media_file_id = transactional_increment_generic_synthetic_id( &user_token, - IdCategory::MediaFile, + args.maybe_creator_file_synthetic_id_category, &mut transaction ).await?; - let next_lipsync_id = transactional_increment_generic_synthetic_id( + let category_id = transactional_increment_generic_synthetic_id( &user_token, - IdCategory::ZeroShotTtsResult, + args.maybe_creator_category_synthetic_id_category, &mut transaction ).await?; maybe_creator_file_synthetic_id = Some(next_media_file_id); - maybe_creator_category_synthetic_id = Some(next_lipsync_id); + maybe_creator_category_synthetic_id = Some(category_id); } - // TODO: Ask if it make sense to use ONE function to do the insert * I agree it does .... use insert media file generic and eventually remove all this. - // Since fill this out apart of the funtion this seems like the only thing to change. - const ORIGIN_CATEGORY : MediaFileOriginCategory = MediaFileOriginCategory::Inference; - const ORIGIN_PRODUCT_CATEGORY : MediaFileOriginProductCategory = MediaFileOriginProductCategory::ZeroShotVoice; - const ORIGIN_MODEL_TYPE : MediaFileOriginModelType = MediaFileOriginModelType::VallEX; - const MEDIA_TYPE : MediaFileType = MediaFileType::Audio; + let query_result = sqlx::query!( + r#" + INSERT INTO media_files + SET + token = ?, + media_type = ?, - let record_id = { - let query_result = sqlx::query!( - r#" - INSERT INTO media_files - SET - token = ?, + origin_category = ?, + origin_product_category = ?, + maybe_origin_model_type = ?, + maybe_origin_model_token = ?, + maybe_origin_filename = ?, - origin_category = ?, - origin_product_category = ?, - maybe_origin_model_type = ?, + is_batch_generated = ?, - media_type = ?, - maybe_mime_type = ?, - file_size_bytes = ?, + maybe_mime_type = ?, + file_size_bytes = ?, + maybe_duration_millis = ?, + maybe_audio_encoding = ?, + maybe_video_encoding = ?, + maybe_frame_width = ?, + maybe_frame_height = ?, + checksum_sha2 = ?, - checksum_sha2 = ?, + public_bucket_directory_hash = ?, + maybe_public_bucket_prefix = ?, + maybe_public_bucket_extension = ?, + extra_file_modification_info = ?, - public_bucket_directory_hash = ?, - maybe_public_bucket_prefix = ?, - maybe_public_bucket_extension = ?, + maybe_creator_user_token = ?, + maybe_creator_anonymous_visitor_token = ?, - maybe_creator_user_token = ?, - maybe_creator_anonymous_visitor_token = ?, - creator_ip_address = ?, + creator_ip_address = ?, + creator_set_visibility = ?, - creator_set_visibility = ?, + maybe_creator_file_synthetic_id = ?, + maybe_creator_category_synthetic_id = ?, - maybe_creator_file_synthetic_id = ?, - maybe_creator_category_synthetic_id = ?, + maybe_mod_user_token = ?, + is_generated_on_prem = ?, + generated_by_worker = ?, + generated_by_cluster = ? + "#, + result_token, + args.media_type.to_str(), - is_generated_on_prem = ?, - generated_by_worker = ?, - generated_by_cluster = ? + args.origin_category.to_str(), + args.origin_product_category.to_str(), + args.maybe_origin_model_type.map(|e| e.to_str()), + args.maybe_origin_model_token.map(|e| e.to_string()), + args.maybe_origin_filename, - "#, - result_token.as_str(), + args.is_batch_generated, - ORIGIN_CATEGORY.to_str(), - ORIGIN_PRODUCT_CATEGORY.to_str(), - ORIGIN_MODEL_TYPE.to_str(), + args.maybe_mime_type, + args.file_size_bytes, + args.maybe_duration_millis, + args.maybe_audio_encoding, + args.maybe_video_encoding, + args.maybe_frame_width, + args.maybe_frame_height, + args.checksum_sha2, - MEDIA_TYPE.to_str(), - args.maybe_mime_type, - args.file_size_bytes, + args.public_bucket_directory_hash, + args.maybe_public_bucket_prefix, + args.maybe_public_bucket_extension, + args.extra_file_modification_info, - args.sha256_checksum, + args.maybe_creator_user_token.map(|e| e.to_string()), + args.maybe_creator_anonymous_visitor_token.map(|e| e.to_string()), - args.public_bucket_directory_hash, - args.maybe_public_bucket_prefix, - args.maybe_public_bucket_extension, + args.creator_ip_address, + args.creator_set_visibility.to_str(), - args.job.maybe_creator_user_token, - args.job.maybe_creator_anonymous_visitor_token, - args.job.creator_ip_address, + maybe_creator_file_synthetic_id, + maybe_creator_category_synthetic_id, - args.job.creator_set_visibility.to_str(), + args.maybe_mod_user_token, + args.is_generated_on_prem, + args.generated_by_worker, + args.generated_by_cluster + ).execute(&mut *transaction).await; - maybe_creator_file_synthetic_id, - maybe_creator_category_synthetic_id, - - args.is_on_prem, - args.worker_hostname, - args.worker_cluster - ) - .execute(&mut *transaction) - .await; - - let record_id = match query_result { - Ok(res) => { - res.last_insert_id() - }, - Err(err) => { - // TODO: handle better - //transaction.rollback().await?; - return Err(anyhow!("Mysql error: {:?}", err)); - } - }; - - record_id + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + // TODO: handle better + //transaction.rollback().await?; + return Err(anyhow!("Mysql error: {:?}", err)); + } }; transaction.commit().await?; - Ok((result_token, record_id)) } \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/mod.rs b/crates/lib/mysql_queries/src/queries/media_files/create/mod.rs index 7d1e1cc8af..492b47ff1b 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/create/mod.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/create/mod.rs @@ -6,4 +6,5 @@ pub mod insert_media_file_from_tacotron2; pub mod insert_media_file_from_voice_conversion; pub mod insert_media_file_from_zero_shot_tts; pub mod insert_media_file_from_mocapnet; -pub mod insert_media_file_from_comfy_ui; \ No newline at end of file +pub mod insert_media_file_generic; +pub mod insert_media_file_from_comfy_ui; diff --git a/crates/lib/mysql_queries/src/queries/model_weights/create/create_weight.rs b/crates/lib/mysql_queries/src/queries/model_weights/create/create_weight.rs index ff6f6f5c9e..fbc686a29e 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/create/create_weight.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/create/create_weight.rs @@ -21,21 +21,17 @@ pub struct CreateModelWeightsArgs<'a> { pub maybe_last_update_user_token: Option, pub original_download_url: Option, pub original_filename: Option, - pub file_size_bytes: i32, + pub file_size_bytes: u64, pub file_checksum_sha2: String, pub public_bucket_hash: String, pub maybe_public_bucket_prefix: Option, pub maybe_public_bucket_extension: Option, - //pub cached_user_ratings_total_count: u32, - //pub cached_user_ratings_positive_count: u32, - //pub cached_user_ratings_negative_count: u32, - //pub maybe_cached_user_ratings_ratio: Option, - pub version: i32, + pub version: u32, pub mysql_pool: &'a MySqlPool, } pub async fn create_weight(args: CreateModelWeightsArgs<'_>) -> AnyhowResult { - let model_weights_token = ModelWeightToken::generate(); + let transaction = args.mysql_pool.begin().await?; let query_result = sqlx ::query!( @@ -78,16 +74,12 @@ pub async fn create_weight(args: CreateModelWeightsArgs<'_>) -> AnyhowResult { Ok(model_weights_token) } + Ok(_) => { Ok(args.token.clone()) } Err(err) => { transaction.rollback().await?; warn!("Transaction failure: {:?}", err); diff --git a/crates/service/job/inference_job/config/inference-job.development.env b/crates/service/job/inference_job/config/inference-job.development.env index 36231d9bcb..3e4ec3202b 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/crates/service/job/inference_job/config/inference-job.development.env @@ -28,6 +28,10 @@ VALL_E_X_INFERENCE_MAYBE_VENV_COMMAND="source .venv/bin/activate" #VALL_E_X_INFERENCE_MAYBE_VENV_COMMAND="source .env/bin/activate" #VALL_E_X_INFERENCE_MAYBE_DOCKER_IMAGE="" # fill this out for dev? +STYLE_TTS2_INFERENCE_ROOT_DIRECTORY="" +STYLE_TTS2_INFERENCE_COMMAND = "" +STYLE_TTS2_INFERENCE_EXECUTABLE = "" + # ================================= # # # # RVC (v2) # @@ -142,3 +146,4 @@ HF_DATASETS_CACHE="/tmp/huggingface" # nltk also caches data. NLTK_DATA="/tmp/nltk_data" + diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs index 0e9170d4d3..cb0a4118fc 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs @@ -1,46 +1,94 @@ -use std::thread; -use std::time::{Duration, Instant}; - +use std::convert::TryInto; +use std::path::PathBuf; +use std::{ fs, thread }; +use std::time::{ Duration, Instant }; +use actix_web::dev::ResourcePath; +use cloud_storage::remote_file_manager::weights_descriptor::{WeightsLoRADescriptor, WeightsSD15Descriptor}; +use filesys::file_exists::file_exists; +use serde_json; use anyhow::anyhow; +use composite_identifiers::by_table::batch_generations::batch_generation_entity::BatchGenerationEntity; +use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use enums::by_table::generic_synthetic_ids::id_category::IdCategory; +use enums::by_table::media_files::media_file_origin_category::MediaFileOriginCategory; +use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType; +use enums::by_table::media_files::media_file_origin_product_category::MediaFileOriginProductCategory; use log::info; use cloud_storage::remote_file_manager::media_descriptor::MediaImagePngDescriptor; use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails; -use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; +use cloud_storage::remote_file_manager::remote_cloud_file_manager::{self, RemoteCloudFileClient}; use enums::by_table::media_files::media_file_type::MediaFileType; + use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; + use filesys::path_to_string::path_to_string; use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand; use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; +use mysql_queries::queries::batch_generations::insert_batch_generation_records::{ + InsertBatchArgs, + self, +}; +use mysql_queries::queries::batch_generations::insert_batch_generation_records::insert_batch_generation_records; use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; -use mysql_queries::queries::media_files::create::insert_media_file_from_file_upload::{insert_media_file_from_file_upload, InsertMediaFileFromUploadArgs, UploadType}; -use mysql_queries::queries::model_weights::create::create_weight::{create_weight, CreateModelWeightsArgs}; +use mysql_queries::queries::media_files::create::insert_media_file_from_file_upload::{ + insert_media_file_from_file_upload, + InsertMediaFileFromUploadArgs, + UploadType, +}; + +use mysql_queries::queries::media_files::create::insert_media_file_generic::insert_media_file_generic; +use mysql_queries::queries::media_files::create::insert_media_file_generic::InsertArgs; + +use mysql_queries::queries::model_categories::create_category; +use mysql_queries::queries::model_weights::create::create_weight::{ + create_weight, + CreateModelWeightsArgs, +}; use mysql_queries::queries::model_weights::get::get_weight::get_weight_by_token; +use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken; +use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::users::UserToken; -use crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::job_success_result::{ JobSuccessResult, ResultEntity }; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::job::job_types::image_generation::sd::sd_inference_command::InferenceArgs; use crate::job_dependencies::JobDependencies; +#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq)] +pub struct InferenceValues { + pub prompt: String, + pub cfg_scale: u32, + pub negative_prompt: Option, + pub lora_model_weight_token: Option, + pub lora_name: Option, + pub sampler: String, + pub width: u32, + pub height: u32, + pub seed: i64, + pub number_of_samples: u32, +} + pub struct StableDiffusionProcessArgs<'a> { pub job_dependencies: &'a JobDependencies, pub job: &'a AvailableInferenceJob, } - // run inference // insert record into the db with the inference job token complete. - pub async fn sd_args_from_job( args: &StableDiffusionProcessArgs<'_> ) -> Result { let inference_args = args.job.maybe_inference_args .as_ref() - .map(|args| args.args.as_ref()) + .map( + | + args: &mysql_queries::payloads::generic_inference_args::generic_inference_args::GenericInferenceArgs + | args.args.as_ref() + ) .flatten(); let polymorphic_args = match inference_args { Some(args) => args, @@ -50,6 +98,7 @@ pub async fn sd_args_from_job( ); } }; + let sd_args = match polymorphic_args { PolymorphicInferenceArgs::Ig(args) => args, _ => { @@ -63,44 +112,436 @@ pub async fn sd_args_from_job( } // store the prompt and cluster them today. - -pub async fn process_job_selection (args: StableDiffusionProcessArgs<'_> +pub async fn process_job_selection( + args: StableDiffusionProcessArgs<'_> ) -> Result { let sd_args = sd_args_from_job(&args).await?; + if sd_args.type_of_inference == "inference" { - process_job_inference(&args).await - } - else if sd_args.type_of_inference == "lora" { - process_job_lora(&args).await - } - else if sd_args.type_of_inference == "checkpoint"{ - process_job_inference(&args).await - } - else { + return process_job_inference(&args).await + } else if sd_args.type_of_inference == "lora" { + return process_job_lora(&args).await + } else if sd_args.type_of_inference == "model" { + return process_job_sd(&args).await + } else { Err(ProcessSingleJobError::Other(anyhow!("inference type doesn't exist!"))) } } -pub async fn process_job_checkpoint(_args: &StableDiffusionProcessArgs<'_>) -> Result { +pub async fn process_job_sd( + args: &StableDiffusionProcessArgs<'_> +) -> Result { - // insert model record in create model record + let job = args.job; + let deps = args.job_dependencies; + let mysql_pool = &deps.db.mysql_pool; + + let sd_args = sd_args_from_job(&args).await?; + + let sd_deps = match + &args.job_dependencies.job.job_specific_dependencies.maybe_stable_diffusion_dependencies + { + None => { + return Err(ProcessSingleJobError::Other(anyhow!("Missing Job Specific Dependencies"))); + } + Some(val) => { val } + }; + + let creator_ip_address = &job.creator_ip_address; + let creator_user_token: UserToken; + + match &job.maybe_creator_user_token { + Some(token) => { + creator_user_token = UserToken::new_from_str(token); + } + None => { + return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Creator User Token"))); + } + } + + // The parameters will be updated on another screen perhaps? + // so right now it will fill with the availible values. + + let work_temp_dir = format!("temp_stable_diffusion_inference_{}", job.id.0); + let work_temp_dir = args.job_dependencies.fs.scoped_temp_dir_creator_for_work + .new_tempdir(&work_temp_dir) + .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + + let sd_checkpoint_path = work_temp_dir.path().join("sd_checkpoint.safetensors"); + let vae_path = work_temp_dir.path().join("vae.safetensors"); + let output_path = work_temp_dir.path().join("output"); + + info!("Paths to download to:"); + info!("sd_checkpoint_path:{}", sd_checkpoint_path.display()); + info!("vae_path:{}", vae_path.display()); + info!("output_path:{}", output_path.display()); + + let download_url = match sd_args.maybe_upload_path { + Some(val) => { val } + None => { "".to_string() } + }; + + if download_url.len() == 0 { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Download URL Missing")) + ); + } + + let file_name = "sd_checkpoint.safetensors"; + let download_script = easyenv::get_env_string_or_default( + "DOWNLOAD_SCRIPT", + "download_internet_file.py" + ); + // Download + let google_drive_downloader = GoogleDriveDownloadCommand::new(&download_script, + None, + None, + None); + info!("Downloading {}", download_url); + let download_filename = match + google_drive_downloader.download_file_with_file_name( + &download_url, + &work_temp_dir, + file_name + ).await + { + Ok(filename) => filename, + Err(_e) => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Download"))); + } + }; + let download_file_path = work_temp_dir.path().join(download_filename); + info!("File Retrieved at {}", download_file_path.display()); + + // download vae for model + + // use this vae doesn't matter though + // VAE token for now + let vae_token = String::from("weight_rb0959wfzjhk3d1k93hr3s0qw"); + let model_weight_vae = ModelWeightToken(vae_token); + let vae_weight_record = get_weight_by_token( + &model_weight_vae, + false, + &deps.db.mysql_pool + ).await?; + let vae_weight_record = match vae_weight_record { + Some(val) => val, + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem.")) + ); + } + }; + + let vae_details = RemoteCloudBucketDetails::new( + vae_weight_record.public_bucket_hash.clone(), + vae_weight_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()), + vae_weight_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string()) + ); + + let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?; + remote_cloud_file_client.download_file(vae_details, path_to_string(vae_path.clone())).await?; + + let stderr_output_file = work_temp_dir.path().join("sd_err.txt"); + let stdout_output_file = work_temp_dir.path().join("sd_out.txt"); + + // run inference on loRA downloaded + sd_deps.inference_command.execute_inference(InferenceArgs { + work_dir: work_temp_dir.path().to_path_buf(), + output_file: output_path.clone(), + stderr_output_file, + stdout_output_file, + prompt:String::from("This is a green sign that says go this is a test prompt to test the model."), + negative_prompt: sd_args.maybe_n_prompt.clone().unwrap_or_default(), + number_of_samples:20, + samplers: sd_args.maybe_sampler.clone().unwrap_or(String::from("Euler a")), + width: sd_args.maybe_width.unwrap_or(512), + height: sd_args.maybe_height.unwrap_or(512), + cfg_scale: sd_args.maybe_cfg_scale.unwrap_or(7), + seed: sd_args.maybe_seed.unwrap_or(1), + lora_path: PathBuf::from(""), + checkpoint_path: sd_checkpoint_path.clone(), + vae: vae_path.clone(), + batch_count: sd_args.maybe_batch_count.unwrap_or(1), + }); + + // check if file exists if it does not error out.... + let path = output_path.clone(); + let file_path = format!("{}_{}.png", path_to_string(path), 0); + + if file_exists(file_path.path()) == false { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Upload Not a SD Model")) + ); + } + + // If it worked and didn't fail! then we should save and create the weight. + // upload and create weights for loRA... + let weights_sd_descriptor = Box::new(WeightsSD15Descriptor{}); + let metadata = remote_cloud_file_client.upload_file(weights_sd_descriptor,sd_checkpoint_path.to_str().unwrap_or_default()).await?; + // chekc the model hash for duplicated models. + let bucket_details = match metadata.bucket_details { + Some(metadata) => metadata, + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to generate bucket details!"))); + } + }; + + let model_weight_token = &ModelWeightToken::generate(); + let model_weight_token_result = create_weight(CreateModelWeightsArgs { + token: &model_weight_token, + weights_type: WeightsType::StableDiffusion15, + weights_category: WeightsCategory::ImageGeneration, + title: sd_args.maybe_name.unwrap_or(String::from("")), + maybe_description_markdown: sd_args.maybe_description, + maybe_description_rendered_html: None, + creator_user_token: Some(&creator_user_token), + creator_ip_address, + creator_set_visibility: Default::default(), + maybe_last_update_user_token: None, + original_download_url: Some(download_url), + original_filename: None, + file_size_bytes: metadata.file_size_bytes, + file_checksum_sha2: metadata.sha256_checksum, + public_bucket_hash: bucket_details.object_hash, + maybe_public_bucket_prefix: Some(bucket_details.prefix), + maybe_public_bucket_extension:Some(bucket_details.suffix), + version: sd_args.maybe_version.unwrap_or(0), + mysql_pool, + }).await?; - // let maybe_result_entity = ResultEntity { - // entity_type: InferenceResultType::UploadModel, - // entity_token: - // }; Ok(JobSuccessResult { - maybe_result_entity: None, + maybe_result_entity: Some(ResultEntity { + entity_type: InferenceResultType::UploadModel, + entity_token: model_weight_token_result.to_string(), + }), inference_duration: Duration::from_secs(0), }) } -pub async fn process_job_lora(_args: &StableDiffusionProcessArgs<'_>) -> Result { + +pub async fn process_job_lora( + args: &StableDiffusionProcessArgs<'_> +) -> Result { + + let job = args.job; + let deps = args.job_dependencies; + let mysql_pool = &deps.db.mysql_pool; + + let sd_args = sd_args_from_job(&args).await?; + + let sd_deps = match + &args.job_dependencies.job.job_specific_dependencies.maybe_stable_diffusion_dependencies + { + None => { + return Err(ProcessSingleJobError::Other(anyhow!("Missing Job Specific Dependencies"))); + } + Some(val) => { val } + }; + + let creator_ip_address = &job.creator_ip_address; + let creator_user_token: UserToken; + + match &job.maybe_creator_user_token { + Some(token) => { + creator_user_token = UserToken::new_from_str(token); + } + None => { + return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Creator User Token"))); + } + } + + + let work_temp_dir = format!("temp_stable_diffusion_inference_{}", job.id.0); + let work_temp_dir = args.job_dependencies.fs.scoped_temp_dir_creator_for_work + .new_tempdir(&work_temp_dir) + .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + + let sd_checkpoint_path = work_temp_dir.path().join("sd_checkpoint.safetensors"); + let lora_path = work_temp_dir.path().join("lora.safetensors"); // input path into execution + let vae_path = work_temp_dir.path().join("vae.safetensors"); + let output_path = work_temp_dir.path().join("output"); + + info!("Paths to download to:"); + info!("sd_checkpoint_path:{}", sd_checkpoint_path.display()); + info!("lora_path:{}", lora_path.display()); + info!("vae_path:{}", vae_path.display()); + info!("output_path:{}", output_path.display()); + + let download_url = match sd_args.maybe_lora_upload_path { + Some(val) => { val } + None => { "".to_string() } + }; + + if download_url.len() == 0 { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Download URL Missing")) + ); + } + + let file_name = "lora.safetensors"; + + let download_script = easyenv::get_env_string_or_default( + "DOWNLOAD_SCRIPT", + "download_internet_file.py" + ); + + // Download + let google_drive_downloader = GoogleDriveDownloadCommand::new(&download_script, + None, + None, + None); + + info!("Downloading {}", download_url); + + let download_filename = match + google_drive_downloader.download_file_with_file_name( + &download_url, + &work_temp_dir, + file_name + ).await + { + Ok(filename) => filename, + Err(_e) => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Download"))); + } + }; + + let download_file_path = work_temp_dir.path().join(download_filename); + info!("File Retrieved at {}", download_file_path.display()); + + // download model + vae for lora + // use lora with whatever checkpoint make it fixed id if it works submit it + let weight_token_sd_model_token = ModelWeightToken(String::from("weight_dmmthavhawqc2hj7yqyemcbf8")); // any generic model. + let sd_weight_record = get_weight_by_token( + &weight_token_sd_model_token, + false, + &deps.db.mysql_pool + ).await?; + let sd_weight_record = match sd_weight_record { + Some(val) => val, + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("No SD weight baked in for loRA upload inference check? thats a problem.")) + ); + } + }; + + let sd_weight_details = RemoteCloudBucketDetails::new( + sd_weight_record.public_bucket_hash.clone(), + sd_weight_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()), + sd_weight_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string()) + ); + + let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?; + remote_cloud_file_client.download_file(sd_weight_details, path_to_string(sd_checkpoint_path.clone())).await?; + + + // use this vae doesn't matter though + // VAE token for now + let vae_token = String::from("weight_rb0959wfzjhk3d1k93hr3s0qw"); + let model_weight_vae = ModelWeightToken(vae_token); + let vae_weight_record = get_weight_by_token( + &model_weight_vae, + false, + &deps.db.mysql_pool + ).await?; + let vae_weight_record = match vae_weight_record { + Some(val) => val, + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem.")) + ); + } + }; + + let vae_details = RemoteCloudBucketDetails::new( + vae_weight_record.public_bucket_hash.clone(), + vae_weight_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()), + vae_weight_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string()) + ); + + let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?; + remote_cloud_file_client.download_file(vae_details, path_to_string(vae_path.clone())).await?; + + let stderr_output_file = work_temp_dir.path().join("sd_err.txt"); + let stdout_output_file = work_temp_dir.path().join("sd_out.txt"); + + // run inference on loRA downloaded + sd_deps.inference_command.execute_inference(InferenceArgs { + work_dir: work_temp_dir.path().to_path_buf(), + output_file: output_path.clone(), + stderr_output_file, + stdout_output_file, + prompt:String::from("This is a green sign that says go this is a test prompt to test the model."), + negative_prompt: sd_args.maybe_n_prompt.clone().unwrap_or_default(), + number_of_samples:20, + samplers: sd_args.maybe_sampler.clone().unwrap_or(String::from("Euler a")), + width: sd_args.maybe_width.unwrap_or(512), + height: sd_args.maybe_height.unwrap_or(512), + cfg_scale: sd_args.maybe_cfg_scale.unwrap_or(7), + seed: sd_args.maybe_seed.unwrap_or(1), + lora_path: lora_path.clone(), + checkpoint_path: sd_checkpoint_path.clone(), + vae: vae_path.clone(), + batch_count: sd_args.maybe_batch_count.unwrap_or(1), + }); + + // check if file exists if it does not error out.... + let path = output_path.clone(); + let file_path = format!("{}_{}.png", path_to_string(path), 0); + + if file_exists(file_path.path()) == false { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Upload Not a LoRA")) + ); + } + + // If it worked and didn't fail! then we should save and create the weight. + // upload and create weights for loRA... + let lora_descriptor = Box::new(WeightsLoRADescriptor{}); + let metadata = remote_cloud_file_client.upload_file(lora_descriptor,lora_path.to_str().unwrap_or_default()).await?; + + let bucket_details = match metadata.bucket_details { + Some(metadata) => metadata, + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to generate bucket details!"))); + } + }; + + let model_weight_token = &ModelWeightToken::generate(); + let model_weight_token_result = create_weight(CreateModelWeightsArgs { + token: &model_weight_token, + weights_type: WeightsType::LoRA, + weights_category: WeightsCategory::ImageGeneration, + title: sd_args.maybe_name.unwrap_or(String::from("")), + maybe_description_markdown: sd_args.maybe_description, + maybe_description_rendered_html: None, + creator_user_token: Some(&creator_user_token), + creator_ip_address, + creator_set_visibility: Default::default(), + maybe_last_update_user_token: None, + original_download_url: Some(download_url), + original_filename: None, + file_size_bytes: metadata.file_size_bytes, + file_checksum_sha2: metadata.sha256_checksum, + public_bucket_hash: bucket_details.object_hash, + maybe_public_bucket_prefix: Some(bucket_details.prefix), + maybe_public_bucket_extension:Some(bucket_details.suffix), + version: sd_args.maybe_version.unwrap_or(0), + mysql_pool, + }).await?; + Ok(JobSuccessResult { - maybe_result_entity: None, + maybe_result_entity: Some(ResultEntity { + entity_type: InferenceResultType::UploadModel, + entity_token: model_weight_token_result.to_string(), + }), inference_duration: Duration::from_secs(0), }) + } + pub async fn process_job_inference( args: &StableDiffusionProcessArgs<'_> ) -> Result { @@ -108,23 +549,37 @@ pub async fn process_job_inference( let deps = args.job_dependencies; let mysql_pool = &deps.db.mysql_pool; - let sd_args = sd_args_from_job(&args).await?; - let sd_deps = match &args.job_dependencies.job.job_specific_dependencies.maybe_stable_diffusion_dependencies { + let sd_deps: &crate::job::job_types::image_generation::sd::stable_diffusion_dependencies::StableDiffusionDependencies = match + &args.job_dependencies.job.job_specific_dependencies.maybe_stable_diffusion_dependencies + { None => { return Err(ProcessSingleJobError::Other(anyhow!("Missing Job Specific Dependencies"))); } - Some(val) => { - val - } + Some(val) => { val } }; - let creator_user_token = match &job.maybe_creator_user_token { - Some(token) => UserToken::new_from_str(token), + let creator_user_token: UserToken; + let anon_user_token: Option; + + match &job.maybe_creator_anonymous_visitor_token { + Some(token) => { + let anonymous_visitor_token = AnonymousVisitorTrackingToken::new_from_str(token); + anon_user_token = Some(anonymous_visitor_token); + } + None => { + anon_user_token = None; + } + } + + match &job.maybe_creator_user_token { + Some(token) => { + creator_user_token = UserToken::new_from_str(token); + } None => { return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Creator User Token"))); } - }; + } let _job_progress_reporter = args.job_dependencies.clients.job_progress_reporter .new_generic_inference(job.inference_job_token.as_str()) @@ -134,72 +589,31 @@ pub async fn process_job_inference( let work_temp_dir = format!("temp_stable_diffusion_inference_{}", job.id.0); - // thread::sleep(seconds) to check the directory - //NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem. let work_temp_dir = args.job_dependencies.fs.scoped_temp_dir_creator_for_work .new_tempdir(&work_temp_dir) .map_err(|e| ProcessSingleJobError::from_io_error(e))?; let sd_checkpoint_path = work_temp_dir.path().join("sd_checkpoint.safetensors"); - let lora_path = work_temp_dir.path().join("lora.safetensors"); + let mut lora_path = work_temp_dir.path().join("lora.safetensors"); let vae_path = work_temp_dir.path().join("vae.safetensors"); let output_path = work_temp_dir.path().join("output"); - let g_drive_path = work_temp_dir.path().join("gdrive"); info!("Paths to download to:"); - info!("sd_checkpoint_path:{}",sd_checkpoint_path.display()); - info!("lora_path:{}",lora_path.display()); - info!("vae_path:{}",vae_path.display()); - info!("output_path:{}",output_path.display()); - info!("tmp_google_drive_path:{}",g_drive_path.display()); + info!("sd_checkpoint_path:{}", sd_checkpoint_path.display()); + info!("lora_path:{}", lora_path.display()); + info!("vae_path:{}", vae_path.display()); + info!("output_path:{}", output_path.display()); - // If it is not inference then it must be loRA or checkpoint downloading - if sd_args.type_of_inference != "inference" { - let mut download_url= String::from(""); - - if sd_args.type_of_inference == "lora" { - download_url = match sd_args.maybe_lora_upload_path { - Some(val) => { val }, - None => { "".to_string() } - }; - } - else if sd_args.type_of_inference == "checkpoint" { - download_url = match sd_args.maybe_upload_path { - Some(val) => { val }, - None => { "".to_string() } - } - } - - if download_url.len() == 0 { - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Download URL Missing"))); - } - - let file_name = "download.safetensors"; - let download_script = easyenv::get_env_string_or_default( - "DOWNLOAD_SCRIPT", - "./scripts/download_internet_file.py"); - let google_drive_downloader = - GoogleDriveDownloadCommand::new_production(&download_script); - info!("Downloading {}",download_url); - - let _download_filename = match google_drive_downloader.download_file_with_file_name(&download_url, &work_temp_dir,file_name).await { - Ok(filename) => filename, - Err(_e) => { - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Download"))); - } - }; - let download_file_path = work_temp_dir.path().join(file_name); - info!("File Retrieved at {}",download_file_path.display()); - } + // thread::sleep(seconds) to check the directory // // Unpack loRA and Checkpoint // // run inference by downloading from google drive. let lora_token = sd_args.maybe_lora_model_token; - let weight_token = sd_args.maybe_sd_model_token; + let weight_token = sd_args.maybe_sd_model_token.clone(); let retrieved_sd_record = match weight_token { - Some(token) => { + Some(ref token) => { let retrieved_sd_record = get_weight_by_token( &token, false, @@ -220,7 +634,7 @@ pub async fn process_job_inference( ); } }; - + // origin file name needs to be just the file name /tmp/downloads_long_lived/temp_stable_diffusion_inference_32.8qJJljxWWZeD/output_0.png // ignore if no lora token let retrieved_loRA_record = match lora_token { Some(token) => { @@ -229,26 +643,33 @@ pub async fn process_job_inference( false, &deps.db.mysql_pool ).await?; + Some(retrieved_loRA_record) } None => None, }; let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?; - + // Details for SD checkpoint let details = RemoteCloudBucketDetails::new( retrieved_sd_record.public_bucket_hash.clone(), retrieved_sd_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()), retrieved_sd_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string()) ); - remote_cloud_file_client.download_file(details, path_to_string(sd_checkpoint_path.clone())).await?; + + remote_cloud_file_client.download_file( + details, + path_to_string(sd_checkpoint_path.clone()) + ).await?; + + let mut lora_name = String::from(""); + let mut lora_token = String::from(""); match retrieved_loRA_record { Some(record) => { match record { Some(model_weight_record) => { - let lora_details = RemoteCloudBucketDetails::new( model_weight_record.public_bucket_hash.clone(), model_weight_record.maybe_public_bucket_prefix @@ -258,23 +679,29 @@ pub async fn process_job_inference( .clone() .unwrap_or_else(|| "".to_string()) ); + + lora_name = model_weight_record.title; + lora_token = model_weight_record.token.to_string(); remote_cloud_file_client.download_file( lora_details, path_to_string(lora_path.clone()) ).await?; - } - None => {} + None => { + lora_path.clear(); + } } } - None => {} + None => { + lora_path.clear(); + } } // VAE token for now let vae_token = String::from("weight_rb0959wfzjhk3d1k93hr3s0qw"); let model_weight_vae = ModelWeightToken(vae_token); - let vae_weight_record = get_weight_by_token( + let vae_weight_record = get_weight_by_token( &model_weight_vae, false, &deps.db.mysql_pool @@ -283,32 +710,29 @@ pub async fn process_job_inference( let vae_weight_record = match vae_weight_record { Some(val) => val, None => { - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))) + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem.")) + ); } }; - let vae_details = RemoteCloudBucketDetails::new( vae_weight_record.public_bucket_hash.clone(), - vae_weight_record.maybe_public_bucket_prefix - .clone() - .unwrap_or_else(|| "".to_string()), - vae_weight_record.maybe_public_bucket_extension - .clone() - .unwrap_or_else(|| "".to_string()) + vae_weight_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()), + vae_weight_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string()) ); - remote_cloud_file_client.download_file( - vae_details, - path_to_string(vae_path.clone()) - ).await?; + remote_cloud_file_client.download_file(vae_details, path_to_string(vae_path.clone())).await?; let prompt = match sd_args.maybe_prompt { Some(val) => val, - None => return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("No Prompt provided!"))) + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("No Prompt provided!"))); + } }; let stderr_output_file = work_temp_dir.path().join("sd_err.txt"); + let stdout_output_file = work_temp_dir.path().join("sd_out.txt"); let number_of_samples = match sd_args.maybe_number_of_samples { Some(val) => val, @@ -321,10 +745,11 @@ pub async fn process_job_inference( work_dir: work_temp_dir.path().to_path_buf(), output_file: output_path.clone(), stderr_output_file, - prompt, - negative_prompt:sd_args.maybe_n_prompt.unwrap_or_default(), + stdout_output_file, + prompt: prompt.clone(), + negative_prompt: sd_args.maybe_n_prompt.clone().unwrap_or_default(), number_of_samples, - samplers: sd_args.maybe_sampler.unwrap_or(String::from("Euler a")), + samplers: sd_args.maybe_sampler.clone().unwrap_or(String::from("Euler a")), width: sd_args.maybe_width.unwrap_or(512), height: sd_args.maybe_height.unwrap_or(512), cfg_scale: sd_args.maybe_cfg_scale.unwrap_or(7), @@ -335,91 +760,131 @@ pub async fn process_job_inference( batch_count: sd_args.maybe_batch_count.unwrap_or(1), }); + + + + // hack to check the directory before clean up. + // let thirtyMinutes = 1800; + // thread::sleep(Duration::from_secs(thirtyMinutes)); + // upload media and create a record. + let inference_duration = Instant::now().duration_since(inference_start_time); let creator_ip_address = &job.creator_ip_address; // run a for loop for output images output_0 in the folder then use upload media. // pngs - if sd_args.type_of_inference == "inference" { - // gather all the images into a single batch + let mut entries = vec![]; + let inputs = InferenceValues { + prompt: prompt.clone(), + cfg_scale: sd_args.maybe_cfg_scale.unwrap_or(7), + negative_prompt: sd_args.maybe_n_prompt, + lora_model_weight_token: Some(lora_token), + lora_name: Some(lora_name), + sampler: sd_args.maybe_sampler.unwrap_or(String::from("Euler a")), + width: sd_args.maybe_width.unwrap_or(512), + height: sd_args.maybe_height.unwrap_or(512), + seed: sd_args.maybe_seed.unwrap_or(1), + number_of_samples: number_of_samples, + }; - for i in 0..sd_args.maybe_batch_count.unwrap_or(1) { - let path = output_path.clone(); - let file_path = format!("{}_{}.png", path_to_string(path), i); - println!("Upload File Path:{}", file_path); - let metadata = remote_cloud_file_client.upload_file(Box::new(MediaImagePngDescriptor {}), file_path.as_ref()).await?; - - let bucket_details = match metadata.bucket_details { - Some(val) => { val } - None => return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))) - }; - - // TODO: Generic Insert with batch id. - // insert media file generic should run here. - insert_media_file_from_file_upload(InsertMediaFileFromUploadArgs { - pool: mysql_pool, - maybe_creator_user_token: Some(&creator_user_token), - maybe_creator_anonymous_visitor_token: None, - creator_ip_address, - creator_set_visibility: Default::default(), - upload_type: UploadType::Filesystem, - media_file_type: MediaFileType::Image, - maybe_mime_type: Some(metadata.mimetype.as_ref()), - file_size_bytes: metadata.file_size_bytes, - duration_millis: inference_duration.as_secs() * 1000, - sha256_checksum: metadata.sha256_checksum.as_str(), - public_bucket_directory_hash: bucket_details.object_hash.as_str(), - maybe_public_bucket_prefix: Some(bucket_details.prefix.as_str()), - maybe_public_bucket_extension: Some(bucket_details.suffix.as_str()), - }).await?; + let inputs = match serde_json::to_string(&inputs) { + Ok(result) => result, + Err(err) => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("couldn't serialize metadata.")) + ); } - - // hack to check the directory before clean up. - let thirtyMinutes = 1800; - thread::sleep(Duration::from_secs(thirtyMinutes)); - // upload media and create a record. - - Ok(JobSuccessResult {// TODO: batch token needs to go here - maybe_result_entity: None, - inference_duration: Duration::from_secs(0), - }) - } - else { // must be an upload and need upload token + }; - // The parameters will be updated on another screen perhaps? - // so right now it will fill with the availible values. - create_weight(CreateModelWeightsArgs { - token: &ModelWeightToken::generate(), - weights_type: WeightsType::StableDiffusion15, - weights_category: WeightsCategory::ImageGeneration, - title: "".to_string(), - maybe_description_markdown: None, - maybe_description_rendered_html: None, - creator_user_token: Some(&creator_user_token), - creator_ip_address, - creator_set_visibility: Default::default(), - maybe_last_update_user_token: None, - original_download_url: None, - original_filename: None, - file_size_bytes: 0, - file_checksum_sha2: "".to_string(), - public_bucket_hash: "".to_string(), - maybe_public_bucket_prefix: None, - maybe_public_bucket_extension: None, - version: 0, - mysql_pool, + for i in 0..sd_args.maybe_batch_count.unwrap_or(1) { + let path = output_path.clone(); + + let file_path = format!("{}_{}.png", path_to_string(path), i); + + println!("Upload File Path:{}", file_path); + + let metadata = remote_cloud_file_client.upload_file( + Box::new(MediaImagePngDescriptor {}), + file_path.as_ref() + ).await?; + + let bucket_details = match metadata.bucket_details { + Some(val) => { val } + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem.")) + ); + } + }; + + // extra_file_modification_info: todo!(), // JSON ENCODED STRUCT + let media_file_token = insert_media_file_generic(InsertArgs { + pool: mysql_pool, + job: job, + media_type: MediaFileType::Image, + origin_category: MediaFileOriginCategory::Upload, + origin_product_category: MediaFileOriginProductCategory::ImageGeneration, + maybe_origin_model_type: Some(MediaFileOriginModelType::StableDiffusion15), + maybe_origin_model_token: weight_token.clone(), + maybe_origin_filename: Some(file_path), + is_batch_generated: true, + maybe_mime_type: Some(metadata.mimetype.as_ref()), + file_size_bytes: metadata.file_size_bytes, + maybe_duration_millis: Some(inference_duration.as_millis() as u64), + maybe_audio_encoding: None, + maybe_video_encoding: None, + maybe_frame_width: Some(sd_args.maybe_width.unwrap_or(512) as u32), + maybe_frame_height: Some(sd_args.maybe_height.unwrap_or(512) as u32), + checksum_sha2: metadata.sha256_checksum.as_str(), + public_bucket_directory_hash: bucket_details.object_hash.as_str(), + maybe_public_bucket_prefix: Some(bucket_details.prefix.as_str()), + maybe_public_bucket_extension: Some(bucket_details.suffix.as_str()), + extra_file_modification_info: Some(&inputs), + maybe_creator_user_token: Some(&creator_user_token), + maybe_creator_anonymous_visitor_token: anon_user_token.as_ref(), + creator_ip_address: creator_ip_address, + creator_set_visibility: args.job.creator_set_visibility, + maybe_creator_file_synthetic_id_category: IdCategory::MediaFile, + maybe_creator_category_synthetic_id_category: IdCategory::ModelWeights, + maybe_mod_user_token: None, + is_generated_on_prem: args.job_dependencies.job.info.container.is_on_prem, + generated_by_worker: Some(&args.job_dependencies.job.info.container.hostname), + generated_by_cluster: Some(&args.job_dependencies.job.info.container.cluster_name), }).await?; - Ok(JobSuccessResult {// TODO: batch token needs to go here - maybe_result_entity: None, - inference_duration: Duration::from_secs(0), - }) + let batch_generation_entity: BatchGenerationEntity = BatchGenerationEntity::MediaFile( + media_file_token.0 + ); + entries.push(batch_generation_entity); } + let mut transaction = mysql_pool.begin().await.unwrap(); + + let batch_token_result = insert_batch_generation_records(InsertBatchArgs { + entries, + transaction: &mut transaction, + }).await; + let batch_token = match batch_token_result { + Ok(v) => { v.to_string() } + Err(err) => { + return Err( + ProcessSingleJobError::from_anyhow_error( + anyhow!("No batch token? something has failed.") + ) + ); + } + }; + Ok(JobSuccessResult { // TODO: batch token needs to go here + maybe_result_entity: Some(ResultEntity { + entity_type: InferenceResultType::MediaFile, + entity_token: batch_token, + }), + inference_duration: inference_duration, + }) } #[ignore] @@ -435,15 +900,16 @@ mod tests { }; use errors::AnyhowResult; + #[ignore] + #[tokio::test] + async fn test_json_extra_info() -> AnyhowResult<()> {} #[ignore] #[tokio::test] async fn test_seed_weights_files() -> AnyhowResult<()> { let seed_path = PathBuf::from("/storyteller/root/custom-seed-tool-data"); let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await; let remote_cloud_file_client = match remote_cloud_file_client { - Ok(res) => { - res - } + Ok(res) => { res } Err(_) => { return Err(anyhow!("failed to get remote cloud file client")); } @@ -466,9 +932,42 @@ mod tests { suffix: String::from("safetensors"), }; - remote_cloud_file_client.download_file(bucket_details1, String::from("./checkpoint")).await?; + remote_cloud_file_client.download_file( + bucket_details1, + String::from("./checkpoint") + ).await?; remote_cloud_file_client.download_file(bucket_details2, String::from("./loRA")).await?; Ok(()) } } + +// /** +// * +// * +// * @CrossProduct this shows the batch table usage: https://github.com/storytold/storyteller-rust/blob/master/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs#L83-L96 +// I wrote a test to show how you can insert multiple media file tokens into the table at once +// I tested and it works: + +// mysql> select * from batch_generations; +// +----+----------------------------------+-------------+--------------+---------------------+---------------------+ +// | id | token | entity_type | entity_token | created_at | updated_at | +// +----+----------------------------------+-------------+--------------+---------------------+---------------------+ +// | 4 | batch_g_25dgzw53jdwgvesgdb48bam6 | media_file | media_foo | 2024-01-17 02:18:52 | 2024-01-17 02:18:52 | +// | 5 | batch_g_25dgzw53jdwgvesgdb48bam6 | media_file | media_bar | 2024-01-17 02:18:52 | 2024-01-17 02:18:52 | +// | 6 | batch_g_25dgzw53jdwgvesgdb48bam6 | media_file | media_baz | 2024-01-17 02:18:52 | 2024-01-17 02:18:52 | +// | 7 | batch_g_79a2evktmnqbn9sq3k3eybze | media_file | media_foo | 2024-01-17 02:20:32 | 2024-01-17 02:20:32 | +// | 8 | batch_g_79a2evktmnqbn9sq3k3eybze | media_file | media_bar | 2024-01-17 02:20:32 | 2024-01-17 02:20:32 | +// | 9 | batch_g_79a2evktmnqbn9sq3k3eybze | media_file | media_baz | 2024-01-17 02:20:32 | 2024-01-17 02:20:32 | +// +----+----------------------------------+-------------+--------------+---------------------+---------------------+ +// 6 rows in set (0.00 sec) + +// (I ran the test twice, which is why there are two separate batches of the same media file tokens) +// I'll write an endpoint, probably something like: + +// GET api.fakeyou.com/v1/batches/{batch_token} + +// That will return a list of all the media files in the batch. + +// We can create a page around this. ALso, we can have a separate endpoint that looks up all the other files in a batch (if you're starting on a given media file page), so that you can show thumbnails of related generations on a media file page. +// */ diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/sd_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/sd_inference_command.rs index 265b35f8f7..32b97d7927 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/sd_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/sd_inference_command.rs @@ -72,19 +72,20 @@ pub struct InferenceArgs { pub work_dir: PathBuf, /// --result_file: path to final file output pub output_file: PathBuf, + pub stdout_output_file: PathBuf, pub stderr_output_file: PathBuf, pub prompt: String, pub negative_prompt:String, - pub number_of_samples:i32, + pub number_of_samples:u32, pub samplers:String, - pub width:i32, - pub height:i32, - pub cfg_scale:i32, + pub width:u32, + pub height:u32, + pub cfg_scale:u32, pub seed:i64, pub lora_path:PathBuf, pub checkpoint_path:PathBuf, pub vae:PathBuf, - pub batch_count:i32, + pub batch_count:u32, } impl StableDiffusionInferenceCommand { @@ -228,8 +229,11 @@ impl StableDiffusionInferenceCommand { command.push_str(" --seed "); command.push_str(args.seed.to_string().as_str()); - command.push_str(" --loRA-path "); - command.push_str(&path_to_string(args.lora_path)); + // TODO ensure lora path is not empty ... + if args.lora_path.as_os_str().is_empty() == false { + command.push_str(" --loRA-path "); + command.push_str(&path_to_string(args.lora_path)); + } command.push_str(" --check-point "); command.push_str(&path_to_string(args.checkpoint_path)); @@ -271,7 +275,9 @@ impl StableDiffusionInferenceCommand { info!("stderr will be written to file: {}", path_to_string(args.stderr_output_file.clone())); let stderr_file = File::create(&args.stderr_output_file)?; + let stdout_file = File::create(&args.stdout_output_file)?; config.stderr = Redirection::File(stderr_file); + config.stdout = Redirection::File(stdout_file); if !env_vars.is_empty() { config.env = Some(env_vars); diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs index 2fe137a43f..5fdbce00d2 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs @@ -1,12 +1,12 @@ // #![forbid(unused_imports)] #![forbid(unused_mut)] -#![forbid(unused_variables)] - -use std::fmt::{Display, Formatter}; +// #![forbid(unused_variables)] +use log::info; +use std::fmt::{ Display, Formatter }; use std::fmt::Debug; use std::sync::Arc; -use actix_web::{HttpRequest, HttpResponse, web}; +use actix_web::{ HttpRequest, HttpResponse, web }; use actix_web::error::ResponseError; use actix_web::http::StatusCode; use log::warn; @@ -44,26 +44,25 @@ const DEBUG_HEADER_NAME: &str = "enable-debug-mode"; /// This is useful for catching the live logs or intercepting the job. const ROUTING_TAG_HEADER_NAME: &str = "routing-tag"; -#[derive(Debug,Deserialize, Clone, Copy, Eq, PartialEq)] +#[derive(Debug, Deserialize, Clone, Copy, Eq, PartialEq)] pub enum TypeOfInference { #[serde(rename = "lora")] - UploadLoRA, - #[serde(rename = "checkpoint")] - CheckPoint, + Lora, + #[serde(rename = "model")] + Model, #[serde(rename = "inference")] - Inference, + Inference } impl Display for TypeOfInference { fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { match self { - TypeOfInference::UploadLoRA => write!(f, "lora"), - TypeOfInference::CheckPoint => write!(f, "checkpoint"), + TypeOfInference::Lora => write!(f, "lora"), + TypeOfInference::Model => write!(f, "model"), TypeOfInference::Inference => write!(f, "inference"), } } } - pub fn is_valid_string(input: &str) -> bool { let valid_strings = [ "DPM++ 2M Karras", @@ -91,41 +90,40 @@ pub fn is_valid_string(input: &str) -> bool { "LMS Karras", "DPM2 Karras", "DPM2 a Karras", - "DPM++ 2S a Karras" + "DPM++ 2S a Karras", ]; valid_strings.contains(&input) } - -#[derive(Deserialize,ToSchema)] +#[derive(Deserialize, ToSchema)] pub struct EnqueueImageGenRequest { uuid_idempotency_token: String, - type_of_inference: TypeOfInference, // upload loRA / check point / standard inference maybe_image_source: Option, maybe_sd_model_token: Option, maybe_lora_model_token: Option, maybe_prompt: Option, maybe_n_prompt: Option, maybe_seed: Option, - maybe_width: Option, - maybe_height: Option, + maybe_width: Option, + maybe_height: Option, maybe_sampler: Option, maybe_upload_path: Option, maybe_lora_upload_path: Option, - maybe_cfg_scale: Option, - maybe_number_of_samples: Option, - maybe_batch_count: Option, - name: String, - description: String, + maybe_cfg_scale: Option, + maybe_number_of_samples: Option, + maybe_batch_count: Option, + maybe_name: Option, + maybe_description: Option, + maybe_version: Option } -#[derive(Serialize,ToSchema)] +#[derive(Serialize, ToSchema)] pub struct EnqueueImageGenRequestSuccessResponse { pub success: bool, pub inference_job_token: InferenceJobToken, } -#[derive(Debug,ToSchema)] +#[derive(Debug, ToSchema)] pub enum EnqueueImageGenRequestError { BadInput(String), NotAuthorized, @@ -166,18 +164,20 @@ impl std::fmt::Display for EnqueueImageGenRequestError { // insert generic inference job.rs // Need to convert it to generic inference job. #[utoipa::path( -post, -path = "/inference/enqueue_image_gen/", -responses( -(status = 200, description = "Enqueue TTS generically", body = EnqueueImageGenRequestSuccessResponse), -(status = 400, description = "Bad input", body = EnqueueImageGenRequestError), -(status = 401, description = "Not authorized", body = EnqueueImageGenRequestError), -(status = 429, description = "Rate limited", body = EnqueueImageGenRequestError), -(status = 500, description = "Server error", body = EnqueueImageGenRequestError) -), -params( -("request" = EnqueueImageGenRequest, description = "Payload for TTS Request") -) + post, + path = "/inference/enqueue_image_gen/", + responses( + ( + status = 200, + description = "Enqueue TTS generically", + body = EnqueueImageGenRequestSuccessResponse, + ), + (status = 400, description = "Bad input", body = EnqueueImageGenRequestError), + (status = 401, description = "Not authorized", body = EnqueueImageGenRequestError), + (status = 429, description = "Rate limited", body = EnqueueImageGenRequestError), + (status = 500, description = "Server error", body = EnqueueImageGenRequestError) + ), + params(("request" = EnqueueImageGenRequest, description = "Payload for TTS Request")) )] pub async fn enqueue_image_generation_request( http_request: HttpRequest, @@ -185,10 +185,59 @@ pub async fn enqueue_image_generation_request( server_state: web::Data> ) -> Result { + // TODO:I know ill fix this later + let path = http_request.path(); + let segments: Vec<&str> = path.split('/').collect(); + let last_segment = segments.last().unwrap_or(&""); + + println!("Last segment: {}", last_segment); + + // use segment to determine what to do. + let mode; + if last_segment.to_string() == "lora" { + mode = "lora"; + } else if last_segment.to_string() == "model" { + mode = "model"; + } else { + mode = "inference"; + } + + // check and ensure we also pass description and the name + if mode == "lora" || mode == "model" { + match request.maybe_name { + Some(_) => {}, + None => return Err(EnqueueImageGenRequestError::BadInput("Missing Model / Lora Name".to_string())) + } + + match request.maybe_description { + Some(_) => {} + None => return Err(EnqueueImageGenRequestError::BadInput("Missing Model / Lora Description".to_string())) + } + } else { + match request.maybe_sd_model_token { + Some(_) => {} + None => return Err(EnqueueImageGenRequestError::BadInput("Missing Model Token".to_string())) + } + } + + if mode == "lora" { + match request.maybe_lora_upload_path { + Some(_) => {} + None => return Err(EnqueueImageGenRequestError::BadInput("Missing Lora Upload Path".to_string())) + } + } + + if mode == "model" { + match request.maybe_upload_path { + Some(_) => {} + None => return Err(EnqueueImageGenRequestError::BadInput("Missing Model Upload Path".to_string())) + } + } + // TODO: Brandon need to figure out premium vs not premium let mut maybe_user_token: Option = None; - let visbility = enums::common::visibility::Visibility::Public; + let visbility = enums::common::visibility::Visibility::Public; let mut mysql_connection = server_state.mysql_pool.acquire().await.map_err(|err| { warn!("MySql pool error: {:?}", err); @@ -204,9 +253,20 @@ pub async fn enqueue_image_generation_request( EnqueueImageGenRequestError::ServerError })?; - + if let Some(user_session) = maybe_user_session.as_ref() { + maybe_user_token = Some(UserToken::new_from_str(&user_session.user_token)); + match maybe_user_token { + Some(ref val) => { + println!("User Token: {}",val.clone().to_string()); + }, + None => { + println!("User Token check? Failed to find!"); + } + } + } else { + println!("Session isn't there"); } // Plan should handle "first anonymous use" and "investor" cases. @@ -226,6 +286,7 @@ pub async fn enqueue_image_generation_request( |routing_tag| routing_tag.trim().to_string() ); + // ==================== BANNED USERS ==================== // if let Some(ref user) = maybe_user_session { @@ -254,25 +315,6 @@ pub async fn enqueue_image_generation_request( // ==================== INFERENCE ARGS ==================== // - if request.type_of_inference == TypeOfInference::Inference { - if request.maybe_sd_model_token.is_none() { - return Err(EnqueueImageGenRequestError::BadInput("No sd model token provided".to_string())); - } - - if request.maybe_prompt.is_none() { - return Err(EnqueueImageGenRequestError::BadInput("No prompt provided".to_string())); - } - - } else if request.type_of_inference == TypeOfInference::UploadLoRA { - if request.maybe_lora_upload_path.is_none() { - return Err(EnqueueImageGenRequestError::BadInput("No lora upload path provided".to_string())); - } - } else if request.type_of_inference == TypeOfInference::CheckPoint { - if request.maybe_upload_path.is_none() { - return Err(EnqueueImageGenRequestError::BadInput("No upload path provided".to_string())); - } - } - let width = request.maybe_width.unwrap_or(512); let height = request.maybe_height.unwrap_or(512); @@ -291,32 +333,31 @@ pub async fn enqueue_image_generation_request( }; let cfg_scale = match request.maybe_cfg_scale { - Some(val) => if val > 32 { 32 } else { val }, - None => 7 + Some(val) => if val > 32 { 32 } else { val } + None => 7, }; let number_of_samples = match request.maybe_number_of_samples { - Some(val) => if val < 0 { 20 } else if val > 128 { 128 } else { val }, + Some(val) => if val < 0 { 20 } else if val > 128 { 128 } else { val } None => 20, }; - let batch_count = match request.maybe_batch_count { - Some(val) => if val > 4 { 4 } else { val }, + Some(val) => if val > 4 { 4 } else { val } None => 1, }; let sampler = match request.maybe_sampler.clone() { - Some(val) => if is_valid_string(&val) { val.clone() } else { String::from("Euler a") }, - None =>String::from( "Euler a"), + Some(val) => if is_valid_string(&val) { val.clone() } else { String::from("Euler a") } + None => String::from("Euler a"), }; - let sd_weight_token = ModelWeightToken(request.maybe_sd_model_token.clone().unwrap_or_default()); + let sd_weight_token = ModelWeightToken( + request.maybe_sd_model_token.clone().unwrap_or_default() + ); let lora_token = ModelWeightToken(request.maybe_lora_model_token.clone().unwrap_or_default()); - let n_prompt = request.maybe_n_prompt.clone().unwrap_or_default(); - // we can only do 1 upload type at a time. // if we are uploading a model. let upload_path = request.maybe_upload_path.clone().unwrap_or_default(); @@ -329,9 +370,15 @@ pub async fn enqueue_image_generation_request( seed = s; } - let type_of_inference = request.type_of_inference.to_string().clone(); - let description = request.description.to_string().clone(); - let name = request.name.to_string().clone(); + let mut version:u32 = 0; + let type_of_inference = mode.to_string().clone(); + let description = request.maybe_description.clone(); + let name = request.maybe_name.clone(); + + if let Some(s) = request.maybe_version { + version = s; + } + let inference_args = StableDiffusionArgs { maybe_sd_model_token: Some(sd_weight_token), maybe_lora_model_token: Some(lora_token), @@ -347,8 +394,9 @@ pub async fn enqueue_image_generation_request( maybe_width: Some(width), maybe_height: Some(height), maybe_sampler: Some(sampler), - description: Some(description), - name: Some(name), + maybe_description: description, + maybe_name: name, + maybe_version: Some(version), }; // create the inference args here @@ -373,7 +421,7 @@ pub async fn enqueue_image_generation_request( creator_ip_address: &ip_address, creator_set_visibility: visbility, priority_level, - requires_keepalive: true, + requires_keepalive: false, //reverse ... TODO fix this. we set it base on account is premium or not ... is_debug_request, maybe_routing_tag: maybe_routing_tag.as_deref(), mysql_pool: &server_state.mysql_pool, @@ -392,8 +440,67 @@ pub async fn enqueue_image_generation_request( inference_job_token: job_token, }; - let body = serde_json::to_string(&response).map_err(|_e| EnqueueImageGenRequestError::ServerError)?; + let body = serde_json + ::to_string(&response) + .map_err(|_e| EnqueueImageGenRequestError::ServerError)?; // Error handling 101 rust result type returned like so. Ok(HttpResponse::Ok().content_type("application/json").body(body)) } + +// with LoRA +// http://127.0.0.1:12345/v1/image_gen/inference/enqueue_image_gen +// { +// "uuid_idempotency_token": "12", +// "maybe_sd_model_token": "weight_dmmthavhawqc2hj7yqyemcbf8", +// "maybe_lora_model_token": "weight_t7gz78fjg27m0wtw6r33gafxs", +// "maybe_prompt": "raiden mei a very well drawn, anime girl, with pink and purple hair sitting down on a chair relaxed, highest quality, semi nude, masterpiece, painted", +// "maybe_a_prompt": "a anime girl, with pink and purple hair sitting down on a chair relaxed.", +// "maybe_n_prompt": "nsfw, black and white, low quality, pixelated", +// "maybe_seed": -1, +// "maybe_width": 1024, +// "maybe_height": 1024, +// "maybe_sampler": "DPM++ 2M SDE Heun", +// "maybe_cfg_scale": 7, +// "maybe_number_of_samples": 64, +// "maybe_batch_count": 4, +// } + +// without LoRA +// { +// "uuid_idempotency_token": "132141", +// "maybe_sd_model_token": "weight_dmmthavhawqc2hj7yqyemcbf8", +// "maybe_prompt": "raiden mei a very well drawn, anime girl, with pink and purple hair sitting down on a chair relaxed, highest quality, semi nude, masterpiece, painted", +// "maybe_a_prompt": "a anime girl, with pink and purple hair sitting down on a chair relaxed.", +// "maybe_n_prompt": "nsfw, black and white, low quality, pixelated", +// "maybe_seed": -1, +// "maybe_width": 512, +// "maybe_height": 512, +// "maybe_sampler": "DPM++ 2M SDE Heun", +// "maybe_cfg_scale": 8, +// "maybe_number_of_samples": 64, +// "maybe_batch_count": 4 +// } + +// http://127.0.0.1:12345/v1/image_gen/upload/lora +// { +// "uuid_idempotency_token": "12", +// "maybe_lora_upload_path": "https://drive.google.com/file/d/1WRgR2pn0Ky8ls5_9Zq6tQlHTBvyWeach/view?usp=sharing", +// "maybe_name":"some_name", +// "maybe_description":"some_description" +// } + +// http://127.0.0.1:12345/v1/image_gen/upload/model +// { +// "uuid_idempotency_token": "13", +// "maybe_upload_path": "https://drive.google.com/file/d/1WRgR2pn0Ky8ls5_9Zq6tQlHTBvyWeach/view?usp=sharing", +// "maybe_name":"some_name", +// "maybe_description":"some_description" +// } + +// { +// "uuid_idempotency_token": "123", +// "maybe_lora_upload_path": "https://drive.google.com/file/d/1WRgR2pn0Ky8ls5_9Zq6tQlHTBvyWeach/view?usp=sharing", +// "maybe_name":"some_name", +// "maybe_description":"some_description" +// } \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/routes.rs b/crates/service/web/storyteller_web/src/routes.rs index 0dc5ece407..862cc8ea07 100644 --- a/crates/service/web/storyteller_web/src/routes.rs +++ b/crates/service/web/storyteller_web/src/routes.rs @@ -1296,12 +1296,13 @@ fn add_image_gen_routes (app:App)-> App app.service( web::scope("/v1/image_gen") .service( - web::scope("/model") - .route("/upload", web::post().to(enqueue_image_generation_request)) + web::scope("/upload") + .route("/lora", web::post().to(enqueue_image_generation_request)) + .route("/model", web::post().to(enqueue_image_generation_request)) ) .service( - web::scope("/inference") - .route("/enqueue_image_gen", web::post().to(enqueue_image_generation_request)) + web::scope("/enqueue") + .route("/inference", web::post().to(enqueue_image_generation_request)) ) ) } @@ -1338,4 +1339,4 @@ fn add_weights_routes(app: App) -> App .route("/list", web::get().to(list_available_weights_handler)) .route("/list_featured", web::get().to(list_featured_weights_handler)) ) -} +} \ No newline at end of file