From 733c24aebb3e0517fcd72bb1c3c8a76cfa3e15b3 Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Tue, 2 Jan 2024 10:50:24 -0500 Subject: [PATCH 01/54] Working on enqueuing image generation --- .../generic_inference_args.rs | 7 + .../image_generation_payload.rs | 54 ++++ .../payloads/generic_inference_args/mod.rs | 3 +- .../image_gen/enqueue_image_generation.rs | 232 ++++++++++++++++++ .../http_server/endpoints/image_gen/mod.rs | 1 + .../src/http_server/endpoints/mod.rs | 3 +- .../service/web/storyteller_web/src/routes.rs | 27 +- 7 files changed, 324 insertions(+), 3 deletions(-) create mode 100644 crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs 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 a67dd0bd91..49da8fb98c 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 @@ -4,6 +4,7 @@ use errors::AnyhowResult; use crate::payloads::generic_inference_args::lipsync_payload::LipsyncArgs; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; use crate::payloads::generic_inference_args::videofilter_payload::{RerenderArgs}; +use crate::payloads::generic_inference_args::image_generation_payload::{SDArgs}; /// Used to encode extra state for the `generic_inference_jobs` table. @@ -40,6 +41,10 @@ pub enum InferenceCategoryAbbreviated { #[serde(rename = "vf")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. #[serde(alias = "video_filter")] VideoFilter, + + #[serde(rename = "ig")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + #[serde(alias = "image_generation")] + ImageGeneration } #[derive(Serialize, Deserialize, Clone, Copy, Debug)] @@ -113,6 +118,7 @@ impl InferenceCategoryAbbreviated { InferenceCategory::TextToSpeech => Self::TextToSpeech, InferenceCategory::VoiceConversion => Self::VoiceConversion, InferenceCategory::VideoFilter => Self::VideoFilter, + InferenceCategory::ImageGeneration => Self::ImageGeneration, } } @@ -122,6 +128,7 @@ impl InferenceCategoryAbbreviated { Self::TextToSpeech => InferenceCategory::TextToSpeech, Self::VoiceConversion => InferenceCategory::VoiceConversion, Self::VideoFilter => InferenceCategory::VideoFilter, + Self::ImageGeneration =>InferenceCategory::ImageGeneration } } } 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 new file mode 100644 index 0000000000..47c1efa11b --- /dev/null +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs @@ -0,0 +1,54 @@ +use tokens::tokens::model_weights::ModelWeightToken; + +/// Video sources can be one of several: +/// - F: media_files (todo) +/// - U: media_uploads (legacy) + +/// For image to image we probably want images +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub enum SDSource { + // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + /// Media File Token (media_files table) + /// Serde cannot yet rename enum variants. + F(String), + + // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + /// Media Upload Token (media_uploads table) + /// Serde cannot yet rename enum variants. + U(String), +} + +#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct SDArgs { + #[serde(rename = "vs")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_video_source: Option, + + #[serde(rename = "sd")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_sd_model_token: Option, + + #[serde(rename = "lora")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_lora_model_token: Option, + + #[serde(rename = "p")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_prompt: Option, + + #[serde(rename = "ap")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_a_prompt: Option, + + #[serde(rename = "np")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_n_prompt: Option, + + #[serde(rename = "se")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_seed: Option, + + #[serde(rename = "upload_path")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_upload_path: Option +} diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs index e68fc5a71c..17b564fc16 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs @@ -1,4 +1,5 @@ pub mod generic_inference_args; pub mod lipsync_payload; pub mod tts_payload; -pub mod videofilter_payload; \ No newline at end of file +pub mod videofilter_payload; +mod image_generation_payload; \ No newline at end of file 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 new file mode 100644 index 0000000000..c0655f3071 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs @@ -0,0 +1,232 @@ +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +use std::fmt::Debug; +use std::sync::Arc; + +use actix_web::{HttpRequest, HttpResponse, web}; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use log::warn; +use serde::Deserialize; +use serde::Serialize; +use utoipa::ToSchema; + +use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; +use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; +use http_server_common::request::get_request_header_optional::get_request_header_optional; +use http_server_common::request::get_request_ip::get_request_ip; +use mysql_queries::payloads::generic_inference_args::generic_inference_args::{ + GenericInferenceArgs, + InferenceCategoryAbbreviated, + PolymorphicInferenceArgs, +}; + +use mysql_queries::payloads::generic_inference_args::tts_payload::TTSArgs; +use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{ + insert_generic_inference_job, + InsertGenericInferenceArgs, +}; + +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::users::UserToken; + +use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; + +/// Debug requests can get routed to special "debug-only" workers, which can +/// be used to trial new code, run debugging, etc. +const DEBUG_HEADER_NAME: &str = "enable-debug-mode"; + +/// The routing tag header can send workloads to particular k8s hosts. +/// This is useful for catching the live logs or intercepting the job. +const ROUTING_TAG_HEADER_NAME: &str = "routing-tag"; + + +#[derive(Deserialize,ToSchema)] +pub struct EnqueueImageGenRequest { + uuid_idempotency_token: String, + text: String, + voice_token: String, +} + +#[derive(Serialize,ToSchema)] +pub struct EnqueueImageGenRequestSuccessResponse { + pub success: bool, + pub inference_job_token: InferenceJobToken, +} + +#[derive(Debug,ToSchema)] +pub enum EnqueueImageGenRequestError { + BadInput(String), + NotAuthorized, + ServerError, + RateLimited, +} + +impl ResponseError for EnqueueImageGenRequestError { + fn status_code(&self) -> StatusCode { + match *self { + EnqueueImageGenRequestError::BadInput(_) => StatusCode::BAD_REQUEST, + EnqueueImageGenRequestError::NotAuthorized => StatusCode::UNAUTHORIZED, + EnqueueImageGenRequestError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + EnqueueImageGenRequestError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + EnqueueImageGenRequestError::BadInput(reason) => reason.to_string(), + EnqueueImageGenRequestError::NotAuthorized => "unauthorized".to_string(), + EnqueueImageGenRequestError::ServerError => "server error".to_string(), + EnqueueImageGenRequestError::RateLimited => "rate limited".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +impl std::fmt::Display for EnqueueImageGenRequestError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self) + } +} + +// Implementation for enqueuing a TTS request +// Reference enqueue_infer_tts_handler.rs for checks: rate limiting / user sessions +// 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") +) +)] +pub async fn enqueue_image_generation_request( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data> +) -> Result { + + let mut maybe_user_token: Option = None; + + let mut mysql_connection = server_state.mysql_pool.acquire().await.map_err(|err| { + warn!("MySql pool error: {:?}", err); + EnqueueImageGenRequestError::ServerError + })?; + + // ==================== USER SESSION ==================== // + + let maybe_user_session = server_state.session_checker + .maybe_get_user_session_extended_from_connection(&http_request, &mut mysql_connection).await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EnqueueImageGenRequestError::ServerError + })?; + + if let Some(user_session) = maybe_user_session.as_ref() { + maybe_user_token = Some(UserToken::new_from_str(&user_session.user_token)); + } + + // Plan should handle "first anonymous use" and "investor" cases. + let plan = get_correct_plan_for_session( + server_state.server_environment, + maybe_user_session.as_ref() + ); + + // Separate priority for animation. + let priority_level = plan.web_vc_base_priority_level(); + + // ==================== DEBUG MODE + ROUTING TAG ==================== // + + let is_debug_request = get_request_header_optional(&http_request, DEBUG_HEADER_NAME).is_some(); + + let maybe_routing_tag = get_request_header_optional(&http_request, ROUTING_TAG_HEADER_NAME).map( + |routing_tag| routing_tag.trim().to_string() + ); + + // ==================== BANNED USERS ==================== // + + if let Some(ref user) = maybe_user_session { + if user.role.is_banned { + return Err(EnqueueImageGenRequestError::NotAuthorized); + } + } + + // ==================== RATE LIMIT ==================== // + + let rate_limiter = match maybe_user_session { + None => &server_state.redis_rate_limiters.logged_out, + Some(ref _user) => &server_state.redis_rate_limiters.logged_in, + }; + + if let Err(_err) = rate_limiter.rate_limit_request(&http_request) { + return Err(EnqueueImageGenRequestError::RateLimited); + } + + // Get up IP address + let ip_address = get_request_ip(&http_request); + + // package as larger component args should always have an embedding token .. + let inference_args = TTSArgs { + voice_token: Some(request.voice_token.clone()), + dataset_token: None, + }; + + // create the inference args here + // enqueue a zero shot tts request here... + let maybe_avt_token = server_state.avt_cookie_manager.get_avt_token_from_request(&http_request); + + // create the job record here! + let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { + uuid_idempotency_token: &request.uuid_idempotency_token, + inference_category: InferenceCategory::TextToSpeech, + maybe_model_type: Some(InferenceModelType::VallEX), // NB: Model is static during inference + maybe_model_token: None, // NB: Model is static during inference + maybe_input_source_token: None, + maybe_input_source_token_type: None, + maybe_raw_inference_text: Some(&request.text), + maybe_max_duration_seconds: None, + maybe_inference_args: Some(GenericInferenceArgs { + inference_category: Some(InferenceCategoryAbbreviated::TextToSpeech), + args: Some(PolymorphicInferenceArgs::Tts(inference_args)), + }), + maybe_creator_user_token: maybe_user_token.as_ref(), + maybe_avt_token: maybe_avt_token.as_ref(), + creator_ip_address: &ip_address, + creator_set_visibility: enums::common::visibility::Visibility::Public, + priority_level, + requires_keepalive: true, // do we need this? I think so + is_debug_request, + maybe_routing_tag: maybe_routing_tag.as_deref(), + mysql_pool: &server_state.mysql_pool, + }).await; + + let job_token = match query_result { + Ok((job_token, _id)) => job_token, + Err(err) => { + warn!("New generic inference job creation DB error: {:?}", err); + return Err(EnqueueImageGenRequestError::ServerError); + } + }; + + let response: EnqueueImageGenRequestSuccessResponse = EnqueueImageGenRequestSuccessResponse { + success: true, + inference_job_token: job_token, + }; + + 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)) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs new file mode 100644 index 0000000000..fd510c2b58 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs @@ -0,0 +1 @@ +pub mod enqueue_image_generation; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/mod.rs index 0a6ed9375d..5b92cab8d4 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/mod.rs @@ -26,4 +26,5 @@ pub mod voice_conversion; pub mod voice_designer; pub mod w2l; pub mod user_bookmarks; -pub mod weights; \ No newline at end of file +pub mod weights; +pub mod image_gen; \ 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 6e85b5e805..ff66d62d46 100644 --- a/crates/service/web/storyteller_web/src/routes.rs +++ b/crates/service/web/storyteller_web/src/routes.rs @@ -162,7 +162,7 @@ use crate::http_server::endpoints::weights::list_featured_weights::list_featured use crate::http_server::endpoints::weights::list_weights_by_user::list_weights_by_user_handler; use crate::http_server::endpoints::weights::set_model_weight_cover_image::set_model_weight_cover_image_handler; use crate::http_server::endpoints::weights::update_weight::update_weight_handler; - +use crate::http_server::endpoints::image_gen::enqueue_image_generation:: pub fn add_routes (app: App, server_environment: ServerEnvironment) -> App where B: MessageBody, @@ -1257,6 +1257,31 @@ fn add_voice_designer_routes (app:App)-> App ) } + +fn add_image_gen_routes (app:App)-> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + app.service( + web::scope("/v1/image_gen") + .service( + web::scope("/model") + .route("/upload", web::post().to(upload_zs_sample_handler)) + ) + .service( + web::scope("/inference") + .route("/enqueue_image_gen", web::post().to(enqueue_tts_request)) + ) + ) +} + // ==================== Weights ROUTES ==================== fn add_weights_routes(app: App) -> App where From 7d858d2c196d65b1031cb17181f1ef773a4dea52 Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Tue, 2 Jan 2024 12:24:13 -0500 Subject: [PATCH 02/54] Payload almost ready need to migrate the db and write unit tests. --- .../inference_category.rs | 10 ++++ .../generic_inference_args.rs | 55 +++++++++++++++++-- .../image_generation_payload.rs | 20 +++++-- 3 files changed, 75 insertions(+), 10 deletions(-) diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs index 19bf5a74c7..99c146683f 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs @@ -28,6 +28,9 @@ pub enum InferenceCategory { #[serde(rename = "video_filter")] VideoFilter, + + #[serde(rename = "image_generation")] + ImageGeneration } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -42,6 +45,7 @@ impl InferenceCategory { Self::TextToSpeech => "text_to_speech", Self::VoiceConversion => "voice_conversion", Self::VideoFilter => "video_filter", + Self::ImageGeneration => "image_generation", } } @@ -51,6 +55,7 @@ impl InferenceCategory { "text_to_speech" => Ok(Self::TextToSpeech), "voice_conversion" => Ok(Self::VoiceConversion), "video_filter" => Ok(Self::VideoFilter), + "image_generation" => Ok(Self::ImageGeneration), _ => Err(format!("invalid value: {:?}", value)), } } @@ -63,6 +68,7 @@ impl InferenceCategory { Self::TextToSpeech, Self::VoiceConversion, Self::VideoFilter, + Self::ImageGeneration, ]) } } @@ -78,6 +84,7 @@ mod tests { assert_serialization(InferenceCategory::TextToSpeech, "text_to_speech"); assert_serialization(InferenceCategory::VoiceConversion, "voice_conversion"); assert_serialization(InferenceCategory::VideoFilter, "video_filter"); + assert_serialization(InferenceCategory::ImageGeneration, "image_generation"); } #[test] @@ -86,6 +93,7 @@ mod tests { assert_eq!(InferenceCategory::TextToSpeech.to_str(), "text_to_speech"); assert_eq!(InferenceCategory::VoiceConversion.to_str(), "voice_conversion"); assert_eq!(InferenceCategory::VideoFilter.to_str(), "video_filter"); + assert_eq!(InferenceCategory::ImageGeneration.to_str(), "image_generation"); } #[test] @@ -94,6 +102,7 @@ mod tests { assert_eq!(InferenceCategory::from_str("text_to_speech").unwrap(), InferenceCategory::TextToSpeech); assert_eq!(InferenceCategory::from_str("voice_conversion").unwrap(), InferenceCategory::VoiceConversion); assert_eq!(InferenceCategory::from_str("video_filter").unwrap(), InferenceCategory::VideoFilter); + assert_eq!(InferenceCategory::from_str("image_generation").unwrap(), InferenceCategory::ImageGeneration); } #[test] @@ -105,6 +114,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(InferenceCategory::TextToSpeech)); assert_eq!(variants.pop_first(), Some(InferenceCategory::VoiceConversion)); assert_eq!(variants.pop_first(), Some(InferenceCategory::VideoFilter)); + assert_eq!(variants.pop_first(), Some(InferenceCategory::ImageGeneration)); assert_eq!(variants.pop_first(), None); // Generated check 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 49da8fb98c..90bc808db2 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 @@ -3,10 +3,10 @@ use errors::AnyhowResult; use crate::payloads::generic_inference_args::lipsync_payload::LipsyncArgs; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; -use crate::payloads::generic_inference_args::videofilter_payload::{RerenderArgs}; -use crate::payloads::generic_inference_args::image_generation_payload::{SDArgs}; - +use crate::payloads::generic_inference_args::videofilter_payload::RerenderArgs; +use crate::payloads::generic_inference_args::image_generation_payload::SDArgs; +use tokens::tokens::{model_weights::ModelWeightToken, media_files::MediaFileToken}; /// Used to encode extra state for the `generic_inference_jobs` table. /// This should act somewhat like a serialized protobuf stored inside a record. #[derive(Clone, Debug, Serialize, Deserialize)] @@ -96,7 +96,10 @@ pub enum PolymorphicInferenceArgs { }, /// Rerender a video. (Short name to save space when serializing.) - Rr(RerenderArgs) + Rr(RerenderArgs), + + /// Image generation. (Short name to save space when serializing.) + Ig(SDArgs) } @@ -135,9 +138,13 @@ impl InferenceCategoryAbbreviated { #[cfg(test)] mod tests { - use crate::payloads::generic_inference_args::generic_inference_args::{FundamentalFrequencyMethodForJob, GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; + use enums::by_table::generic_inference_jobs::inference_category; + +use crate::payloads::generic_inference_args::generic_inference_args::{FundamentalFrequencyMethodForJob, GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; use crate::payloads::generic_inference_args::lipsync_payload::{LipsyncAnimationAudioSource, LipsyncAnimationImageSource, LipsyncArgs}; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; + use crate::payloads::generic_inference_args::image_generation_payload::SDArgs; + use tokens::tokens::{model_weights::ModelWeightToken, media_files::MediaFileToken}; #[test] fn typical_lipsync_animation_args_serialize() { @@ -184,6 +191,44 @@ mod tests { assert!(json.len() < 1000); } + #[test] + fn typical_image_gen_args_serialize() { + + let video_media_token = MediaFileToken("video_media_token".to_string()); + let image_media_token = MediaFileToken("image_media_token".to_string()); + let sd_model_token = ModelWeightToken("sd_model_token".to_string()); + let lora_model_token = ModelWeightToken("lora_model_token".to_string()); + + let prompt = "prompt".to_string(); + let a_prompt = "a_prompt".to_string(); + let n_prompt = "n_prompt".to_string(); + let seed = 1; + let upload_path = "upload_path".to_string(); + let lora_upload_path = "lora_upload_path".to_string(); + + let args = GenericInferenceArgs { + inference_category: Some(InferenceCategoryAbbreviated::ImageGeneration), + args: Some(PolymorphicInferenceArgs::Ig(SDArgs { + maybe_video_source: Some(video_media_token), + maybe_image_source: Some(image_media_token), + maybe_sd_model_token: Some(sd_model_token), + maybe_lora_model_token: Some(lora_model_token), + maybe_prompt: Some(prompt), + maybe_a_prompt: Some(a_prompt), + maybe_n_prompt: Some(n_prompt), + maybe_seed: Some(seed), + maybe_upload_path: Some(upload_path), + maybe_lora_upload_path: Some(lora_upload_path) + })), + }; + + let json = serde_json::ser::to_string(&args).unwrap(); + + assert_eq!(json, r#"{"cat":"ig","args":{"Ig":{"v":"video_media_token","i":"image_media_token","s":"sd_model_token","l":"lora_model_token","p":"prompt","a":"a_prompt","n":"n_prompt","e":1,"u":"upload_path","o":"lora_upload_path"}}}"#.to_string()); + + + } + #[test] fn typical_voice_conversion_args_serialize() { let args = GenericInferenceArgs { 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 47c1efa11b..87a23e9969 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 @@ -1,4 +1,4 @@ -use tokens::tokens::model_weights::ModelWeightToken; +use tokens::tokens::{model_weights::ModelWeightToken, media_files::MediaFileToken}; /// Video sources can be one of several: /// - F: media_files (todo) @@ -22,7 +22,11 @@ pub enum SDSource { pub struct SDArgs { #[serde(rename = "vs")] #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_video_source: Option, + pub maybe_video_source: Option, + + #[serde(rename = "is")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_image_source: Option, #[serde(rename = "sd")] #[serde(skip_serializing_if = "Option::is_none")] @@ -47,8 +51,14 @@ pub struct SDArgs { #[serde(rename = "se")] #[serde(skip_serializing_if = "Option::is_none")] pub maybe_seed: Option, - - #[serde(rename = "upload_path")] + + #[serde(rename = "mup")] #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_upload_path: Option + pub maybe_upload_path: Option, + + #[serde(rename = "lup")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_lora_upload_path: Option } + + From 87b2f4a113c7a21e7a769e46555c9ee9805aa3ef Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Tue, 2 Jan 2024 12:55:34 -0500 Subject: [PATCH 03/54] Args are tested --- .../generic_inference_args/generic_inference_args.rs | 2 +- .../generic_inference_args/image_generation_payload.rs | 6 +++--- .../generic_inference/web/get_inference_job_status.rs | 2 ++ crates/lib/mysql_queries/src/queries/model_weights/tests.rs | 2 ++ 4 files changed, 8 insertions(+), 4 deletions(-) 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 90bc808db2..eadffa7e5b 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 @@ -224,7 +224,7 @@ use crate::payloads::generic_inference_args::generic_inference_args::{Fundamenta let json = serde_json::ser::to_string(&args).unwrap(); - assert_eq!(json, r#"{"cat":"ig","args":{"Ig":{"v":"video_media_token","i":"image_media_token","s":"sd_model_token","l":"lora_model_token","p":"prompt","a":"a_prompt","n":"n_prompt","e":1,"u":"upload_path","o":"lora_upload_path"}}}"#.to_string()); + assert_eq!(json, r#"{"cat":"ig","args":{"Ig":{"vs":"video_media_token","is":"image_media_token","sd":"sd_model_token","lm":"lora_model_token","p":"prompt","ap":"a_prompt","np":"n_prompt","se":1,"mu":"upload_path","lu":"lora_upload_path"}}}"#.to_string()); } 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 87a23e9969..ce2b0a4d22 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 @@ -32,7 +32,7 @@ pub struct SDArgs { #[serde(skip_serializing_if = "Option::is_none")] pub maybe_sd_model_token: Option, - #[serde(rename = "lora")] + #[serde(rename = "lm")] #[serde(skip_serializing_if = "Option::is_none")] pub maybe_lora_model_token: Option, @@ -52,11 +52,11 @@ pub struct SDArgs { #[serde(skip_serializing_if = "Option::is_none")] pub maybe_seed: Option, - #[serde(rename = "mup")] + #[serde(rename = "mu")] #[serde(skip_serializing_if = "Option::is_none")] pub maybe_upload_path: Option, - #[serde(rename = "lup")] + #[serde(rename = "lu")] #[serde(skip_serializing_if = "Option::is_none")] pub maybe_lora_upload_path: Option } diff --git a/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs b/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs index e2a9b7be9c..04c14aab2f 100644 --- a/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs +++ b/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs @@ -179,6 +179,7 @@ fn raw_record_to_public_result(record: RawGenericInferenceJobStatus) -> GenericI InferenceCategory::TextToSpeech => record.maybe_tts_model_title.as_deref(), InferenceCategory::VoiceConversion => record.maybe_voice_conversion_model_title.as_deref(), InferenceCategory::VideoFilter => Some("Video Filter"), + InferenceCategory::ImageGeneration => Some("Image Generation"), }; // NB: A bit of a hack. We store TTS results with a full path. @@ -188,6 +189,7 @@ fn raw_record_to_public_result(record: RawGenericInferenceJobStatus) -> GenericI InferenceCategory::TextToSpeech => (false, record.maybe_tts_public_bucket_path.as_deref()), InferenceCategory::VoiceConversion => (true, record.maybe_voice_conversion_public_bucket_hash.as_deref()), InferenceCategory::VideoFilter => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), + InferenceCategory::ImageGeneration => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), }; // NB: We've moved voice conversion out of their own table and into media_files diff --git a/crates/lib/mysql_queries/src/queries/model_weights/tests.rs b/crates/lib/mysql_queries/src/queries/model_weights/tests.rs index 788f03cd78..09f10b570e 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/tests.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/tests.rs @@ -295,6 +295,8 @@ mod tests { sort_ascending: false, view_as, mysql_pool: &pool, + maybe_scoped_weight_type: None, + maybe_scoped_weight_category: None, }).await?; for weight in weights_by_username.records.iter() { // print weight From 20b95c090eb9f2e495dc6959956866824903911c Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Tue, 2 Jan 2024 16:16:51 -0500 Subject: [PATCH 04/54] Now testing the end point. --- .../inference_model_type.rs | 13 ++- .../generic_inference_args.rs | 4 - .../image_generation_payload.rs | 3 + .../payloads/generic_inference_args/mod.rs | 2 +- .../image_gen/enqueue_generation_type.rs | 76 +++++++++++++ .../image_gen/enqueue_image_generation.rs | 101 +++++++++++++++--- .../http_server/endpoints/image_gen/mod.rs | 1 + .../service/web/storyteller_web/src/routes.rs | 8 +- 8 files changed, 182 insertions(+), 26 deletions(-) create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_generation_type.rs diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs index a3ea0bcd29..823995a59e 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs @@ -33,6 +33,8 @@ pub enum InferenceModelType { VallEX, #[serde(rename = "rerender_a_video")] RerenderAVideo, + #[serde(rename = "stable_diffusion")] + StableDiffusion, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -50,6 +52,7 @@ impl InferenceModelType { Self::Vits => "vits", Self::VallEX => "vall_e_x", Self::RerenderAVideo => "rerender_a_video", + Self::StableDiffusion => "stable_diffusion" } } @@ -62,6 +65,7 @@ impl InferenceModelType { "vits" => Ok(Self::Vits), "vall_e_x" => Ok(Self::VallEX), "rerender_a_video" => Ok(Self::RerenderAVideo), + "stable_diffusion" => Ok(Self::StableDiffusion), _ => Err(format!("invalid value: {:?}", value)), } } @@ -76,7 +80,8 @@ impl InferenceModelType { InferenceModelType::Tacotron2, InferenceModelType::Vits, InferenceModelType::VallEX, - InferenceModelType::RerenderAVideo + InferenceModelType::RerenderAVideo, + InferenceModelType::StableDiffusion ]) } } @@ -95,6 +100,7 @@ mod tests { assert_serialization(InferenceModelType::Vits, "vits"); assert_serialization(InferenceModelType::VallEX, "vall_e_x"); assert_serialization(InferenceModelType::RerenderAVideo, "rerender_a_video"); + assert_serialization(InferenceModelType::StableDiffusion, "stable_diffusion"); } #[test] @@ -106,6 +112,7 @@ mod tests { assert_eq!(InferenceModelType::Vits.to_str(), "vits"); assert_eq!(InferenceModelType::VallEX.to_str(), "vall_e_x"); assert_eq!(InferenceModelType::RerenderAVideo.to_str(), "rerender_a_video"); + assert_eq!(InferenceModelType::StableDiffusion.to_str(), "stable_diffusion"); } #[test] @@ -117,13 +124,14 @@ mod tests { assert_eq!(InferenceModelType::from_str("vits").unwrap(), InferenceModelType::Vits); assert_eq!(InferenceModelType::from_str("vall_e_x").unwrap(), InferenceModelType::VallEX); assert_eq!(InferenceModelType::from_str("rerender_a_video").unwrap(), InferenceModelType::RerenderAVideo); + assert_eq!(InferenceModelType::from_str("stable_diffusion").unwrap(), InferenceModelType::StableDiffusion); } #[test] fn all_variants() { // Static check let mut variants = InferenceModelType::all_variants(); - assert_eq!(variants.len(), 7); + assert_eq!(variants.len(), 8); assert_eq!(variants.pop_first(), Some(InferenceModelType::RvcV2)); assert_eq!(variants.pop_first(), Some(InferenceModelType::SadTalker)); assert_eq!(variants.pop_first(), Some(InferenceModelType::SoVitsSvc)); @@ -131,6 +139,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(InferenceModelType::Vits)); assert_eq!(variants.pop_first(), Some(InferenceModelType::VallEX)); assert_eq!(variants.pop_first(), Some(InferenceModelType::RerenderAVideo)); + assert_eq!(variants.pop_first(), Some(InferenceModelType::StableDiffusion)); assert_eq!(variants.pop_first(), None); // Generated check 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 eadffa7e5b..a4fb7f1df2 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 @@ -6,7 +6,6 @@ use crate::payloads::generic_inference_args::tts_payload::TTSArgs; use crate::payloads::generic_inference_args::videofilter_payload::RerenderArgs; use crate::payloads::generic_inference_args::image_generation_payload::SDArgs; -use tokens::tokens::{model_weights::ModelWeightToken, media_files::MediaFileToken}; /// Used to encode extra state for the `generic_inference_jobs` table. /// This should act somewhat like a serialized protobuf stored inside a record. #[derive(Clone, Debug, Serialize, Deserialize)] @@ -223,10 +222,7 @@ use crate::payloads::generic_inference_args::generic_inference_args::{Fundamenta }; let json = serde_json::ser::to_string(&args).unwrap(); - assert_eq!(json, r#"{"cat":"ig","args":{"Ig":{"vs":"video_media_token","is":"image_media_token","sd":"sd_model_token","lm":"lora_model_token","p":"prompt","ap":"a_prompt","np":"n_prompt","se":1,"mu":"upload_path","lu":"lora_upload_path"}}}"#.to_string()); - - } #[test] 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 ce2b0a4d22..c3385bdcb3 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 @@ -59,6 +59,9 @@ pub struct SDArgs { #[serde(rename = "lu")] #[serde(skip_serializing_if = "Option::is_none")] pub maybe_lora_upload_path: Option + + #[serde(rename = "t")] + pub inference_type: Option } diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs index 17b564fc16..510e503861 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs @@ -2,4 +2,4 @@ pub mod generic_inference_args; pub mod lipsync_payload; pub mod tts_payload; pub mod videofilter_payload; -mod image_generation_payload; \ No newline at end of file +pub mod image_generation_payload; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_generation_type.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_generation_type.rs new file mode 100644 index 0000000000..da85580539 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_generation_type.rs @@ -0,0 +1,76 @@ +use std::collections::BTreeSet; + +#[cfg(test)] +use strum::EnumCount; +#[cfg(test)] +use strum::EnumIter; +use utoipa::ToSchema; + + +#[cfg_attr(test, derive(EnumIter, EnumCount))] +#[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Deserialize, Serialize, ToSchema)] +pub enum EnqueueImageGenType { + #[serde(rename = "upload_lora")] + UploadLoRA, + #[serde(rename = "upload_sd")] + Checkpoint, + #[serde(rename = "inference")] + Inference, +} + +impl EnqueueImageGenType { + pub fn to_str(&self) -> &'static str { + match self { + Self::UploadLoRA => "upload_lora", + Self::Checkpoint => "upload_sd", + Self::Inference => "inference", + } + } + + pub fn from_str(value: &str) -> Result { + match value { + "upload_lora" => Ok(Self::UploadLoRA), + "upload_sd" => Ok(Self::Checkpoint), + "inference" => Ok(Self::Inference), + _ => Err(format!("invalid value: {:?}", value)), + } + } + + pub fn all_variants() -> BTreeSet { + BTreeSet::from([ + Self::UploadLoRA, + Self::Checkpoint, + Self::Inference + ]) + } +} + + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_to_str() { + assert_eq!(EnqueueImageGenType::UploadLoRA.to_str(), "upload_lora"); + assert_eq!(EnqueueImageGenType::Checkpoint.to_str(), "upload_sd"); + assert_eq!(EnqueueImageGenType::Inference.to_str(), "inference"); + } + + #[test] + fn test_from_str() { + assert_eq!(EnqueueImageGenType::from_str("upload_lora").unwrap(), EnqueueImageGenType::UploadLoRA); + assert_eq!(EnqueueImageGenType::from_str("upload_sd").unwrap(), EnqueueImageGenType::Checkpoint); + assert_eq!(EnqueueImageGenType::from_str("inference").unwrap(), EnqueueImageGenType::Inference); + assert!(EnqueueImageGenType::from_str("invalid").is_err()); + } + + #[test] + fn test_all_variants() { + let variants = EnqueueImageGenType::all_variants(); + assert_eq!(variants.len(), 3); + assert!(variants.contains(&EnqueueImageGenType::UploadLoRA)); + assert!(variants.contains(&EnqueueImageGenType::Checkpoint)); + assert!(variants.contains(&EnqueueImageGenType::Inference)); + } +} \ No newline at end of file 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 c0655f3071..417887fd1a 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 @@ -8,9 +8,12 @@ use std::sync::Arc; use actix_web::{HttpRequest, HttpResponse, web}; use actix_web::error::ResponseError; use actix_web::http::StatusCode; +use elasticsearch::http::request; use log::warn; use serde::Deserialize; use serde::Serialize; +use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::model_weights::ModelWeightToken; use utoipa::ToSchema; use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; @@ -23,16 +26,19 @@ use mysql_queries::payloads::generic_inference_args::generic_inference_args::{ PolymorphicInferenceArgs, }; -use mysql_queries::payloads::generic_inference_args::tts_payload::TTSArgs; +use mysql_queries::payloads::generic_inference_args::image_generation_payload::SDArgs; use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{ insert_generic_inference_job, InsertGenericInferenceArgs, }; + use tokens::tokens::generic_inference_jobs::InferenceJobToken; use tokens::tokens::users::UserToken; use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session; +use crate::http_server::endpoints::media_files::drain_multipart_request::MediaSource; +use crate::http_server::endpoints::weights::list_weights_by_user::Weight; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::server_state::ServerState; @@ -44,12 +50,30 @@ 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(Deserialize,ToSchema)] +pub enum TypeOfInference { + #[serde(rename = "lora")] + UploadLoRA, + #[serde(rename = "checkpoint")] + CheckPoint, + #[serde(rename = "standard_inference")] + StandardInference, +} #[derive(Deserialize,ToSchema)] pub struct EnqueueImageGenRequest { uuid_idempotency_token: String, - text: String, - voice_token: String, + type_of_inference: TypeOfInference, // upload loRA / check point / standard inference + maybe_video_source: Option, + maybe_image_source: Option, + maybe_sd_model_token: Option, + maybe_lora_model_token: Option, + maybe_prompt: Option, + maybe_a_prompt: Option, + maybe_n_prompt: Option, + maybe_seed: Option, + maybe_upload_path: Option, + maybe_lora_upload_path: Option, } #[derive(Serialize,ToSchema)] @@ -119,7 +143,7 @@ pub async fn enqueue_image_generation_request( ) -> Result { let mut maybe_user_token: Option = None; - + let mut visbility = enums::common::visibility::Visibility::Public; let mut mysql_connection = server_state.mysql_pool.acquire().await.map_err(|err| { warn!("MySql pool error: {:?}", err); EnqueueImageGenRequestError::ServerError @@ -133,7 +157,8 @@ pub async fn enqueue_image_generation_request( warn!("Session checker error: {:?}", e); EnqueueImageGenRequestError::ServerError })?; - + + if let Some(user_session) = maybe_user_session.as_ref() { maybe_user_token = Some(UserToken::new_from_str(&user_session.user_token)); } @@ -177,36 +202,80 @@ pub async fn enqueue_image_generation_request( // Get up IP address let ip_address = get_request_ip(&http_request); - // package as larger component args should always have an embedding token .. - let inference_args = TTSArgs { - voice_token: Some(request.voice_token.clone()), - dataset_token: None, + // Check the inference args to make sure everything is all there for upload loRA / model or standard inference + + // ==================== INFERENCE ARGS ==================== // + + if request.type_of_inference == EnqueueImageGenType::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 == EnqueueImageGenRequest::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 == EnqueueImageGenRequest::CheckPoint { + if request.maybe_upload_path.is_none() { + return Err(EnqueueImageGenRequestError::BadInput("No upload path provided".to_string())); + } + } + + let video_token = MediaFileToken(request.maybe_video_source.clone().unwrap_or_default()); + let image_source_token = MediaFileToken(request.maybe_image_source.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 a_prompt = request.maybe_a_prompt.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(); + // if we are uploading a lora model. + let lora_upload_path = request.maybe_lora_upload_path.clone().unwrap_or_default(); + + let inference_args = SDArgs { + maybe_video_source: Some(video_token), + maybe_image_source: Some(image_source_token), + maybe_sd_model_token: Some(sd_weight_token), + maybe_lora_model_token: Some(lora_token), + maybe_prompt: Some(request.maybe_prompt.clone().unwrap_or_default()), + maybe_a_prompt: Some(a_prompt), + maybe_n_prompt: Some(n_prompt), + maybe_seed: Some(request.maybe_seed.clone().unwrap_or_default()), + maybe_upload_path: Some(upload_path), + maybe_lora_upload_path: Some(lora_upload_path), + inference_type: Some(request.type_of_inference.clone()), }; // create the inference args here - // enqueue a zero shot tts request here... let maybe_avt_token = server_state.avt_cookie_manager.get_avt_token_from_request(&http_request); // create the job record here! let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, - inference_category: InferenceCategory::TextToSpeech, - maybe_model_type: Some(InferenceModelType::VallEX), // NB: Model is static during inference + inference_category: InferenceCategory::ImageGeneration, + maybe_model_type: Some(InferenceModelType::StableDiffusion), // NB: Model is static during inference maybe_model_token: None, // NB: Model is static during inference maybe_input_source_token: None, maybe_input_source_token_type: None, maybe_raw_inference_text: Some(&request.text), maybe_max_duration_seconds: None, maybe_inference_args: Some(GenericInferenceArgs { - inference_category: Some(InferenceCategoryAbbreviated::TextToSpeech), - args: Some(PolymorphicInferenceArgs::Tts(inference_args)), + inference_category: Some(InferenceCategoryAbbreviated::ImageGeneration), + args: Some(PolymorphicInferenceArgs::Ig(inference_args)), }), maybe_creator_user_token: maybe_user_token.as_ref(), maybe_avt_token: maybe_avt_token.as_ref(), creator_ip_address: &ip_address, - creator_set_visibility: enums::common::visibility::Visibility::Public, + creator_set_visibility: visbility, priority_level, - requires_keepalive: true, // do we need this? I think so + requires_keepalive: true, is_debug_request, maybe_routing_tag: maybe_routing_tag.as_deref(), mysql_pool: &server_state.mysql_pool, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs index fd510c2b58..460b32b57c 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs @@ -1 +1,2 @@ pub mod enqueue_image_generation; +pub mod enqueue_generation_type; \ 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 ff66d62d46..633b0e6bf1 100644 --- a/crates/service/web/storyteller_web/src/routes.rs +++ b/crates/service/web/storyteller_web/src/routes.rs @@ -162,7 +162,8 @@ use crate::http_server::endpoints::weights::list_featured_weights::list_featured use crate::http_server::endpoints::weights::list_weights_by_user::list_weights_by_user_handler; use crate::http_server::endpoints::weights::set_model_weight_cover_image::set_model_weight_cover_image_handler; use crate::http_server::endpoints::weights::update_weight::update_weight_handler; -use crate::http_server::endpoints::image_gen::enqueue_image_generation:: +use crate::http_server::endpoints::image_gen::enqueue_image_generation::enqueue_image_generation_request; + pub fn add_routes (app: App, server_environment: ServerEnvironment) -> App where B: MessageBody, @@ -1269,15 +1270,16 @@ fn add_image_gen_routes (app:App)-> App InitError = (), >, { + // app.service( web::scope("/v1/image_gen") .service( web::scope("/model") - .route("/upload", web::post().to(upload_zs_sample_handler)) + .route("/upload", web::post().to(enqueue_image_generation_request)) ) .service( web::scope("/inference") - .route("/enqueue_image_gen", web::post().to(enqueue_tts_request)) + .route("/enqueue_image_gen", web::post().to(enqueue_image_generation_request)) ) ) } From bcd806e7f0b1f4095d16bb3601b45ae269586f4c Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Tue, 2 Jan 2024 17:50:22 -0500 Subject: [PATCH 05/54] Some boilerplate done. --- .../download_job/src/job_loop/process_single_job.rs | 2 ++ .../src/job/job_loop/process_single_job.rs | 4 ++++ .../job/inference_job/src/job/job_types/mod.rs | 3 ++- .../rerender_a_video/download_model_file.rs | 2 +- .../rerender_a_video/download_video_file.rs | 2 +- .../inference_job/src/job_specific_dependencies.rs | 11 +++++++++++ 6 files changed, 21 insertions(+), 3 deletions(-) diff --git a/crates/service/job/download_job/src/job_loop/process_single_job.rs b/crates/service/job/download_job/src/job_loop/process_single_job.rs index ba69531e70..4a133216d9 100644 --- a/crates/service/job/download_job/src/job_loop/process_single_job.rs +++ b/crates/service/job/download_job/src/job_loop/process_single_job.rs @@ -76,6 +76,8 @@ pub async fn do_process_single_job(job_state: &JobState, job: &AvailableDownload return Err(anyhow!("Bad download URL: `{}`", &job.download_url)); } + // TODO USE THIS DOWNLOADER + let download_filename = match job_state.sidecar_configs.google_drive_downloader.download_file(&job.download_url, &temp_dir).await { Ok(filename) => filename, Err(e) => { diff --git a/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs b/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs index 1aab84a477..a9896d8a7b 100644 --- a/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs +++ b/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs @@ -17,6 +17,7 @@ use crate::job::job_types::lipsync::process_single_lipsync_job::process_single_l use crate::job::job_types::tts::process_single_tts_job::process_single_tts_job; use crate::job::job_types::vc::process_single_vc_job::process_single_vc_job; use crate::job::job_types::videofilter::process_single_vf_job::process_single_vf_job; +use crate::job::job_types::imagegen::process_single_ig_job::process_single_ig_job; use crate::job_dependencies::JobDependencies; pub async fn process_single_job( @@ -186,6 +187,9 @@ async fn do_process_single_job( InferenceCategory::VideoFilter => { process_single_vf_job(job_dependencies, job).await? } + InferenceCategory::ImageGeneration => { + process_single_ig_job(job_dependencies, job).await? + } }; let maybe_entity_type = job_success_result.maybe_result_entity diff --git a/crates/service/job/inference_job/src/job/job_types/mod.rs b/crates/service/job/inference_job/src/job/job_types/mod.rs index a688e281ba..d73c394afb 100644 --- a/crates/service/job/inference_job/src/job/job_types/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/mod.rs @@ -1,4 +1,5 @@ pub mod lipsync; pub mod tts; pub mod vc; -pub(crate) mod videofilter; \ No newline at end of file +pub(crate) mod videofilter; +pub mod imagegen; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_model_file.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_model_file.rs index bc5ded36a7..0cd5fa52f4 100644 --- a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_model_file.rs +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_model_file.rs @@ -16,7 +16,7 @@ use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; pub struct ModelFile { pub filesystem_path: PathBuf, } - +// TODO replace this with remote cloud file manager. pub async fn download_model_file( model_record: &Option, public_bucket_client: &BucketClient, diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs index 6ec512d1ea..86d73379a8 100644 --- a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs @@ -22,7 +22,7 @@ use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; pub struct VideoFile { pub filesystem_path: PathBuf, } - +// TODO replace this with remote cloud file manager. pub async fn download_video_file( video_source: &VideofilterVideoSource, public_bucket_client: &BucketClient, diff --git a/crates/service/job/inference_job/src/job_specific_dependencies.rs b/crates/service/job/inference_job/src/job_specific_dependencies.rs index a863b3872c..f70052dc2c 100644 --- a/crates/service/job/inference_job/src/job_specific_dependencies.rs +++ b/crates/service/job/inference_job/src/job_specific_dependencies.rs @@ -1,5 +1,6 @@ use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; use errors::AnyhowResult; +use sqlx::pool::maybe; use crate::job::job_types::lipsync::sad_talker::sad_talker_dependencies::SadTalkerDependencies; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_dependencies::Tacotron2Dependencies; @@ -8,6 +9,8 @@ use crate::job::job_types::tts::vits::vits_dependencies::VitsDependencies; use crate::job::job_types::vc::rvc_v2::rvc_v2_dependencies::RvcV2Dependencies; use crate::job::job_types::vc::so_vits_svc::svc_dependencies::SvcDependencies; use crate::job::job_types::videofilter::rerender_a_video::rerender_dependencies::RerenderDependencies; +use crate::job::job_types::imagegen::sd::stable_diffusion_dependencies::StableDiffusionDependencies; + use crate::util::scoped_execution::ScopedExecution; pub struct JobSpecificDependencies { @@ -18,6 +21,7 @@ pub struct JobSpecificDependencies { pub maybe_vall_e_x_dependencies: Option, pub maybe_vits_dependencies: Option, pub maybe_rerender_dependencies: Option, + pub maybe_stable_diffusion_dependencies: Option, } impl JobSpecificDependencies { @@ -30,6 +34,7 @@ impl JobSpecificDependencies { let mut maybe_vall_e_x_dependencies = None; let mut maybe_vits_dependencies = None; let mut maybe_rerender_dependencies = None; + let mut maybe_stable_diffusion_dependencies = None; if scoped_execution.can_run_job(InferenceModelType::RvcV2) { print_with_space("Setting RVCv2 dependencies..."); @@ -65,6 +70,11 @@ impl JobSpecificDependencies { print_with_space("Setting Rerender dependencies..."); maybe_rerender_dependencies = Some(RerenderDependencies::setup()?); } + + if scoped_execution.can_run_job(job_model_type::StableDiffusion) { + print_with_space("Setting Stable Diffusion dependencies..."); + maybe_stable_diffusion_dependencies = Some(StableDiffusionDependencies::setup()?); + } Ok(JobSpecificDependencies { maybe_rvc_v2_dependencies, @@ -74,6 +84,7 @@ impl JobSpecificDependencies { maybe_vall_e_x_dependencies, maybe_vits_dependencies, maybe_rerender_dependencies, + maybe_stable_diffusion_dependencies }) } } From 6e5e326d0ed47b571c66a3939dbe77d9cb2a46e8 Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Tue, 2 Jan 2024 23:06:12 -0500 Subject: [PATCH 06/54] Ok it builds time to fillout the inference command. --- .../image_generation_payload.rs | 4 +- .../src/job/job_loop/process_single_job.rs | 2 +- .../src/job/job_types/image_generation/mod.rs | 2 + .../image_generation/process_single_ig_job.rs | 34 ++ .../job/job_types/image_generation/sd/mod.rs | 4 + .../image_generation/sd/model_downloaders.rs | 20 ++ .../image_generation/sd/process_job.rs | 17 + .../sd/sd_inference_command.rs | 320 ++++++++++++++++++ .../sd/stable_diffusion_dependencies.rs | 19 ++ .../inference_job/src/job/job_types/mod.rs | 2 +- .../job/job_types/vc/rvc_v2/process_job.rs | 2 + .../job_types/vc/so_vits_svc/process_job.rs | 2 + .../rerender_a_video/process_job.rs | 1 - .../src/job_specific_dependencies.rs | 4 +- .../image_gen/enqueue_image_generation.rs | 41 ++- .../inference_job/get_inference_job_status.rs | 5 +- 16 files changed, 457 insertions(+), 22 deletions(-) create mode 100644 crates/service/job/inference_job/src/job/job_types/image_generation/mod.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/image_generation/sd/model_downloaders.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/image_generation/sd/sd_inference_command.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs 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 c3385bdcb3..ee06b830d4 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 @@ -50,7 +50,7 @@ pub struct SDArgs { #[serde(rename = "se")] #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_seed: Option, + pub maybe_seed: Option, #[serde(rename = "mu")] #[serde(skip_serializing_if = "Option::is_none")] @@ -58,7 +58,7 @@ pub struct SDArgs { #[serde(rename = "lu")] #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_lora_upload_path: Option + pub maybe_lora_upload_path: Option, #[serde(rename = "t")] pub inference_type: Option diff --git a/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs b/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs index a9896d8a7b..5f2e9fd064 100644 --- a/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs +++ b/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs @@ -17,7 +17,7 @@ use crate::job::job_types::lipsync::process_single_lipsync_job::process_single_l use crate::job::job_types::tts::process_single_tts_job::process_single_tts_job; use crate::job::job_types::vc::process_single_vc_job::process_single_vc_job; use crate::job::job_types::videofilter::process_single_vf_job::process_single_vf_job; -use crate::job::job_types::imagegen::process_single_ig_job::process_single_ig_job; +use crate::job::job_types::image_generation::process_single_ig_job::process_single_ig_job; use crate::job_dependencies::JobDependencies; pub async fn process_single_job( diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/mod.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/mod.rs new file mode 100644 index 0000000000..fa1c0e36ec --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/mod.rs @@ -0,0 +1,2 @@ +pub mod process_single_ig_job; +pub mod sd; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs new file mode 100644 index 0000000000..e93d0cda62 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs @@ -0,0 +1,34 @@ +use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::image_generation::sd::process_job::process_job; +use crate::job_dependencies::JobDependencies; +use crate::job::job_types::image_generation::sd::process_job::SDProcessArgs; +use errors::anyhow; + +pub async fn process_single_ig_job(job_dependencies: &JobDependencies, job: &AvailableInferenceJob) -> Result { + let job = job; + match job.maybe_model_type { + Some(InferenceModelType::StableDiffusion) => { + dispatch_sd_job( + job_dependencies, + job + ).await + }, + Some(other_model_type) => { + Err(ProcessSingleJobError::Other(anyhow!("Wrong model type for SD: {:?}", other_model_type))) + } + None => { + Err(ProcessSingleJobError::Other(anyhow!("SD model type not set"))) + } + } +} + +pub async fn dispatch_sd_job(job_dependencies: &JobDependencies, job: &AvailableInferenceJob) -> Result { + let args = SDProcessArgs { + job_dependencies, + job, + }; + process_job(args).await +} \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs new file mode 100644 index 0000000000..ff0f816bda --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs @@ -0,0 +1,4 @@ +pub mod process_job; +pub mod model_downloaders; +pub mod sd_inference_command; +pub mod stable_diffusion_dependencies; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/model_downloaders.rs new file mode 100644 index 0000000000..38e480227c --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/model_downloaders.rs @@ -0,0 +1,20 @@ + +crate::impl_model_downloader!( + SDVAEEncoder, + "SD Variation Auto Encoder", + "SD_VARIATIONAL_AUTO_ENCODER_BUCKET_PATH", + "/animation_sadtalker/enhancer/parsing_parsenet.pth", + "SD_VARIATIONAL_AUTO_ENCODER_FILESYSTEM_PATH", + "/tmp/downloads/sadtalker/parsing_parsenet.pth" + ); + +pub struct SDDownloaders { + sd_vae_encoder: SDVAEEncoder +} +impl SDDownloaders { + pub fn build_all_from_env() -> Self { + Self { + sd_vae_encoder: SDVAEEncoder::from_env(), + } + } +} \ No newline at end of file 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 new file mode 100644 index 0000000000..b67082a33c --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs @@ -0,0 +1,17 @@ + +use std::time::Duration; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job_dependencies::JobDependencies; + +pub struct SDProcessArgs<'a> { + pub job_dependencies: &'a JobDependencies, + pub job: &'a AvailableInferenceJob, +} +pub async fn process_job(args: SDProcessArgs<'_>) -> Result { + Ok(JobSuccessResult { + maybe_result_entity: None, + inference_duration: Duration::from_secs(0), + }) +} \ No newline at end of file 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 new file mode 100644 index 0000000000..1d05017faf --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/sd_inference_command.rs @@ -0,0 +1,320 @@ + +use std::collections::HashSet; +use std::env; +use std::ffi::OsString; +use std::fs::File; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use anyhow::anyhow; +use log::info; +use once_cell::sync::Lazy; +use subprocess::{Popen, PopenConfig, Redirection}; + +use errors::AnyhowResult; +use filesys::path_to_string::path_to_string; +use subprocess_common::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions}; + +use crate::job::job_loop::command_exit_status::CommandExitStatus; + +// These environment vars are not copied over to the subprocess +// TODO/FIXME(bt, 2023-05-28): This is horrific security! +static IGNORED_ENVIRONMENT_VARS : Lazy> = Lazy::new(|| { + let env_var_names= [ + "MYSQL_URL", + "ACCESS_KEY", + "SECRET_KEY", + "NEWRELIC_API_KEY", + ]; + + env_var_names.iter() + .map(|value| value.to_string()) + .collect::>() +}); + +#[derive(Clone)] +pub struct StableDiffusionInferenceCommand { + /// Where the code lives + stable_diffusion_root_code_directory: PathBuf, + + /// A single executable script or a much larger bash command. + executable_or_command: ExecutableOrCommand, + + /// eg. `source python/bin/activate` + maybe_virtual_env_activation_command: Option, + + /// Optional default config file to use + maybe_default_config_path: Option, + + /// If this is run under Docker (eg. in development), these are the options. + maybe_docker_options: Option, + + /// If the execution should be ended after a certain point. + maybe_execution_timeout: Option, + + /// Inference arg. + /// --checkpoint_dir: optional location for checkpoints directory + pub alternate_checkpoint_dir: Option, +} + +#[derive(Clone)] +pub enum ExecutableOrCommand { + /// Eg. `inference.py` + Executable(PathBuf), + + /// Eg. `python3 inference.py` + Command(String), +} + +pub struct InferenceArgs<'s, P: AsRef> { + /// --source_image: path to the input image (or video) + pub input_image: P, + + /// --result_dir: path to directory work is performed + pub work_dir: P, + + /// --result_file: path to final file output + pub output_file: P, + + pub stderr_output_file: P, + + /// --enhancer: "gfpgan" + pub maybe_enhancer: Option<&'s str>, + + /// --preprocess: "crop", etc. + pub maybe_preprocess: Option<&'s str>, +} + +impl StableDiffusionInferenceCommand { + pub fn new>( + stable_diffusion_root_code_directory: P, + executable_or_command: ExecutableOrCommand, + maybe_virtual_env_activation_command: Option<&str>, + maybe_default_config_path: Option

, + maybe_docker_options: Option, + maybe_execution_timeout: Option, + alternate_checkpoint_dir: Option + ) -> Self { + Self { + stable_diffusion_root_code_directory: stable_diffusion_root_code_directory.as_ref().to_path_buf(), + executable_or_command, + maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()), + maybe_default_config_path: maybe_default_config_path.map(|p| p.as_ref().to_path_buf()), + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + } + } + + pub fn from_env() -> AnyhowResult { + let stable_diffusion_root_code_directory = easyenv::get_env_pathbuf_required( + "STABLE_DIFFUSION_INFERENCE_ROOT_DIRECTORY")?; + + let maybe_inference_command = easyenv::get_env_string_optional( + "STABLE_DIFFUSION_INFERENCE_COMMAND"); + + // Optional, eg. `./infer.py`. Typically we'll use the command form instead. + let maybe_inference_executable = easyenv::get_env_pathbuf_optional( + "STABLE_DIFFUSION_INFERENCE_EXECUTABLE"); + + let executable_or_command = match maybe_inference_command { + Some(command) => ExecutableOrCommand::Command(command), + None => match maybe_inference_executable { + Some(executable) => ExecutableOrCommand::Executable(executable), + None => return Err(anyhow!("neither command nor executable passed")), + }, + }; + + let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( + "STABLE_DIFFUSION_INFERENCE_MAYBE_VENV_COMMAND"); + + let maybe_default_config_path = easyenv::get_env_pathbuf_optional( + "STABLE_DIFFUSION_INFERENCE_MAYBE_DEFAULT_CONFIG_PATH"); + + let maybe_execution_timeout = + easyenv::get_env_duration_seconds_optional("STABLE_DIFFUSION_TIMEOUT_SECONDS"); + + let maybe_docker_options = easyenv::get_env_string_optional( + "STABLE_DIFFUSION_INFERENCE_MAYBE_DOCKER_IMAGE") + .map(|image_name| { + DockerOptions { + image_name, + maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()), + maybe_environment_variables: None, + maybe_gpu: Some(DockerGpu::All), + } + }); + + // Override for --checkpoint_dir at inference time + let alternate_checkpoint_dir = easyenv::get_env_pathbuf_optional( + "STABLE_DIFFUSION_ALTERNATE_CHECKPOINT_PATH"); + + Ok(Self { + stable_diffusion_root_code_directory, + executable_or_command, + maybe_virtual_env_activation_command, + maybe_default_config_path, + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + }) + } + + pub fn execute_inference>( + &self, + args: InferenceArgs

, + ) -> CommandExitStatus { + match self.do_execute_inference(args) { + Ok(exit_status) => exit_status, + Err(error) => CommandExitStatus::FailureWithReason { reason: format!("error: {:?}", error) }, + } + } + + fn do_execute_inference>( + &self, + args: InferenceArgs

, + ) -> AnyhowResult { + + let mut command = String::new(); + command.push_str(&format!("cd {}", path_to_string(&self.stable_diffusion_root_code_directory))); + + if let Some(venv_command) = self.maybe_virtual_env_activation_command.as_deref() { + command.push_str(" && "); + command.push_str(venv_command); + command.push_str(" "); + } + + command.push_str(" && "); + + match self.executable_or_command { + ExecutableOrCommand::Executable(ref executable) => { + command.push_str(&path_to_string(executable)); + command.push_str(" "); + } + ExecutableOrCommand::Command(ref cmd) => { + command.push_str(cmd); + command.push_str(" "); + } + } + + // ===== Begin Python Args ===== + + // command.push_str(" --driven_audio "); + // command.push_str(&path_to_string(args.input_audio)); + + // command.push_str(" --source_image "); + // command.push_str(&path_to_string(args.input_image)); + + // command.push_str(" --result_dir "); + // command.push_str(&path_to_string(args.work_dir)); + + // command.push_str(" --result_file "); + // command.push_str(&path_to_string(args.output_file)); + + // if let Some(preprocess) = args.maybe_preprocess.as_deref() { + // command.push_str(" --preprocess "); + // command.push_str(preprocess); + // command.push_str(" "); + // } + // if let Some(enhancer) = args.maybe_enhancer.as_deref() { + // command.push_str(" --enhancer "); + // command.push_str(enhancer); + // command.push_str(" "); + // } + + // if args.make_still { + // command.push_str(" --still "); + // } + + if let Some(dir) = self.alternate_checkpoint_dir.as_ref() { + command.push_str(" --checkpoint_dir "); + command.push_str(&path_to_string(dir)); + } + + // ===== End Python Args ===== + + if let Some(docker_options) = self.maybe_docker_options.as_ref() { + command = docker_options.to_command_string(&command); + } + + info!("Command: {:?}", command); + + let command_parts = [ + "bash", + "-c", + &command + ]; + + let mut env_vars = Vec::new(); + + // Copy all environment variables from the parent process. + // This is necessary to send all the kubernetes settings for Nvidia / CUDA. + for (env_key, env_value) in env::vars() { + if IGNORED_ENVIRONMENT_VARS.contains(&env_key) { + continue; + } + env_vars.push(( + OsString::from(env_key), + OsString::from(env_value), + )); + } + + let mut config = PopenConfig::default(); + + info!("stderr will be written to file: {:?}", args.stderr_output_file.as_ref()); + + let stderr_file = File::create(&args.stderr_output_file)?; + config.stderr = Redirection::File(stderr_file); + + if !env_vars.is_empty() { + config.env = Some(env_vars); + } + + let mut p = Popen::create(&command_parts, config)?; + + info!("Subprocess PID: {:?}", p.pid()); + + match self.maybe_execution_timeout { + None => { + let exit_status = p.wait()?; + info!("Subprocess exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + Some(timeout) => { + info!("Executing with timeout: {:?}", &timeout); + let exit_status = p.wait_timeout(timeout.clone())?; + + match exit_status { + None => { + // NB: If the program didn't successfully terminate, kill it. + info!("Subprocess didn't end after timeout: {:?}; terminating...", &timeout); + let _r = p.terminate()?; + Ok(CommandExitStatus::Timeout) + } + Some(exit_status) => { + info!("Subprocess timed wait exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + } + } + } + } +} + + + + + +// pub struct StableDiffusionInferenceCommand { +// pub command: String, +// pub args: Vec +// } + +// impl StableDiffusionInferenceCommand { +// pub from_env() -> Self { +// Self { +// command: "python".to_string(), +// args: vec![], +// } +// } +// } \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs new file mode 100644 index 0000000000..4668917834 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs @@ -0,0 +1,19 @@ +use anyhow::{Ok}; +use errors::AnyhowResult; + +use crate::job::job_types::image_generation::sd::model_downloaders::SDDownloaders; +use crate::job::job_types::image_generation::sd::sd_inference_command::StableDiffusionInferenceCommand; + +pub struct StableDiffusionDependencies { + pub downloaders: SDDownloaders, + pub inference_command: StableDiffusionInferenceCommand, +} + +impl StableDiffusionDependencies { + pub fn setup() -> AnyhowResult { + Ok(Self { + downloaders: SDDownloaders::build_all_from_env(), + inference_command: StableDiffusionInferenceCommand::from_env()?, + }) + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/mod.rs b/crates/service/job/inference_job/src/job/job_types/mod.rs index d73c394afb..e67f85ba66 100644 --- a/crates/service/job/inference_job/src/job/job_types/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/mod.rs @@ -2,4 +2,4 @@ pub mod lipsync; pub mod tts; pub mod vc; pub(crate) mod videofilter; -pub mod imagegen; \ No newline at end of file +pub mod image_generation; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs b/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs index 0d2f22e34b..cfce433e25 100644 --- a/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs @@ -260,6 +260,7 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::Vc { override_f0_method, .. } => *override_f0_method, PolymorphicInferenceArgs::Rr(_) => None, + PolymorphicInferenceArgs::Ig(_) => None, }) .flatten(); @@ -269,6 +270,7 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::Vc { transpose, .. } => *transpose, PolymorphicInferenceArgs::Rr(_) => None, + PolymorphicInferenceArgs::Ig(_) => None, }) .flatten(); diff --git a/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs b/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs index a46e8c89f5..239063ef6e 100644 --- a/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs @@ -180,6 +180,7 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::Vc { auto_predict_f0, .. } => *auto_predict_f0, PolymorphicInferenceArgs::Rr(_) => None, + PolymorphicInferenceArgs::Ig(_) => None, }) .flatten() .unwrap_or(false); @@ -190,6 +191,7 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::Vc { transpose, .. } => *transpose, PolymorphicInferenceArgs::Rr(_) => None, + PolymorphicInferenceArgs::Ig(_) => None, }) .flatten(); diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs index 7de58ac925..f6e556dddf 100644 --- a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs @@ -35,7 +35,6 @@ const BUCKET_FILE_EXTENSION: &str = ".mp4"; pub struct RerenderProcessJobArgs<'a> { pub job_dependencies: &'a JobDependencies, pub job: &'a AvailableInferenceJob, - // pub media_file: MediaFile, } pub async fn process_job(args: RerenderProcessJobArgs<'_>) -> Result { diff --git a/crates/service/job/inference_job/src/job_specific_dependencies.rs b/crates/service/job/inference_job/src/job_specific_dependencies.rs index f70052dc2c..ec1bccd7db 100644 --- a/crates/service/job/inference_job/src/job_specific_dependencies.rs +++ b/crates/service/job/inference_job/src/job_specific_dependencies.rs @@ -9,7 +9,7 @@ use crate::job::job_types::tts::vits::vits_dependencies::VitsDependencies; use crate::job::job_types::vc::rvc_v2::rvc_v2_dependencies::RvcV2Dependencies; use crate::job::job_types::vc::so_vits_svc::svc_dependencies::SvcDependencies; use crate::job::job_types::videofilter::rerender_a_video::rerender_dependencies::RerenderDependencies; -use crate::job::job_types::imagegen::sd::stable_diffusion_dependencies::StableDiffusionDependencies; +use crate::job::job_types::image_generation::sd::stable_diffusion_dependencies::StableDiffusionDependencies; use crate::util::scoped_execution::ScopedExecution; @@ -71,7 +71,7 @@ impl JobSpecificDependencies { maybe_rerender_dependencies = Some(RerenderDependencies::setup()?); } - if scoped_execution.can_run_job(job_model_type::StableDiffusion) { + if scoped_execution.can_run_job(InferenceModelType::StableDiffusion) { print_with_space("Setting Stable Diffusion dependencies..."); maybe_stable_diffusion_dependencies = Some(StableDiffusionDependencies::setup()?); } 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 417887fd1a..ae193d0c86 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 @@ -8,7 +8,6 @@ use std::sync::Arc; use actix_web::{HttpRequest, HttpResponse, web}; use actix_web::error::ResponseError; use actix_web::http::StatusCode; -use elasticsearch::http::request; use log::warn; use serde::Deserialize; use serde::Serialize; @@ -37,11 +36,10 @@ use tokens::tokens::generic_inference_jobs::InferenceJobToken; use tokens::tokens::users::UserToken; use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session; -use crate::http_server::endpoints::media_files::drain_multipart_request::MediaSource; -use crate::http_server::endpoints::weights::list_weights_by_user::Weight; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::server_state::ServerState; +use std::fmt::{Display, Formatter}; /// Debug requests can get routed to special "debug-only" workers, which can /// be used to trial new code, run debugging, etc. const DEBUG_HEADER_NAME: &str = "enable-debug-mode"; @@ -50,14 +48,23 @@ 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(Deserialize,ToSchema)] +#[derive(Debug,Deserialize, Clone, Copy, Eq, PartialEq)] pub enum TypeOfInference { #[serde(rename = "lora")] UploadLoRA, #[serde(rename = "checkpoint")] CheckPoint, - #[serde(rename = "standard_inference")] - StandardInference, + #[serde(rename = "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::Inference => write!(f, "inference"), + } + } } #[derive(Deserialize,ToSchema)] @@ -71,7 +78,7 @@ pub struct EnqueueImageGenRequest { maybe_prompt: Option, maybe_a_prompt: Option, maybe_n_prompt: Option, - maybe_seed: Option, + maybe_seed: Option, maybe_upload_path: Option, maybe_lora_upload_path: Option, } @@ -143,7 +150,7 @@ pub async fn enqueue_image_generation_request( ) -> Result { let mut maybe_user_token: Option = None; - let mut 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); EnqueueImageGenRequestError::ServerError @@ -206,7 +213,7 @@ pub async fn enqueue_image_generation_request( // ==================== INFERENCE ARGS ==================== // - if request.type_of_inference == EnqueueImageGenType::Inference { + 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())); } @@ -215,11 +222,11 @@ pub async fn enqueue_image_generation_request( return Err(EnqueueImageGenRequestError::BadInput("No prompt provided".to_string())); } - } else if request.type_of_inference == EnqueueImageGenRequest::UploadLoRA { + } 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 == EnqueueImageGenRequest::CheckPoint { + } 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())); } @@ -239,6 +246,12 @@ pub async fn enqueue_image_generation_request( // if we are uploading a lora model. let lora_upload_path = request.maybe_lora_upload_path.clone().unwrap_or_default(); + let mut seed = -1; + + if let Some(s) = request.maybe_seed { + seed = s; + } + let inference_args = SDArgs { maybe_video_source: Some(video_token), maybe_image_source: Some(image_source_token), @@ -247,10 +260,10 @@ pub async fn enqueue_image_generation_request( maybe_prompt: Some(request.maybe_prompt.clone().unwrap_or_default()), maybe_a_prompt: Some(a_prompt), maybe_n_prompt: Some(n_prompt), - maybe_seed: Some(request.maybe_seed.clone().unwrap_or_default()), + maybe_seed: Some(seed), maybe_upload_path: Some(upload_path), maybe_lora_upload_path: Some(lora_upload_path), - inference_type: Some(request.type_of_inference.clone()), + inference_type: Some(request.type_of_inference.to_string()), }; // create the inference args here @@ -264,7 +277,7 @@ pub async fn enqueue_image_generation_request( maybe_model_token: None, // NB: Model is static during inference maybe_input_source_token: None, maybe_input_source_token_type: None, - maybe_raw_inference_text: Some(&request.text), + maybe_raw_inference_text: None, maybe_max_duration_seconds: None, maybe_inference_args: Some(GenericInferenceArgs { inference_category: Some(InferenceCategoryAbbreviated::ImageGeneration), diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs index 071b3795a6..3e058ab60b 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs @@ -302,8 +302,11 @@ fn record_to_payload( } } } + InferenceCategory::ImageGeneration => { + "".to_string() + } }; - + ResultDetailsResponse { entity_type: result_details.entity_type, entity_token: result_details.entity_token, From 74db1683b16abae9d816fc1df5ad9883b263675e Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Wed, 3 Jan 2024 11:01:37 -0500 Subject: [PATCH 07/54] unwraped the args --- .../generic_inference_args.rs | 8 ++--- .../image_generation_payload.rs | 6 ++-- .../image_generation/sd/process_job.rs | 29 ++++++++++++++++++- .../image_gen/enqueue_image_generation.rs | 6 ++-- 4 files changed, 38 insertions(+), 11 deletions(-) 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 a4fb7f1df2..e9f6b37b5b 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 @@ -4,7 +4,7 @@ use errors::AnyhowResult; use crate::payloads::generic_inference_args::lipsync_payload::LipsyncArgs; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; use crate::payloads::generic_inference_args::videofilter_payload::RerenderArgs; -use crate::payloads::generic_inference_args::image_generation_payload::SDArgs; +use crate::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; /// Used to encode extra state for the `generic_inference_jobs` table. /// This should act somewhat like a serialized protobuf stored inside a record. @@ -98,7 +98,7 @@ pub enum PolymorphicInferenceArgs { Rr(RerenderArgs), /// Image generation. (Short name to save space when serializing.) - Ig(SDArgs) + Ig(StableDiffusionArgs) } @@ -142,7 +142,7 @@ mod tests { use crate::payloads::generic_inference_args::generic_inference_args::{FundamentalFrequencyMethodForJob, GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; use crate::payloads::generic_inference_args::lipsync_payload::{LipsyncAnimationAudioSource, LipsyncAnimationImageSource, LipsyncArgs}; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; - use crate::payloads::generic_inference_args::image_generation_payload::SDArgs; + use crate::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; use tokens::tokens::{model_weights::ModelWeightToken, media_files::MediaFileToken}; #[test] @@ -207,7 +207,7 @@ use crate::payloads::generic_inference_args::generic_inference_args::{Fundamenta let args = GenericInferenceArgs { inference_category: Some(InferenceCategoryAbbreviated::ImageGeneration), - args: Some(PolymorphicInferenceArgs::Ig(SDArgs { + args: Some(PolymorphicInferenceArgs::Ig(StableDiffusionArgs { maybe_video_source: Some(video_media_token), maybe_image_source: Some(image_media_token), maybe_sd_model_token: Some(sd_model_token), 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 ee06b830d4..7c0ea03136 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 @@ -6,7 +6,7 @@ use tokens::tokens::{model_weights::ModelWeightToken, media_files::MediaFileToke /// For image to image we probably want images #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] -pub enum SDSource { +pub enum StableDiffusionSource { // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. /// Media File Token (media_files table) /// Serde cannot yet rename enum variants. @@ -19,7 +19,7 @@ pub enum SDSource { } #[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)] -pub struct SDArgs { +pub struct StableDiffusionArgs { #[serde(rename = "vs")] #[serde(skip_serializing_if = "Option::is_none")] pub maybe_video_source: Option, @@ -61,7 +61,7 @@ pub struct SDArgs { pub maybe_lora_upload_path: Option, #[serde(rename = "t")] - pub inference_type: Option + pub inference_type: String, } 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 b67082a33c..9dba7756c3 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,15 +1,42 @@ use std::time::Duration; +use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; use crate::job::job_loop::job_success_result::JobSuccessResult; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::job_dependencies::JobDependencies; - +use anyhow::anyhow; +use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; pub struct SDProcessArgs<'a> { pub job_dependencies: &'a JobDependencies, pub job: &'a AvailableInferenceJob, } + pub async fn process_job(args: SDProcessArgs<'_>) -> Result { + + // validation + let inference_args = args.job.maybe_inference_args + .as_ref() + .map(|args| args.args.as_ref()) + .flatten(); + + let args = match inference_args { + Some(args) => args, + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no inference args for job!"))); + } + }; + + let args = match args { + PolymorphicInferenceArgs::Ig(args) => args, + _ => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!"))); + } + }; + + + let stable_diffusion_args: StableDiffusionArgs = StableDiffusionArgs::from(args.clone()); + Ok(JobSuccessResult { maybe_result_entity: None, inference_duration: Duration::from_secs(0), 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 ae193d0c86..57fdf2a705 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 @@ -25,7 +25,7 @@ use mysql_queries::payloads::generic_inference_args::generic_inference_args::{ PolymorphicInferenceArgs, }; -use mysql_queries::payloads::generic_inference_args::image_generation_payload::SDArgs; +use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{ insert_generic_inference_job, InsertGenericInferenceArgs, @@ -71,7 +71,6 @@ impl Display for TypeOfInference { pub struct EnqueueImageGenRequest { uuid_idempotency_token: String, type_of_inference: TypeOfInference, // upload loRA / check point / standard inference - maybe_video_source: Option, maybe_image_source: Option, maybe_sd_model_token: Option, maybe_lora_model_token: Option, @@ -81,6 +80,7 @@ pub struct EnqueueImageGenRequest { maybe_seed: Option, maybe_upload_path: Option, maybe_lora_upload_path: Option, + maybe_cfg_scale: Option } #[derive(Serialize,ToSchema)] @@ -252,7 +252,7 @@ pub async fn enqueue_image_generation_request( seed = s; } - let inference_args = SDArgs { + let inference_args = StableDiffusionArgs { maybe_video_source: Some(video_token), maybe_image_source: Some(image_source_token), maybe_sd_model_token: Some(sd_weight_token), From a2cad0b289502b3d62424b9e7f5351531d10869c Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Wed, 3 Jan 2024 14:57:49 -0500 Subject: [PATCH 08/54] getting the args for inference working. --- .../remote_cloud_file_manager.rs | 2 +- .../generic_inference_args.rs | 11 +- .../image_generation/process_single_ig_job.rs | 4 +- .../image_generation/sd/model_downloaders.rs | 4 +- .../image_generation/sd/process_job.rs | 107 ++++++++++++++++-- .../sd/sd_inference_command.rs | 10 +- .../sd/stable_diffusion_dependencies.rs | 6 +- 7 files changed, 117 insertions(+), 27 deletions(-) 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 513c0dfe98..5cfd9ed402 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 @@ -11,7 +11,7 @@ use crate::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDe use super::file_descriptor::FileDescriptor; use super::file_meta_data::FileMetaData; -struct RemoteCloudFileClient { +pub struct RemoteCloudFileClient { bucket_orchestration_client: Box } 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 e9f6b37b5b..cc88ebd3c9 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 @@ -101,7 +101,6 @@ pub enum PolymorphicInferenceArgs { Ig(StableDiffusionArgs) } - impl GenericInferenceArgs { pub fn from_json(json: &str) -> AnyhowResult { @@ -137,9 +136,7 @@ impl InferenceCategoryAbbreviated { #[cfg(test)] mod tests { - use enums::by_table::generic_inference_jobs::inference_category; - -use crate::payloads::generic_inference_args::generic_inference_args::{FundamentalFrequencyMethodForJob, GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; + use crate::payloads::generic_inference_args::generic_inference_args::{FundamentalFrequencyMethodForJob, GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; use crate::payloads::generic_inference_args::lipsync_payload::{LipsyncAnimationAudioSource, LipsyncAnimationImageSource, LipsyncArgs}; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; use crate::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; @@ -204,7 +201,8 @@ use crate::payloads::generic_inference_args::generic_inference_args::{Fundamenta let seed = 1; let upload_path = "upload_path".to_string(); let lora_upload_path = "lora_upload_path".to_string(); - + let checkpoint = "checkpoint".to_string(); + let args = GenericInferenceArgs { inference_category: Some(InferenceCategoryAbbreviated::ImageGeneration), args: Some(PolymorphicInferenceArgs::Ig(StableDiffusionArgs { @@ -217,7 +215,8 @@ use crate::payloads::generic_inference_args::generic_inference_args::{Fundamenta maybe_n_prompt: Some(n_prompt), maybe_seed: Some(seed), maybe_upload_path: Some(upload_path), - maybe_lora_upload_path: Some(lora_upload_path) + maybe_lora_upload_path: Some(lora_upload_path), + inference_type: Some(checkpoint) })), }; diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs index e93d0cda62..959642865f 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs @@ -4,7 +4,7 @@ use crate::job::job_loop::job_success_result::JobSuccessResult; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::job::job_types::image_generation::sd::process_job::process_job; use crate::job_dependencies::JobDependencies; -use crate::job::job_types::image_generation::sd::process_job::SDProcessArgs; +use crate::job::job_types::image_generation::sd::process_job::StableDiffusionProcessArgs; use errors::anyhow; pub async fn process_single_ig_job(job_dependencies: &JobDependencies, job: &AvailableInferenceJob) -> Result { @@ -26,7 +26,7 @@ pub async fn process_single_ig_job(job_dependencies: &JobDependencies, job: &Ava } pub async fn dispatch_sd_job(job_dependencies: &JobDependencies, job: &AvailableInferenceJob) -> Result { - let args = SDProcessArgs { + let args = StableDiffusionProcessArgs { job_dependencies, job, }; diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/model_downloaders.rs index 38e480227c..0dd95e7cfb 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/model_downloaders.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/model_downloaders.rs @@ -8,10 +8,10 @@ crate::impl_model_downloader!( "/tmp/downloads/sadtalker/parsing_parsenet.pth" ); -pub struct SDDownloaders { +pub struct StableDiffusionDownloaders { sd_vae_encoder: SDVAEEncoder } -impl SDDownloaders { +impl StableDiffusionDownloaders { pub fn build_all_from_env() -> Self { Self { sd_vae_encoder: SDVAEEncoder::from_env(), 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 9dba7756c3..ae70b22e88 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 @@ -7,18 +7,50 @@ use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::job_dependencies::JobDependencies; use anyhow::anyhow; use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; -pub struct SDProcessArgs<'a> { +use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; + +pub struct StableDiffusionProcessArgs<'a> { pub job_dependencies: &'a JobDependencies, pub job: &'a AvailableInferenceJob, } -pub async fn process_job(args: SDProcessArgs<'_>) -> Result { +pub async fn prepare_inputs() { + // download the required files and weights using the tokens +} - // validation +pub async fn validate_inputs(args: StableDiffusionProcessArgs<'_>) -> Result<(), ProcessSingleJobError> { + // validate the inputs + let model_dependencies = args.job_dependencies + .job + .job_specific_dependencies + .maybe_stable_diffusion_dependencies + .as_ref() + .ok_or_else(|| ProcessSingleJobError::JobSystemMisconfiguration(Some("missing stable diffusion dependencies".to_string())))?; + + let inference_args = args.job.maybe_inference_args - .as_ref() - .map(|args| args.args.as_ref()) - .flatten(); + .as_ref() + .map(|args| args.args.as_ref()) + .flatten(); + + let args = match args { + PolymorphicInferenceArgs::Ig(args) => args, + _ => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!"))); + } + }; + + let stable_diffusion_args: StableDiffusionArgs = StableDiffusionArgs::from(args.clone()); + + if stable_diffusion_args.inference_type == "checkpoint" { + + } else if stable_diffusion_args.inference_type == "lora" { + + } else if stable_diffusion_args.inference_type == "inference" { + + } else { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inference type for job!"))); + } let args = match inference_args { Some(args) => args, @@ -33,9 +65,68 @@ pub async fn process_job(args: SDProcessArgs<'_>) -> Result) -> Result { + let job = args.job; + let deps = args.job_dependencies; + + let inference_args = args.job.maybe_inference_args + .as_ref() + .map(|args| args.args.as_ref()) + .flatten(); + + let polymorphic_args = match inference_args { + Some(args) => args, + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no inference args for job!"))); + } + }; + + let polymorphic_args: StableDiffusionArgs = match polymorphic_args { + PolymorphicInferenceArgs::Ig(args) => args, + _ => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!"))); + } + }; + + // let mut job_progress_reporter = args.job_dependencies + // .clients + // .job_progress_reporter + // .new_generic_inference(job.inference_job_token.as_str()) + // .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + + // ==================== UNPACK + VALIDATE INFERENCE ARGS ==================== // + let job_args = validate_inputs(args).await?; + + // ==================== TEMP DIR ==================== // + let work_temp_dir = format!("temp_rerender_inference_{}", job.id.0); + + // 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))?; + + // validation + let stable_diffusion_args: StableDiffusionArgs = StableDiffusionArgs::from(polymorphic_args.clone()); + + if stable_diffusion_args.inference_type == "checkpoint" { + // run inference with checkpoint and upload + } else if stable_diffusion_args.inference_type == "lora" { + // run inference with the lora + random checkpoint + } else if stable_diffusion_args.inference_type == "inference" { + // run inference with or without lora + } else { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inference type for job!"))); + } + + // run inference + + // insert record into the db with the inference job token complete. Ok(JobSuccessResult { maybe_result_entity: None, 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 1d05017faf..94aa981249 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 @@ -131,7 +131,7 @@ impl StableDiffusionInferenceCommand { let maybe_default_config_path = easyenv::get_env_pathbuf_optional( "STABLE_DIFFUSION_INFERENCE_MAYBE_DEFAULT_CONFIG_PATH"); - let maybe_execution_timeout = + let maybe_execution_timeout: Option = easyenv::get_env_duration_seconds_optional("STABLE_DIFFUSION_TIMEOUT_SECONDS"); let maybe_docker_options = easyenv::get_env_string_optional( @@ -226,10 +226,10 @@ impl StableDiffusionInferenceCommand { // command.push_str(" --still "); // } - if let Some(dir) = self.alternate_checkpoint_dir.as_ref() { - command.push_str(" --checkpoint_dir "); - command.push_str(&path_to_string(dir)); - } + // if let Some(dir) = self.alternate_checkpoint_dir.as_ref() { + // command.push_str(" --checkpoint_dir "); + // command.push_str(&path_to_string(dir)); + // } // ===== End Python Args ===== diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs index 4668917834..bf14dfbdc8 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs @@ -1,18 +1,18 @@ use anyhow::{Ok}; use errors::AnyhowResult; -use crate::job::job_types::image_generation::sd::model_downloaders::SDDownloaders; +use crate::job::job_types::image_generation::sd::model_downloaders::StableDiffusionDownloaders; use crate::job::job_types::image_generation::sd::sd_inference_command::StableDiffusionInferenceCommand; pub struct StableDiffusionDependencies { - pub downloaders: SDDownloaders, + pub downloaders: StableDiffusionDownloaders, pub inference_command: StableDiffusionInferenceCommand, } impl StableDiffusionDependencies { pub fn setup() -> AnyhowResult { Ok(Self { - downloaders: SDDownloaders::build_all_from_env(), + downloaders: StableDiffusionDownloaders::build_all_from_env(), inference_command: StableDiffusionInferenceCommand::from_env()?, }) } From dbe989c5a12ea790f5af4f407170f8019ea9aea9 Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Wed, 3 Jan 2024 22:32:25 -0500 Subject: [PATCH 09/54] rebuilds on the routes in storyteller web. --- .../image_generation_payload.rs | 4 - .../job/job_types/image_generation/sd/mod.rs | 3 +- .../image_generation/sd/process_job.rs | 127 ++++++------------ .../sd/sd_inference_command.rs | 114 +++++++++------- .../image_generation/sd/validate_inputs.rs | 42 ++++++ .../image_gen/enqueue_image_generation.rs | 7 +- 6 files changed, 152 insertions(+), 145 deletions(-) create mode 100644 crates/service/job/inference_job/src/job/job_types/image_generation/sd/validate_inputs.rs 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 7c0ea03136..9a94b3a7e9 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 @@ -20,10 +20,6 @@ pub enum StableDiffusionSource { #[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)] pub struct StableDiffusionArgs { - #[serde(rename = "vs")] - #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_video_source: Option, - #[serde(rename = "is")] #[serde(skip_serializing_if = "Option::is_none")] pub maybe_image_source: Option, diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs index ff0f816bda..17702a9f82 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs @@ -1,4 +1,5 @@ pub mod process_job; pub mod model_downloaders; pub mod sd_inference_command; -pub mod stable_diffusion_dependencies; \ No newline at end of file +pub mod stable_diffusion_dependencies; +pub mod validate_inputs; \ No newline at end of file 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 ae70b22e88..4a2314eec6 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 @@ -8,125 +8,76 @@ use crate::job_dependencies::JobDependencies; use anyhow::anyhow; use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; +use crate::job::job_types::image_generation::sd::validate_inputs::validate_inputs; pub struct StableDiffusionProcessArgs<'a> { pub job_dependencies: &'a JobDependencies, pub job: &'a AvailableInferenceJob, } -pub async fn prepare_inputs() { - // download the required files and weights using the tokens -} +// fill out sd diffusion dependencies +// get specific weight +// download weight / lora +// run inference +// upload inference result +// upload model checkpoint or loRA +// create record in db + // if stable_diffusion_args.inference_type == "checkpoint" { + // // run inference with checkpoint and upload + // } else if stable_diffusion_args.inference_type == "lora" { + // // run inference with the lora + random checkpoint + // } else if stable_diffusion_args.inference_type == "inference" { + // // run inference with or without lora + // } else { + // return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inference type for job!"))); + // } + // run inference + // insert record into the db with the inference job token complete. -pub async fn validate_inputs(args: StableDiffusionProcessArgs<'_>) -> Result<(), ProcessSingleJobError> { - // validate the inputs - let model_dependencies = args.job_dependencies - .job - .job_specific_dependencies - .maybe_stable_diffusion_dependencies - .as_ref() - .ok_or_else(|| ProcessSingleJobError::JobSystemMisconfiguration(Some("missing stable diffusion dependencies".to_string())))?; - - +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()) .flatten(); - - let args = match args { - PolymorphicInferenceArgs::Ig(args) => args, - _ => { - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!"))); - } - }; - - let stable_diffusion_args: StableDiffusionArgs = StableDiffusionArgs::from(args.clone()); - - if stable_diffusion_args.inference_type == "checkpoint" { - - } else if stable_diffusion_args.inference_type == "lora" { - - } else if stable_diffusion_args.inference_type == "inference" { - - } else { - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inference type for job!"))); - } - - let args = match inference_args { - Some(args) => args, - None => { - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no inference args for job!"))); - } - }; - - let args = match args { - PolymorphicInferenceArgs::Ig(args) => args, - _ => { - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!"))); - } - }; - - Ok(()) -} - -pub async fn process_job(args: StableDiffusionProcessArgs<'_>) -> Result { - let job = args.job; - let deps = args.job_dependencies; - - let inference_args = args.job.maybe_inference_args - .as_ref() - .map(|args| args.args.as_ref()) - .flatten(); - let polymorphic_args = match inference_args { Some(args) => args, None => { return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no inference args for job!"))); } }; - - let polymorphic_args: StableDiffusionArgs = match polymorphic_args { + let sd_args = match polymorphic_args { PolymorphicInferenceArgs::Ig(args) => args, _ => { return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!"))); } }; - - // let mut job_progress_reporter = args.job_dependencies - // .clients - // .job_progress_reporter - // .new_generic_inference(job.inference_job_token.as_str()) - // .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + let stable_diffusion_args: StableDiffusionArgs = StableDiffusionArgs::from(sd_args.clone()); + Ok(stable_diffusion_args) +} - // ==================== UNPACK + VALIDATE INFERENCE ARGS ==================== // - let job_args = validate_inputs(args).await?; - - // ==================== TEMP DIR ==================== // - let work_temp_dir = format!("temp_rerender_inference_{}", job.id.0); +pub async fn process_job(args: StableDiffusionProcessArgs<'_>) -> Result { + let job = args.job; + let deps = args.job_dependencies; - // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem. + let sd_args = sd_args_from_job(&args).await?; + + let mut job_progress_reporter = args.job_dependencies + .clients + .job_progress_reporter + .new_generic_inference(job.inference_job_token.as_str()) + .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + + //==================== TEMP DIR ==================== // + let work_temp_dir = format!("temp_stable_diffusion_inference_{}", job.id.0); + + //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))?; - // validation - let stable_diffusion_args: StableDiffusionArgs = StableDiffusionArgs::from(polymorphic_args.clone()); - if stable_diffusion_args.inference_type == "checkpoint" { - // run inference with checkpoint and upload - } else if stable_diffusion_args.inference_type == "lora" { - // run inference with the lora + random checkpoint - } else if stable_diffusion_args.inference_type == "inference" { - // run inference with or without lora - } else { - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inference type for job!"))); - } - - // run inference - - // insert record into the db with the inference job token complete. Ok(JobSuccessResult { maybe_result_entity: None, 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 94aa981249..2c5b37c62d 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 @@ -66,23 +66,27 @@ pub enum ExecutableOrCommand { Command(String), } -pub struct InferenceArgs<'s, P: AsRef> { +pub struct InferenceArgs> { /// --source_image: path to the input image (or video) - pub input_image: P, - + pub input_image: P, // Not availbile rn /// --result_dir: path to directory work is performed pub work_dir: P, - /// --result_file: path to final file output pub output_file: P, - pub stderr_output_file: P, - - /// --enhancer: "gfpgan" - pub maybe_enhancer: Option<&'s str>, - - /// --preprocess: "crop", etc. - pub maybe_preprocess: Option<&'s str>, + pub prompt: P, + pub negative_prompt:P, + pub number_of_samples:i32, + pub samplers:String, + pub width:i32, + pub height:i32, + pub cfg_scale:i32, + pub seed:i32, + pub lora_path:P, + pub check_point:P, + pub vae:P, + pub batch_size:i32, + pub batch_count:i32, } impl StableDiffusionInferenceCommand { @@ -198,44 +202,56 @@ impl StableDiffusionInferenceCommand { } // ===== Begin Python Args ===== - - // command.push_str(" --driven_audio "); - // command.push_str(&path_to_string(args.input_audio)); - - // command.push_str(" --source_image "); - // command.push_str(&path_to_string(args.input_image)); - - // command.push_str(" --result_dir "); - // command.push_str(&path_to_string(args.work_dir)); - - // command.push_str(" --result_file "); - // command.push_str(&path_to_string(args.output_file)); - - // if let Some(preprocess) = args.maybe_preprocess.as_deref() { - // command.push_str(" --preprocess "); - // command.push_str(preprocess); - // command.push_str(" "); - // } - // if let Some(enhancer) = args.maybe_enhancer.as_deref() { - // command.push_str(" --enhancer "); - // command.push_str(enhancer); - // command.push_str(" "); - // } - - // if args.make_still { - // command.push_str(" --still "); - // } - - // if let Some(dir) = self.alternate_checkpoint_dir.as_ref() { - // command.push_str(" --checkpoint_dir "); - // command.push_str(&path_to_string(dir)); - // } - - // ===== End Python Args ===== - - if let Some(docker_options) = self.maybe_docker_options.as_ref() { - command = docker_options.to_command_string(&command); - } + command.push_str("--prompt "); + command.push_str(&args.prompt.as_ref().to_string_lossy()); + + command.push_str(" --result_dir "); + command.push_str(&args.work_dir.as_ref().to_string_lossy()); + + command.push_str(" --result_file "); + command.push_str(&args.output_file.as_ref().to_string_lossy()); + + command.push_str(" --stderr_output_file "); + command.push_str(&args.stderr_output_file.as_ref().to_string_lossy()); + + command.push_str(" --prompt "); + command.push_str(&args.prompt.as_ref().to_string_lossy()); + + command.push_str(" --negative_prompt "); + command.push_str(&args.negative_prompt.as_ref().to_string_lossy()); + + command.push_str(" --number_of_samples "); + command.push_str(&args.number_of_samples.to_string()); + + command.push_str(" --samplers "); + command.push_str(&args.samplers); + + command.push_str(" --width "); + command.push_str(&args.width.to_string()); + + command.push_str(" --height "); + command.push_str(&args.height.to_string()); + + command.push_str(" --cfg_scale "); + command.push_str(&args.cfg_scale.to_string()); + + command.push_str(" --seed "); + command.push_str(&args.seed.to_string()); + + command.push_str(" --lora_path "); + command.push_str(&args.lora_path.as_ref().to_string_lossy()); + + command.push_str(" --check_point "); + command.push_str(&args.check_point.as_ref().to_string_lossy()); + + command.push_str(" --vae "); + command.push_str(&args.vae.as_ref().to_string_lossy()); + + command.push_str(" --batch_size "); + command.push_str(&args.batch_size.to_string()); + + command.push_str(" --batch_count "); + command.push_str(&args.batch_count.to_string()); info!("Command: {:?}", command); diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/validate_inputs.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/validate_inputs.rs new file mode 100644 index 0000000000..129ee19461 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/validate_inputs.rs @@ -0,0 +1,42 @@ +use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use anyhow::anyhow; +use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; +use crate::job::job_types::image_generation::sd::process_job::StableDiffusionProcessArgs; + + +pub async fn validate_inputs(args: StableDiffusionProcessArgs<'_>) -> Result<(), ProcessSingleJobError> { + let inference_args = args.job.maybe_inference_args + .as_ref() + .map(|args| args.args.as_ref()) + .flatten(); + + let polymorphic_args = match inference_args { + Some(args) => args, + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no inference args for job!"))); + } + }; + + let sd_args = match polymorphic_args { + PolymorphicInferenceArgs::Ig(args) => args, + _ => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!"))); + } + }; + + let stable_diffusion_args: StableDiffusionArgs = StableDiffusionArgs::from(sd_args.clone()); + if stable_diffusion_args.inference_type == "checkpoint" { + + } else if stable_diffusion_args.inference_type == "lora" { + + } else if stable_diffusion_args.inference_type == "inference" { + + } else { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inference type for job!"))); + } + + + + Ok(()) +} 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 57fdf2a705..6b03b384fd 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 @@ -232,7 +232,7 @@ pub async fn enqueue_image_generation_request( } } - let video_token = MediaFileToken(request.maybe_video_source.clone().unwrap_or_default()); + let image_source_token = MediaFileToken(request.maybe_image_source.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()); @@ -252,8 +252,9 @@ pub async fn enqueue_image_generation_request( seed = s; } + let inference_type = request.type_of_inference.to_string().clone(); + let inference_args = StableDiffusionArgs { - maybe_video_source: Some(video_token), maybe_image_source: Some(image_source_token), maybe_sd_model_token: Some(sd_weight_token), maybe_lora_model_token: Some(lora_token), @@ -263,7 +264,7 @@ pub async fn enqueue_image_generation_request( maybe_seed: Some(seed), maybe_upload_path: Some(upload_path), maybe_lora_upload_path: Some(lora_upload_path), - inference_type: Some(request.type_of_inference.to_string()), + inference_type: inference_type, }; // create the inference args here From 1614cba2f86b8a308a19a21ef76d80460a3b9f4e Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Wed, 3 Jan 2024 23:16:00 -0500 Subject: [PATCH 10/54] properties added --- .../image_generation_payload.rs | 31 ++++++++++++++++--- .../image_gen/enqueue_image_generation.rs | 10 +++++- 2 files changed, 36 insertions(+), 5 deletions(-) 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 9a94b3a7e9..9ef9da3481 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 @@ -20,9 +20,6 @@ pub enum StableDiffusionSource { #[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)] pub struct StableDiffusionArgs { - #[serde(rename = "is")] - #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_image_source: Option, #[serde(rename = "sd")] #[serde(skip_serializing_if = "Option::is_none")] @@ -32,6 +29,18 @@ pub struct StableDiffusionArgs { #[serde(skip_serializing_if = "Option::is_none")] pub maybe_lora_model_token: Option, + #[serde(rename = "w")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_width: Option, + + #[serde(rename = "h")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_height: Option, + + #[serde(rename = "s")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_sampler: Option, + #[serde(rename = "p")] #[serde(skip_serializing_if = "Option::is_none")] pub maybe_prompt: Option, @@ -51,13 +60,27 @@ pub struct StableDiffusionArgs { #[serde(rename = "mu")] #[serde(skip_serializing_if = "Option::is_none")] pub maybe_upload_path: Option, + + #[serde(rename = "cf")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_cfg_scale: Option, #[serde(rename = "lu")] #[serde(skip_serializing_if = "Option::is_none")] pub maybe_lora_upload_path: Option, + #[serde(rename = "sa")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_number_of_samples: Option, + + #[serde(rename = "bz")] + pub maybe_batch_size: Option, + + #[serde(rename = "bc")] + pub maybe_batch_count: Option, + #[serde(rename = "t")] - pub inference_type: String, + pub type_of_inference: String, } 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 6b03b384fd..c6960a7eb4 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 @@ -67,6 +67,7 @@ impl Display for TypeOfInference { } } + #[derive(Deserialize,ToSchema)] pub struct EnqueueImageGenRequest { uuid_idempotency_token: String, @@ -78,9 +79,15 @@ pub struct EnqueueImageGenRequest { maybe_a_prompt: Option, maybe_n_prompt: Option, maybe_seed: Option, + maybe_width: Option, + maybe_height: Option, + maybe_sampler: Option, maybe_upload_path: Option, maybe_lora_upload_path: Option, - maybe_cfg_scale: Option + maybe_cfg_scale: Option, + maybe_number_of_samples: Option, + maybe_batch_size: Option, + maybe_batch_count: Option, } #[derive(Serialize,ToSchema)] @@ -236,6 +243,7 @@ pub async fn enqueue_image_generation_request( let image_source_token = MediaFileToken(request.maybe_image_source.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 a_prompt = request.maybe_a_prompt.clone().unwrap_or_default(); let n_prompt = request.maybe_n_prompt.clone().unwrap_or_default(); From 796e8b5aee938214a21e30cef8f4497a785f582e Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Thu, 4 Jan 2024 09:57:45 -0500 Subject: [PATCH 11/54] Compile fix switching to rust rover since less mapping is required now. --- .../bucket_orchestration.rs | 1 - .../remote_cloud_file_manager.rs | 16 +++++ .../image_generation/sd/process_job.rs | 69 ++++++++++++++++++- .../image_generation/sd/validate_inputs.rs | 6 +- 4 files changed, 87 insertions(+), 5 deletions(-) diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/bucket_orchestration.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/bucket_orchestration.rs index 7b790fc432..dc7ef315df 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/bucket_orchestration.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/bucket_orchestration.rs @@ -19,7 +19,6 @@ pub trait BucketOrchestrationCore { bytes: &[u8], content_type: &str, is_public: bool) -> AnyhowResult<()>; - } 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 5cfd9ed402..1954459e9e 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 @@ -1,3 +1,4 @@ +use anyhow::anyhow; use errors::AnyhowResult; use filesys::file_read_bytes::file_read_bytes; use filesys::file_size::file_size; @@ -5,6 +6,7 @@ use hashing::sha256::sha256_hash_file::sha256_hash_file; use mimetypes::mimetype_for_bytes::get_mimetype_for_bytes; use crate::remote_file_manager::bucket_orchestration::BucketOrchestrationCore; +use crate::remote_file_manager::bucket_orchestration::BucketOrchestration; use crate::remote_file_manager::file_directory::FileBucketDirectory; use crate::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails; @@ -16,6 +18,20 @@ pub struct RemoteCloudFileClient { } impl RemoteCloudFileClient { + + pub async fn get_remote_cloud_file_client() -> AnyhowResult { + let bucket_orchestration = match BucketOrchestration::new_bucket_client_from_existing_env() { + Ok(client) => client, + Err(e) => { + return Err(anyhow!("Error creating bucket orchestration client: {:?}", e)); + } + }; + + Ok(Self { + bucket_orchestration_client: Box::new(bucket_orchestration) + }) + } + pub fn new(bucket_orchestration_client: Box) -> Self { Self { bucket_orchestration_client 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 4a2314eec6..a041be97d0 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,13 +1,16 @@ use std::time::Duration; +use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails; use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use mysql_queries::queries::model_weights::get_weight::get_weight_by_token; use crate::job::job_loop::job_success_result::JobSuccessResult; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::job_dependencies::JobDependencies; use anyhow::anyhow; use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; +use cloud_storage::remote_file_manager::bucket_orchestration::BucketOrchestration; use crate::job::job_types::image_generation::sd::validate_inputs::validate_inputs; pub struct StableDiffusionProcessArgs<'a> { @@ -15,6 +18,16 @@ pub struct StableDiffusionProcessArgs<'a> { pub job: &'a AvailableInferenceJob, } +async fn get_remote_cloud_file_client() -> RemoteCloudFileClient { + let bucket_orchestration = match BucketOrchestration::new_bucket_client_from_existing_env() { + Ok(client) => client, + Err(e) => { + panic!("Error creating bucket orchestration client: {:?}", e); + } + }; + RemoteCloudFileClient::new(Box::new(bucket_orchestration)) +} + // fill out sd diffusion dependencies // get specific weight // download weight / lora @@ -66,7 +79,6 @@ pub async fn process_job(args: StableDiffusionProcessArgs<'_>) -> Result) -> Result { + let retrieved_sd_record = get_weight_by_token( + &token, + false, + &deps.db.mysql_pool + ).await?; + match retrieved_sd_record { + Some(record) => record, + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no record of model!"))) + } + } + }, + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no sd model token for job!"))) + } + }; + // ignore if no lora token + let retrieved_loRA_record = match lora_token { + Some(token) => { + let retrieved_loRA_record = get_weight_by_token( + &token, + false, + &deps.db.mysql_pool + ).await?; + Some(retrieved_loRA_record) + }, + None => None + }; + + 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()) + ); + + let remote_cloud_file_client = get_remote_cloud_file_client().await; + remote_cloud_file_client.download_file(remote_cloud_bucket_details, to_system_file_path); + + match retrieved_loRA_record { + Some(record) => { + + let lora_details = RemoteCloudBucketDetails::new( + record.public_bucket_hash.clone(), + record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()), + record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string()) + ); + remote_cloud_file_client.download_file(lora_details, work_temp_dir.to_string()); + }, + None => {} + } Ok(JobSuccessResult { maybe_result_entity: None, diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/validate_inputs.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/validate_inputs.rs index 129ee19461..75cf105f6f 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/validate_inputs.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/validate_inputs.rs @@ -26,11 +26,11 @@ pub async fn validate_inputs(args: StableDiffusionProcessArgs<'_>) -> Result<(), }; let stable_diffusion_args: StableDiffusionArgs = StableDiffusionArgs::from(sd_args.clone()); - if stable_diffusion_args.inference_type == "checkpoint" { + if stable_diffusion_args.type_of_inference == "checkpoint" { - } else if stable_diffusion_args.inference_type == "lora" { + } else if stable_diffusion_args.type_of_inference == "lora" { - } else if stable_diffusion_args.inference_type == "inference" { + } else if stable_diffusion_args.type_of_inference == "inference" { } else { return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inference type for job!"))); From 1b4e3cccb59bc88d80faac8097cc9178e95ad3c4 Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Thu, 4 Jan 2024 17:28:00 -0500 Subject: [PATCH 12/54] Fixed over sight in the cloud file remote_file_manager, where the metadata didn't include the bucket path for the file. Then we didn't pass the correct upload path with the file name --- .../inference_result_type.rs | 3 + .../src/remote_file_manager/file_meta_data.rs | 5 +- .../remote_cloud_bucket_details.rs | 2 +- .../remote_cloud_file_manager.rs | 20 ++++- .../image_generation/sd/process_job.rs | 77 ++++++++++++------- .../rerender_a_video/download_video_file.rs | 32 ++++---- 6 files changed, 91 insertions(+), 48 deletions(-) diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs index 4a7d879c21..73befb6186 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs @@ -29,6 +29,9 @@ pub enum InferenceResultType { #[serde(rename = "zs_voice_embedding")] ZeroShotVoiceEmbedding, + + #[serde(renane = "upload_model")] + UploadModel } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/file_meta_data.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/file_meta_data.rs index 8066d28659..b7f80df895 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/file_meta_data.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/file_meta_data.rs @@ -1,6 +1,9 @@ +use crate::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails; + #[derive(Debug, Clone)] pub struct FileMetaData { pub file_size_bytes: u64, pub sha256_checksum: String, - pub mimetype: String + pub mimetype: String, + pub bucket_details: Option } 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 21b3a3aee1..68627337ff 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,7 +1,7 @@ use crate::remote_file_manager::weights_descriptor::{self}; - use super::file_descriptor::FileDescriptor; +#[derive(Debug, Clone)] pub struct RemoteCloudBucketDetails { pub object_hash: String, pub prefix: String, 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 1954459e9e..55f7bccb04 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 @@ -13,6 +13,7 @@ use crate::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDe use super::file_descriptor::FileDescriptor; use super::file_meta_data::FileMetaData; + pub struct RemoteCloudFileClient { bucket_orchestration_client: Box } @@ -55,17 +56,26 @@ impl RemoteCloudFileClient { println!("Reading media file: {:?}", from_system_file_path); // get meta data let bytes = file_read_bytes(from_system_file_path)?; - let result = Self::get_file_meta_data(from_system_file_path)?; + let mut result = Self::get_file_meta_data(from_system_file_path)?; let is_public = file_descriptor.is_public(); + let suffix = file_descriptor.get_suffix().clone(); + let prefix = file_descriptor.get_prefix().clone(); + let directory = FileBucketDirectory::generate_new( file_descriptor ); + + result.bucket_details = Some(RemoteCloudBucketDetails { + object_hash: directory.get_file_object_hash().to_string(), + suffix:suffix, + prefix:prefix, + }); println!("Uploading media file to bucket path: {:?}",directory.get_full_remote_cloud_file_path()); self.bucket_orchestration_client.upload_file_with_content_type_process( - &directory.get_remote_cloud_base_directory(), + &directory.get_full_remote_cloud_file_path(), bytes.as_ref(), result.mimetype.as_ref(), is_public @@ -85,7 +95,8 @@ impl RemoteCloudFileClient { Ok(FileMetaData { file_size_bytes, sha256_checksum, - mimetype: mimetype.to_string() + mimetype: mimetype.to_string(), + bucket_details: None }) } } @@ -127,7 +138,8 @@ mod tests { is_public: bool) -> AnyhowResult<()> { println!("Upload File to Disk"); println!("{}",object_name); - assert_eq!(object_name,String::from("/weights/2/y/q/m/2/2yqm2f1bamh88seyd690h9v24apgezhr/loRA_2yqm2f1bamh88seyd690h9v24apgezhr.safetensors")); + // this is random you have to just check the outputs + //assert_eq!(object_name,String::from("/weights/2/y/q/m/2/2yqm2f1bamh88seyd690h9v24apgezhr/loRA_2yqm2f1bamh88seyd690h9v24apgezhr.safetensors")); println!("ContentType:{}",content_type); println!("{}",is_public); assert_eq!(is_public,true); 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 a041be97d0..adef75aff9 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 @@ -4,13 +4,13 @@ use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloud use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; use mysql_queries::queries::model_weights::get_weight::get_weight_by_token; -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_dependencies::JobDependencies; use anyhow::anyhow; use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; -use cloud_storage::remote_file_manager::bucket_orchestration::BucketOrchestration; +use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType::UploadModel; use crate::job::job_types::image_generation::sd::validate_inputs::validate_inputs; pub struct StableDiffusionProcessArgs<'a> { @@ -18,22 +18,10 @@ pub struct StableDiffusionProcessArgs<'a> { pub job: &'a AvailableInferenceJob, } -async fn get_remote_cloud_file_client() -> RemoteCloudFileClient { - let bucket_orchestration = match BucketOrchestration::new_bucket_client_from_existing_env() { - Ok(client) => client, - Err(e) => { - panic!("Error creating bucket orchestration client: {:?}", e); - } - }; - RemoteCloudFileClient::new(Box::new(bucket_orchestration)) -} - -// fill out sd diffusion dependencies -// get specific weight -// download weight / lora // run inference // upload inference result -// upload model checkpoint or loRA +// upload model checkpoint or loRA + // create record in db // if stable_diffusion_args.inference_type == "checkpoint" { // // run inference with checkpoint and upload @@ -47,6 +35,10 @@ async fn get_remote_cloud_file_client() -> RemoteCloudFileClient { // run inference // insert record into the db with the inference job token complete. + +pub async fn download_from_google_link() { + +} pub async fn sd_args_from_job(args: &StableDiffusionProcessArgs<'_>) -> Result { let inference_args = args.job.maybe_inference_args .as_ref() @@ -82,6 +74,7 @@ pub async fn process_job(args: StableDiffusionProcessArgs<'_>) -> Result) -> Result { + value + }, + None => { + + } + }; + + let lora_path = work_temp_dir.path().join("lora").to_str(); + let lora_path = match lora_path { + Some(value) => { + value + }, + None => { + } + }; + + // 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 retrieved_sd_record = match weight_token { Some(token) => { let retrieved_sd_record = get_weight_by_token( @@ -130,22 +144,33 @@ pub async fn process_job(args: StableDiffusionProcessArgs<'_>) -> Result { - - let lora_details = RemoteCloudBucketDetails::new( - record.public_bucket_hash.clone(), - record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()), - record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string()) - ); - remote_cloud_file_client.download_file(lora_details, work_temp_dir.to_string()); + 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.clone().unwrap_or_else(|| "".to_string()), + model_weight_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string()) + ); + remote_cloud_file_client.download_file(lora_details, lora_path.to_string()).await?; + }, + None => {} + } }, None => {} } + // insert model record in create model record + + // let maybe_result_entity = ResultEntity { + // entity_type: InferenceResultType::UploadModel, + // entity_token: + // }; + Ok(JobSuccessResult { maybe_result_entity: None, inference_duration: Duration::from_secs(0), diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs index 86d73379a8..597757985a 100644 --- a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs @@ -63,23 +63,23 @@ pub async fn download_video_file( } } } - VideofilterVideoSource::U(media_upload_token) => { - let media_upload_token = MediaUploadToken::new_from_str(media_upload_token); - let media_upload_result = get_media_upload_for_inference( - &media_upload_token, - mysql_pool - ).await; + VideofilterVideoSource::U(media_upload_token) => { + let media_upload_token = MediaUploadToken::new_from_str(media_upload_token); + let media_upload_result = get_media_upload_for_inference( + &media_upload_token, + mysql_pool + ).await; - let media_upload_result = match media_upload_result { - Ok(Some(result)) => result, - Ok(None) => { - warn!("could not find media upload: {:?}", media_upload_token); - return Err(ProcessSingleJobError::from_anyhow_error( - anyhow!("could not find media upload: {:?}", media_upload_token))) - } - Err(e) => { - error!("could not query media upload: {:?}", e); - return Err(ProcessSingleJobError::from_anyhow_error(e)) + let media_upload_result = match media_upload_result { + Ok(Some(result)) => result, + Ok(None) => { + warn!("could not find media upload: {:?}", media_upload_token); + return Err(ProcessSingleJobError::from_anyhow_error( + anyhow!("could not find media upload: {:?}", media_upload_token))) + } + Err(e) => { + error!("could not query media upload: {:?}", e); + return Err(ProcessSingleJobError::from_anyhow_error(e)) } }; From aa202172a2948809b442ea0082f97c43d15c7366 Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Thu, 4 Jan 2024 21:46:25 -0500 Subject: [PATCH 13/54] uploading models and then reseed. --- .../src/seeding/model_weights.rs | 291 ++++++++++++++++++ .../service/job/download_job/src/job_state.rs | 2 +- .../image_generation/sd/process_job.rs | 1 + 3 files changed, 293 insertions(+), 1 deletion(-) diff --git a/crates/cli/dev-database-seed/src/seeding/model_weights.rs b/crates/cli/dev-database-seed/src/seeding/model_weights.rs index c7b9704bb4..03054774f5 100644 --- a/crates/cli/dev-database-seed/src/seeding/model_weights.rs +++ b/crates/cli/dev-database-seed/src/seeding/model_weights.rs @@ -2,6 +2,10 @@ use std::vec; use log::info; use sqlx::{MySql, Pool}; +use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; +use cloud_storage::remote_file_manager::weights_descriptor; +use cloud_storage::remote_file_manager::weights_descriptor::{WeightsLoRADescriptor, WeightsSD15Descriptor}; +use easyenv::from_filename; use enums::by_table::model_weights::{ weights_category::WeightsCategory, @@ -11,6 +15,7 @@ use enums::common::visibility::Visibility; use errors::{anyhow, AnyhowResult}; use mysql_queries::queries::model_weights::create::create_weight::{create_weight, CreateModelWeightsArgs}; 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::model_weights::ModelWeightToken; use tokens::tokens::users::UserToken; @@ -330,6 +335,291 @@ pub async fn seed_weights_for_paging(mysql_pool: &Pool, user_token: UserT } Ok(()) } + +pub async fn seed_weights_for_testing_inference(mysql_pool: &Pool, user_token: UserToken) -> AnyhowResult<()>{ + let miyhoyo_description = r#"一个面向米哈游角色的模型合集~A collection for MIHOYO Characters~ +对应人物tag记录在版本信息里，可在右侧“About this version”选项中查看。The corresponding character prompts are recorded in the version information and can be viewed in the About this version option on the right. +打开目录寻找最爱~（附带链接，点击直达）~Open the catalog, find your favorites~ (with links, click to go directly): +崩坏三 Honkai Impact 3rd： + +梅比乌斯- Mobius + +琪亚娜-Kiana：琪亚娜 合集-Collections for Kiana (Honkai Impact 3rd) + +布洛妮娅-bronya：布洛妮娅 合集-Collections for Bronya (Honkai Impact 3rd | Honkai Star Rail) + +芽衣：雷电芽衣 合集-Collections for Raiden Mei (Honkai Impact 3rd) + +西琳-Sirin + +魔法少女-Miracle ☆ Magical Girl + +赤鸢-Jingwei + +丹朱苍玄-Nuwa/Fuxi + +李素裳-Li Sushang + +玉骑士·月痕/锦岁玉团-Jade Knight/Auspicious Dazzle + +卡莲-kallen + +第六夜想曲-Sixth Serenade + +符华-Fu hua + +崩落 + +云墨丹心-Azure Empyrea + +八重樱-Yae Sakura + +逆神巫女-Gyakushinn Miko + +御神装·勿忘-Goushinnso Memento + +信花舞伎-Blooming Maiko + +希尔-Seele + +生死律者-Herrscher of Rebirth + +栖于永夜-Eternal Night's Embrace + +德莉莎-Theresa：德莉莎 合集-Collections for Theresa (Honkai Impact 3rd) + +月下誓约-Lunar Vow: Crimson Love + +月下初拥-Luna Kindred + +花裳月纱-Rosy Bridesmaid + +格蕾修-Griseo + +繁星·绘世之卷-Starry Impression + +天行·绘星之卷-Cosmic Expression + +画中的童话-Maroon Riding Hood + +原神 Genshin： + +凝光-Ningguang + +申鹤-Shenhe + +妮露 + +诺艾尔 + +神里绫华（包含礼装） + +甘雨(包含喜茶联动礼装)-Ganyu + +胡桃(包含森罗联动礼装)-Hutao + +琳妮特-Lynette + +绮良良-Kirara + +坎蒂丝-Candace + +芙宁娜-Furina + +芙宁娜/芙卡洛斯-Furina/Focalors + +克洛琳德-Clorinde + +娜维娅-Navia + +丝柯克-Skirk + +夏洛蒂-Charlotte + +归终-Guizhong + +闲云-Xianyun + +千织-Chiori + +夏沃蕾-Chevreuse + +崩坏：星穹铁道 Honkai Star Rand： + +驭空-Yukong + +银狼-Silver Wolf + +符玄-Fu Xuan + +李素裳-Sushang + +三月七(礼装)-March 7th (New Outfit) + +停云-Tingyun + +青雀-Qingque + +布洛妮娅·兰德-Bronya Rand + +镜流-Jingliu + +卡芙卡-Kafka + +托帕-Topaz + +桂乃芬-Guinaifen + +玲可-Lynx + +佩拉-Peia + +寒鸦-Hanya + +藿藿-Huohuo + +阮·梅-Ruan Mei + +雪衣-Xueyi + +希露瓦-Serval + +绝区零 Zenless-zone-zero： + +星见雅-Miyabi + +猫宫又奈-Nekomata + +艾莲·乔-Ellen + +妮可-Nicole + +苍角-Soukaku + +格莉丝-Grace + +珂蕾妲-koleda + +可琳-Corin + +其他 Others： + +鹿鸣-Luming + """#; + + let sd1_5_markdown_description = r#" +00025-2161235528-1girl,(orange_hair_1.1),(zentangle, mandala, tangle, entangle_0.6),(fractal art),the most beautiful form of chaos,brutalist desi.png + + +00010-1626070972-elf portrait,enchanting beauty,fantasy,ethereal glow,pointed ears,delicate facial features,long elegant hair,nature-themed attir.png + + +00001-301370162-girl in a snowy landscape,winter coat,fluffy hood,white fur trim,breath visible in cold air,soft falling snowflakes,glistening s.png + + +00002-3996269517-venomous sorcerer,ominous dark robes,glowing green eyes,holding a staff with a serpent design,mystical green flames,poisonous mi.png + + +00003-3817469732-Thunder Deity,immense power,hovering above the sea,stormy ocean waves below,crackling lightning in fingertips,dramatic stormy cl.png + +if you want to remake my front image, check this out: Lightflow | SD webUI Extension - workflow restore & save in one click - majicMIXfantasy example | Stable Diffusion Other | Civitai + +v3 is here. Shouldn't be difficult to use. Simple and complex prompts can be used, unleash your imagination. + +v3来了，简单复杂的提示词都可以用。 + +I developed the v3 version towards 2D in order to make more room for 2.5D models like majicmix lux. + +v3版本我往2D发展是为了给majicmix lux等2.5D模型让出更多空间。 + +merged from融了: +Jack Of All by GeneratedJudge + +AWPainting by DynamicWang + +majicMIX fantasy v2 = +Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix +"#; + + info!("Seeding weights for inference..."); + ModelWeightToken::reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(54321); + let model_weight_token1 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let model_weight_token2 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + + let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?; + + let sd_15_weights_descriptor = Box::new(WeightsSD15Descriptor {}); + let lora_descriptor = Box::new(WeightsLoRADescriptor{}); + + + let path_object_SD = get_seed_tool_data_root().await?; + path_object_SD.push("models/imagegen/sd15/majicmixFantasy_v30Vae.safetensors"); + + let path_object_loRA = get_seed_tool_data_root().await?; + path_object_loRA.push("model/imagegen/loRA/xiawolei-v100-000019.safetensors"); + + let metadata1 = remote_cloud_file_client.upload_file(sd_15_weights_descriptor,path_object_SD.string()).await?; + let metadata2 = remote_cloud_file_client.upload_file(lora_descriptor,path_object_loRA.string()).await?; + + let weights1 = CreateModelWeightsArgs { + token: &model_weight_token1, // replace with actual ModelWeightToken + weights_type: WeightsType::StableDiffusion15, // replace with actual WeightsType + weights_category: WeightsCategory::ImageGeneration, // replace with actual WeightsCategory + title: "Dragonfruit".to_string(), + maybe_description_markdown: Some(sd1_5_markdown_description.to_string()), + maybe_description_rendered_html: Some("

Description

".to_string()), + creator_user_token: Some(&user_token), // replace with actual UserToken + creator_ip_address: "192.168.1.1", + creator_set_visibility: Visibility::Public, + maybe_last_update_user_token: Some("Last Update User Token".to_string()), + original_download_url: Some("https://civitai.com/models/41865/majicmix-fantasy".to_string()), + original_filename: Some("majicmixFantasy_v30Vae.safetensors".to_string()), + file_size_bytes: metadata1.file_size_bytes as i32, + file_checksum_sha2: metadata1.sha256_checksum.to_string(), + public_bucket_hash: metadata1.bucket_details.unwrap().object_hash, + maybe_public_bucket_prefix: Some(metadata1.bucket_details.unwrap().prefix), + maybe_public_bucket_extension: Some(metadata1.bucket_details.unwrap().suffix), + cached_user_ratings_total_count: 10, + cached_user_ratings_positive_count: 9, + cached_user_ratings_negative_count: 1, + maybe_cached_user_ratings_ratio: Some(0.9), + version: 1, + mysql_pool: &mysql_pool, // replace with actual MySqlPool + }; + + let weights2 = CreateModelWeightsArgs { + token: &model_weight_token2, // replace with actual ModelWeightToken + weights_type: WeightsType::LoRA, // replace with actual WeightsType + weights_category: WeightsCategory::ImageGeneration, // replace with actual WeightsCategory + title: "MIHOYO Collection 米家全家桶 (Honkai Impact 3rd | Honkai Star Rail | Genshin Impact | Zenless Zone Zero)".to_string(), + maybe_description_markdown: Some(miyhoyo_description.to_string()), + maybe_description_rendered_html: Some(miyhoyo_description.to_string()), + creator_user_token: Some(&user_token), // replace with actual UserToken + creator_ip_address: "292.268.2.2", + creator_set_visibility: Visibility::Public, + maybe_last_update_user_token: Some("

Honkai

".to_string()), + original_download_url: Some("https://civitai.com/models/95243/mihoyo-collection-honkai-impact-3rd-or-honkai-star-rail-or-genshin-impact-or-zenless-zone-zero".to_string()), + original_filename: Some("xiawolei-v100-000019.safetensors".to_string()), + file_size_bytes: metadata2.file_size_bytes as i32, + file_checksum_sha2: metadata2.sha256_checksum.to_string(), + public_bucket_hash: metadata2.bucket_details.unwrap().object_hash, + maybe_public_bucket_prefix: Some(metadata2.bucket_details.unwrap().prefix), + maybe_public_bucket_extension: Some(metadata2.bucket_details.unwrap().suffix), + cached_user_ratings_total_count: 20, + cached_user_ratings_positive_count: 9, + cached_user_ratings_negative_count: 2, + maybe_cached_user_ratings_ratio: Some(0.9), + version: 2, + mysql_pool: &mysql_pool, // replace with actual MySqlPool + }; + + create_weight(weights1).await?; + + create_weight(weights2).await?; + + + Ok(()) +} pub async fn seed_weights_for_user_token( mysql_pool: &Pool, user_token: UserToken @@ -630,6 +920,7 @@ pub async fn original_seed_weights(mysql_pool: &Pool, _user_token: UserTo Ok(()) } + pub async fn seed_weights(mysql_pool: &Pool) -> AnyhowResult<()> { info!("Seeding weights..."); diff --git a/crates/service/job/download_job/src/job_state.rs b/crates/service/job/download_job/src/job_state.rs index 14b476a1f3..26b3e90a94 100644 --- a/crates/service/job/download_job/src/job_state.rs +++ b/crates/service/job/download_job/src/job_state.rs @@ -70,7 +70,7 @@ pub struct PretrainedModels { /// Configurations and interfaces to code deployed as sidecars or container mounts. pub struct SidecarConfigs { - pub google_drive_downloader: GoogleDriveDownloadCommand, + pub(crate) google_drive_downloader: GoogleDriveDownloadCommand, pub rvc_v2_model_check_command: RvcV2ModelCheckCommand, pub softvc_model_check_command: SoftVcModelCheckCommand, pub so_vits_svc_model_check_command: SoVitsSvcModelCheckCommand, 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 adef75aff9..ea7d09d396 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 @@ -64,6 +64,7 @@ pub async fn process_job(args: StableDiffusionProcessArgs<'_>) -> Result Date: Fri, 5 Jan 2024 16:01:42 -0500 Subject: [PATCH 14/54] Fix compile issues --- .../inference_result_type.rs | 9 +- .../src/seeding/model_weights.rs | 28 +- .../src/seeding/seed_model_weights.rs | 2 + .../image_generation/sd/process_job.rs | 274 ++++++++++-------- 4 files changed, 180 insertions(+), 133 deletions(-) diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs index 73befb6186..0b2b7ffaaa 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs @@ -30,7 +30,7 @@ pub enum InferenceResultType { #[serde(rename = "zs_voice_embedding")] ZeroShotVoiceEmbedding, - #[serde(renane = "upload_model")] + #[serde(rename = "upload_model")] UploadModel } @@ -46,6 +46,7 @@ impl InferenceResultType { Self::TextToSpeech => "text_to_speech", Self::VoiceConversion => "voice_conversion", Self::ZeroShotVoiceEmbedding => "zs_voice_embedding", + Self::UploadModel => "upload_model", } } @@ -55,6 +56,7 @@ impl InferenceResultType { "text_to_speech" => Ok(Self::TextToSpeech), "voice_conversion" => Ok(Self::VoiceConversion), "zs_voice_embedding" => Ok(Self::ZeroShotVoiceEmbedding), + "upload_model" => Ok(Self::UploadModel), _ => Err(format!("invalid value: {:?}", value)), } } @@ -67,6 +69,7 @@ impl InferenceResultType { InferenceResultType::TextToSpeech, InferenceResultType::VoiceConversion, InferenceResultType::ZeroShotVoiceEmbedding, + InferenceResultType::UploadModel, ]) } } @@ -82,6 +85,7 @@ mod tests { assert_serialization(InferenceResultType::TextToSpeech, "text_to_speech"); assert_serialization(InferenceResultType::VoiceConversion, "voice_conversion"); assert_serialization(InferenceResultType::ZeroShotVoiceEmbedding, "zs_voice_embedding"); + assert_serialization(InferenceResultType::UploadModel, "upload_model"); } #[test] @@ -90,6 +94,7 @@ mod tests { assert_eq!(InferenceResultType::TextToSpeech.to_str(), "text_to_speech"); assert_eq!(InferenceResultType::VoiceConversion.to_str(), "voice_conversion"); assert_eq!(InferenceResultType::ZeroShotVoiceEmbedding.to_str(), "zs_voice_embedding"); + assert_eq!(InferenceResultType::UploadModel.to_str(), "upload_model"); } #[test] @@ -98,6 +103,7 @@ mod tests { assert_eq!(InferenceResultType::from_str("text_to_speech").unwrap(), InferenceResultType::TextToSpeech); assert_eq!(InferenceResultType::from_str("voice_conversion").unwrap(), InferenceResultType::VoiceConversion); assert_eq!(InferenceResultType::from_str("zs_voice_embedding").unwrap(), InferenceResultType::ZeroShotVoiceEmbedding); + assert_eq!(InferenceResultType::from_str("upload_model").unwrap(), InferenceResultType::UploadModel); } #[test] @@ -109,6 +115,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(InferenceResultType::TextToSpeech)); assert_eq!(variants.pop_first(), Some(InferenceResultType::VoiceConversion)); assert_eq!(variants.pop_first(), Some(InferenceResultType::ZeroShotVoiceEmbedding)); + assert_eq!(variants.pop_first(), Some(InferenceResultType::UploadModel)); assert_eq!(variants.pop_first(), None); // Generated check diff --git a/crates/cli/dev-database-seed/src/seeding/model_weights.rs b/crates/cli/dev-database-seed/src/seeding/model_weights.rs index 03054774f5..e6136837c0 100644 --- a/crates/cli/dev-database-seed/src/seeding/model_weights.rs +++ b/crates/cli/dev-database-seed/src/seeding/model_weights.rs @@ -551,15 +551,14 @@ Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix let sd_15_weights_descriptor = Box::new(WeightsSD15Descriptor {}); let lora_descriptor = Box::new(WeightsLoRADescriptor{}); - - let path_object_SD = get_seed_tool_data_root().await?; + let mut path_object_SD = get_seed_tool_data_root(); path_object_SD.push("models/imagegen/sd15/majicmixFantasy_v30Vae.safetensors"); - let path_object_loRA = get_seed_tool_data_root().await?; - path_object_loRA.push("model/imagegen/loRA/xiawolei-v100-000019.safetensors"); + let mut path_object_loRA = get_seed_tool_data_root(); + path_object_loRA.push("models/imagegen/loRA/xiawolei-v100-000019.safetensors"); - let metadata1 = remote_cloud_file_client.upload_file(sd_15_weights_descriptor,path_object_SD.string()).await?; - let metadata2 = remote_cloud_file_client.upload_file(lora_descriptor,path_object_loRA.string()).await?; + let metadata1 = remote_cloud_file_client.upload_file(sd_15_weights_descriptor,path_object_SD.as_path().to_str().unwrap()).await?; + let metadata2 = remote_cloud_file_client.upload_file(lora_descriptor,path_object_loRA.as_path().to_str().unwrap()).await?; let weights1 = CreateModelWeightsArgs { token: &model_weight_token1, // replace with actual ModelWeightToken @@ -576,9 +575,9 @@ Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix original_filename: Some("majicmixFantasy_v30Vae.safetensors".to_string()), file_size_bytes: metadata1.file_size_bytes as i32, file_checksum_sha2: metadata1.sha256_checksum.to_string(), - public_bucket_hash: metadata1.bucket_details.unwrap().object_hash, - maybe_public_bucket_prefix: Some(metadata1.bucket_details.unwrap().prefix), - maybe_public_bucket_extension: Some(metadata1.bucket_details.unwrap().suffix), + public_bucket_hash: metadata1.bucket_details.clone().unwrap().object_hash, + maybe_public_bucket_prefix: Some(metadata1.bucket_details.clone().unwrap().prefix), + maybe_public_bucket_extension: Some(metadata1.bucket_details.clone().unwrap().suffix), cached_user_ratings_total_count: 10, cached_user_ratings_positive_count: 9, cached_user_ratings_negative_count: 1, @@ -602,9 +601,9 @@ Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix original_filename: Some("xiawolei-v100-000019.safetensors".to_string()), file_size_bytes: metadata2.file_size_bytes as i32, file_checksum_sha2: metadata2.sha256_checksum.to_string(), - public_bucket_hash: metadata2.bucket_details.unwrap().object_hash, - maybe_public_bucket_prefix: Some(metadata2.bucket_details.unwrap().prefix), - maybe_public_bucket_extension: Some(metadata2.bucket_details.unwrap().suffix), + public_bucket_hash: metadata2.bucket_details.clone().unwrap().object_hash.clone(), + maybe_public_bucket_prefix: Some(metadata2.bucket_details.clone().unwrap().prefix), + maybe_public_bucket_extension: Some(metadata2.bucket_details.clone().unwrap().suffix), cached_user_ratings_total_count: 20, cached_user_ratings_positive_count: 9, cached_user_ratings_negative_count: 2, @@ -614,7 +613,6 @@ Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix }; create_weight(weights1).await?; - create_weight(weights2).await?; @@ -933,7 +931,7 @@ pub async fn seed_weights(mysql_pool: &Pool) -> AnyhowResult<()> { //original_seed_weights(mysql_pool,user_token).await?; //seed_weights_for_user_token(mysql_pool, user_token).await?; - seed_weights_for_paging(mysql_pool,user_token).await?; - + //seed_weights_for_paging(mysql_pool,user_token).await?; + seed_weights_for_testing_inference(mysql_pool,user_token).await?; Ok(()) } diff --git a/crates/cli/dev-database-seed/src/seeding/seed_model_weights.rs b/crates/cli/dev-database-seed/src/seeding/seed_model_weights.rs index 0704288417..631daeb9e9 100644 --- a/crates/cli/dev-database-seed/src/seeding/seed_model_weights.rs +++ b/crates/cli/dev-database-seed/src/seeding/seed_model_weights.rs @@ -226,3 +226,5 @@ async fn seed_file_to_bucket( maybe_bucket_extension: maybe_bucket_extension.map(|s| s.to_string()), }) } + + 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 ea7d09d396..075e3e7a47 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,10 +1,9 @@ - use std::time::Duration; use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails; use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; use mysql_queries::queries::model_weights::get_weight::get_weight_by_token; -use crate::job::job_loop::job_success_result::{JobSuccessResult, ResultEntity}; +use crate::job::job_loop::job_success_result::{ JobSuccessResult, ResultEntity }; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::job_dependencies::JobDependencies; use anyhow::anyhow; @@ -23,147 +22,142 @@ pub struct StableDiffusionProcessArgs<'a> { // upload model checkpoint or loRA // create record in db - // if stable_diffusion_args.inference_type == "checkpoint" { - // // run inference with checkpoint and upload - // } else if stable_diffusion_args.inference_type == "lora" { - // // run inference with the lora + random checkpoint - // } else if stable_diffusion_args.inference_type == "inference" { - // // run inference with or without lora - // } else { - // return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inference type for job!"))); - // } - // run inference - // insert record into the db with the inference job token complete. +// if stable_diffusion_args.inference_type == "checkpoint" { +// // run inference with checkpoint and upload +// } else if stable_diffusion_args.inference_type == "lora" { +// // run inference with the lora + random checkpoint +// } else if stable_diffusion_args.inference_type == "inference" { +// // run inference with or without lora +// } else { +// return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inference type for job!"))); +// } +// run inference +// insert record into the db with the inference job token complete. +pub async fn download_from_google_link() {} -pub async fn download_from_google_link() { - -} -pub async fn sd_args_from_job(args: &StableDiffusionProcessArgs<'_>) -> Result { +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()) - .flatten(); + .as_ref() + .map(|args| args.args.as_ref()) + .flatten(); let polymorphic_args = match inference_args { Some(args) => args, None => { - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no inference args for job!"))); + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no inference args for job!")) + ); } }; let sd_args = match polymorphic_args { PolymorphicInferenceArgs::Ig(args) => args, _ => { - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!"))); + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!")) + ); } }; let stable_diffusion_args: StableDiffusionArgs = StableDiffusionArgs::from(sd_args.clone()); Ok(stable_diffusion_args) } -pub async fn process_job(args: StableDiffusionProcessArgs<'_>) -> Result { +pub async fn process_job( + args: StableDiffusionProcessArgs<'_> +) -> Result { let job = args.job; let deps = args.job_dependencies; + // let sd_args = sd_args_from_job(&args).await?; - let sd_args = sd_args_from_job(&args).await?; - - let mut job_progress_reporter = args.job_dependencies - .clients - .job_progress_reporter - .new_generic_inference(job.inference_job_token.as_str()) - .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; - //==================== TEMP DIR ==================== // - let work_temp_dir = format!("temp_stable_diffusion_inference_{}", job.id.0); + // let mut job_progress_reporter = args.job_dependencies.clients.job_progress_reporter + // .new_generic_inference(job.inference_job_token.as_str()) + // .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + // //==================== TEMP DIR ==================== // + // let work_temp_dir = format!("temp_stable_diffusion_inference_{}", job.id.0); + // //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))?; - //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"); + // let lora_path = work_temp_dir.path().join("lora"); - let sd_checkpoint_path = work_temp_dir.path().join("sd_checkpoint").to_str(); - let sd_checkpoint_path = match sd_checkpoint_path { - Some(value) => { - value - }, - None => { + // // 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 retrieved_sd_record = match weight_token { + // Some(token) => { + // let retrieved_sd_record = get_weight_by_token( + // &token, + // false, + // &deps.db.mysql_pool + // ).await?; + // match retrieved_sd_record { + // Some(record) => record, + // None => { + // return Err( + // ProcessSingleJobError::from_anyhow_error(anyhow!("no record of model!")) + // ); + // } + // } + // } + // None => { + // return Err( + // ProcessSingleJobError::from_anyhow_error(anyhow!("no sd model token for job!")) + // ); + // } + // }; - let lora_path = work_temp_dir.path().join("lora").to_str(); - let lora_path = match lora_path { - Some(value) => { - value - }, - None => { - } - }; + // // ignore if no lora token + // let retrieved_loRA_record = match lora_token { + // Some(token) => { + // let retrieved_loRA_record = get_weight_by_token( + // &token, + // false, + // &deps.db.mysql_pool + // ).await?; + // Some(retrieved_loRA_record) + // } + // None => None, + // }; - // 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 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()) + // ); + // let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?; + // remote_cloud_file_client.download_file(details, sd_checkpoint_path.to_string()).await?; - let retrieved_sd_record = match weight_token { - Some(token) => { - let retrieved_sd_record = get_weight_by_token( - &token, - false, - &deps.db.mysql_pool - ).await?; - match retrieved_sd_record { - Some(record) => record, - None => { - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no record of model!"))) - } - } - }, - None => { - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no sd model token for job!"))) - } - }; - // ignore if no lora token - let retrieved_loRA_record = match lora_token { - Some(token) => { - let retrieved_loRA_record = get_weight_by_token( - &token, - false, - &deps.db.mysql_pool - ).await?; - Some(retrieved_loRA_record) - }, - None => None - }; - - 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()) - ); - - let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?; - remote_cloud_file_client.download_file(details, sd_checkpoint_path.to_string()).await?; - - 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.clone().unwrap_or_else(|| "".to_string()), - model_weight_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string()) - ); - remote_cloud_file_client.download_file(lora_details, lora_path.to_string()).await?; - }, - None => {} - } - }, - None => {} - } + // 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 + // .clone() + // .unwrap_or_else(|| "".to_string()), + // model_weight_record.maybe_public_bucket_extension + // .clone() + // .unwrap_or_else(|| "".to_string()) + // ); + // remote_cloud_file_client.download_file( + // lora_details, + // lora_path.to_str + // ).await?; + // } + // None => {} + // } + // } + // None => {} + // } // insert model record in create model record @@ -173,7 +167,53 @@ pub async fn process_job(args: StableDiffusionProcessArgs<'_>) -> Result 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 + } + Err(_) => { + return Err(anyhow!("failed to get remote cloud file client")); + } + }; + + let mut path_dl1 = seed_path.clone(); + path_dl1.push("downloads/loRA"); + let mut path_dl2 = seed_path.clone(); + path_dl2.push("downloads/checkpoint"); + + let bucket_details1 = RemoteCloudBucketDetails { + object_hash: String::from("apa0ej6es8d3ss2gwtf1cghge35qn9tn"), + prefix: String::from("sd15"), + suffix: String::from("safetensors"), + }; + + let bucket_details2 = RemoteCloudBucketDetails { + object_hash: String::from("27kz11et18fargyyxbj66ntfn621k9d3"), + prefix: String::from("loRA"), + suffix: String::from("safetensors"), + }; + + 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(()) + } +} From e5ed4d986fa7fcc381b55a51054cc7c66139c4dc Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Fri, 5 Jan 2024 17:20:21 -0500 Subject: [PATCH 15/54] downloading works. --- .env | 2 +- crates/cli/dev-database-seed/src/main.rs | 9 ++-- .../src/seeding/model_weights.rs | 43 ++++++++++++++++++- .../image_generation/sd/process_job.rs | 16 +++---- 4 files changed, 56 insertions(+), 14 deletions(-) diff --git a/.env b/.env index a39f1b1d8d..c41492d472 100644 --- a/.env +++ b/.env @@ -76,4 +76,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/cli/dev-database-seed/src/main.rs b/crates/cli/dev-database-seed/src/main.rs index c4e3f6b459..020f93015a 100644 --- a/crates/cli/dev-database-seed/src/main.rs +++ b/crates/cli/dev-database-seed/src/main.rs @@ -9,6 +9,7 @@ use crate::bucket_clients::get_bucket_clients; use crate::cli_args::parse_cli_args; 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::model_weights::test_seed_weights_files; use crate::seeding::seed_media_files::seed_media_files; use crate::seeding::tts_tacotron2::seed_tts_tacotron2; use crate::seeding::users::seed_user_accounts; @@ -59,13 +60,13 @@ pub async fn main() -> AnyhowResult<()> { seed_voice_conversion(&pool).await?; 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!"); + //seed_weights_files(&pool, maybe_bucket_clients.as_ref()).await?; + //println!("TESTING DOWLOAD"); + //test_seed_weights_files().await?; + //info!("Done!"); Ok(()) } diff --git a/crates/cli/dev-database-seed/src/seeding/model_weights.rs b/crates/cli/dev-database-seed/src/seeding/model_weights.rs index e6136837c0..3a78ed932b 100644 --- a/crates/cli/dev-database-seed/src/seeding/model_weights.rs +++ b/crates/cli/dev-database-seed/src/seeding/model_weights.rs @@ -1,8 +1,11 @@ +use std::path::PathBuf; use std::vec; use log::info; use sqlx::{MySql, Pool}; use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; +use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails; + use cloud_storage::remote_file_manager::weights_descriptor; use cloud_storage::remote_file_manager::weights_descriptor::{WeightsLoRADescriptor, WeightsSD15Descriptor}; use easyenv::from_filename; @@ -21,6 +24,42 @@ use tokens::tokens::users::UserToken; use crate::seeding::users::HANASHI_USERNAME; +pub 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 + } + Err(_) => { + return Err(anyhow!("failed to get remote cloud file client")); + } + }; + + let mut path_dl1 = seed_path.clone(); + path_dl1.push("downloads/loRA"); + let mut path_dl2 = seed_path.clone(); + path_dl2.push("downloads/checkpoint"); + + let bucket_details1 = RemoteCloudBucketDetails { + object_hash: String::from("apa0ej6es8d3ss2gwtf1cghge35qn9tn"), + prefix: String::from("sd15"), + suffix: String::from("safetensors"), + }; + + let bucket_details2 = RemoteCloudBucketDetails { + object_hash: String::from("27kz11et18fargyyxbj66ntfn621k9d3"), + prefix: String::from("loRA"), + suffix: String::from("safetensors"), + }; + + 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(()) +} + pub async fn seed_weights_for_paging(mysql_pool: &Pool, user_token: UserToken) -> AnyhowResult<()> { let sd1_5_markdown_description = r#" # Dragonfruit AI Models Update and Workflow @@ -932,6 +971,8 @@ pub async fn seed_weights(mysql_pool: &Pool) -> AnyhowResult<()> { //original_seed_weights(mysql_pool,user_token).await?; //seed_weights_for_user_token(mysql_pool, user_token).await?; //seed_weights_for_paging(mysql_pool,user_token).await?; - seed_weights_for_testing_inference(mysql_pool,user_token).await?; + //seed_weights_for_testing_inference(mysql_pool,user_token).await?; + println!("TESTING DOWLOAD"); + test_seed_weights_files().await?; Ok(()) } 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 075e3e7a47..0d133492f1 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 @@ -174,14 +174,14 @@ pub async fn process_job( #[cfg(test)] mod tests { - use anyhow::anyhow; - use std::path::PathBuf; - use cloud_storage::remote_file_manager::{ - remote_cloud_file_manager::{RemoteCloudFileClient, self}, - remote_cloud_bucket_details::RemoteCloudBucketDetails, - }; - use errors::{AnyhowError, AnyhowResult}; - #[tokio::test] +use anyhow::anyhow; +use std::path::PathBuf; +use cloud_storage::remote_file_manager::{ +remote_cloud_file_manager::{RemoteCloudFileClient, self}, +remote_cloud_bucket_details::RemoteCloudBucketDetails, +}; +use errors::{AnyhowError, AnyhowResult}; +#[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; From 4250469fd3275b72c34a78ed605ba41b04d94edc Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Wed, 10 Jan 2024 20:35:53 -0500 Subject: [PATCH 16/54] enqueue_image_generation fixes added checks and changes to the interface. Will document the api as well. --- Cargo.lock | 1 + .../image_generation_payload.rs | 4 +- crates/service/job/inference_job/Cargo.toml | 1 + .../image_gen/enqueue_image_generation.rs | 95 +++++++++++++++++-- 4 files changed, 92 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 706c846123..61c758ad67 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2977,6 +2977,7 @@ dependencies = [ "enums", "errors", "filesys", + "google_drive_common", "hashing", "hostname", "hyper", 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 9ef9da3481..d0d298201b 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 @@ -39,7 +39,7 @@ pub struct StableDiffusionArgs { #[serde(rename = "s")] #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_sampler: Option, + pub maybe_sampler: Option, #[serde(rename = "p")] #[serde(skip_serializing_if = "Option::is_none")] @@ -71,7 +71,7 @@ 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 = "bz")] pub maybe_batch_size: Option, diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml index 19e17e422b..bc9f8d5e2f 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/crates/service/job/inference_job/Cargo.toml @@ -19,6 +19,7 @@ actix_helpers = { path = "../../../lib/actix_helpers" } bootstrap = { path = "../../../lib/bootstrap" } buckets = { path = "../../../_types/buckets" } cloud_storage = { path = "../../../lib/api_clients/cloud_storage" } +google_drive_common = { path = "../../../lib/api_clients/google_drive_common" } collections = { path = "../../../lib/collections" } concurrency = { path = "../../../lib/concurrency" } config = { path = "../../../lib/deprecated/config" } 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 c6960a7eb4..e2a60ad23f 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 @@ -11,7 +11,7 @@ use actix_web::http::StatusCode; use log::warn; use serde::Deserialize; use serde::Serialize; -use tokens::tokens::media_files::MediaFileToken; + use tokens::tokens::model_weights::ModelWeightToken; use utoipa::ToSchema; @@ -68,6 +68,39 @@ impl Display for TypeOfInference { } +pub fn is_valid_string(input: &str) -> bool { + let valid_strings = [ + "DPM++ 2M Karras", + "DPM++ SDE Karras", + "DPM++ 2M SDE Exponential", + "DPM++ 2M SDE Karras", + "Euler a", + "Euler", + "LMS", + "Heun", + "DPM2", + "DPM2 a", + "DPM++ 2S a", + "DPM++ 2M", + "DPM++ SDE", + "DPM++ 2M SDE", + "DPM++ 2M SDE Heun", + "DPM++ 2M SDE Heun Karras", + "DPM++ 2M SDE Heun Exponential", + "DPM++ 3M SDE", + "DPM++ 3M SDE Karras", + "DPM++ 3M SDE Exponential", + "DPM fast", + "DPM adaptive", + "LMS Karras", + "DPM2 Karras", + "DPM2 a Karras", + "DPM++ 2S a Karras" + ]; + valid_strings.contains(&input) +} + + #[derive(Deserialize,ToSchema)] pub struct EnqueueImageGenRequest { uuid_idempotency_token: String, @@ -85,7 +118,7 @@ pub struct EnqueueImageGenRequest { maybe_upload_path: Option, maybe_lora_upload_path: Option, maybe_cfg_scale: Option, - maybe_number_of_samples: Option, + maybe_number_of_samples: Option, maybe_batch_size: Option, maybe_batch_count: Option, } @@ -239,8 +272,48 @@ pub async fn enqueue_image_generation_request( } } - - let image_source_token = MediaFileToken(request.maybe_image_source.clone().unwrap_or_default()); + let width = request.maybe_width.unwrap_or(512); + let height = request.maybe_height.unwrap_or(512); + + let width = match width { + 0..=512 => 512, + 513..=768 => 768, + 769..=1024 => 1024, + _ => 1024, + }; + + let height = match height { + 0..=512 => 512, + 513..=768 => 768, + 769..=1024 => 1024, + _ => 1024, + }; + + let cfg_scale = match request.maybe_cfg_scale { + 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 }, + None => 20, + }; + + let batch_size = match request.maybe_batch_size { + Some(val) => if val > 4 { 4 } else { val }, + None => 1, + }; + + let batch_count = match request.maybe_batch_count { + Some(val) => if val > 2 { 2 } 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"), + }; + 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()); @@ -260,10 +333,11 @@ pub async fn enqueue_image_generation_request( seed = s; } - let inference_type = request.type_of_inference.to_string().clone(); + let type_of_inference = request.type_of_inference.to_string().clone(); + + let inference_args = StableDiffusionArgs { - maybe_image_source: Some(image_source_token), maybe_sd_model_token: Some(sd_weight_token), maybe_lora_model_token: Some(lora_token), maybe_prompt: Some(request.maybe_prompt.clone().unwrap_or_default()), @@ -272,7 +346,14 @@ pub async fn enqueue_image_generation_request( maybe_seed: Some(seed), maybe_upload_path: Some(upload_path), maybe_lora_upload_path: Some(lora_upload_path), - inference_type: inference_type, + type_of_inference: type_of_inference, + maybe_cfg_scale: Some(cfg_scale), + maybe_number_of_samples: Some(number_of_samples), + maybe_batch_count: Some(batch_count), + maybe_batch_size: Some(batch_size), + maybe_width: Some(width), + maybe_height: Some(height), + maybe_sampler: Some(sampler) }; // create the inference args here From 8709fa9c8a1a3e5b0abe4117452dcd3b9220a400 Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Thu, 11 Jan 2024 16:43:04 -0500 Subject: [PATCH 17/54] enqueue is there. testing inference side. --- .../src/seeding/model_weights.rs | 6 +++--- .../image_generation_payload.rs | 3 --- .../config/inference-job.development.env | 8 ++++++++ .../image_gen/enqueue_image_generation.rs | 12 ++++-------- .../http_server/endpoints/image_gen/json.txt | 18 ++++++++++++++++++ .../service/web/storyteller_web/src/routes.rs | 10 ++++++---- 6 files changed, 39 insertions(+), 18 deletions(-) create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/json.txt diff --git a/crates/cli/dev-database-seed/src/seeding/model_weights.rs b/crates/cli/dev-database-seed/src/seeding/model_weights.rs index 3a78ed932b..af7f744545 100644 --- a/crates/cli/dev-database-seed/src/seeding/model_weights.rs +++ b/crates/cli/dev-database-seed/src/seeding/model_weights.rs @@ -971,8 +971,8 @@ pub async fn seed_weights(mysql_pool: &Pool) -> AnyhowResult<()> { //original_seed_weights(mysql_pool,user_token).await?; //seed_weights_for_user_token(mysql_pool, user_token).await?; //seed_weights_for_paging(mysql_pool,user_token).await?; - //seed_weights_for_testing_inference(mysql_pool,user_token).await?; - println!("TESTING DOWLOAD"); - test_seed_weights_files().await?; + seed_weights_for_testing_inference(mysql_pool,user_token).await?; + //println!("TESTING DOWLOAD"); + //test_seed_weights_files().await?; Ok(()) } 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 d0d298201b..83415daea4 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 @@ -73,9 +73,6 @@ pub struct StableDiffusionArgs { #[serde(skip_serializing_if = "Option::is_none")] pub maybe_number_of_samples: Option, - #[serde(rename = "bz")] - pub maybe_batch_size: Option, - #[serde(rename = "bc")] pub maybe_batch_count: Option, 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 b911d80f6e..58472916ad 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/crates/service/job/inference_job/config/inference-job.development.env @@ -104,6 +104,14 @@ VITS_INFERENCE_MAYBE_DOCKER_IMAGE_SHA="03017896367b" RERENDER_INFERENCE_ROOT_DIRECTORY="/home/salt/clone/storyteller-ml/animation/rerender-lora/" RERENDER_INFERENCE_EXECUTABLE_OR_COMMAND="venv/bin/python rerender.py" +# ================================= # +# # +# STABLE DIFFUSION # +# # +# ================================= # + +STABLE_DIFFUSION_INFERENCE_COMMAND="/home/tensor/code/storyteller/storyteller-ml/image-gen" + # ================================= # # # # CACHE SETTINGS # 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 e2a60ad23f..a84efdffcc 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 @@ -119,7 +119,6 @@ pub struct EnqueueImageGenRequest { maybe_lora_upload_path: Option, maybe_cfg_scale: Option, maybe_number_of_samples: Option, - maybe_batch_size: Option, maybe_batch_count: Option, } @@ -189,6 +188,7 @@ pub async fn enqueue_image_generation_request( server_state: web::Data> ) -> Result { + // 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 mut mysql_connection = server_state.mysql_pool.acquire().await.map_err(|err| { @@ -235,6 +235,8 @@ pub async fn enqueue_image_generation_request( } } + // DETECT premium user and queue + // ==================== RATE LIMIT ==================== // let rate_limiter = match maybe_user_session { @@ -299,13 +301,9 @@ pub async fn enqueue_image_generation_request( None => 20, }; - let batch_size = match request.maybe_batch_size { - Some(val) => if val > 4 { 4 } else { val }, - None => 1, - }; let batch_count = match request.maybe_batch_count { - Some(val) => if val > 2 { 2 } else { val }, + Some(val) => if val > 4 { 4 } else { val }, None => 1, }; @@ -336,7 +334,6 @@ pub async fn enqueue_image_generation_request( let type_of_inference = request.type_of_inference.to_string().clone(); - let inference_args = StableDiffusionArgs { maybe_sd_model_token: Some(sd_weight_token), maybe_lora_model_token: Some(lora_token), @@ -350,7 +347,6 @@ pub async fn enqueue_image_generation_request( maybe_cfg_scale: Some(cfg_scale), maybe_number_of_samples: Some(number_of_samples), maybe_batch_count: Some(batch_count), - maybe_batch_size: Some(batch_size), maybe_width: Some(width), maybe_height: Some(height), maybe_sampler: Some(sampler) diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/json.txt b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/json.txt new file mode 100644 index 0000000000..67177d418f --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/json.txt @@ -0,0 +1,18 @@ +{ + "uuid_idempotency_token": "123", + "type_of_inference": "checkpoint", + "maybe_sd_model_token": "weight_dmmthavhawqc2hj7yqyemcbf8", + "maybe_lora_model_token": "weight_t7gz78fjg27m0wtw6r33gafxs", + "maybe_prompt": "a very well drawn, anime girl, with pink and purple hair sitting down on a chair relaxed.", + "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_upload_path": "https://drive.google.com/file/d/1WRgR2pn0Ky8ls5_9Zq6tQlHTBvyWeach/view?usp=sharing", + "maybe_lora_upload_path": "https://drive.google.com/file/d/1WRgR2pn0Ky8ls5_9Zq6tQlHTBvyWeach/view?usp=sharing", + "maybe_cfg_scale": "7", + "maybe_number_of_samples": "32", + "maybe_batch_count": "4" +} \ 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 633b0e6bf1..a3210d7849 100644 --- a/crates/service/web/storyteller_web/src/routes.rs +++ b/crates/service/web/storyteller_web/src/routes.rs @@ -195,10 +195,11 @@ pub fn add_routes (app: App, server_environment: ServerEnvironment) -> app = add_user_rating_routes(app); /* /v1/user_rating/... */ app = add_subscription_routes(app); /* /v1/subscriptions/... */ app = add_voice_designer_routes(app); /* /v1/voice_designer */ - - if server_environment == ServerEnvironment::Development { - app = add_weights_routes(app); /* /v1/stubs/... */ - } + app = add_image_gen_routes(app); + app = add_weights_routes(app); + // if server_environment == ServerEnvironment::Development { + // + // } // ==================== Comments ==================== let mut app = RouteBuilder::from_app(app) @@ -217,6 +218,7 @@ pub fn add_routes (app: App, server_environment: ServerEnvironment) -> .add_get("/v1/user_bookmarks/list/entity/{entity_type}/{entity_token}", list_user_bookmarks_for_entity_handler) .into_app(); + // ==================== Animations ==================== let mut app = RouteBuilder::from_app(app) From 2d52ac0fa9691f490355bf6260a116618329ae9a Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Thu, 11 Jan 2024 21:24:38 -0500 Subject: [PATCH 18/54] Removed some components, also added the command env variables, now going to run inference. --- .../generic_inference_args.rs | 13 ++++--- .../image_generation_payload.rs | 4 --- .../config/inference-job.development.env | 4 +-- .../image_generation/process_single_ig_job.rs | 4 +-- .../job/job_types/image_generation/sd/mod.rs | 1 + .../image_generation/sd/process_job.rs | 34 ++++++++++++++++++- .../sd/sd_inference_command.rs | 27 +++++---------- .../image_gen/enqueue_image_generation.rs | 6 ++-- 8 files changed, 56 insertions(+), 37 deletions(-) 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 cc88ebd3c9..ed4aa3e961 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 @@ -202,21 +202,24 @@ mod tests { let upload_path = "upload_path".to_string(); let lora_upload_path = "lora_upload_path".to_string(); let checkpoint = "checkpoint".to_string(); - + let type_of_inference = "type_of_inference".to_string(); let args = GenericInferenceArgs { inference_category: Some(InferenceCategoryAbbreviated::ImageGeneration), args: Some(PolymorphicInferenceArgs::Ig(StableDiffusionArgs { - maybe_video_source: Some(video_media_token), - maybe_image_source: Some(image_media_token), maybe_sd_model_token: Some(sd_model_token), maybe_lora_model_token: Some(lora_model_token), + maybe_height: Some(512), + maybe_width: Some(512), + maybe_cfg_scale: Some(7), maybe_prompt: Some(prompt), - maybe_a_prompt: Some(a_prompt), maybe_n_prompt: Some(n_prompt), + maybe_batch_count: Some(1), + maybe_batch_size: Some(1), + maybe_number_of_samples: Some(25), maybe_seed: Some(seed), maybe_upload_path: Some(upload_path), maybe_lora_upload_path: Some(lora_upload_path), - inference_type: Some(checkpoint) + type_of_inference: type_of_inference })), }; 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 83415daea4..233f802533 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 @@ -45,10 +45,6 @@ pub struct StableDiffusionArgs { #[serde(skip_serializing_if = "Option::is_none")] pub maybe_prompt: Option, - #[serde(rename = "ap")] - #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_a_prompt: Option, - #[serde(rename = "np")] #[serde(skip_serializing_if = "Option::is_none")] pub maybe_n_prompt: Option, 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 58472916ad..bc9a8165ff 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/crates/service/job/inference_job/config/inference-job.development.env @@ -109,8 +109,8 @@ RERENDER_INFERENCE_EXECUTABLE_OR_COMMAND="venv/bin/python rerender.py" # STABLE DIFFUSION # # # # ================================= # - -STABLE_DIFFUSION_INFERENCE_COMMAND="/home/tensor/code/storyteller/storyteller-ml/image-gen" +STABLE_DIFFUSION_INFERENCE_ROOT_DIRECTORY="/home/tensor/code/storyteller/storyteller-ml/image_generation/Storyteller-Studio-Filter/" +STABLE_DIFFUSION_INFERENCE_COMMAND="venv/bin/python main.py" # ================================= # # # diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs index 959642865f..a5f09ef9f9 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs @@ -2,7 +2,7 @@ use enums::by_table::generic_inference_jobs::inference_model_type::InferenceMode use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; use crate::job::job_loop::job_success_result::JobSuccessResult; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; -use crate::job::job_types::image_generation::sd::process_job::process_job; +use crate::job::job_types::image_generation::sd::process_job::process_job_selection; use crate::job_dependencies::JobDependencies; use crate::job::job_types::image_generation::sd::process_job::StableDiffusionProcessArgs; use errors::anyhow; @@ -30,5 +30,5 @@ pub async fn dispatch_sd_job(job_dependencies: &JobDependencies, job: &Available job_dependencies, job, }; - process_job(args).await + process_job_selection(args).await } \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs index 17702a9f82..bf85551a02 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs @@ -1,4 +1,5 @@ pub mod process_job; + pub mod model_downloaders; pub mod sd_inference_command; pub mod stable_diffusion_dependencies; 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 0d133492f1..c6b765888e 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 @@ -63,7 +63,39 @@ pub async fn sd_args_from_job( Ok(stable_diffusion_args) } -pub async fn process_job( + +pub async fn process_job_selection (args: StableDiffusionProcessArgs<'_> +) -> Result { + let job = args.job; + 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 { + Err(ProcessSingleJobError::Other(anyhow!("inference type doesn't exist!"))) + } +} + +pub async fn process_job_checkpoint(args: StableDiffusionProcessArgs<'_>) -> Result { + Ok(JobSuccessResult { + maybe_result_entity: None, + inference_duration: Duration::from_secs(0), + }) +} +pub async fn process_job_lora(args: StableDiffusionProcessArgs<'_>) -> Result { + Ok(JobSuccessResult { + maybe_result_entity: None, + inference_duration: Duration::from_secs(0), + }) +} + +pub async fn process_job_inference( args: StableDiffusionProcessArgs<'_> ) -> Result { let job = args.job; 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 2c5b37c62d..5014b00602 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 @@ -83,9 +83,8 @@ pub struct InferenceArgs> { pub cfg_scale:i32, pub seed:i32, pub lora_path:P, - pub check_point:P, + pub checkpoint_path:P, pub vae:P, - pub batch_size:i32, pub batch_count:i32, } @@ -117,16 +116,12 @@ impl StableDiffusionInferenceCommand { let maybe_inference_command = easyenv::get_env_string_optional( "STABLE_DIFFUSION_INFERENCE_COMMAND"); - // Optional, eg. `./infer.py`. Typically we'll use the command form instead. - let maybe_inference_executable = easyenv::get_env_pathbuf_optional( - "STABLE_DIFFUSION_INFERENCE_EXECUTABLE"); let executable_or_command = match maybe_inference_command { Some(command) => ExecutableOrCommand::Command(command), - None => match maybe_inference_executable { - Some(executable) => ExecutableOrCommand::Executable(executable), - None => return Err(anyhow!("neither command nor executable passed")), - }, + None => { + return Err(anyhow!("neither command nor executable passed")); + } }; let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( @@ -204,19 +199,16 @@ impl StableDiffusionInferenceCommand { // ===== Begin Python Args ===== command.push_str("--prompt "); command.push_str(&args.prompt.as_ref().to_string_lossy()); - + command.push_str(" --result_dir "); command.push_str(&args.work_dir.as_ref().to_string_lossy()); - + command.push_str(" --result_file "); command.push_str(&args.output_file.as_ref().to_string_lossy()); command.push_str(" --stderr_output_file "); command.push_str(&args.stderr_output_file.as_ref().to_string_lossy()); - - command.push_str(" --prompt "); - command.push_str(&args.prompt.as_ref().to_string_lossy()); - + command.push_str(" --negative_prompt "); command.push_str(&args.negative_prompt.as_ref().to_string_lossy()); @@ -242,14 +234,11 @@ impl StableDiffusionInferenceCommand { command.push_str(&args.lora_path.as_ref().to_string_lossy()); command.push_str(" --check_point "); - command.push_str(&args.check_point.as_ref().to_string_lossy()); + command.push_str(&args.checkpoint_path.as_ref().to_string_lossy()); command.push_str(" --vae "); command.push_str(&args.vae.as_ref().to_string_lossy()); - command.push_str(" --batch_size "); - command.push_str(&args.batch_size.to_string()); - command.push_str(" --batch_count "); command.push_str(&args.batch_count.to_string()); 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 a84efdffcc..9f09505293 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 @@ -109,7 +109,6 @@ pub struct EnqueueImageGenRequest { maybe_sd_model_token: Option, maybe_lora_model_token: Option, maybe_prompt: Option, - maybe_a_prompt: Option, maybe_n_prompt: Option, maybe_seed: Option, maybe_width: Option, @@ -189,8 +188,10 @@ pub async fn enqueue_image_generation_request( ) -> Result { // 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 mut mysql_connection = server_state.mysql_pool.acquire().await.map_err(|err| { warn!("MySql pool error: {:?}", err); EnqueueImageGenRequestError::ServerError @@ -316,7 +317,6 @@ pub async fn enqueue_image_generation_request( let lora_token = ModelWeightToken(request.maybe_lora_model_token.clone().unwrap_or_default()); - let a_prompt = request.maybe_a_prompt.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. @@ -333,12 +333,10 @@ pub async fn enqueue_image_generation_request( let type_of_inference = request.type_of_inference.to_string().clone(); - let inference_args = StableDiffusionArgs { maybe_sd_model_token: Some(sd_weight_token), maybe_lora_model_token: Some(lora_token), maybe_prompt: Some(request.maybe_prompt.clone().unwrap_or_default()), - maybe_a_prompt: Some(a_prompt), maybe_n_prompt: Some(n_prompt), maybe_seed: Some(seed), maybe_upload_path: Some(upload_path), From c01ffe33a54678e69454dc0f4a0f3eb2bb01357d Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 11 Jan 2024 23:32:52 -0500 Subject: [PATCH 19/54] cover image details on all informational weight endpoints --- .../http_server/endpoints/weights/get_weight.rs | 13 +++++++++++++ .../endpoints/weights/list_featured_weights.rs | 12 ++++++++++++ .../endpoints/weights/list_weights_by_user.rs | 15 ++++++++++++++- .../weights/search_model_weights_handler.rs | 13 +++++++++++++ 4 files changed, 52 insertions(+), 1 deletion(-) diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/get_weight.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/get_weight.rs index eb8efa27e5..4197638c5b 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/get_weight.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/get_weight.rs @@ -15,6 +15,7 @@ use enums::by_table::model_weights::weights_types::WeightsType; use enums::common::visibility::Visibility; use mysql_queries::queries::model_weights::get_weight::get_weight_by_token; use tokens::tokens::model_weights::ModelWeightToken; +use crate::http_server::common_responses::cover_image_details::CoverImageDetails; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; use crate::http_server::common_responses::user_details_lite::UserDetailsLight; @@ -40,8 +41,12 @@ pub struct GetWeightResponse { file_size_bytes: i32, file_checksum_sha2: String, + /// Information about the cover image. + pub cover_image: CoverImageDetails, + /// Cover images are small descriptive images that can be set for any model. /// If a cover image is set, this is the path to the asset. + #[deprecated(note="switch to CoverImageDetails")] maybe_cover_image_public_bucket_path: Option, /// Statistics about the weights @@ -149,6 +154,13 @@ pub async fn get_weight_handler( } } + let cover_image_details = CoverImageDetails::from_optional_db_fields( + &weight.token, + weight.maybe_cover_image_public_bucket_hash.as_deref(), + weight.maybe_cover_image_public_bucket_prefix.as_deref(), + weight.maybe_cover_image_public_bucket_extension.as_deref(), + ); + let maybe_cover_image = weight.maybe_cover_image_public_bucket_hash .as_deref() .map(|hash| { @@ -176,6 +188,7 @@ pub async fn get_weight_handler( // TODO(bt,2023-12-24): Migrated the column. We should return nullable fields, but I don't want to break the frontend description_markdown: weight.maybe_description_markdown.unwrap_or_else(|| "".to_string()), description_rendered_html: weight.maybe_description_rendered_html.unwrap_or_else(|| "".to_string()), + cover_image: cover_image_details, maybe_cover_image_public_bucket_path: maybe_cover_image, creator, creator_set_visibility: weight.creator_set_visibility, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_featured_weights.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_featured_weights.rs index 91c86066e5..fa2a42e6c3 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_featured_weights.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_featured_weights.rs @@ -15,6 +15,7 @@ use enums::by_table::model_weights::{ }; use mysql_queries::queries::model_weights::list::list_weights_by_tokens::list_weights_by_tokens; use tokens::tokens::model_weights::ModelWeightToken; +use crate::http_server::common_responses::cover_image_details::CoverImageDetails; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; use crate::http_server::common_responses::user_details_lite::UserDetailsLight; @@ -37,8 +38,12 @@ pub struct ModelWeightForList { pub creator: UserDetailsLight, + /// Information about the cover image. + pub cover_image: CoverImageDetails, + /// Cover images are small descriptive images that can be set for any model. /// If a cover image is set, this is the path to the asset. + #[deprecated(note="switch to CoverImageDetails")] pub maybe_cover_image_public_bucket_path: Option, /// Statistics about the weights @@ -128,6 +133,12 @@ pub async fn list_featured_weights_handler( success: true, results: weights.into_iter() .map(|w| { + let cover_image_details = CoverImageDetails::from_optional_db_fields( + &w.token, + w.maybe_cover_image_public_bucket_hash.as_deref(), + w.maybe_cover_image_public_bucket_prefix.as_deref(), + w.maybe_cover_image_public_bucket_extension.as_deref(), + ); let maybe_cover_image = w.maybe_cover_image_public_bucket_hash .as_deref() @@ -145,6 +156,7 @@ pub async fn list_featured_weights_handler( title: w.title, weight_type: w.weights_type, weight_category: w.weights_category, + cover_image: cover_image_details, maybe_cover_image_public_bucket_path: maybe_cover_image, creator: UserDetailsLight::from_db_fields( &w.creator_user_token, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_weights_by_user.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_weights_by_user.rs index 6347b81d27..08ec39718b 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_weights_by_user.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_weights_by_user.rs @@ -16,6 +16,7 @@ use enums::common::view_as::ViewAs; use enums::common::visibility::Visibility; use mysql_queries::queries::model_weights::list::list_weights_by_user::{list_weights_by_creator_username, ListWeightsForUserArgs}; use tokens::tokens::model_weights::ModelWeightToken; +use crate::http_server::common_responses::cover_image_details::CoverImageDetails; use crate::http_server::common_responses::pagination_page::PaginationPage; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; @@ -39,8 +40,12 @@ pub struct Weight { file_size_bytes: i32, file_checksum_sha2: String, + /// Information about the cover image. + pub cover_image: CoverImageDetails, + /// Cover images are small descriptive images that can be set for any model. /// If a cover image is set, this is the path to the asset. + #[deprecated(note="switch to CoverImageDetails")] maybe_cover_image_public_bucket_path: Option, /// Statistics about the weights @@ -179,7 +184,14 @@ pub async fn list_weights_by_user_handler( } }; - let weights:Vec = results_page.records.into_iter().map(|weight| { + let weights: Vec = results_page.records.into_iter().map(|weight| { + + let cover_image_details = CoverImageDetails::from_optional_db_fields( + &weight.token, + weight.maybe_cover_image_public_bucket_hash.as_deref(), + weight.maybe_cover_image_public_bucket_prefix.as_deref(), + weight.maybe_cover_image_public_bucket_extension.as_deref(), + ); let maybe_cover_image = weight.maybe_cover_image_public_bucket_hash .as_deref() @@ -203,6 +215,7 @@ pub async fn list_weights_by_user_handler( &weight.creator_display_name, &weight.creator_email_gravatar_hash, ), + cover_image: cover_image_details, maybe_cover_image_public_bucket_path: maybe_cover_image, file_size_bytes: weight.file_size_bytes, file_checksum_sha2: weight.file_checksum_sha2, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs index f8fb80f89f..fb880b43ab 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs @@ -20,6 +20,7 @@ use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; use enums::common::visibility::Visibility; use tokens::tokens::model_weights::ModelWeightToken; +use crate::http_server::common_responses::cover_image_details::CoverImageDetails; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; use crate::http_server::common_responses::user_details_lite::UserDetailsLight; @@ -44,6 +45,10 @@ pub struct ModelWeightSearchResult { pub creator: UserDetailsLight, + /// Information about the cover image. + pub cover_image: CoverImageDetails, + + #[deprecated(note="switch to CoverImageDetails")] pub maybe_cover_image_public_bucket_path: Option, pub stats: SimpleEntityStats, @@ -117,6 +122,13 @@ pub async fn search_model_weights_handler( let results = results.into_iter() .map(|result| { + let cover_image_details = CoverImageDetails::from_optional_db_fields( + &result.token, + result.maybe_cover_image_public_bucket_hash.as_deref(), + result.maybe_cover_image_public_bucket_prefix.as_deref(), + result.maybe_cover_image_public_bucket_extension.as_deref(), + ); + let maybe_cover_image = result.maybe_cover_image_public_bucket_hash .as_deref() .map(|hash| { @@ -139,6 +151,7 @@ pub async fn search_model_weights_handler( &result.creator_display_name, &result.creator_gravatar_hash, ), + cover_image: cover_image_details, maybe_cover_image_public_bucket_path: maybe_cover_image, stats: SimpleEntityStats { positive_rating_count: result.ratings_positive_count, From 5949d1d4e79dcd565e349b1b675d44585ef48f3e Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 12 Jan 2024 00:18:44 -0500 Subject: [PATCH 20/54] fix names of search fields --- .../endpoints/weights/search_model_weights_handler.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs index fb880b43ab..e27ece3566 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs @@ -36,8 +36,8 @@ pub struct SearchModelWeightsRequest { pub struct ModelWeightSearchResult { pub weight_token: ModelWeightToken, - pub weights_type: WeightsType, - pub weights_category: WeightsCategory, + pub weight_type: WeightsType, + pub weight_category: WeightsCategory, pub creator_set_visibility: Visibility, @@ -142,8 +142,8 @@ pub async fn search_model_weights_handler( ModelWeightSearchResult { weight_token: result.token, - weights_type: result.weights_type, - weights_category: result.weights_category, + weight_type: result.weights_type, + weight_category: result.weights_category, title: result.title, creator: UserDetailsLight::from_db_fields( &result.creator_user_token, From c4681937b389bb8ebdc6634475d618331402b85a Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Fri, 12 Jan 2024 00:19:09 -0500 Subject: [PATCH 21/54] Image generation on it's way... --- .../image_generation/sd/process_job.rs | 157 +++++++++++------- .../sd/sd_inference_command.rs | 19 --- .../vall_e_x/vall_e_x_inference_command.rs | 1 + 3 files changed, 94 insertions(+), 83 deletions(-) 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 c6b765888e..83e9fda3b3 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 @@ -8,8 +8,11 @@ use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::job_dependencies::JobDependencies; use anyhow::anyhow; use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; + use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType::UploadModel; +use crate::job::job_types::image_generation::sd::sd_inference_command::InferenceArgs; +use crate::job::job_types::image_generation::sd::stable_diffusion_dependencies::StableDiffusionDependencies; use crate::job::job_types::image_generation::sd::validate_inputs::validate_inputs; pub struct StableDiffusionProcessArgs<'a> { @@ -21,16 +24,6 @@ pub struct StableDiffusionProcessArgs<'a> { // upload inference result // upload model checkpoint or loRA -// create record in db -// if stable_diffusion_args.inference_type == "checkpoint" { -// // run inference with checkpoint and upload -// } else if stable_diffusion_args.inference_type == "lora" { -// // run inference with the lora + random checkpoint -// } else if stable_diffusion_args.inference_type == "inference" { -// // run inference with or without lora -// } else { -// return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inference type for job!"))); -// } // run inference // insert record into the db with the inference job token complete. @@ -63,32 +56,33 @@ pub async fn sd_args_from_job( Ok(stable_diffusion_args) } +// store the prompt and cluster them today. pub async fn process_job_selection (args: StableDiffusionProcessArgs<'_> ) -> Result { let job = args.job; let sd_args = sd_args_from_job(&args).await?; if sd_args.type_of_inference == "inference" { - process_job_inference(args).await + process_job_inference(&args).await } else if sd_args.type_of_inference == "lora" { - process_job_lora(args).await + process_job_lora(&args).await } else if sd_args.type_of_inference == "checkpoint"{ - process_job_inference(args).await + process_job_inference(&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_checkpoint(args: &StableDiffusionProcessArgs<'_>) -> Result { Ok(JobSuccessResult { maybe_result_entity: None, inference_duration: Duration::from_secs(0), }) } -pub async fn process_job_lora(args: StableDiffusionProcessArgs<'_>) -> Result { +pub async fn process_job_lora(args: &StableDiffusionProcessArgs<'_>) -> Result { Ok(JobSuccessResult { maybe_result_entity: None, inference_duration: Duration::from_secs(0), @@ -96,67 +90,82 @@ pub async fn process_job_lora(args: StableDiffusionProcessArgs<'_>) -> Result + args: &StableDiffusionProcessArgs<'_> ) -> Result { let job = args.job; let deps = args.job_dependencies; - // let sd_args = sd_args_from_job(&args).await?; + let sd_args = sd_args_from_job(&args).await?; - // let mut job_progress_reporter = args.job_dependencies.clients.job_progress_reporter - // .new_generic_inference(job.inference_job_token.as_str()) - // .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; - // //==================== TEMP DIR ==================== // - // let work_temp_dir = format!("temp_stable_diffusion_inference_{}", job.id.0); + 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 mut job_progress_reporter = args.job_dependencies.clients.job_progress_reporter + .new_generic_inference(job.inference_job_token.as_str()) + .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + //==================== TEMP DIR ==================== // + + 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"); + let lora_path = work_temp_dir.path().join("lora"); - // //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"); - // let lora_path = work_temp_dir.path().join("lora"); // // 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 lora_token = sd_args.maybe_lora_model_token; + let weight_token = sd_args.maybe_sd_model_token; - // let retrieved_sd_record = match weight_token { - // Some(token) => { - // let retrieved_sd_record = get_weight_by_token( - // &token, - // false, - // &deps.db.mysql_pool - // ).await?; - // match retrieved_sd_record { - // Some(record) => record, - // None => { - // return Err( - // ProcessSingleJobError::from_anyhow_error(anyhow!("no record of model!")) - // ); - // } - // } - // } - // None => { - // return Err( - // ProcessSingleJobError::from_anyhow_error(anyhow!("no sd model token for job!")) - // ); - // } - // }; - // // ignore if no lora token - // let retrieved_loRA_record = match lora_token { - // Some(token) => { - // let retrieved_loRA_record = get_weight_by_token( - // &token, - // false, - // &deps.db.mysql_pool - // ).await?; - // Some(retrieved_loRA_record) - // } - // None => None, - // }; + let retrieved_sd_record = match weight_token { + Some(token) => { + let retrieved_sd_record = get_weight_by_token( + &token, + false, + &deps.db.mysql_pool + ).await?; + match retrieved_sd_record { + Some(record) => record, + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no record of model!")) + ); + } + } + } + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no sd model token for job!")) + ); + } + }; + + // ignore if no lora token + let retrieved_loRA_record = match lora_token { + Some(token) => { + let retrieved_loRA_record = get_weight_by_token( + &token, + false, + &deps.db.mysql_pool + ).await?; + Some(retrieved_loRA_record) + } + None => None, + }; // let details = RemoteCloudBucketDetails::new( // retrieved_sd_record.public_bucket_hash.clone(), @@ -198,6 +207,26 @@ pub async fn process_job_inference( // entity_token: // }; + + sd_deps.inference_command.execute_inference(InferenceArgs { + work_dir: (), + output_file: (), + stderr_output_file: (), + prompt: (), + negative_prompt: (), + number_of_samples: 0, + samplers: "".to_string(), + width: 0, + height: 0, + cfg_scale: 0, + seed: 0, + lora_path: (), + checkpoint_path: (), + vae: (), + batch_count: 0, + }); + + Ok(JobSuccessResult { maybe_result_entity: None, inference_duration: Duration::from_secs(0), 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 5014b00602..3c64d5e316 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 @@ -68,7 +68,6 @@ pub enum ExecutableOrCommand { pub struct InferenceArgs> { /// --source_image: path to the input image (or video) - pub input_image: P, // Not availbile rn /// --result_dir: path to directory work is performed pub work_dir: P, /// --result_file: path to final file output @@ -305,21 +304,3 @@ impl StableDiffusionInferenceCommand { } } } - - - - - -// pub struct StableDiffusionInferenceCommand { -// pub command: String, -// pub args: Vec -// } - -// impl StableDiffusionInferenceCommand { -// pub from_env() -> Self { -// Self { -// command: "python".to_string(), -// args: vec![], -// } -// } -// } \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs index 09721a2a22..46516f93d7 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs @@ -72,6 +72,7 @@ pub struct VallEXInferenceCommand { pub alternate_checkpoint_dir: Option, } + impl VallEXInferenceCommand { pub fn new>( root_code_directory: P, From 7f65014d6d0c21076403f0b314032f4b49d808b2 Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Fri, 12 Jan 2024 02:33:36 -0500 Subject: [PATCH 22/54] Almost done inference. --- .../src/seeding/model_weights.rs | 90 +++++++--- .../remote_cloud_file_manager.rs | 1 + .../image_generation/sd/process_job.rs | 156 ++++++++++++------ .../sd/sd_inference_command.rs | 4 +- 4 files changed, 174 insertions(+), 77 deletions(-) diff --git a/crates/cli/dev-database-seed/src/seeding/model_weights.rs b/crates/cli/dev-database-seed/src/seeding/model_weights.rs index af7f744545..2ce8693b9a 100644 --- a/crates/cli/dev-database-seed/src/seeding/model_weights.rs +++ b/crates/cli/dev-database-seed/src/seeding/model_weights.rs @@ -17,6 +17,7 @@ use enums::by_table::model_weights::{ use enums::common::visibility::Visibility; use errors::{anyhow, AnyhowResult}; use mysql_queries::queries::model_weights::create::create_weight::{create_weight, CreateModelWeightsArgs}; +use mysql_queries::queries::users::user::create_account::create_account; 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::model_weights::ModelWeightToken; @@ -376,6 +377,8 @@ pub async fn seed_weights_for_paging(mysql_pool: &Pool, user_token: UserT } pub async fn seed_weights_for_testing_inference(mysql_pool: &Pool, user_token: UserToken) -> AnyhowResult<()>{ + + let miyhoyo_description = r#"一个面向米哈游角色的模型合集~A collection for MIHOYO Characters~ 对应人物tag记录在版本信息里，可在右侧“About this version”选项中查看。The corresponding character prompts are recorded in the version information and can be viewed in the About this version option on the right. 打开目录寻找最爱~（附带链接，点击直达）~Open the catalog, find your favorites~ (with links, click to go directly): @@ -580,24 +583,35 @@ majicMIX fantasy v2 = Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix "#; + + let sd_vae_description = r"This is an earlier version of a stable VAE. Compared to other VAEs, it has a higher level of stability. I am not the creator, but I could not find this VAE on the website, so I am sharing it here. +这是较早时候的稳定VAE，与其他VAE相比具有较高的稳定性，不容易坏图。我不是作者，但是站里没找着所以搬运。#"; + + info!("Seeding weights for inference..."); ModelWeightToken::reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(54321); let model_weight_token1 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); let model_weight_token2 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let model_weight_token3 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?; let sd_15_weights_descriptor = Box::new(WeightsSD15Descriptor {}); let lora_descriptor = Box::new(WeightsLoRADescriptor{}); - + let sd_vae_15_weights_descriptor = Box::new(WeightsSD15Descriptor {}); let mut path_object_SD = get_seed_tool_data_root(); path_object_SD.push("models/imagegen/sd15/majicmixFantasy_v30Vae.safetensors"); let mut path_object_loRA = get_seed_tool_data_root(); path_object_loRA.push("models/imagegen/loRA/xiawolei-v100-000019.safetensors"); + let mut path_to_VAE = get_seed_tool_data_root(); + path_to_VAE.push("models/imagegen/vae/zVae_v20.safetensors"); + let metadata1 = remote_cloud_file_client.upload_file(sd_15_weights_descriptor,path_object_SD.as_path().to_str().unwrap()).await?; let metadata2 = remote_cloud_file_client.upload_file(lora_descriptor,path_object_loRA.as_path().to_str().unwrap()).await?; + let metadata3 = remote_cloud_file_client.upload_file(sd_vae_15_weights_descriptor,path_to_VAE.as_path().to_str().unwrap()).await?; let weights1 = CreateModelWeightsArgs { token: &model_weight_token1, // replace with actual ModelWeightToken @@ -626,34 +640,60 @@ Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix }; let weights2 = CreateModelWeightsArgs { - token: &model_weight_token2, // replace with actual ModelWeightToken - weights_type: WeightsType::LoRA, // replace with actual WeightsType - weights_category: WeightsCategory::ImageGeneration, // replace with actual WeightsCategory - title: "MIHOYO Collection 米家全家桶 (Honkai Impact 3rd | Honkai Star Rail | Genshin Impact | Zenless Zone Zero)".to_string(), - maybe_description_markdown: Some(miyhoyo_description.to_string()), - maybe_description_rendered_html: Some(miyhoyo_description.to_string()), - creator_user_token: Some(&user_token), // replace with actual UserToken - creator_ip_address: "292.268.2.2", - creator_set_visibility: Visibility::Public, - maybe_last_update_user_token: Some("

Honkai

".to_string()), - original_download_url: Some("https://civitai.com/models/95243/mihoyo-collection-honkai-impact-3rd-or-honkai-star-rail-or-genshin-impact-or-zenless-zone-zero".to_string()), - original_filename: Some("xiawolei-v100-000019.safetensors".to_string()), - file_size_bytes: metadata2.file_size_bytes as i32, - file_checksum_sha2: metadata2.sha256_checksum.to_string(), - public_bucket_hash: metadata2.bucket_details.clone().unwrap().object_hash.clone(), - maybe_public_bucket_prefix: Some(metadata2.bucket_details.clone().unwrap().prefix), - maybe_public_bucket_extension: Some(metadata2.bucket_details.clone().unwrap().suffix), - cached_user_ratings_total_count: 20, - cached_user_ratings_positive_count: 9, - cached_user_ratings_negative_count: 2, - maybe_cached_user_ratings_ratio: Some(0.9), - version: 2, - mysql_pool: &mysql_pool, // replace with actual MySqlPool + token: &model_weight_token2, // replace with actual ModelWeightToken + weights_type: WeightsType::LoRA, // replace with actual WeightsType + weights_category: WeightsCategory::ImageGeneration, // replace with actual WeightsCategory + title: "MIHOYO Collection 米家全家桶 (Honkai Impact 3rd | Honkai Star Rail | Genshin Impact | Zenless Zone Zero)".to_string(), + maybe_description_markdown: Some(miyhoyo_description.to_string()), + maybe_description_rendered_html: Some(miyhoyo_description.to_string()), + creator_user_token: Some(&user_token), // replace with actual UserToken + creator_ip_address: "292.268.2.2", + creator_set_visibility: Visibility::Public, + maybe_last_update_user_token: Some("

Honkai

".to_string()), + original_download_url: Some("https://civitai.com/models/95243/mihoyo-collection-honkai-impact-3rd-or-honkai-star-rail-or-genshin-impact-or-zenless-zone-zero".to_string()), + original_filename: Some("xiawolei-v100-000019.safetensors".to_string()), + file_size_bytes: metadata2.file_size_bytes as i32, + file_checksum_sha2: metadata2.sha256_checksum.to_string(), + public_bucket_hash: metadata2.bucket_details.clone().unwrap().object_hash.clone(), + maybe_public_bucket_prefix: Some(metadata2.bucket_details.clone().unwrap().prefix), + maybe_public_bucket_extension: Some(metadata2.bucket_details.clone().unwrap().suffix), + cached_user_ratings_total_count: 20, + cached_user_ratings_positive_count: 9, + cached_user_ratings_negative_count: 2, + maybe_cached_user_ratings_ratio: Some(0.9), + version: 2, + mysql_pool: &mysql_pool, // replace with actual MySqlPool + }; + + let weights3 = CreateModelWeightsArgs { + token: &model_weight_token3, // replace with actual ModelWeightToken + weights_type: WeightsType::StableDiffusion15, // replace with actual WeightsType + weights_category: WeightsCategory::ImageGeneration, // replace with actual WeightsCategory + title: "z-vae".to_string(), + maybe_description_markdown: Some(sd_vae_description.to_string()), + maybe_description_rendered_html: Some("

Description

".to_string()), + creator_user_token: Some(&user_token), // replace with actual UserToken + creator_ip_address: "192.168.1.1", + creator_set_visibility: Visibility::Public, + maybe_last_update_user_token: Some("Last Update User Token".to_string()), + original_download_url: Some("https://civitai.com/models/97653/z-vae".to_string()), + original_filename: Some("zVae_v20.safetensors".to_string()), + file_size_bytes: metadata3.file_size_bytes as i32, + file_checksum_sha2: metadata3.sha256_checksum.to_string(), + public_bucket_hash: metadata3.bucket_details.clone().unwrap().object_hash, + maybe_public_bucket_prefix: Some(metadata3.bucket_details.clone().unwrap().prefix), + maybe_public_bucket_extension: Some(metadata3.bucket_details.clone().unwrap().suffix), + cached_user_ratings_total_count: 10, + cached_user_ratings_positive_count: 9, + cached_user_ratings_negative_count: 1, + maybe_cached_user_ratings_ratio: Some(0.9), + version: 1, + mysql_pool: &mysql_pool, // replace with actual MySqlPool }; create_weight(weights1).await?; create_weight(weights2).await?; - + create_weight(weights3).await?; Ok(()) } 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 55f7bccb04..e11507ddbd 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 @@ -39,6 +39,7 @@ impl RemoteCloudFileClient { } } + 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); 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 83e9fda3b3..a2e0abdd31 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 @@ -11,6 +11,7 @@ use mysql_queries::payloads::generic_inference_args::generic_inference_args::Pol use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType::UploadModel; +use tokens::tokens::model_weights::ModelWeightToken; use crate::job::job_types::image_generation::sd::sd_inference_command::InferenceArgs; use crate::job::job_types::image_generation::sd::stable_diffusion_dependencies::StableDiffusionDependencies; use crate::job::job_types::image_generation::sd::validate_inputs::validate_inputs; @@ -77,6 +78,14 @@ pub async fn process_job_selection (args: StableDiffusionProcessArgs<'_> } pub async fn process_job_checkpoint(args: &StableDiffusionProcessArgs<'_>) -> Result { + + // insert model record in create model record + + // let maybe_result_entity = ResultEntity { + // entity_type: InferenceResultType::UploadModel, + // entity_token: + // }; + Ok(JobSuccessResult { maybe_result_entity: None, inference_duration: Duration::from_secs(0), @@ -96,7 +105,6 @@ pub async fn process_job_inference( let deps = args.job_dependencies; 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"))); @@ -122,8 +130,14 @@ pub async fn process_job_inference( let sd_checkpoint_path = work_temp_dir.path().join("sd_checkpoint"); let lora_path = work_temp_dir.path().join("lora"); + let vae_path = work_temp_dir.path().join("vae"); + let output_path = work_temp_dir.path().join("output"); - + println!("Pathes to download to"); + println!("sd_checkpoint_path:{}",sd_checkpoint_path.to_string()); + println!("lora_path:{}",lora_path.to_string()); + println!("vae_path:{}",vae_path.to_string()); + println!("output_path:{}",output_path.to_string()); // // Unpack loRA and Checkpoint // // run inference by downloading from google drive. @@ -167,65 +181,107 @@ pub async fn process_job_inference( None => None, }; - // 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()) - // ); + let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?; - // let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?; - // remote_cloud_file_client.download_file(details, sd_checkpoint_path.to_string()).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, sd_checkpoint_path.to_string()).await?; - // 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 - // .clone() - // .unwrap_or_else(|| "".to_string()), - // model_weight_record.maybe_public_bucket_extension - // .clone() - // .unwrap_or_else(|| "".to_string()) - // ); - // remote_cloud_file_client.download_file( - // lora_details, - // lora_path.to_str - // ).await?; - // } - // None => {} - // } - // } - // None => {} - // } + match retrieved_loRA_record { + Some(record) => { + match record { + Some(model_weight_record) => { - // insert model record in create model record + let lora_details = RemoteCloudBucketDetails::new( + model_weight_record.public_bucket_hash.clone(), + model_weight_record.maybe_public_bucket_prefix + .clone() + .unwrap_or_else(|| "".to_string()), + model_weight_record.maybe_public_bucket_extension + .clone() + .unwrap_or_else(|| "".to_string()) + ); + remote_cloud_file_client.download_file( + lora_details, + lora_path.to_string() + ).await?; - // let maybe_result_entity = ResultEntity { - // entity_type: InferenceResultType::UploadModel, - // entity_token: - // }; + } + None => {} + } + } + None => {} + } + + // 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()) + ); + + remote_cloud_file_client.download_file( + vae_details, + vae_path.to_string() + ).await?; + + let prompt = match sd_args.maybe_prompt { + Some(val) => { + val + } + None => Err(ProcessSingleJobError::from_anyhow_error(anyhow!("No Prompt provided!"))) + }; + + // #TODO allow for batch and saving later. output files sd_deps.inference_command.execute_inference(InferenceArgs { - work_dir: (), - output_file: (), - stderr_output_file: (), - prompt: (), - negative_prompt: (), + work_dir: work_temp_dir, + output_file: output_path, + stderr_output_file:, + prompt: prompt, + negative_prompt:sd_args.maybe_n_prompt.unwrap_or_default(), number_of_samples: 0, samplers: "".to_string(), - width: 0, - height: 0, - cfg_scale: 0, - seed: 0, - lora_path: (), - checkpoint_path: (), - vae: (), - batch_count: 0, + 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().as_path(), + checkpoint_path: sd_checkpoint_path.clone(), + vae: vae_path.as_path(), + batch_count: sd_args.maybe_batch_count.unwrap_or(1), }); + // upload media and create a record. + Ok(JobSuccessResult { maybe_result_entity: None, 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 3c64d5e316..899f79eb3e 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 @@ -73,8 +73,8 @@ pub struct InferenceArgs> { /// --result_file: path to final file output pub output_file: P, pub stderr_output_file: P, - pub prompt: P, - pub negative_prompt:P, + pub prompt: String, + pub negative_prompt:String, pub number_of_samples:i32, pub samplers:String, pub width:i32, From c6a6676f088d946df3874f1798d410983dc520ab Mon Sep 17 00:00:00 2001 From: Kasisnu <5044120+kasisnu@users.noreply.github.com> Date: Fri, 5 Jan 2024 18:11:43 +0530 Subject: [PATCH 23/54] adds styletts2 as voice designer model --- .../src/private/zs_voices/bucket_directory.rs | 2 + .../inference_model_type.rs | 9 + .../src/by_table/zs_voices/model_type.rs | 10 +- .../src/job/job_types/tts/mod.rs | 1 + .../src/job/job_types/tts/styletts2/mod.rs | 7 + .../tts/styletts2/model_downloaders.rs | 94 +++ .../job_types/tts/styletts2/process_job.rs | 45 ++ .../tts/styletts2/process_job_create_voice.rs | 370 ++++++++++++ .../styletts2/process_job_inference_voice.rs | 318 ++++++++++ .../tts/styletts2/styletts2_dependencies.rs | 20 + .../styletts2/styletts2_inference_command.rs | 549 ++++++++++++++++++ .../job_types/tts/styletts2/validate_job.rs | 84 +++ .../src/job_specific_dependencies.rs | 11 +- 13 files changed, 1518 insertions(+), 2 deletions(-) create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/mod.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_dependencies.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/tts/styletts2/validate_job.rs diff --git a/crates/_types/buckets/src/private/zs_voices/bucket_directory.rs b/crates/_types/buckets/src/private/zs_voices/bucket_directory.rs index 9d4ecee304..0c95cf153a 100644 --- a/crates/_types/buckets/src/private/zs_voices/bucket_directory.rs +++ b/crates/_types/buckets/src/private/zs_voices/bucket_directory.rs @@ -26,12 +26,14 @@ impl ModelCategory { #[derive(Debug, Clone, Copy, Eq, PartialEq)] pub enum ModelType { VallEx, + StyleTTS2, } impl ModelType { pub fn to_str(&self) -> &'static str { match self { ModelType::VallEx => "vallex", + ModelType::StyleTTS2 => "styletts2", } } } diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs index cd71c3246a..48d7de9dce 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs @@ -35,6 +35,8 @@ pub enum InferenceModelType { RerenderAVideo, #[serde(rename = "mocap_net")] MocapNet, + #[serde(rename = "styletts2")] + StyleTTS2, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -53,6 +55,7 @@ impl InferenceModelType { Self::VallEX => "vall_e_x", Self::RerenderAVideo => "rerender_a_video", Self::MocapNet => "mocap_net", + Self::StyleTTS2 => "styletts2", } } @@ -66,6 +69,7 @@ impl InferenceModelType { "vall_e_x" => Ok(Self::VallEX), "rerender_a_video" => Ok(Self::RerenderAVideo), "mocap_net" => Ok(Self::MocapNet), + "styletts2" => Ok(Self::StyleTTS2), _ => Err(format!("invalid value: {:?}", value)), } } @@ -82,6 +86,7 @@ impl InferenceModelType { InferenceModelType::VallEX, InferenceModelType::RerenderAVideo, InferenceModelType::MocapNet, + InferenceModelType::StyleTTS2, ]) } } @@ -101,6 +106,7 @@ mod tests { assert_serialization(InferenceModelType::VallEX, "vall_e_x"); assert_serialization(InferenceModelType::RerenderAVideo, "rerender_a_video"); assert_serialization(InferenceModelType::MocapNet, "mocap_net"); + assert_serialization(InferenceModelType::StyleTTS2, "styletts2"); } #[test] @@ -113,6 +119,7 @@ mod tests { assert_eq!(InferenceModelType::VallEX.to_str(), "vall_e_x"); assert_eq!(InferenceModelType::RerenderAVideo.to_str(), "rerender_a_video"); assert_eq!(InferenceModelType::MocapNet.to_str(), "mocap_net"); + assert_eq!(InferenceModelType::StyleTTS2.to_str(), "styletts2"); } #[test] @@ -125,6 +132,7 @@ mod tests { assert_eq!(InferenceModelType::from_str("vall_e_x").unwrap(), InferenceModelType::VallEX); assert_eq!(InferenceModelType::from_str("rerender_a_video").unwrap(), InferenceModelType::RerenderAVideo); assert_eq!(InferenceModelType::from_str("mocap_net").unwrap(), InferenceModelType::MocapNet); + assert_eq!(InferenceModelType::from_str("styletts2").unwrap(), InferenceModelType::StyleTTS2); } #[test] @@ -140,6 +148,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(InferenceModelType::VallEX)); assert_eq!(variants.pop_first(), Some(InferenceModelType::RerenderAVideo)); assert_eq!(variants.pop_first(), Some(InferenceModelType::MocapNet)); + assert_eq!(variants.pop_first(), Some(InferenceModelType::StyleTTS2)); assert_eq!(variants.pop_first(), None); // Generated check diff --git a/crates/_types/enums/src/by_table/zs_voices/model_type.rs b/crates/_types/enums/src/by_table/zs_voices/model_type.rs index d8f022035d..a5e4dca813 100644 --- a/crates/_types/enums/src/by_table/zs_voices/model_type.rs +++ b/crates/_types/enums/src/by_table/zs_voices/model_type.rs @@ -16,6 +16,8 @@ pub enum ZsVoiceModelType { /// TTS-type zero shot models #[serde(rename = "vall-e-x")] VallEX, + #[serde(rename = "styletts2")] + StyleTTS2, } // TODO(bt, 2023-01-17): This desperately needs MySQL integration tests! @@ -27,12 +29,14 @@ impl ZsVoiceModelType { pub fn to_str(&self) -> &'static str { match self { Self::VallEX=> "vall-e-x", + Self::StyleTTS2 => "styletts2", } } pub fn from_str(value: &str) -> Result { match value { "vall-e-x" => Ok(Self::VallEX), + "styletts2" => Ok(Self::StyleTTS2), _ => Err(format!("invalid value: {:?}", value)), } } @@ -42,6 +46,7 @@ impl ZsVoiceModelType { // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. BTreeSet::from([ Self::VallEX, + Self::StyleTTS2, ]) } } @@ -66,11 +71,13 @@ mod tests { #[test] fn to_str() { assert_eq!(ZsVoiceModelType::VallEX.to_str(), "vall-e-x"); + assert_eq!(ZsVoiceModelType::StyleTTS2.to_str(), "styletts2"); } #[test] fn from_str() { assert_eq!(ZsVoiceModelType::from_str("vall-e-x").unwrap(), ZsVoiceModelType::VallEX); + assert_eq!(ZsVoiceModelType::from_str("styletts2").unwrap(), ZsVoiceModelType::StyleTTS2); assert!(ZsVoiceModelType::from_str("foo").is_err()); } } @@ -81,8 +88,9 @@ mod tests { #[test] fn all_variants() { let mut variants = ZsVoiceModelType::all_variants(); - assert_eq!(variants.len(), 1); + assert_eq!(variants.len(), 2); assert_eq!(variants.pop_first(), Some(ZsVoiceModelType::VallEX)); + assert_eq!(variants.pop_first(), Some(ZsVoiceModelType::StyleTTS2)); assert_eq!(variants.pop_first(), None); } } diff --git a/crates/service/job/inference_job/src/job/job_types/tts/mod.rs b/crates/service/job/inference_job/src/job/job_types/tts/mod.rs index 26d581e703..137d21bf61 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/mod.rs @@ -1,4 +1,5 @@ pub mod process_single_tts_job; +pub mod styletts2; pub mod tacotron2_v2_early_fakeyou; pub mod vits; pub mod vall_e_x; diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/mod.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/mod.rs new file mode 100644 index 0000000000..5623a593ff --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/mod.rs @@ -0,0 +1,7 @@ +pub mod model_downloaders; +pub mod process_job; +pub mod process_job_create_voice; +pub mod process_job_inference_voice; +pub mod validate_job; +pub mod styletts2_dependencies; +pub mod styletts2_inference_command; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs new file mode 100644 index 0000000000..18bcc1c931 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs @@ -0,0 +1,94 @@ +use crate::util::model_downloader::ModelDownloader; + +pub struct StyleTTS2Downloaders { + pub styletts2: StyleTTS2ModelMapping, + pub whisper_lg: WhisperLargeModelMapping, + pub whisper_md: WhisperMediumModelMapping, + pub vocos_encodec_24khz: VocosEncodecModelMapping, + pub vocos_config: VocosEncodecConfigMapping +} + +crate::impl_model_downloader!( + StyleTTS2ModelMapping, + "Vall-e-x Model Mapping", + "VALL_E_X_MAPPING_BUCKET_PATH", + "/dependencies/zero_shot_tts/vall-e-x_1.0/StyleTTS2-checkpoint.pt", + "VALL_E_X_MAPPING_FILESYSTEM_PATH", + "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/StyleTTS2-checkpoint.pt" +); + +crate::impl_model_downloader!( + WhisperLargeModelMapping, + "whisper Large Model Mapping", + "WHISPER_LARGE_MAPPING_BUCKET_PATH", + "/dependencies/zero_shot_tts/vall-e-x_1.0/whisper-large/model.safetensors", + "WHISPER_LARGE_MAPPING_FILESYSTEM_PATH", + "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper-large/model.safetensors" +); + +crate::impl_model_downloader!( + WhisperMediumModelMapping, + "whisper Medium Model Mapping", + "WHISPER_MEDIUM_TENSOR_BUCKET_PATH", + "/dependencies/zero_shot_tts/vall-e-x_1.0/whisper/medium.pt", + "WHISPER_MEDIUM_TENSOR_FILESYSTEM_PATH", + "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper/medium.pt" +); + +crate::impl_model_downloader!( + VocosEncodecModelMapping, + "Vocos Encodec Model Mapping", + "VOCOS_ENCODEC_TENSOR_BUCKET_PATH", + "/dependencies/zero_shot_tts/vall-e-x_1.0/vocos-encodec/encodec_pytorch_model.bin", + "VOCOS_ENCODEC_TENSOR_FILESYSTEM_PATH", + "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/encodec_pytorch_model.bin" +); + +crate::impl_model_downloader!( + VocosEncodecConfigMapping, + "Vocos Config Mapping", + "VOCOS_ENCODEC_CONFIG_BUCKET_PATH", + "/dependencies/zero_shot_tts/vall-e-x_1.0/vocos-encodec/config.yaml", + "VOCOS_ENCODEC_CONFIG_FILESYSTEM_PATH", + "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/config.yaml" +); + +impl StyleTTS2Downloaders { + pub fn build_all_from_env() -> Self { + Self { + styletts2: StyleTTS2ModelMapping::from_env(), + whisper_lg: WhisperLargeModelMapping::from_env(), + whisper_md: WhisperMediumModelMapping::from_env(), + vocos_encodec_24khz: VocosEncodecModelMapping::from_env(), + vocos_config: VocosEncodecConfigMapping::from_env(), + } + } + + pub fn all_downloaders(&self) -> Vec<&dyn ModelDownloader> { + vec![&self.styletts2, &self.whisper_lg, &self.whisper_md, &self.vocos_encodec_24khz, &self.vocos_config] + } +} + +mod test { + // use std::path::Path; + + // use crate::job::job_types::lipsync::sad_talker::model_downloaders::{SadTalkerModelFirstMapping, SadTalkerModelSecondMapping}; + // use crate::util::model_downloader::ModelDownloader; + + // #[test] + // fn test_downloader_model_first_mapping() { + // let model = SadTalkerModelFirstMapping::default(); + // assert_eq!(model.get_model_name(), "sad talker model first mapping"); + // assert_eq!(model.get_cloud_bucket_path(), "/animation_sadtalker/model/mapping_00109-model.pth.tar"); + // assert_eq!(model.get_filesystem_path(), Path::new("/tmp/downloads/sadtalker/mapping_00109-model.pth.tar")); + // } + + // #[test] + // fn test_downloader_model_second_mapping() { + // let model = SadTalkerModelSecondMapping::default(); + // assert_eq!(model.get_model_name(), "sad talker model second mapping"); + // assert_eq!(model.get_cloud_bucket_path(), "/animation_sadtalker/model/mapping_00229-model.pth.tar"); + // assert_eq!(model.get_filesystem_path(), Path::new("/tmp/downloads/sadtalker/mapping_00229-model.pth.tar")); + // } + +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job.rs new file mode 100644 index 0000000000..e98c004db0 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job.rs @@ -0,0 +1,45 @@ +use anyhow::anyhow; + +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; + +use crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::tts::styletts2::process_job_create_voice::process_create_voice; +use crate::job::job_types::tts::styletts2::process_job_inference_voice::process_inference_voice; +use crate::job::job_types::tts::styletts2::validate_job::validate_job; +use crate::job_dependencies::JobDependencies; + +use super::validate_job::JobType; + +// This will download everything get into the root host OS then ... will invoke inference using the pathes from the files invoked +pub struct StyleTTS2ProcessJobArgs<'a> { + pub job_dependencies: &'a JobDependencies, + pub job: &'a AvailableInferenceJob, +} + +// query using the token then grab the bucket hash +pub async fn process_job( + args: StyleTTS2ProcessJobArgs<'_> +) -> Result { + let job = args.job; + + // get args token + let jobArgs = validate_job(&job)?; // bubbles error up + + match jobArgs.job_type { + JobType::Create => { + if let Some(voice_dataset_token) = jobArgs.voice_dataset_token { + process_create_voice(args, voice_dataset_token).await + } else { + Err(ProcessSingleJobError::Other(anyhow!("Missing Dataset Token?"))) + } + } + JobType::Inference => { + if let Some(voice_token) = jobArgs.voice_token { + process_inference_voice(args, voice_token).await + } else { + Err(ProcessSingleJobError::Other(anyhow!("Missing Voice Token?"))) + } + } + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs new file mode 100644 index 0000000000..d64cc4f809 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs @@ -0,0 +1,370 @@ +use std::fs::read_to_string; +use std::path::PathBuf; +use std::time::Instant; + +use anyhow::anyhow; +use log::{error, info, warn}; + +use buckets::private::zs_voices::bucket_directory::{ModelCategory, ModelType}; +use buckets::private::zs_voices::bucket_file_path::ZeroShotVoiceEmbeddingBucketPath; +use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; +use cloud_storage::bucket_client::BucketClient; +use cloud_storage::bucket_path_unifier::BucketPathUnifier; +use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use enums::by_table::zs_voices::encoding_type::ZsVoiceEncodingType; +use enums::by_table::zs_voices::model_category::ZsVoiceModelCategory; +use enums::by_table::zs_voices::model_type::ZsVoiceModelType; +use enums::common::visibility::Visibility; +use mysql_queries::queries::voice_designer::datasets::get_dataset::get_dataset_by_token; +use mysql_queries::queries::voice_designer::datasets::get_dataset::ZsDataset; +use mysql_queries::queries::voice_designer::voice_samples::list_dataset_samples_for_dataset_token::DatasetSampleRecordForList; +use mysql_queries::queries::voice_designer::voice_samples::list_dataset_samples_for_dataset_token::list_dataset_samples_for_dataset_token; +use mysql_queries::queries::voice_designer::voices::create_voice::create_voice; +use mysql_queries::queries::voice_designer::voices::create_voice::CreateVoiceArgs; +use tokens::tokens::users::UserToken; + +use crate::job; +use crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::job_success_result::ResultEntity; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::tts::styletts2::process_job::StyleTTS2ProcessJobArgs; + +const BUCKET_FILE_PREFIX_CREATE: &str = "fakeyou_"; +const BUCKET_FILE_EXTENSION_CREATE: &str = ".bin"; +const MIME_TYPE_CREATE: &str = "application/x-binary"; + +/// We'll increment this if we change the underlying StyleTTS2 code in an incompatible way. +const CURRENT_STYLE_TTS2_MODEL_VERSION: u64 = 0; + +pub async fn process_create_voice( + args: StyleTTS2ProcessJobArgs<'_>, + dataset_token: String +) -> Result { + let deps = args.job_dependencies; + let job = args.job; + let mysql_pool = &deps.db.mysql_pool; + + let model_dependencies = deps + .job + .job_specific_dependencies + .maybe_styletts2_dependencies + .as_ref() + .ok_or_else(|| ProcessSingleJobError::JobSystemMisconfiguration(Some("missing StyleTTS2 dependencies".to_string())))?; + + // get some globals + let mut job_progress_reporter = deps.clients.job_progress_reporter + .new_generic_inference(job.inference_job_token.as_str()) + .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + + info!("token! {}", dataset_token); + let voice_dataset_token = tokens::tokens::zs_voice_datasets::ZsVoiceDatasetToken(dataset_token); + + let creator_ip_address = &job.creator_ip_address; + + let creator_user_token = match &job.maybe_creator_user_token { + Some(token) => { + UserToken::new_from_str(token) + }, + None => { + return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Creator User Token"))); + } + }; + + // STEP 1. SETUP A TEMP DIRECTORY + let work_temp_dir = format!("/tmp/temp_zeroshot_create_voice_{}", job.id.0); + // 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 workdir = work_temp_dir.path().to_path_buf(); + + // STEP 2. Get dataset for the title for the voice + let voice_dataset = get_dataset_by_token(&voice_dataset_token, false, &mysql_pool).await; + let single_dataset: ZsDataset; + match voice_dataset { + Ok(val) => { + match val { + Some(val) => { + single_dataset = val; + } + None => { + return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Dataset"))); + } + } + } + Err(e) => { + return Err(ProcessSingleJobError::from_anyhow_error(e)); + } + } + info!("Title:{} Token:{}", &single_dataset.title, &single_dataset.token); + + // STEP 3. Download dataset each audio file + let result = list_dataset_samples_for_dataset_token( + &voice_dataset_token, + false, + &mysql_pool + ).await; + + let dataset: Vec = match result { + Ok(val) => val, + Err(e) => return Err(ProcessSingleJobError::from_anyhow_error(e)), + }; + + info!("Dataset length info: {}", dataset.len()); + +// let temp_extension = String::from(".bin"); +// let temp_prefix: String; +// +// if !deps.job.info.container.is_on_prem { +// temp_prefix = String::from("sample_"); // this is for seed in local dev to download the samples +// } else { +// temp_prefix = String::from(BUCKET_FILE_PREFIX_CREATE); +// } + + let mut downloaded_dataset: Vec = Vec::new(); + + for (index, record) in dataset.iter().enumerate() { + //https://storage.googleapis.com/dev-vocodes-public/media/5/3/3/w/8/533w8zs0fy11nv7gkcna7p7vt03h8nda/dev_zs_533w8zs0fy11nv7gkcna7p7vt03h8nda.bin <-- where the file actually is + + let prefix: Option<&str> = record.maybe_public_bucket_prefix.as_ref().map(|s| s.as_str()); + let extension: Option<&str> = record.maybe_public_bucket_extension.as_ref().map(|s| s.as_str()); + + info!("Record=> hash:{} prefix:{:?} extension:{:?}", + record.public_bucket_directory_hash, + prefix, + extension + ); + + let audio_media_file = MediaFileBucketPath::from_object_hash( + &record.public_bucket_directory_hash, + prefix, + extension + ); + + info!("Download using audio_media_file_path: {:?}", audio_media_file.to_full_object_pathbuf()); + + let file_name_wav = format!("{}.wav", index); + + let mut file_path = PathBuf::new(); + file_path.push(workdir.clone()); + file_path.push(file_path.clone()); + file_path.push(file_name_wav); + + info!("Downloading to path: {:?}", file_path); + + // TODO: we might want to catch the error and not include the pathes into download dataset? + let result = deps.buckets.public_bucket_client.download_file_to_disk( + audio_media_file.to_full_object_pathbuf(), + &file_path + ).await; + + if let Err(err) = result { + error!("could not download sample: {:?}", err); + return Err(ProcessSingleJobError::from_anyhow_error(err)); + } + + info!("FilePath to clone voice: {}", file_path.to_string_lossy()); + downloaded_dataset.push(file_path.clone()); + } + + info!("Dataset Length {}", downloaded_dataset.len()); + + // STEP 4 Download the models + info!("Download models (if not present)..."); + for downloader in model_dependencies.downloaders.all_downloaders() { + let result = downloader.download_if_not_on_filesystem( + &args.job_dependencies.buckets.private_bucket_client, + &args.job_dependencies.fs.scoped_temp_dir_creator_for_short_lived_downloads + ).await; + if let Err(err) = result { + error!("could not download: {:?}", err); + return Err(err); + } + } + + job_progress_reporter + .log_status("running inference") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // Command line arg for a list of paths to insert the container + let audio_files = join_paths(downloaded_dataset); + + info!("Files to process: {:?}", audio_files); + + // Name of the output file + // NB: don't use the extension... for the inference since the container will add the extension. + let output_file_name = PathBuf::from("temp"); + + let stderr_output_file = work_temp_dir.path().join("zero_shot_create_voice_err.txt"); + + let inference_start_time = Instant::now(); + + // Run Inference + let command_exit_status = + model_dependencies.create_embedding_command.execute_inference( + job::job_types::tts::styletts2::styletts2_inference_command::CreateVoiceInferenceArgs { + output_embedding_path: &workdir, + output_embedding_name: &output_file_name, + audio_files, + stderr_output_file: &stderr_output_file, + } + ); + + let inference_duration = Instant::now().duration_since(inference_start_time); + + if !command_exit_status.is_success() { + error!("Create Embedding Failed: {:?}", command_exit_status); + + let error = ProcessSingleJobError::Other(anyhow!("CommandExitStatus: {:?}", command_exit_status)); + + if let Ok(contents) = read_to_string(&stderr_output_file) { + warn!("Captured stderr output: {}", contents); + + // Re-categorize error? + //match categorize_error(&contents) { + // Some(ProcessSingleJobError::FaceDetectionFailure) => { + // warn!("Face not detected in source image"); + // error = ProcessSingleJobError::FaceDetectionFailure; + // } + // _ => {} + //} + } + + //thread::sleep(Duration::from_secs(300)); + + // Clean up temp files + //safe_delete_temp_file(&audio_path.filesystem_path); + //safe_delete_temp_file(&image_path.filesystem_path); + //safe_delete_temp_file(&usable_image_path); + //safe_delete_temp_file(&output_video_fs_path); + //safe_delete_temp_file(&stderr_output_file); + //safe_delete_temp_directory(&work_temp_dir); + + return Err(error); + } + + info!("Inference success!"); + + //info!("Success; waiting..."); + //thread::sleep(Duration::from_secs(300)); + + + // STEP 4. Download dataset each audio file + info!("Uploading Media ..."); + + let embedding_bucket_location = ZeroShotVoiceEmbeddingBucketPath::generate_new( + ModelCategory::Tts, + ModelType::StyleTTS2, + CURRENT_STYLE_TTS2_MODEL_VERSION + ); + + let embedding_bucket_object_pathbuf = embedding_bucket_location.to_full_object_pathbuf(); + + // Get Finished File + let mut finished_file = work_temp_dir.path().to_path_buf(); + + let output_bucket_file_name = String::from("temp.npz"); // use extension for bucket upload. + finished_file.push(&output_bucket_file_name); + + info!("Upload File Path: {:?}", finished_file); + info!("Upload Bucket Path: {:?}", embedding_bucket_object_pathbuf); + + args.job_dependencies.buckets.private_bucket_client + .upload_filename_with_content_type( + &embedding_bucket_object_pathbuf, + &finished_file, + &MIME_TYPE_CREATE + ) + .await // TODO: We should check the mimetype to make sure bad payloads can't get uploaded + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // CLEARIFY! these items + // 1.Should this be object hash? + // 2.As well as this what should the voice name be? + let voice_name = single_dataset.title; + + // insert record + let voice_token = create_voice(CreateVoiceArgs { + dataset_token: &voice_dataset_token, + model_category: ZsVoiceModelCategory::Tts, + model_type: ZsVoiceModelType::StyleTTS2, + model_version: CURRENT_STYLE_TTS2_MODEL_VERSION, + model_encoding_type: ZsVoiceEncodingType::Encodec, + voice_title: &voice_name, + bucket_hash: embedding_bucket_location.get_object_hash().clone(), + maybe_creator_user_token: Some(&creator_user_token), + creator_ip_address: &creator_ip_address, + creator_set_visibility: Visibility::Public, + mysql_pool, + }).await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + Ok(JobSuccessResult { + maybe_result_entity: Some(ResultEntity { + entity_type: InferenceResultType::ZeroShotVoiceEmbedding, + entity_token: voice_token.to_string(), + }), + inference_duration, + }) +} + +pub struct AudioFile { + pub filesystem_path: PathBuf, +} + +pub async fn download_audio_from_hash( + bucket_hash: &str, + name: &str, + private_bucket_client: &BucketClient, + path: &PathBuf +) -> Result { + let unifer = BucketPathUnifier::default_paths(); + let object_path = unifer.zero_shot_tts_speaker_encoding(bucket_hash, 0); + + let mut path = path.clone(); + + let file_name = format!("{}", name); + path.push(&file_name); + + private_bucket_client.download_file_to_disk(object_path, &path) + .await + .map_err(|err| { + error!("Could not download audio file: {err}"); + ProcessSingleJobError::from_anyhow_error(anyhow!("Could not download audio file: {err}")) + })?; + + let audio_file = AudioFile { + filesystem_path: PathBuf::from(&path.clone()), + }; + + Ok(audio_file) +} + +fn join_paths(paths: Vec) -> String { + paths + .into_iter() + .map(|p| format!("\"{}\"", p.display())) + .collect::>() + .join(" ") +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use crate::job::job_types::tts::styletts2::process_job_create_voice::join_paths; + + #[test] + fn test_path_build() { + let paths = vec![ + PathBuf::from("/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/20.wav"), + PathBuf::from("/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/21.wav") + ]; + + let value = join_paths(paths); + let expected = + "\"/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/20.wav\" \"/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/21.wav\""; + assert_eq!(value, expected); + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs new file mode 100644 index 0000000000..cb8d476954 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs @@ -0,0 +1,318 @@ +use std::fs::read_to_string; +use std::path::PathBuf; +use std::time::Instant; + +use anyhow::anyhow; +use log::{error, info, warn}; + +use buckets::private::zs_voices::bucket_directory::{ModelCategory, ModelType}; +use buckets::private::zs_voices::bucket_file_path::ZeroShotVoiceEmbeddingBucketPath; +use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; +use cloud_storage::bucket_client::BucketClient; +use cloud_storage::bucket_path_unifier::BucketPathUnifier; +use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use filesys::file_size::file_size; +use hashing::sha256::sha256_hash_file::sha256_hash_file; +use mysql_queries::queries::media_files::create::insert_media_file_from_zero_shot_tts::insert_media_file_from_zero_shot; +use mysql_queries::queries::media_files::create::insert_media_file_from_zero_shot_tts::InsertArgs; +use mysql_queries::queries::voice_designer::voices::get_voice::{get_voice_by_token, ZsVoice}; + +use crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::job_success_result::ResultEntity; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::tts::styletts2::process_job::StyleTTS2ProcessJobArgs; +use crate::job::job_types::tts::styletts2::styletts2_inference_command::InferenceArgs; + +const BUCKET_FILE_PREFIX: &str = "fakeyou_"; +const BUCKET_FILE_EXTENSION: &str = ".wav"; +const MIME_TYPE: &str = "audio/wav"; + + +pub async fn process_inference_voice( + args: StyleTTS2ProcessJobArgs<'_>, + voice_token: String +) -> Result { + let deps = args.job_dependencies; + let job = args.job; + let mysql_pool = &deps.db.mysql_pool; + + let model_dependencies = deps + .job + .job_specific_dependencies + .maybe_styletts2_dependencies + .as_ref() + .ok_or_else(|| ProcessSingleJobError::JobSystemMisconfiguration(Some("missing StyleTTS2 dependencies".to_string())))?; + + // get some globals + let mut job_progress_reporter = deps + .clients + .job_progress_reporter + .new_generic_inference(job.inference_job_token.as_str()) + .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + + // get job args + let text = match job.maybe_raw_inference_text.clone() { + Some(value) => { value } + None => { + return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Text for Inference"))); + } + }; + + let voice_token = tokens::tokens::zs_voices::ZsVoiceToken(voice_token); + + // Get voice bucket hash - from voice token + let voice_lookup_result = get_voice_by_token(&voice_token, false, &mysql_pool).await; + + let voice = match voice_lookup_result { + Ok(Some(voice)) => voice, + Ok(None) => { + warn!("Voice not found: {:?}", voice_token); + return Err(ProcessSingleJobError::Other(anyhow!("Voice not found: {:?}", voice_token))); + } + Err(err) => { + warn!("Error looking up voice: {:?}", err); + return Err(ProcessSingleJobError::Other(anyhow!("Error looking up voice: {:?}", err))); + } + }; + + // Need to download the models + info!("Download models (if not present)..."); + + for downloader in model_dependencies.downloaders.all_downloaders() { + let result = downloader.download_if_not_on_filesystem( + &args.job_dependencies.buckets.private_bucket_client, + &args.job_dependencies.fs.scoped_temp_dir_creator_for_short_lived_downloads + ).await; + + if let Err(err) = result { + error!("could not download: {:?}", err); + return Err(err); + } + } + + // Might not need this for inference. + // 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"))); + // } + // } + + // run inference + let work_temp_dir = format!("/tmp/temp_zeroshot_inference_{}", job.id.0); + + // 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 workdir = work_temp_dir.path().to_path_buf(); + let filename = "weights.npz".to_string(); + + let mut downloaded_weights_path = work_temp_dir.path().to_path_buf(); + downloaded_weights_path.push(&filename); + + let voice_file = download_voice_embedding( + &voice, + &deps.buckets.private_bucket_client, + &downloaded_weights_path + ).await?; + + println!("voice file path! {}", voice_file.filesystem_path.to_string_lossy()); + + // Download embeddings file using embedding token + // Create a temp dir to download things to + job_progress_reporter + .log_status("running inference") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + let inference_start_time = Instant::now(); + + let output_file_name = String::from("output.wav"); + + let stderr_output_file = work_temp_dir.path().join("zero_shot_inference.txt"); + + // Run Inference + let command_exit_status = + model_dependencies.inference_command.execute_inference( + InferenceArgs { + input_embedding_path: &workdir, + input_embedding_name: filename, + input_text: String::from(text), // text + output_file_name: output_file_name.clone(), // output file name in the output folder + stderr_output_file: &stderr_output_file, + } + ); + + let inference_duration = Instant::now().duration_since(inference_start_time); + + if !command_exit_status.is_success() { + error!("Zero shot TTS failed: {:?}", command_exit_status); + + let error = ProcessSingleJobError::Other(anyhow!("CommandExitStatus: {:?}", command_exit_status)); + + if let Ok(contents) = read_to_string(&stderr_output_file) { + warn!("Captured stderr output: {}", contents); + + // Re-categorize error? + //match categorize_error(&contents) { + // Some(ProcessSingleJobError::FaceDetectionFailure) => { + // warn!("Face not detected in source image"); + // error = ProcessSingleJobError::FaceDetectionFailure; + // } + // _ => {} + //} + } + + //thread::sleep(Duration::from_secs(300)); + + // Clean up temp files + //safe_delete_temp_file(&audio_path.filesystem_path); + //safe_delete_temp_file(&image_path.filesystem_path); + //safe_delete_temp_file(&usable_image_path); + //safe_delete_temp_file(&output_video_fs_path); + //safe_delete_temp_file(&stderr_output_file); + //safe_delete_temp_directory(&work_temp_dir); + + return Err(error); + } + + // upload audio to bucket + info!("Uploading media ..."); + + let result_bucket_location = MediaFileBucketPath::generate_new( + Some(BUCKET_FILE_PREFIX), + Some(BUCKET_FILE_EXTENSION) + ); + + let result_bucket_object_pathbuf = result_bucket_location.to_full_object_pathbuf(); + + // Finished file path + let mut finished_file = work_temp_dir.path().to_path_buf(); + finished_file.push(&output_file_name); + + info!("Upload Bucket Path: {:?}", result_bucket_object_pathbuf); + info!("Upload File Path: {:?}", finished_file); + + args.job_dependencies.buckets.public_bucket_client + .upload_filename_with_content_type( + &result_bucket_object_pathbuf, + &finished_file, + &MIME_TYPE + ) + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // ==================== UPLOAD AUDIO TO BUCKET ==================== + + info!("Calculating sha256..."); + + let file_checksum = sha256_hash_file(&finished_file).map_err(|err| { + ProcessSingleJobError::Other(anyhow!("Error hashing file: {:?}", err)) + })?; + + let file_size_bytes = file_size(&finished_file).map_err(|err| + ProcessSingleJobError::Other(err) + )?; + + job_progress_reporter.log_status("done").map_err(|e| ProcessSingleJobError::Other(e))?; + + // insert into db the record + let (media_file_token, id) = insert_media_file_from_zero_shot(InsertArgs { + pool: &args.job_dependencies.db.mysql_pool, + job: &job, + maybe_mime_type: Some(&MIME_TYPE), + file_size_bytes, + sha256_checksum: &file_checksum, + public_bucket_directory_hash: result_bucket_location.get_object_hash(), + maybe_public_bucket_prefix: Some(BUCKET_FILE_PREFIX), + maybe_public_bucket_extension: Some(BUCKET_FILE_EXTENSION), + is_on_prem: args.job_dependencies.job.info.container.is_on_prem, + worker_hostname: &args.job_dependencies.job.info.container.hostname, + worker_cluster: &args.job_dependencies.job.info.container.cluster_name, + }).await.map_err(|e| ProcessSingleJobError::Other(e))?; + + info!( + "Job {:?} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}, Result Token: {}", + job.id, + id, + &media_file_token + ); + + Ok(JobSuccessResult { + maybe_result_entity: Some(ResultEntity { + entity_type: InferenceResultType::MediaFile, + entity_token: media_file_token.to_string(), + }), + inference_duration, + }) +} + +pub struct VoiceFile { + pub filesystem_path: PathBuf, +} + +pub async fn download_voice_embedding_from_hash( + bucket_hash: &str, + name: &str, + private_bucket_client: &BucketClient, + path: &PathBuf +) -> Result { + let unifer = BucketPathUnifier::default_paths(); + let object_path = unifer.zero_shot_tts_speaker_encoding(bucket_hash, 0); + + let mut path = path.clone(); + + let file_name = format!("{}", name); + path.push(&file_name); + + private_bucket_client.download_file_to_disk(object_path, &path) + .await + .map_err(|err| { + error!("Could not download embedding file: {err}"); + ProcessSingleJobError::from_anyhow_error(anyhow!("Could not download embedding file: {err}")) + })?; + + let voice_file = VoiceFile { + filesystem_path: PathBuf::from(&path.clone()), + }; + + Ok(voice_file) +} + +pub async fn download_voice_embedding( + voice: &ZsVoice, + private_bucket_client: &BucketClient, + download_path: &PathBuf, +) -> Result { + + let embedding_bucket_location = ZeroShotVoiceEmbeddingBucketPath::from_object_hash( + &voice.bucket_hash, + ModelCategory::Tts, + ModelType::StyleTTS2, + voice.model_version + ); + + info!("Downloading embedding from: {:?}", &embedding_bucket_location.to_full_object_pathbuf()); + info!("Downloading to filesystem location: {:?}", download_path); + + let result = private_bucket_client.download_file_to_disk( + embedding_bucket_location.to_full_object_pathbuf(), + &download_path + ).await; + + if let Err(err) = result { + error!("could not download embedding file: {:?}", err); + return Err(ProcessSingleJobError::from_anyhow_error(err)); + } + + let voice_file = VoiceFile { + filesystem_path: PathBuf::from(&download_path.clone()), + }; + + Ok(voice_file) +} + diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_dependencies.rs new file mode 100644 index 0000000000..eb96c87c7a --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_dependencies.rs @@ -0,0 +1,20 @@ +use errors::AnyhowResult; + +use crate::job::job_types::tts::styletts2::model_downloaders::StyleTTS2Downloaders; +use crate::job::job_types::tts::styletts2::styletts2_inference_command::{StyleTTS2CreateEmbeddingCommand, StyleTTS2InferenceCommand}; + +pub struct StyleTTS2Dependencies { + pub create_embedding_command: StyleTTS2CreateEmbeddingCommand, + pub downloaders: StyleTTS2Downloaders, + pub inference_command: StyleTTS2InferenceCommand, +} + +impl StyleTTS2Dependencies { + pub fn setup() -> AnyhowResult { + Ok(Self { + create_embedding_command: StyleTTS2CreateEmbeddingCommand::from_env()?, + downloaders: StyleTTS2Downloaders::build_all_from_env(), + inference_command: StyleTTS2InferenceCommand::from_env()?, + }) + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs new file mode 100644 index 0000000000..f23fe05c09 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs @@ -0,0 +1,549 @@ +use std::collections::HashSet; +use std::env; +use std::ffi::OsString; +use std::fs::File; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use anyhow::anyhow; +use log::info; +use once_cell::sync::Lazy; +use subprocess::{Popen, PopenConfig, Redirection}; + +use errors::AnyhowResult; +use filesys::path_to_string::path_to_string; +use subprocess_common::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions}; + +use crate::job::job_loop::command_exit_status::CommandExitStatus; + +#[derive(Clone)] +pub enum ExecutableOrCommand { + /// Eg. `inference.py` + Executable(PathBuf), + /// Eg. `python3 inference.py` + Command(String), +} + +// These environment vars are not copied over to the subprocess +// TODO/FIXME(bt, 2023-05-28): This is horrific security! +static IGNORED_ENVIRONMENT_VARS: Lazy> = Lazy::new(|| { + let env_var_names = ["MYSQL_URL", "ACCESS_KEY", "SECRET_KEY", "NEWRELIC_API_KEY"]; + + env_var_names + .iter() + .map(|value| value.to_string()) + .collect::>() +}); + +// The inference command assumes everything has been setup to run inference +// so we have all the files downloaded for the inputs. +pub struct InferenceArgs> { + /// --driven_audio: path to the input audio + pub input_embedding_path: P, // name of the embedding.npz in the work dir + pub input_embedding_name: String, + pub input_text: String, // name of the text + pub output_file_name: String, // output file name in the output folder + pub stderr_output_file: P, +} + +pub struct CreateVoiceInferenceArgs> { + pub output_embedding_path: P, + pub output_embedding_name: P, + pub audio_files: String, + pub stderr_output_file: P, +} + +#[derive(Clone)] +pub struct StyleTTS2InferenceCommand { + /// Where the code lives + root_code_directory: PathBuf, + /// A single executable script or a much larger bash command. + executable_or_command: ExecutableOrCommand, + /// eg. `source python/bin/activate` + maybe_virtual_env_activation_command: Option, + /// Optional default config file to use + maybe_default_config_path: Option, + /// If this is run under Docker (eg. in development), these are the options. + maybe_docker_options: Option, + /// If the execution should be ended after a certain point. + maybe_execution_timeout: Option, + /// Inference arg. + /// --checkpoint_dir: optional location for checkpoints directory + pub alternate_checkpoint_dir: Option, +} + +impl StyleTTS2InferenceCommand { + pub fn new>( + root_code_directory: P, + executable_or_command: ExecutableOrCommand, + maybe_virtual_env_activation_command: Option<&str>, + maybe_default_config_path: Option

, + maybe_docker_options: Option, + maybe_execution_timeout: Option, + alternate_checkpoint_dir: Option + ) -> Self { + Self { + root_code_directory: root_code_directory.as_ref().to_path_buf(), + executable_or_command, + maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| + s.to_string() + ), + maybe_default_config_path: maybe_default_config_path.map(|p| p.as_ref().to_path_buf()), + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + } + } + + pub fn from_env() -> AnyhowResult { + let root_code_directory = easyenv::get_env_pathbuf_required( + "STYLE_TTS2_INFERENCE_ROOT_DIRECTORY" + )?; + + let maybe_inference_command = easyenv::get_env_string_optional( + "STYLE_TTS2_INFERENCE_COMMAND" + ); + + // EVENTUALLY should remove the check + // Optional, eg. `./infer.py`. Typically we'll use the command form instead. + let maybe_inference_executable = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_INFERENCE_EXECUTABLE" + ); + + let executable_or_command = match maybe_inference_command { + Some(command) => ExecutableOrCommand::Command(command), + None => + match maybe_inference_executable { + Some(executable) => ExecutableOrCommand::Executable(executable), + None => { + return Err(anyhow!("neither command nor executable passed")); + } + } + }; + + let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( + "STYLE_TTS2_INFERENCE_MAYBE_VENV_COMMAND" + ); + + let maybe_default_config_path = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_INFERENCE_MAYBE_DEFAULT_CONFIG_PATH" + ); + + let maybe_execution_timeout = easyenv::get_env_duration_seconds_optional("TIMEOUT_SECONDS"); + + // Probably for local + let maybe_docker_options = easyenv + ::get_env_string_optional("STYLE_TTS2_INFERENCE_MAYBE_DOCKER_IMAGE") + .map(|image_name| { + DockerOptions { + image_name, + maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()), + maybe_environment_variables: None, + maybe_gpu: Some(DockerGpu::All), + } + }); + + // Override for --checkpoint_dir at inference time + let alternate_checkpoint_dir = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_ALTERNATE_CHECKPOINT_PATH" + ); + + Ok(Self { + root_code_directory, + executable_or_command, + maybe_virtual_env_activation_command, + maybe_default_config_path, + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + }) + } + + pub fn execute_inference>(&self, args: InferenceArgs

) -> CommandExitStatus { + match self.do_execute_inference(args) { + Ok(exit_status) => exit_status, + Err(error) => + CommandExitStatus::FailureWithReason { reason: format!("error: {:?}", error) }, + } + } + + fn do_execute_inference>( + &self, + args: InferenceArgs

+ ) -> AnyhowResult { + let mut command = String::new(); + command.push_str(&format!("cd {}", path_to_string(&self.root_code_directory))); + + if let Some(venv_command) = self.maybe_virtual_env_activation_command.as_deref() { + command.push_str(" && "); + command.push_str(venv_command); + command.push_str(" "); + } + + command.push_str(" && "); + + match self.executable_or_command { + ExecutableOrCommand::Executable(ref executable) => { + command.push_str(&path_to_string(executable)); + command.push_str(" "); + } + ExecutableOrCommand::Command(ref cmd) => { + command.push_str(cmd); + command.push_str(" "); + } + } + + // ===== Begin Python Args ===== + + command.push_str(" --text "); + command.push_str(&format!("\"{}\"", &args.input_text)); + + command.push_str(" --prompt-name "); + command.push_str(&path_to_string(args.input_embedding_name)); + + command.push_str(" --prompt-path "); + command.push_str(&path_to_string(&args.input_embedding_path)); + + command.push_str(" --audio-name "); + command.push_str(&path_to_string(args.output_file_name)); + + command.push_str(" --audio-path "); + command.push_str(&path_to_string(&args.input_embedding_path)); + + command.push_str(" --mode "); + command.push_str(&path_to_string("0")); + // TODO improve and use a better model for premium users + command.push_str(" --whisper-model "); + command.push_str(&path_to_string(Path::new("medium"))); + + command.push_str(" --whisper-folder-path "); + command.push_str( + &path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper/")) + ); + + command.push_str(" --vocos-folder-path "); + command.push_str( + &path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/")) + ); + + command.push_str(" --StyleTTS2-path "); + command.push_str( + &path_to_string( + Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/StyleTTS2-checkpoint.pt") + ) + ); + + command.push_str(" --tmp-work-dir "); + command.push_str(&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/"))); + + // ===== End Python Args ===== + + if let Some(docker_options) = self.maybe_docker_options.as_ref() { + command = docker_options.to_command_string(&command); + } + + info!("Command: {:?}", command); + + let command_parts = ["bash", "-c", &command]; + + let mut env_vars = Vec::new(); + + // Copy all environment variables from the parent process. + // This is necessary to send all the kubernetes settings for Nvidia / CUDA. + for (env_key, env_value) in env::vars() { + if IGNORED_ENVIRONMENT_VARS.contains(&env_key) { + continue; + } + env_vars.push((OsString::from(env_key), OsString::from(env_value))); + } + + let mut config = PopenConfig::default(); + + info!("stderr will be written to file: {:?}", args.stderr_output_file.as_ref()); + + let stderr_file = File::create(&args.stderr_output_file)?; + config.stderr = Redirection::File(stderr_file); + + if !env_vars.is_empty() { + config.env = Some(env_vars); + } + + let mut p = Popen::create(&command_parts, config)?; + + info!("Subprocess PID: {:?}", p.pid()); + + match self.maybe_execution_timeout { + None => { + let exit_status = p.wait()?; + info!("Subprocess exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + Some(timeout) => { + info!("Executing with timeout: {:?}", &timeout); + let exit_status = p.wait_timeout(timeout.clone())?; + + match exit_status { + None => { + // NB: If the program didn't successfully terminate, kill it. + info!( + "Subprocess didn't end after timeout: {:?}; terminating...", + &timeout + ); + let _r = p.terminate()?; + Ok(CommandExitStatus::Timeout) + } + Some(exit_status) => { + info!("Subprocess timed wait exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + } + } + } + } +} + +#[derive(Clone)] +pub struct StyleTTS2CreateEmbeddingCommand { + /// Where the code lives + root_code_directory: PathBuf, + /// A single executable script or a much larger bash command. + executable_or_command: ExecutableOrCommand, + /// eg. `source python/bin/activate` + maybe_virtual_env_activation_command: Option, + /// Optional default config file to use + maybe_default_config_path: Option, + /// If this is run under Docker (eg. in development), these are the options. + maybe_docker_options: Option, + /// If the execution should be ended after a certain point. + maybe_execution_timeout: Option, + /// Inference arg. + /// --checkpoint_dir: optional location for checkpoints directory + pub alternate_checkpoint_dir: Option, +} + +impl StyleTTS2CreateEmbeddingCommand { + pub fn new>( + root_code_directory: P, + executable_or_command: ExecutableOrCommand, + maybe_virtual_env_activation_command: Option<&str>, + maybe_default_config_path: Option

, + maybe_docker_options: Option, + maybe_execution_timeout: Option, + alternate_checkpoint_dir: Option + ) -> Self { + Self { + root_code_directory: root_code_directory.as_ref().to_path_buf(), + executable_or_command, + maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| + s.to_string() + ), + maybe_default_config_path: maybe_default_config_path.map(|p| p.as_ref().to_path_buf()), + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + } + } + + // helper funciton to set up env + pub fn from_env() -> AnyhowResult { + // TODO: fix for now just for compilation + let root_code_directory = easyenv::get_env_pathbuf_required( + "STYLE_TTS2_INFERENCE_ROOT_DIRECTORY" + )?; + + let maybe_inference_command = easyenv::get_env_string_optional( + "STYLE_TTS2_INFERENCE_COMMAND" + ); + + // EVENTUALLY should remove the check + // Optional, eg. `./infer.py`. Typically we'll use the command form instead. + let maybe_inference_executable = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_INFERENCE_EXECUTABLE" + ); + + let executable_or_command = match maybe_inference_command { + Some(command) => ExecutableOrCommand::Command(command), + None => + match maybe_inference_executable { + Some(executable) => ExecutableOrCommand::Executable(executable), + None => { + return Err(anyhow!("neither command nor executable passed")); + } + } + }; + + let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( + "STYLE_TTS2_INFERENCE_MAYBE_VENV_COMMAND" + ); + + let maybe_default_config_path = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_INFERENCE_MAYBE_DEFAULT_CONFIG_PATH" + ); + + let maybe_execution_timeout = easyenv::get_env_duration_seconds_optional("TIMEOUT_SECONDS"); + + // Probably for local + let maybe_docker_options = easyenv + ::get_env_string_optional("STYLE_TTS2_INFERENCE_MAYBE_DOCKER_IMAGE") + .map(|image_name| { + DockerOptions { + image_name, + maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()), + maybe_environment_variables: None, + maybe_gpu: Some(DockerGpu::All), + } + }); + + // Override for --checkpoint_dir at inference time + let alternate_checkpoint_dir = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_CHECKPOINT_PATH" + ); + + Ok(Self { + root_code_directory, + executable_or_command, + maybe_virtual_env_activation_command, + maybe_default_config_path, + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + }) + } + + pub fn execute_inference>(&self, args: CreateVoiceInferenceArgs

) -> CommandExitStatus { + match self.do_execute_inference(args) { + Ok(exit_status) => exit_status, + Err(error) => + CommandExitStatus::FailureWithReason { reason: format!("error: {:?}", error) }, + } + } + + fn do_execute_inference>( + &self, + args: CreateVoiceInferenceArgs

+ ) -> AnyhowResult { + let mut command = String::new(); + command.push_str(&format!("cd {}", path_to_string(&self.root_code_directory))); + + if let Some(venv_command) = self.maybe_virtual_env_activation_command.as_deref() { + command.push_str(" && "); + command.push_str(venv_command); + command.push_str(" "); + } + + command.push_str(" && "); + + match self.executable_or_command { + ExecutableOrCommand::Executable(ref executable) => { + command.push_str(&path_to_string(executable)); + command.push_str(" "); + } + ExecutableOrCommand::Command(ref cmd) => { + command.push_str(cmd); + command.push_str(" "); + } + } + + // ===== Begin Python Args ===== + + command.push_str(" --prompt-name "); + command.push_str(&path_to_string(args.output_embedding_name)); + + command.push_str(" --prompt-path "); + command.push_str(&path_to_string(&args.output_embedding_path)); + + command.push_str(" --audio-wav-files "); + command.push_str(&args.audio_files); + + command.push_str(" --mode "); + command.push_str(&path_to_string("1")); + // TODO improve and use a better model for premium users + command.push_str(" --whisper-model "); + command.push_str(&path_to_string(Path::new("medium"))); + + command.push_str(" --whisper-folder-path "); + command.push_str( + &path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper/")) + ); + + command.push_str(" --vocos-folder-path "); + command.push_str( + &path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/")) + ); + + command.push_str(" --StyleTTS2-path "); + command.push_str( + &path_to_string( + Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/StyleTTS2-checkpoint.pt") + ) + ); + + command.push_str(" --tmp-work-dir "); + command.push_str(&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/"))); + + // ===== End Python Args ===== + + if let Some(docker_options) = self.maybe_docker_options.as_ref() { + command = docker_options.to_command_string(&command); + } + + info!("Command: {:?}", command); + + let command_parts = ["bash", "-c", &command]; + + let mut env_vars = Vec::new(); + + // Copy all environment variables from the parent process. + // This is necessary to send all the kubernetes settings for Nvidia / CUDA. + for (env_key, env_value) in env::vars() { + if IGNORED_ENVIRONMENT_VARS.contains(&env_key) { + continue; + } + env_vars.push((OsString::from(env_key), OsString::from(env_value))); + } + + let mut config = PopenConfig::default(); + + info!("stderr will be written to file: {:?}", args.stderr_output_file.as_ref()); + + let stderr_file = File::create(&args.stderr_output_file)?; + config.stderr = Redirection::File(stderr_file); + + if !env_vars.is_empty() { + config.env = Some(env_vars); + } + + let mut p = Popen::create(&command_parts, config)?; + + info!("Subprocess PID: {:?}", p.pid()); + + match self.maybe_execution_timeout { + None => { + let exit_status = p.wait()?; + info!("Subprocess exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + Some(timeout) => { + info!("Executing with timeout: {:?}", &timeout); + let exit_status = p.wait_timeout(timeout.clone())?; + + match exit_status { + None => { + // NB: If the program didn't successfully terminate, kill it. + info!( + "Subprocess didn't end after timeout: {:?}; terminating...", + &timeout + ); + let _r = p.terminate()?; + Ok(CommandExitStatus::Timeout) + } + Some(exit_status) => { + info!("Subprocess timed wait exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + } + } + } + } + +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/validate_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/validate_job.rs new file mode 100644 index 0000000000..41b76feac1 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/validate_job.rs @@ -0,0 +1,84 @@ +use anyhow::anyhow; + +use mysql_queries::payloads::generic_inference_args::generic_inference_args::{ + InferenceCategoryAbbreviated, + PolymorphicInferenceArgs, +}; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; + +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; + +pub enum JobType { + Create, + Inference +} +// determines the job type voice inference or voice creation +pub struct JobArgs { + pub voice_token: Option, // infers + pub voice_dataset_token: Option, // creates + pub job_type: JobType +} + +pub fn validate_job(job: &AvailableInferenceJob) -> Result { + let inference_category = job.maybe_inference_args + .as_ref() + .map(|args| args.inference_category) + .flatten(); + + match inference_category { + Some(InferenceCategoryAbbreviated::TextToSpeech) => {} + Some(category) => { + return Err( + ProcessSingleJobError::from_anyhow_error( + anyhow!("wrong inference category for job: {:?}", category) + ) + ); + } + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no inference category for job!")) + ); + } + } + + let inference_args = job.maybe_inference_args + .as_ref() + .map(|args| args.args.as_ref()) + .flatten(); + + let inference_args = match inference_args { + Some(args) => args, + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no inference args for job!")) + ); + } + }; + + let ttsArgs = match inference_args { + PolymorphicInferenceArgs::Tts(inference_args) => inference_args, + _ => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!")) + ); + } + }; + + if let Some(voice_token) = &ttsArgs.voice_token { + return Ok::(JobArgs { + voice_token: Some(voice_token.clone()), + voice_dataset_token: None, + job_type: JobType::Inference + }); + } + + if let Some(dataset_token) = &ttsArgs.dataset_token { + return Ok::(JobArgs { + voice_token: None, + voice_dataset_token: Some(dataset_token.clone()), + job_type: JobType::Create + }); + } + + Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Missing Voice Token!"))) +} diff --git a/crates/service/job/inference_job/src/job_specific_dependencies.rs b/crates/service/job/inference_job/src/job_specific_dependencies.rs index a17b4ef472..1263583493 100644 --- a/crates/service/job/inference_job/src/job_specific_dependencies.rs +++ b/crates/service/job/inference_job/src/job_specific_dependencies.rs @@ -2,6 +2,7 @@ use enums::by_table::generic_inference_jobs::inference_model_type::InferenceMode use errors::AnyhowResult; use crate::job::job_types::lipsync::sad_talker::sad_talker_dependencies::SadTalkerDependencies; +use crate::job::job_types::tts::styletts2::styletts2_dependencies::StyleTTS2Dependencies; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_dependencies::Tacotron2Dependencies; use crate::job::job_types::tts::vall_e_x::vall_e_x_dependencies::VallExDependencies; use crate::job::job_types::tts::vits::vits_dependencies::VitsDependencies; @@ -20,6 +21,7 @@ pub struct JobSpecificDependencies { pub maybe_vits_dependencies: Option, pub maybe_rerender_dependencies: Option, pub maybe_mocapnet_dependencies: Option, + pub maybe_styletts2_dependencies: Option, } impl JobSpecificDependencies { @@ -33,6 +35,7 @@ impl JobSpecificDependencies { let mut maybe_vits_dependencies = None; let mut maybe_rerender_dependencies = None; let mut maybe_mocapnet_dependencies = None; + let mut maybe_styletts2_dependencies = None; if scoped_execution.can_run_job(InferenceModelType::RvcV2) { print_with_space("Setting RVCv2 dependencies..."); @@ -59,6 +62,11 @@ impl JobSpecificDependencies { maybe_vall_e_x_dependencies = Some(VallExDependencies::setup()?); } + if scoped_execution.can_run_job(InferenceModelType::StyleTTS2) { + print_with_space("Setting StyleTTS2 dependencies..."); + maybe_styletts2_dependencies = Some(StyleTTS2Dependencies::setup()?); + } + if scoped_execution.can_run_job(InferenceModelType::Vits) { print_with_space("Setting Vits dependencies..."); maybe_vits_dependencies = Some(VitsDependencies::setup()?); @@ -82,7 +90,8 @@ impl JobSpecificDependencies { maybe_vall_e_x_dependencies, maybe_vits_dependencies, maybe_rerender_dependencies, - maybe_mocapnet_dependencies + maybe_mocapnet_dependencies, + maybe_styletts2_dependencies, }) } } From ea3a6d3c0b6ce7a9914b656bda2e215cbf3508eb Mon Sep 17 00:00:00 2001 From: Kasisnu <5044120+kasisnu@users.noreply.github.com> Date: Fri, 12 Jan 2024 17:58:12 +0530 Subject: [PATCH 24/54] pass appropriate arguments to styletts inference command --- .../tts/styletts2/model_downloaders.rs | 15 +--- .../tts/styletts2/process_job_create_voice.rs | 45 +++------- .../styletts2/process_job_inference_voice.rs | 44 +++------- .../styletts2/styletts2_inference_command.rs | 88 ++++--------------- 4 files changed, 45 insertions(+), 147 deletions(-) diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs index 18bcc1c931..c50a6dcfff 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs @@ -1,11 +1,7 @@ use crate::util::model_downloader::ModelDownloader; pub struct StyleTTS2Downloaders { - pub styletts2: StyleTTS2ModelMapping, - pub whisper_lg: WhisperLargeModelMapping, - pub whisper_md: WhisperMediumModelMapping, - pub vocos_encodec_24khz: VocosEncodecModelMapping, - pub vocos_config: VocosEncodecConfigMapping + // pub styletts2: StyleTTS2ModelMapping, } crate::impl_model_downloader!( @@ -54,18 +50,15 @@ crate::impl_model_downloader!( ); impl StyleTTS2Downloaders { + // TODO(KS): All checkpoints are in the container, so we don't need to download them here. + // Maybe they should be here instead pub fn build_all_from_env() -> Self { Self { - styletts2: StyleTTS2ModelMapping::from_env(), - whisper_lg: WhisperLargeModelMapping::from_env(), - whisper_md: WhisperMediumModelMapping::from_env(), - vocos_encodec_24khz: VocosEncodecModelMapping::from_env(), - vocos_config: VocosEncodecConfigMapping::from_env(), } } pub fn all_downloaders(&self) -> Vec<&dyn ModelDownloader> { - vec![&self.styletts2, &self.whisper_lg, &self.whisper_md, &self.vocos_encodec_24khz, &self.vocos_config] + vec![] } } diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs index d64cc4f809..3c890b3cea 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs @@ -113,20 +113,10 @@ pub async fn process_create_voice( info!("Dataset length info: {}", dataset.len()); -// let temp_extension = String::from(".bin"); -// let temp_prefix: String; -// -// if !deps.job.info.container.is_on_prem { -// temp_prefix = String::from("sample_"); // this is for seed in local dev to download the samples -// } else { -// temp_prefix = String::from(BUCKET_FILE_PREFIX_CREATE); -// } - let mut downloaded_dataset: Vec = Vec::new(); for (index, record) in dataset.iter().enumerate() { //https://storage.googleapis.com/dev-vocodes-public/media/5/3/3/w/8/533w8zs0fy11nv7gkcna7p7vt03h8nda/dev_zs_533w8zs0fy11nv7gkcna7p7vt03h8nda.bin <-- where the file actually is - let prefix: Option<&str> = record.maybe_public_bucket_prefix.as_ref().map(|s| s.as_str()); let extension: Option<&str> = record.maybe_public_bucket_extension.as_ref().map(|s| s.as_str()); @@ -166,9 +156,16 @@ pub async fn process_create_voice( info!("FilePath to clone voice: {}", file_path.to_string_lossy()); downloaded_dataset.push(file_path.clone()); + + + // TODO(KS): This is a hack to get only the first file to download + // Ideally we could have improved logic to pick the "best" media file from a dataset + if (index > 0) { + break; + } } - info!("Dataset Length {}", downloaded_dataset.len()); + info!("Dataset Length downloaded: {}", downloaded_dataset.len()); // STEP 4 Download the models info!("Download models (if not present)..."); @@ -193,9 +190,7 @@ pub async fn process_create_voice( info!("Files to process: {:?}", audio_files); // Name of the output file - // NB: don't use the extension... for the inference since the container will add the extension. - let output_file_name = PathBuf::from("temp"); - + let output_embedding_path = work_temp_dir.path().join("style.npz"); let stderr_output_file = work_temp_dir.path().join("zero_shot_create_voice_err.txt"); let inference_start_time = Instant::now(); @@ -204,9 +199,8 @@ pub async fn process_create_voice( let command_exit_status = model_dependencies.create_embedding_command.execute_inference( job::job_types::tts::styletts2::styletts2_inference_command::CreateVoiceInferenceArgs { - output_embedding_path: &workdir, - output_embedding_name: &output_file_name, - audio_files, + output_embedding_path: &output_embedding_path, + audio_file: audio_files, stderr_output_file: &stderr_output_file, } ); @@ -220,15 +214,6 @@ pub async fn process_create_voice( if let Ok(contents) = read_to_string(&stderr_output_file) { warn!("Captured stderr output: {}", contents); - - // Re-categorize error? - //match categorize_error(&contents) { - // Some(ProcessSingleJobError::FaceDetectionFailure) => { - // warn!("Face not detected in source image"); - // error = ProcessSingleJobError::FaceDetectionFailure; - // } - // _ => {} - //} } //thread::sleep(Duration::from_secs(300)); @@ -246,13 +231,7 @@ pub async fn process_create_voice( info!("Inference success!"); - //info!("Success; waiting..."); - //thread::sleep(Duration::from_secs(300)); - - - // STEP 4. Download dataset each audio file info!("Uploading Media ..."); - let embedding_bucket_location = ZeroShotVoiceEmbeddingBucketPath::generate_new( ModelCategory::Tts, ModelType::StyleTTS2, @@ -264,7 +243,7 @@ pub async fn process_create_voice( // Get Finished File let mut finished_file = work_temp_dir.path().to_path_buf(); - let output_bucket_file_name = String::from("temp.npz"); // use extension for bucket upload. + let output_bucket_file_name = String::from("style.npz"); // use extension for bucket upload. finished_file.push(&output_bucket_file_name); info!("Upload File Path: {:?}", finished_file); diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs index cb8d476954..b7fcd141de 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs @@ -16,6 +16,7 @@ use hashing::sha256::sha256_hash_file::sha256_hash_file; use mysql_queries::queries::media_files::create::insert_media_file_from_zero_shot_tts::insert_media_file_from_zero_shot; use mysql_queries::queries::media_files::create::insert_media_file_from_zero_shot_tts::InsertArgs; use mysql_queries::queries::voice_designer::voices::get_voice::{get_voice_by_token, ZsVoice}; +use tts_common::clean_symbols::clean_symbols; use crate::job::job_loop::job_success_result::JobSuccessResult; use crate::job::job_loop::job_success_result::ResultEntity; @@ -52,7 +53,9 @@ pub async fn process_inference_voice( // get job args let text = match job.maybe_raw_inference_text.clone() { - Some(value) => { value } + Some(value) => { + clean_symbols(value.as_str()) + } None => { return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Text for Inference"))); } @@ -90,17 +93,6 @@ pub async fn process_inference_voice( } } - // Might not need this for inference. - // 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"))); - // } - // } - // run inference let work_temp_dir = format!("/tmp/temp_zeroshot_inference_{}", job.id.0); @@ -110,7 +102,7 @@ pub async fn process_inference_voice( .map_err(|e| ProcessSingleJobError::from_io_error(e))?; let workdir = work_temp_dir.path().to_path_buf(); - let filename = "weights.npz".to_string(); + let filename = "style.npz".to_string(); let mut downloaded_weights_path = work_temp_dir.path().to_path_buf(); downloaded_weights_path.push(&filename); @@ -129,20 +121,23 @@ pub async fn process_inference_voice( .log_status("running inference") .map_err(|e| ProcessSingleJobError::Other(e))?; + let text_input_fs_path = work_temp_dir.path().join("inference_input.txt"); + + std::fs::write(&text_input_fs_path, &text) + .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + let inference_start_time = Instant::now(); - let output_file_name = String::from("output.wav"); - + let output_file_name = work_temp_dir.path().join("output.wav"); let stderr_output_file = work_temp_dir.path().join("zero_shot_inference.txt"); // Run Inference let command_exit_status = model_dependencies.inference_command.execute_inference( InferenceArgs { - input_embedding_path: &workdir, - input_embedding_name: filename, - input_text: String::from(text), // text - output_file_name: output_file_name.clone(), // output file name in the output folder + input_embedding_file_path: &workdir, + input_text_file_path: &text_input_fs_path, // text + output_file_name: &output_file_name, // output file name in the output folder stderr_output_file: &stderr_output_file, } ); @@ -156,19 +151,8 @@ pub async fn process_inference_voice( if let Ok(contents) = read_to_string(&stderr_output_file) { warn!("Captured stderr output: {}", contents); - - // Re-categorize error? - //match categorize_error(&contents) { - // Some(ProcessSingleJobError::FaceDetectionFailure) => { - // warn!("Face not detected in source image"); - // error = ProcessSingleJobError::FaceDetectionFailure; - // } - // _ => {} - //} } - //thread::sleep(Duration::from_secs(300)); - // Clean up temp files //safe_delete_temp_file(&audio_path.filesystem_path); //safe_delete_temp_file(&image_path.filesystem_path); diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs index f23fe05c09..bda3b67f5d 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs @@ -39,17 +39,15 @@ static IGNORED_ENVIRONMENT_VARS: Lazy> = Lazy::new(|| { // so we have all the files downloaded for the inputs. pub struct InferenceArgs> { /// --driven_audio: path to the input audio - pub input_embedding_path: P, // name of the embedding.npz in the work dir - pub input_embedding_name: String, - pub input_text: String, // name of the text - pub output_file_name: String, // output file name in the output folder + pub input_embedding_file_path: P, // name of the embedding.npz in the work dir + pub input_text_file_path: P, // name of the text + pub output_file_name: P, // output file name in the output folder pub stderr_output_file: P, } pub struct CreateVoiceInferenceArgs> { pub output_embedding_path: P, - pub output_embedding_name: P, - pub audio_files: String, + pub audio_file: String, pub stderr_output_file: P, } @@ -195,46 +193,19 @@ impl StyleTTS2InferenceCommand { // ===== Begin Python Args ===== - command.push_str(" --text "); - command.push_str(&format!("\"{}\"", &args.input_text)); + command.push_str(" --input "); + command.push_str(&path_to_string(args.input_text_file_path)); - command.push_str(" --prompt-name "); - command.push_str(&path_to_string(args.input_embedding_name)); + command.push_str(" --input-style-npz "); + command.push_str(&path_to_string(args.input_embedding_file_path)); - command.push_str(" --prompt-path "); - command.push_str(&path_to_string(&args.input_embedding_path)); - command.push_str(" --audio-name "); + command.push_str(" --output "); command.push_str(&path_to_string(args.output_file_name)); - command.push_str(" --audio-path "); - command.push_str(&path_to_string(&args.input_embedding_path)); - - command.push_str(" --mode "); - command.push_str(&path_to_string("0")); - // TODO improve and use a better model for premium users - command.push_str(" --whisper-model "); - command.push_str(&path_to_string(Path::new("medium"))); - - command.push_str(" --whisper-folder-path "); - command.push_str( - &path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper/")) - ); - - command.push_str(" --vocos-folder-path "); - command.push_str( - &path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/")) - ); - - command.push_str(" --StyleTTS2-path "); - command.push_str( - &path_to_string( - Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/StyleTTS2-checkpoint.pt") - ) - ); - - command.push_str(" --tmp-work-dir "); - command.push_str(&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/"))); + //TODO(KS): Make this a parameter + command.push_str(" --vcsteps "); + command.push_str(format!("{}", 20).as_str()); // ===== End Python Args ===== @@ -446,40 +417,11 @@ impl StyleTTS2CreateEmbeddingCommand { // ===== Begin Python Args ===== - command.push_str(" --prompt-name "); - command.push_str(&path_to_string(args.output_embedding_name)); - - command.push_str(" --prompt-path "); + command.push_str(" --output-style-npz "); command.push_str(&path_to_string(&args.output_embedding_path)); - command.push_str(" --audio-wav-files "); - command.push_str(&args.audio_files); - - command.push_str(" --mode "); - command.push_str(&path_to_string("1")); - // TODO improve and use a better model for premium users - command.push_str(" --whisper-model "); - command.push_str(&path_to_string(Path::new("medium"))); - - command.push_str(" --whisper-folder-path "); - command.push_str( - &path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper/")) - ); - - command.push_str(" --vocos-folder-path "); - command.push_str( - &path_to_string(Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/")) - ); - - command.push_str(" --StyleTTS2-path "); - command.push_str( - &path_to_string( - Path::new("/tmp/downloads/zero_shot_tts/vall-e-x_1.0/StyleTTS2-checkpoint.pt") - ) - ); - - command.push_str(" --tmp-work-dir "); - command.push_str(&path_to_string(Path::new("/tmp/downloads/zero_shot_tts/"))); + command.push_str(" --voice "); + command.push_str(&args.audio_file); // ===== End Python Args ===== From ce852eda8aeff1d42caa5d0f59b624f008a944cd Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 12 Jan 2024 15:51:50 -0500 Subject: [PATCH 25/54] netlify feature branch cors --- crates/lib/deprecated/http_server_common/src/cors.rs | 2 ++ .../src/http_server/endpoints/weights/get_weight.rs | 2 +- .../src/http_server/endpoints/weights/list_weights_by_user.rs | 2 +- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/crates/lib/deprecated/http_server_common/src/cors.rs b/crates/lib/deprecated/http_server_common/src/cors.rs index 95f51ff0e9..a72c174710 100644 --- a/crates/lib/deprecated/http_server_common/src/cors.rs +++ b/crates/lib/deprecated/http_server_common/src/cors.rs @@ -59,6 +59,8 @@ pub fn add_fakeyou(cors: Cors, is_production: bool) -> Cors { // FakeYou (Staging) .allowed_origin("http://staging.fakeyou.com") .allowed_origin("https://staging.fakeyou.com") + // FakeYou (Netlify Staging / Production) + .allowed_origin("https://feature-mvp--fakeyou.netlify.app") } else { cors // FakeYou (Development) diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/get_weight.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/get_weight.rs index 4197638c5b..28f1500449 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/get_weight.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/get_weight.rs @@ -42,7 +42,7 @@ pub struct GetWeightResponse { file_checksum_sha2: String, /// Information about the cover image. - pub cover_image: CoverImageDetails, + cover_image: CoverImageDetails, /// Cover images are small descriptive images that can be set for any model. /// If a cover image is set, this is the path to the asset. diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_weights_by_user.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_weights_by_user.rs index 08ec39718b..b73e722da8 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_weights_by_user.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_weights_by_user.rs @@ -41,7 +41,7 @@ pub struct Weight { file_checksum_sha2: String, /// Information about the cover image. - pub cover_image: CoverImageDetails, + cover_image: CoverImageDetails, /// Cover images are small descriptive images that can be set for any model. /// If a cover image is set, this is the path to the asset. From 447a457dffcc50d9576054982f67b8e2b1388cb4 Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Sat, 13 Jan 2024 00:26:51 -0500 Subject: [PATCH 26/54] image gen should build. --- .../image_generation/sd/process_job.rs | 34 ++++++++++++------- .../sd/sd_inference_command.rs | 24 ++++++------- 2 files changed, 33 insertions(+), 25 deletions(-) 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 a2e0abdd31..db653be265 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,3 +1,5 @@ +use std::ops::Deref; +use std::path::PathBuf; use std::time::Duration; use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails; use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; @@ -21,9 +23,6 @@ pub struct StableDiffusionProcessArgs<'a> { pub job: &'a AvailableInferenceJob, } -// run inference -// upload inference result -// upload model checkpoint or loRA // run inference // insert record into the db with the inference job token complete. @@ -128,12 +127,13 @@ pub async fn process_job_inference( .new_tempdir(&work_temp_dir) .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + let sd_checkpoint_path = work_temp_dir.path().join("sd_checkpoint"); let lora_path = work_temp_dir.path().join("lora"); let vae_path = work_temp_dir.path().join("vae"); let output_path = work_temp_dir.path().join("output"); - println!("Pathes to download to"); + println!("Pathes to download to:"); println!("sd_checkpoint_path:{}",sd_checkpoint_path.to_string()); println!("lora_path:{}",lora_path.to_string()); println!("vae_path:{}",vae_path.to_string()); @@ -217,8 +217,6 @@ pub async fn process_job_inference( None => {} } - - // VAE token for now let vae_token = String::from("weight_rb0959wfzjhk3d1k93hr3s0qw"); let model_weight_vae = ModelWeightToken(vae_token); @@ -261,28 +259,38 @@ pub async fn process_job_inference( None => Err(ProcessSingleJobError::from_anyhow_error(anyhow!("No Prompt provided!"))) }; + let stderr_output_file = work_temp_dir.path().join("zero_shot_create_voice_err.txt"); + + let number_of_samples = match sd_args.maybe_number_of_samples { + Some(val) => { + val + }, + None => { + 20 + } + }; + // #TODO allow for batch and saving later. output files sd_deps.inference_command.execute_inference(InferenceArgs { - work_dir: work_temp_dir, + work_dir: work_temp_dir.path().to_path_buf(), output_file: output_path, - stderr_output_file:, - prompt: prompt, + stderr_output_file: stderr_output_file, + prompt: prompt, negative_prompt:sd_args.maybe_n_prompt.unwrap_or_default(), - number_of_samples: 0, + number_of_samples: number_of_samples, samplers: "".to_string(), 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().as_path(), + lora_path: Some(lora_path), checkpoint_path: sd_checkpoint_path.clone(), - vae: vae_path.as_path(), + vae: vae_path.clone(), batch_count: sd_args.maybe_batch_count.unwrap_or(1), }); // upload media and create a record. - Ok(JobSuccessResult { maybe_result_entity: None, inference_duration: Duration::from_secs(0), 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 899f79eb3e..a47ccd0f22 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 @@ -66,13 +66,13 @@ pub enum ExecutableOrCommand { Command(String), } -pub struct InferenceArgs> { +pub struct InferenceArgs { /// --source_image: path to the input image (or video) /// --result_dir: path to directory work is performed - pub work_dir: P, + pub work_dir: PathBuf, /// --result_file: path to final file output - pub output_file: P, - pub stderr_output_file: P, + pub output_file: PathBuf, + pub stderr_output_file: PathBuf, pub prompt: String, pub negative_prompt:String, pub number_of_samples:i32, @@ -80,10 +80,10 @@ pub struct InferenceArgs> { pub width:i32, pub height:i32, pub cfg_scale:i32, - pub seed:i32, - pub lora_path:P, - pub checkpoint_path:P, - pub vae:P, + pub seed:i64, + pub lora_path:Option, + pub checkpoint_path:PathBuf, + pub vae:PathBuf, pub batch_count:i32, } @@ -158,9 +158,9 @@ impl StableDiffusionInferenceCommand { }) } - pub fn execute_inference>( + pub fn execute_inference( &self, - args: InferenceArgs

, + args: InferenceArgs, ) -> CommandExitStatus { match self.do_execute_inference(args) { Ok(exit_status) => exit_status, @@ -168,9 +168,9 @@ impl StableDiffusionInferenceCommand { } } - fn do_execute_inference>( + fn do_execute_inference( &self, - args: InferenceArgs

, + args: InferenceArgs, ) -> AnyhowResult { let mut command = String::new(); From 108806cb5231bec4cef6e9008705671e98934d20 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 13 Jan 2024 01:00:55 -0500 Subject: [PATCH 27/54] list media files - MediaFileOriginDetails + title --- _sql/migrations_squashed/media_files.sql | 9 +- .../media_file_origin_product_category.rs | 25 ++--- .../media_files/list/list_media_files.rs | 12 +- .../web/storyteller_web/src/api_doc.rs | 3 + .../media_file_origin_details.rs | 104 ++++++++++++++++++ .../src/http_server/common_responses/mod.rs | 1 + .../endpoints/media_files/list_media_files.rs | 20 +++- 7 files changed, 156 insertions(+), 18 deletions(-) create mode 100644 crates/service/web/storyteller_web/src/http_server/common_responses/media_file_origin_details.rs diff --git a/_sql/migrations_squashed/media_files.sql b/_sql/migrations_squashed/media_files.sql index 90fc13febb..e04cb2de30 100644 --- a/_sql/migrations_squashed/media_files.sql +++ b/_sql/migrations_squashed/media_files.sql @@ -26,7 +26,7 @@ CREATE TABLE media_files ( -- * 'device_api' for direct user uploads recorded using Browser/Device APIs. origin_category VARCHAR(16) NOT NULL, - -- TODO: Remove default value once records backfilled and queries updated. + -- TODO(bt,2024-01-12): Rename to origin_product. -- Product area where the media file originated. -- This is not the *model* that created the thing, this is the *product*. -- (The underlying models can change over time.) @@ -37,8 +37,13 @@ CREATE TABLE media_files ( -- "Voice Designer" dataset samples in a video generation flow.) -- -- Possible values: + -- * 'unknown' for legacy records without an associated product -- * 'face_animator' for uploads and outputs - -- * 'unknown' for legacy records without an associated product (TODO: This should be temporary) + -- * 'tts' for text to speech outputs + -- * 'voice_conversion' for uploads or outputs (RVC or SVC) + -- * 'zs_voice' for uploads or outputs for zero shot voice products + -- * 'video_filter' filters on videos + -- * 'mocap' for files uploaded or processed by motion capture origin_product_category VARCHAR(16) NOT NULL DEFAULT "unknown", -- For inference that can be tied back to a model, the type of model. 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 bc99d59e07..f9d258f356 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 @@ -17,6 +17,11 @@ use utoipa::ToSchema; #[cfg_attr(test, derive(EnumIter, EnumCount))] #[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Deserialize, Serialize, ToSchema)] pub enum MediaFileOriginProductCategory { + /// Unknown which product is attached to the file (generated the file, the file was + /// uploaded on behalf of, etc.) + #[serde(rename = "unknown")] + Unknown, + /// Media files created by (or uploaded for) the Face Animator product. /// The underlying model could be SadTalker, Wav2Lip, or some future model #[serde(rename = "face_animator")] @@ -26,12 +31,6 @@ pub enum MediaFileOriginProductCategory { #[serde(rename = "tts")] TextToSpeech, - // TODO: This should be a temporary category until we migrate the DB to remove this default value - /// Unknown which product is attached to the file (generated the file, the file was - /// uploaded on behalf of, etc.) - #[serde(rename = "unknown")] - Unknown, - /// Voice conversion (either RVC or SVC) #[serde(rename = "voice_conversion")] VoiceConversion, @@ -58,9 +57,9 @@ impl_mysql_from_row!(MediaFileOriginProductCategory); impl MediaFileOriginProductCategory { pub fn to_str(&self) -> &'static str { match self { + Self::Unknown => "unknown", Self::FaceAnimator => "face_animator", Self::TextToSpeech => "tts", - Self::Unknown => "unknown", Self::VoiceConversion => "voice_conversion", Self::ZeroShotVoice => "zs_voice", Self::VideoFilter => "video_filter", @@ -70,9 +69,9 @@ impl MediaFileOriginProductCategory { pub fn from_str(value: &str) -> Result { match value { + "unknown" => Ok(Self::Unknown), "face_animator" => Ok(Self::FaceAnimator), "tts" => Ok(Self::TextToSpeech), - "unknown" => Ok(Self::Unknown), "voice_conversion" => Ok(Self::VoiceConversion), "zs_voice" => Ok(Self::ZeroShotVoice), "video_filter" => Ok(Self::VideoFilter), @@ -85,10 +84,10 @@ impl MediaFileOriginProductCategory { // NB: BTreeSet is sorted // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. BTreeSet::from([ + Self::Unknown, Self::FaceAnimator, Self::TextToSpeech, Self::VideoFilter, - Self::Unknown, Self::VoiceConversion, Self::ZeroShotVoice, Self::Mocap, @@ -106,9 +105,9 @@ mod tests { #[test] fn test_serialization() { + assert_serialization(MediaFileOriginProductCategory::Unknown, "unknown"); assert_serialization(MediaFileOriginProductCategory::FaceAnimator, "face_animator"); assert_serialization(MediaFileOriginProductCategory::TextToSpeech, "tts"); - assert_serialization(MediaFileOriginProductCategory::Unknown, "unknown"); assert_serialization(MediaFileOriginProductCategory::VoiceConversion, "voice_conversion"); assert_serialization(MediaFileOriginProductCategory::ZeroShotVoice, "zs_voice"); assert_serialization(MediaFileOriginProductCategory::VideoFilter, "video_filter"); @@ -117,9 +116,9 @@ mod tests { #[test] fn to_str() { + assert_eq!(MediaFileOriginProductCategory::Unknown.to_str(), "unknown"); assert_eq!(MediaFileOriginProductCategory::FaceAnimator.to_str(), "face_animator"); assert_eq!(MediaFileOriginProductCategory::TextToSpeech.to_str(), "tts"); - assert_eq!(MediaFileOriginProductCategory::Unknown.to_str(), "unknown"); assert_eq!(MediaFileOriginProductCategory::VoiceConversion.to_str(), "voice_conversion"); assert_eq!(MediaFileOriginProductCategory::ZeroShotVoice.to_str(), "zs_voice"); assert_eq!(MediaFileOriginProductCategory::VideoFilter.to_str(), "video_filter"); @@ -128,9 +127,9 @@ mod tests { #[test] fn from_str() { + assert_eq!(MediaFileOriginProductCategory::from_str("unknown").unwrap(), MediaFileOriginProductCategory::Unknown); assert_eq!(MediaFileOriginProductCategory::from_str("face_animator").unwrap(), MediaFileOriginProductCategory::FaceAnimator); assert_eq!(MediaFileOriginProductCategory::from_str("tts").unwrap(), MediaFileOriginProductCategory::TextToSpeech); - assert_eq!(MediaFileOriginProductCategory::from_str("unknown").unwrap(), MediaFileOriginProductCategory::Unknown); assert_eq!(MediaFileOriginProductCategory::from_str("voice_conversion").unwrap(), MediaFileOriginProductCategory::VoiceConversion); assert_eq!(MediaFileOriginProductCategory::from_str("zs_voice").unwrap(), MediaFileOriginProductCategory::ZeroShotVoice); assert_eq!(MediaFileOriginProductCategory::from_str("video_filter").unwrap(), MediaFileOriginProductCategory::VideoFilter); @@ -141,9 +140,9 @@ mod tests { fn all_variants() { let mut variants = MediaFileOriginProductCategory::all_variants(); assert_eq!(variants.len(), 7); + assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Unknown)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::FaceAnimator)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::TextToSpeech)); - assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Unknown)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::VoiceConversion)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::ZeroShotVoice)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::VideoFilter)); diff --git a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files.rs b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files.rs index 447a106c90..1286b5be2b 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files.rs @@ -30,8 +30,11 @@ pub struct MediaFileListItem { pub origin_category: MediaFileOriginCategory, pub origin_product_category: MediaFileOriginProductCategory, + pub maybe_origin_model_type: Option, pub maybe_origin_model_token: Option, + // NB: The title won't be populated for `tts_models` records or non-`model_weights` records. + pub maybe_origin_model_title: Option, pub media_type: MediaFileType, pub public_bucket_directory_hash: String, @@ -98,6 +101,7 @@ pub async fn list_media_files(args: ListMediaFilesArgs<'_>) -> AnyhowResult, maybe_origin_model_token: Option, + maybe_origin_model_title: Option, media_type: MediaFileType, public_bucket_directory_hash: String, @@ -317,6 +326,7 @@ impl FromRow<'_, MySqlRow> for MediaFileListItemInternal { origin_product_category: MediaFileOriginProductCategory::try_from_mysql_row(row, "origin_product_category")?, maybe_origin_model_type: MediaFileOriginModelType::try_from_mysql_row_nullable(row, "maybe_origin_model_type")?, maybe_origin_model_token: row.try_get("maybe_origin_model_token")?, + maybe_origin_model_title: row.try_get("maybe_origin_model_title")?, media_type: MediaFileType::try_from_mysql_row(row, "media_type")?, public_bucket_directory_hash: row.try_get("public_bucket_directory_hash")?, maybe_public_bucket_prefix: row.try_get("maybe_public_bucket_prefix")?, diff --git a/crates/service/web/storyteller_web/src/api_doc.rs b/crates/service/web/storyteller_web/src/api_doc.rs index f8f0a63291..a6648d4983 100644 --- a/crates/service/web/storyteller_web/src/api_doc.rs +++ b/crates/service/web/storyteller_web/src/api_doc.rs @@ -16,6 +16,7 @@ use tokens::tokens::user_bookmarks::*; use tokens::tokens::users::*; use crate::http_server::common_responses::cover_image_details::*; +use crate::http_server::common_responses::media_file_origin_details::*; use crate::http_server::common_responses::media_file_social_meta_lite::MediaFileSocialMetaLight; use crate::http_server::common_responses::pagination_cursors::PaginationCursors; use crate::http_server::common_responses::pagination_page::PaginationPage; @@ -90,6 +91,8 @@ use crate::http_server::web_utils::response_success_helpers::*; CoverImageDetails, DefaultAvatarInfo, DefaultCoverInfo, + MediaFileModelDetails, + MediaFileOriginDetails, MediaFileSocialMetaLight, PaginationCursors, PaginationPage, diff --git a/crates/service/web/storyteller_web/src/http_server/common_responses/media_file_origin_details.rs b/crates/service/web/storyteller_web/src/http_server/common_responses/media_file_origin_details.rs new file mode 100644 index 0000000000..e8f3352760 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/common_responses/media_file_origin_details.rs @@ -0,0 +1,104 @@ +use utoipa::ToSchema; + +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 tokens::tokens::model_weights::ModelWeightToken; + +/// Fields useful for enriching media file listings +#[derive(Clone, Serialize, ToSchema)] +pub struct MediaFileOriginDetails { + /// Where the file came from (broadly) + pub origin_category: MediaFileOriginCategory, + + /// Where the file came from (specifically, a product area, eg. "face_animator") + pub product_category: MediaFileOriginProductCategory, + + /// If the file was produced by a model, the details. + pub maybe_model: Option, +} + +/// Information about the model weights +/// https://serde.rs/enum-representations.html#untagged +#[derive(Clone, Serialize, ToSchema)] +#[serde(untagged)] +pub enum MediaFileModelDetails { + /// A system model is one that we maintain and update, + /// eg. SadTalker, Wav2Lip, MocapNet, etc. + SystemModel { + /// The type of model weight + model_type: MediaFileOriginModelType, + }, + /// A model weight is typically a user-submitted model. + /// We store lots of these. + ModelWeight { + /// The type of model weight + model_type: MediaFileOriginModelType, + /// The model token + token: ModelWeightToken, + /// The model title + title: String, + }, +} + +impl MediaFileOriginDetails { + pub fn from_db_fields_owned( + origin_category: MediaFileOriginCategory, + product_category: MediaFileOriginProductCategory, + maybe_model_type: Option, + maybe_model_token: Option, + maybe_model_title: Option, + ) -> Self { + let maybe_model = match (maybe_model_type, maybe_model_token, maybe_model_title) { + (Some(model_type), Some(model_token), Some(model_title)) => { + Some(MediaFileModelDetails::ModelWeight { + model_type, + token: model_token, + title: model_title, + }) + }, + (Some(model_type), None, None) => { + Some(MediaFileModelDetails::SystemModel { + model_type, + }) + }, + _ => None, + }; + + Self { + origin_category, + product_category, + maybe_model, + } + } + + pub fn from_db_fields( + origin_category: MediaFileOriginCategory, + product_category: MediaFileOriginProductCategory, + maybe_model_type: Option, + maybe_model_token: Option<&ModelWeightToken>, + maybe_model_title: Option<&str>, + ) -> Self { + Self::from_db_fields_owned( + origin_category, + product_category, + maybe_model_type, + maybe_model_token.cloned(), + maybe_model_title.map(|s|s.to_string())) + } + + pub fn from_db_fields_str( + origin_category: MediaFileOriginCategory, + product_category: MediaFileOriginProductCategory, + maybe_model_type: Option, + maybe_model_token: Option<&str>, + maybe_model_title: Option<&str>, + ) -> Self { + Self::from_db_fields_owned( + origin_category, + product_category, + maybe_model_type, + maybe_model_token.map(|s| ModelWeightToken::new_from_str(s)), + maybe_model_title.map(|s| s.to_string())) + } +} diff --git a/crates/service/web/storyteller_web/src/http_server/common_responses/mod.rs b/crates/service/web/storyteller_web/src/http_server/common_responses/mod.rs index 4d3f5bb33e..a46efd6cc9 100644 --- a/crates/service/web/storyteller_web/src/http_server/common_responses/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/common_responses/mod.rs @@ -1,4 +1,5 @@ pub mod cover_image_details; +pub mod media_file_origin_details; pub mod media_file_social_meta_lite; pub mod pagination_cursors; pub mod pagination_page; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files.rs index c999b90112..ddda36c3ff 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files.rs @@ -18,6 +18,7 @@ use enums::common::visibility::Visibility; use mysql_queries::queries::media_files::list::list_media_files::{list_media_files, ListMediaFilesArgs}; use tokens::tokens::media_files::MediaFileToken; +use crate::http_server::common_responses::media_file_origin_details::MediaFileOriginDetails; use crate::http_server::common_responses::pagination_cursors::PaginationCursors; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; use crate::http_server::common_responses::user_details_lite::UserDetailsLight; @@ -44,13 +45,22 @@ pub struct ListMediaFilesSuccessResponse { pub struct MediaFileListItem { pub token: MediaFileToken, + pub media_type: MediaFileType, + + #[deprecated(note="Use MediaFileOriginDetails instead")] pub origin_category: MediaFileOriginCategory, + + #[deprecated(note="Use MediaFileOriginDetails instead")] pub origin_product_category: MediaFileOriginProductCategory, + #[deprecated(note="Use MediaFileOriginDetails instead")] pub maybe_origin_model_type: Option, + + #[deprecated(note="Use MediaFileOriginDetails instead")] pub maybe_origin_model_token: Option, - pub media_type: MediaFileType, + /// Details where the media file came from. + pub origin: MediaFileOriginDetails, /// URL to the media file. pub public_bucket_path: String, @@ -193,11 +203,17 @@ pub async fn list_media_files_handler( let results = results_page.records.into_iter() .map(|record| MediaFileListItem { token: record.token.clone(), + media_type: record.media_type, + origin: MediaFileOriginDetails::from_db_fields_str( + record.origin_category, + record.origin_product_category, + record.maybe_origin_model_type, + record.maybe_origin_model_token.as_deref(), + record.maybe_origin_model_title.as_deref()), origin_category: record.origin_category, origin_product_category: record.origin_product_category, maybe_origin_model_type: record.maybe_origin_model_type, maybe_origin_model_token: record.maybe_origin_model_token, - media_type: record.media_type, public_bucket_path: MediaFileBucketPath::from_object_hash( &record.public_bucket_directory_hash, record.maybe_public_bucket_prefix.as_deref(), From 4c9cb5d9867a1b51e12f888adbfaa3d6e7f97da0 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 13 Jan 2024 01:12:04 -0500 Subject: [PATCH 28/54] add MediaFileOriginDetails to list featured media files --- .../media_files/list/list_media_files.rs | 1 + .../list/list_media_files_by_tokens.rs | 19 +++++++++++++++++++ .../media_files/list_featured_media_files.rs | 17 +++++++++++++++++ 3 files changed, 37 insertions(+) diff --git a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files.rs b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files.rs index 1286b5be2b..0d7a14faef 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files.rs @@ -33,6 +33,7 @@ pub struct MediaFileListItem { pub maybe_origin_model_type: Option, pub maybe_origin_model_token: Option, + // NB: The title won't be populated for `tts_models` records or non-`model_weights` records. pub maybe_origin_model_title: Option, diff --git a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_by_tokens.rs b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_by_tokens.rs index 7adb769287..76a7ab510b 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_by_tokens.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_by_tokens.rs @@ -19,9 +19,13 @@ pub struct MediaFilesByTokensRecord { pub origin_category: MediaFileOriginCategory, pub origin_product_category: MediaFileOriginProductCategory, + pub maybe_origin_model_type: Option, pub maybe_origin_model_token: Option, + // NB: The title won't be populated for `tts_models` records or non-`model_weights` records. + pub maybe_origin_model_title: Option, + pub public_bucket_directory_hash: String, pub maybe_public_bucket_prefix: Option, pub maybe_public_bucket_extension: Option, @@ -68,9 +72,12 @@ async fn get_raw_media_files_by_tokens( m.origin_category, m.origin_product_category, + m.maybe_origin_model_type, m.maybe_origin_model_token, + w.title as maybe_origin_model_title, + users.token as maybe_creator_user_token, users.username as maybe_creator_username, users.display_name as maybe_creator_display_name, @@ -90,6 +97,8 @@ async fn get_raw_media_files_by_tokens( FROM media_files as m LEFT OUTER JOIN users ON users.token = m.maybe_creator_user_token + LEFT OUTER JOIN model_weights as w + ON m.maybe_origin_model_token = w.token LEFT OUTER JOIN entity_stats ON entity_stats.entity_type = "media_file" AND entity_stats.entity_token = m.token @@ -107,9 +116,12 @@ async fn get_raw_media_files_by_tokens( m.origin_category, m.origin_product_category, + m.maybe_origin_model_type, m.maybe_origin_model_token, + w.title as maybe_origin_model_title, + users.token as maybe_creator_user_token, users.username as maybe_creator_username, users.display_name as maybe_creator_display_name, @@ -129,6 +141,8 @@ async fn get_raw_media_files_by_tokens( FROM media_files as m LEFT OUTER JOIN users ON users.token = m.maybe_creator_user_token + LEFT OUTER JOIN model_weights as w + ON m.maybe_origin_model_token = w.token LEFT OUTER JOIN entity_stats ON entity_stats.entity_type = "media_file" AND entity_stats.entity_token = m.token @@ -173,6 +187,7 @@ fn map_to_media_files(dataset:Vec) -> Vec) -> Vec, pub maybe_origin_model_token: Option, + pub maybe_origin_model_title: Option, + pub maybe_creator_user_token: Option, pub maybe_creator_username: Option, pub maybe_creator_display_name: Option, @@ -246,6 +264,7 @@ impl FromRow<'_, MySqlRow> for RawMediaFileJoinUser { origin_product_category: MediaFileOriginProductCategory::try_from_mysql_row(row, "origin_product_category")?, maybe_origin_model_type: MediaFileOriginModelType::try_from_mysql_row_nullable(row, "maybe_origin_model_type")?, maybe_origin_model_token: row.try_get("maybe_origin_model_token")?, + maybe_origin_model_title: row.try_get("maybe_origin_model_title")?, maybe_creator_user_token: row.try_get::, _>("maybe_creator_user_token")? .and_then(|token| Some(UserToken::new_from_str(&token))), maybe_creator_username: row.try_get("maybe_creator_username")?, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_featured_media_files.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_featured_media_files.rs index b5be212aa0..883fc01212 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_featured_media_files.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_featured_media_files.rs @@ -15,6 +15,7 @@ use enums::by_table::media_files::media_file_origin_product_category::MediaFileO use enums::by_table::media_files::media_file_type::MediaFileType; use mysql_queries::queries::media_files::list::list_media_files_by_tokens::list_media_files_by_tokens; use tokens::tokens::media_files::MediaFileToken; +use crate::http_server::common_responses::media_file_origin_details::MediaFileOriginDetails; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; use crate::http_server::common_responses::user_details_lite::UserDetailsLight; @@ -37,11 +38,21 @@ pub struct MediaFile { /// URL to the media file pub public_bucket_path: String, + #[deprecated(note="Use MediaFileOriginDetails instead")] pub origin_category: MediaFileOriginCategory, + + #[deprecated(note="Use MediaFileOriginDetails instead")] pub origin_product_category: MediaFileOriginProductCategory, + + #[deprecated(note="Use MediaFileOriginDetails instead")] pub maybe_origin_model_type: Option, + + #[deprecated(note="Use MediaFileOriginDetails instead")] pub maybe_origin_model_token: Option, + /// Details where the media file came from. + pub origin: MediaFileOriginDetails, + pub maybe_creator: Option, /// Statistics about the media file @@ -142,6 +153,12 @@ pub async fn list_featured_media_files_handler( token: m.token, media_type: m.media_type, public_bucket_path, + origin: MediaFileOriginDetails::from_db_fields_str( + m.origin_category, + m.origin_product_category, + m.maybe_origin_model_type, + m.maybe_origin_model_token.as_deref(), + m.maybe_origin_model_title.as_deref()), origin_category: m.origin_category, origin_product_category: m.origin_product_category, maybe_origin_model_type: m.maybe_origin_model_type, From a76272aeb3e2a78db0003ec4f3b03193fb20f724 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 13 Jan 2024 01:13:21 -0500 Subject: [PATCH 29/54] doc --- .../http_server/common_responses/media_file_origin_details.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/service/web/storyteller_web/src/http_server/common_responses/media_file_origin_details.rs b/crates/service/web/storyteller_web/src/http_server/common_responses/media_file_origin_details.rs index e8f3352760..11aaaccbe1 100644 --- a/crates/service/web/storyteller_web/src/http_server/common_responses/media_file_origin_details.rs +++ b/crates/service/web/storyteller_web/src/http_server/common_responses/media_file_origin_details.rs @@ -36,7 +36,7 @@ pub enum MediaFileModelDetails { model_type: MediaFileOriginModelType, /// The model token token: ModelWeightToken, - /// The model title + /// The model title (typically only populated for `model_weights` models, not legacy tables such as `tts_models`.) title: String, }, } From 929105bb85767a412b8d1926d7534fae66f8669d Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 13 Jan 2024 01:22:31 -0500 Subject: [PATCH 30/54] add color index to cover images --- .../common_responses/cover_image_details.rs | 4 + .../default_cover_image_color_from_token.rs | 73 +++++++++++++++++++ .../default_cover_image_from_token.rs | 4 +- .../placeholder_images/cover_images/mod.rs | 1 + 4 files changed, 79 insertions(+), 3 deletions(-) create mode 100644 crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_color_from_token.rs diff --git a/crates/service/web/storyteller_web/src/http_server/common_responses/cover_image_details.rs b/crates/service/web/storyteller_web/src/http_server/common_responses/cover_image_details.rs index de1a9f7b3b..f0e7faf854 100644 --- a/crates/service/web/storyteller_web/src/http_server/common_responses/cover_image_details.rs +++ b/crates/service/web/storyteller_web/src/http_server/common_responses/cover_image_details.rs @@ -3,6 +3,7 @@ use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::users::UserToken; +use crate::util::placeholder_images::cover_images::default_cover_image_color_from_token::default_cover_image_color_from_token; use crate::util::placeholder_images::cover_images::default_cover_image_from_token::default_cover_image_from_token; use crate::util::placeholder_images::user_avatars::default_avatar_color_from_username::default_avatar_color_from_username; @@ -25,6 +26,8 @@ pub struct CoverImageDetails { pub struct DefaultCoverInfo { /// Which image to show. pub image_index: u8, + /// Which color to use. + pub color_index: u8, } impl CoverImageDetails { @@ -60,6 +63,7 @@ impl DefaultCoverInfo { pub fn from_token(model_weight_token: &ModelWeightToken) -> Self { Self { image_index: default_cover_image_from_token(model_weight_token), + color_index: default_cover_image_color_from_token(model_weight_token), } } } diff --git a/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_color_from_token.rs b/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_color_from_token.rs new file mode 100644 index 0000000000..bac7ff1806 --- /dev/null +++ b/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_color_from_token.rs @@ -0,0 +1,73 @@ +use std::collections::hash_map::DefaultHasher; +use std::hash::{Hash, Hasher}; + +use tokens::tokens::model_weights::ModelWeightToken; + +/// We'll return an index from [0, NUMBER_OF_COLORS) - [inclusive, exclusive) +const NUMBER_OF_COLORS: u64 = 9; + +/// Not that it matters, but this perturbs the hash. +const SALT_LIKE_OFFSET : u8 = 51; + +/// We return an index instead of a filename, that way the frontend can drive. +/// The hash should be stable with respect to username. +pub fn default_cover_image_color_from_token(token: &ModelWeightToken) -> u8 { + // Not as important as with usernames, but seems like a reasonable choice. + let token_string = token.as_str().to_lowercase(); + + let mut hasher = DefaultHasher::new(); + + token_string.hash(&mut hasher); + SALT_LIKE_OFFSET.hash(&mut hasher); + + let hash = hasher.finish(); + + let cover_image_color_index = hash % NUMBER_OF_COLORS; + cover_image_color_index as u8 +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use rand::distributions::{Alphanumeric, DistString}; + + use tokens::tokens::model_weights::ModelWeightToken; + + use crate::util::placeholder_images::cover_images::default_cover_image_color_from_token::default_cover_image_color_from_token; + + #[test] + fn test_stability() { + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("TOKEN")), 3); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("T_TOKEN")), 7); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("foo")), 6); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("bar")), 1); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("baz")), 6); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("bin")), 1); + } + + #[test] + fn test_case_insensitivity() { + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("T_TOKEN")), 7); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("T_TOKen")), 7); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("T_tOkEn")), 7); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("t_ToKeN")), 7); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("t_toKEN")), 7); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("t_token")), 7); + } + + #[test] + fn test_range() { + let mut distribution = HashSet::new(); + for _ in 0..1000 { + let random_token = Alphanumeric.sample_string(&mut rand::thread_rng(), 16); + let random_token= ModelWeightToken::new_from_str(&random_token); + let avatar_id = default_cover_image_color_from_token(&random_token); + distribution.insert(avatar_id); + assert!(avatar_id >= 0); + assert!(avatar_id <= 8); + } + // NB: We could test frequency of the distribution too + assert_eq!(distribution.len(), 9); + } +} diff --git a/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_from_token.rs b/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_from_token.rs index 2cf0b81211..711eeacbf1 100644 --- a/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_from_token.rs +++ b/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_from_token.rs @@ -3,9 +3,7 @@ use std::hash::{Hash, Hasher}; use tokens::tokens::model_weights::ModelWeightToken; -// TODO(bt,2024-01-11): How many cover images are there? - -/// There are currently 25 avatars numbered 0 to 24 (0-indexed). +/// There are currently 25 cover images numbered 0 to 24 (0-indexed). /// The original dataset was numbered 1 - 25, but I renamed 25 to 0. const NUMBER_OF_COVER_IMAGES: u64 = 25; diff --git a/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/mod.rs b/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/mod.rs index 574d4791c4..75c63a82a6 100644 --- a/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/mod.rs +++ b/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/mod.rs @@ -1 +1,2 @@ +pub mod default_cover_image_color_from_token; pub mod default_cover_image_from_token; \ No newline at end of file From 6122f11cf39503d8ee4a991626d8cdf84b71cd2e Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 13 Jan 2024 01:40:34 -0500 Subject: [PATCH 31/54] Add MediaFileOriginDetails to list media files for user --- .../list/list_media_files_for_user.rs | 18 ++++++++++-- .../media_files/list_media_files_for_user.rs | 29 +++++++++++++++---- 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_for_user.rs b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_for_user.rs index 748cb46d98..9c10eb9e61 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_for_user.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_for_user.rs @@ -26,8 +26,12 @@ pub struct MediaFileListItem { pub origin_category: MediaFileOriginCategory, pub origin_product_category: MediaFileOriginProductCategory, + pub maybe_origin_model_type: Option, pub maybe_origin_model_token: Option, + + // NB: The title won't be populated for `tts_models` records or non-`model_weights` records. + pub maybe_origin_model_title: Option, pub media_type: MediaFileType, pub public_bucket_directory_hash: String, @@ -97,6 +101,7 @@ pub async fn list_media_files_for_user(args: ListMediaFileForUserArgs<'_>) -> An origin_product_category: record.origin_product_category, maybe_origin_model_type: record.maybe_origin_model_type, maybe_origin_model_token: record.maybe_origin_model_token, + maybe_origin_model_title: record.maybe_origin_model_title, media_type: record.media_type, public_bucket_directory_hash: record.public_bucket_directory_hash, maybe_public_bucket_prefix: record.maybe_public_bucket_prefix, @@ -125,13 +130,15 @@ fn select_result_fields() -> String { m.id, m.token, + m.media_type, + m.origin_category, m.origin_product_category, m.maybe_origin_model_type, m.maybe_origin_model_token, - - m.media_type, + + w.title as maybe_origin_model_title, m.public_bucket_directory_hash, m.maybe_public_bucket_prefix, @@ -175,6 +182,8 @@ SELECT FROM media_files AS m LEFT OUTER JOIN users AS u ON m.maybe_creator_user_token = u.token +LEFT OUTER JOIN model_weights as w + ON m.maybe_origin_model_token = w.token LEFT OUTER JOIN entity_stats ON entity_stats.entity_type = "media_file" AND entity_stats.entity_token = m.token @@ -227,8 +236,12 @@ struct MediaFileListItemInternal { origin_category: MediaFileOriginCategory, origin_product_category: MediaFileOriginProductCategory, + maybe_origin_model_type: Option, maybe_origin_model_token: Option, + + // NB: The title won't be populated for `tts_models` records or non-`model_weights` records. + maybe_origin_model_title: Option, media_type: MediaFileType, public_bucket_directory_hash: String, @@ -266,6 +279,7 @@ impl FromRow<'_, MySqlRow> for MediaFileListItemInternal { origin_product_category: MediaFileOriginProductCategory::try_from_mysql_row(row, "origin_product_category")?, maybe_origin_model_type: MediaFileOriginModelType::try_from_mysql_row_nullable(row, "maybe_origin_model_type")?, maybe_origin_model_token: row.try_get("maybe_origin_model_token")?, + maybe_origin_model_title: row.try_get("maybe_origin_model_title")?, media_type: MediaFileType::try_from_mysql_row(row, "media_type")?, public_bucket_directory_hash: row.try_get("public_bucket_directory_hash")?, maybe_public_bucket_prefix: row.try_get("maybe_public_bucket_prefix")?, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files_for_user.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files_for_user.rs index 87fcecca7d..845087a2f8 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files_for_user.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files_for_user.rs @@ -17,6 +17,7 @@ use enums::common::view_as::ViewAs; use enums::common::visibility::Visibility; use mysql_queries::queries::media_files::list::list_media_files_for_user::{list_media_files_for_user, ListMediaFileForUserArgs}; use tokens::tokens::media_files::MediaFileToken; +use crate::http_server::common_responses::media_file_origin_details::MediaFileOriginDetails; use crate::http_server::common_responses::pagination_page::PaginationPage; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; @@ -47,13 +48,23 @@ pub struct ListMediaFilesForUserSuccessResponse { pub struct MediaFileForUserListItem { pub token: MediaFileToken, - pub origin_category: MediaFileOriginCategory, - pub origin_product_category: MediaFileOriginProductCategory, - pub maybe_origin_model_type: Option, - pub maybe_origin_model_token: Option, - pub media_type: MediaFileType, + /// Details where the media file came from. + pub origin: MediaFileOriginDetails, + + #[deprecated(note="Use MediaFileOriginDetails instead")] + pub origin_category: MediaFileOriginCategory, + + #[deprecated(note="Use MediaFileOriginDetails instead")] + pub origin_product_category: MediaFileOriginProductCategory, + + #[deprecated(note="Use MediaFileOriginDetails instead")] + pub maybe_origin_model_type: Option, + + #[deprecated(note="Use MediaFileOriginDetails instead")] + pub maybe_origin_model_token: Option, + /// URL to the media file. pub public_bucket_path: String, @@ -168,11 +179,17 @@ pub async fn list_media_files_for_user_handler( let results = results_page.records.into_iter() .map(|record| MediaFileForUserListItem { token: record.token, + media_type: record.media_type, + origin: MediaFileOriginDetails::from_db_fields_str( + record.origin_category, + record.origin_product_category, + record.maybe_origin_model_type, + record.maybe_origin_model_token.as_deref(), + record.maybe_origin_model_title.as_deref()), origin_category: record.origin_category, origin_product_category: record.origin_product_category, maybe_origin_model_type: record.maybe_origin_model_type, maybe_origin_model_token: record.maybe_origin_model_token, - media_type: record.media_type, public_bucket_path: MediaFileBucketPath::from_object_hash( &record.public_bucket_directory_hash, record.maybe_public_bucket_prefix.as_deref(), From f4497a4aeeeb806766be9c67bcba167aceaf46a6 Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Sat, 13 Jan 2024 02:27:12 -0500 Subject: [PATCH 32/54] Fixed build errors inference job should be good to run now. --- .../image_generation/sd/process_job.rs | 25 +++++++------- .../sd/sd_inference_command.rs | 34 +++++++++---------- 2 files changed, 30 insertions(+), 29 deletions(-) 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 db653be265..5b698036e9 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 @@ -13,6 +13,7 @@ use mysql_queries::payloads::generic_inference_args::generic_inference_args::Pol use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType::UploadModel; +use filesys::path_to_string::path_to_string; use tokens::tokens::model_weights::ModelWeightToken; use crate::job::job_types::image_generation::sd::sd_inference_command::InferenceArgs; use crate::job::job_types::image_generation::sd::stable_diffusion_dependencies::StableDiffusionDependencies; @@ -133,11 +134,11 @@ pub async fn process_job_inference( let vae_path = work_temp_dir.path().join("vae"); let output_path = work_temp_dir.path().join("output"); - println!("Pathes to download to:"); - println!("sd_checkpoint_path:{}",sd_checkpoint_path.to_string()); - println!("lora_path:{}",lora_path.to_string()); - println!("vae_path:{}",vae_path.to_string()); - println!("output_path:{}",output_path.to_string()); + println!("Paths to download to:"); + println!("sd_checkpoint_path:{}",sd_checkpoint_path.display()); + println!("lora_path:{}",lora_path.display()); + println!("vae_path:{}",vae_path.display()); + println!("output_path:{}",output_path.display()); // // Unpack loRA and Checkpoint // // run inference by downloading from google drive. @@ -189,7 +190,7 @@ pub async fn process_job_inference( 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, sd_checkpoint_path.to_string()).await?; + remote_cloud_file_client.download_file(details, path_to_string(sd_checkpoint_path.clone())).await?; match retrieved_loRA_record { Some(record) => { @@ -207,7 +208,7 @@ pub async fn process_job_inference( ); remote_cloud_file_client.download_file( lora_details, - lora_path.to_string() + path_to_string(lora_path.clone()) ).await?; } @@ -249,17 +250,17 @@ pub async fn process_job_inference( remote_cloud_file_client.download_file( vae_details, - vae_path.to_string() + path_to_string(vae_path.clone()) ).await?; let prompt = match sd_args.maybe_prompt { Some(val) => { val } - None => 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("zero_shot_create_voice_err.txt"); + let stderr_output_file = work_temp_dir.path().join("sd_err.txt"); let number_of_samples = match sd_args.maybe_number_of_samples { Some(val) => { @@ -278,12 +279,12 @@ pub async fn process_job_inference( prompt: prompt, negative_prompt:sd_args.maybe_n_prompt.unwrap_or_default(), number_of_samples: number_of_samples, - samplers: "".to_string(), + samplers: 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), cfg_scale: sd_args.maybe_cfg_scale.unwrap_or(7), seed: sd_args.maybe_seed.unwrap_or(1), - lora_path: Some(lora_path), + 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), 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 a47ccd0f22..e76ce47cab 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 @@ -81,7 +81,7 @@ pub struct InferenceArgs { pub height:i32, pub cfg_scale:i32, pub seed:i64, - pub lora_path:Option, + pub lora_path:PathBuf, pub checkpoint_path:PathBuf, pub vae:PathBuf, pub batch_count:i32, @@ -197,49 +197,49 @@ impl StableDiffusionInferenceCommand { // ===== Begin Python Args ===== command.push_str("--prompt "); - command.push_str(&args.prompt.as_ref().to_string_lossy()); + command.push_str(&path_to_string(args.prompt)); command.push_str(" --result_dir "); - command.push_str(&args.work_dir.as_ref().to_string_lossy()); + command.push_str(&path_to_string(args.work_dir)); command.push_str(" --result_file "); - command.push_str(&args.output_file.as_ref().to_string_lossy()); + command.push_str(&path_to_string(args.output_file)); command.push_str(" --stderr_output_file "); - command.push_str(&args.stderr_output_file.as_ref().to_string_lossy()); + command.push_str(&path_to_string(args.stderr_output_file.clone())); command.push_str(" --negative_prompt "); - command.push_str(&args.negative_prompt.as_ref().to_string_lossy()); + command.push_str(&path_to_string(args.negative_prompt)); command.push_str(" --number_of_samples "); - command.push_str(&args.number_of_samples.to_string()); + command.push_str(args.number_of_samples.to_string().as_str()); command.push_str(" --samplers "); - command.push_str(&args.samplers); + command.push_str(&path_to_string(args.samplers)); command.push_str(" --width "); - command.push_str(&args.width.to_string()); + command.push_str(args.width.to_string().as_str()); command.push_str(" --height "); - command.push_str(&args.height.to_string()); + command.push_str(args.height.to_string().as_str()); command.push_str(" --cfg_scale "); - command.push_str(&args.cfg_scale.to_string()); + command.push_str(args.cfg_scale.to_string().as_str()); command.push_str(" --seed "); - command.push_str(&args.seed.to_string()); + command.push_str(args.seed.to_string().as_str()); command.push_str(" --lora_path "); - command.push_str(&args.lora_path.as_ref().to_string_lossy()); + command.push_str(&path_to_string(args.lora_path)); command.push_str(" --check_point "); - command.push_str(&args.checkpoint_path.as_ref().to_string_lossy()); + command.push_str(&path_to_string(args.checkpoint_path)); command.push_str(" --vae "); - command.push_str(&args.vae.as_ref().to_string_lossy()); + command.push_str(&path_to_string(args.vae)); command.push_str(" --batch_count "); - command.push_str(&args.batch_count.to_string()); + command.push_str(args.batch_count.to_string().as_str()); info!("Command: {:?}", command); @@ -265,7 +265,7 @@ impl StableDiffusionInferenceCommand { let mut config = PopenConfig::default(); - info!("stderr will be written to file: {:?}", args.stderr_output_file.as_ref()); + info!("stderr will be written to file: {}", path_to_string(args.stderr_output_file.clone())); let stderr_file = File::create(&args.stderr_output_file)?; config.stderr = Redirection::File(stderr_file); From de4548795c05cad47821665e47a49469f199eb3a Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 13 Jan 2024 03:38:17 -0500 Subject: [PATCH 33/54] handle BVH file uploads --- .../media_files/upload/process_upload_media_file.rs | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/process_upload_media_file.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/process_upload_media_file.rs index f63e4eadaa..d8ce495964 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/process_upload_media_file.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/process_upload_media_file.rs @@ -144,7 +144,7 @@ pub async fn process_upload_media_file( info!("Upload maybe filesize: {:?}", maybe_file_size_bytes); - let maybe_mimetype = upload_media_request.file_bytes + let mut maybe_mimetype = upload_media_request.file_bytes .as_ref() .map(|bytes| get_mimetype_for_bytes(bytes.as_ref())) .flatten(); @@ -256,6 +256,16 @@ pub async fn process_upload_media_file( } } + if media_file_type.is_none() && maybe_mimetype.is_none() { + // https://research.cs.wisc.edu/graphics/Courses/cs-838-1999/Jeff/BVH.html + const BVH_HEADER : &[u8] = "HIERARCHY".as_bytes(); + + if bytes.starts_with(BVH_HEADER) { + media_file_type = Some(MediaFileType::Mocap); + maybe_mimetype = Some("application/octet-stream"); + } + } + let media_file_type = match media_file_type { Some(m) => m, None => { From b5009d59856a3751562589506bf1c2541d2e72d6 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 13 Jan 2024 03:42:06 -0500 Subject: [PATCH 34/54] bvh extension --- .../media_files/upload/process_upload_media_file.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/process_upload_media_file.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/process_upload_media_file.rs index d8ce495964..b0e5017e15 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/process_upload_media_file.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/process_upload_media_file.rs @@ -289,9 +289,13 @@ pub async fn process_upload_media_file( }, }; - let extension = mimetype_to_extension(mime_type) + let mut extension = mimetype_to_extension(mime_type) .map(|extension| format!(".{extension}")); + if extension.is_none() && media_file_type == MediaFileType::Mocap { + extension = Some(".bvh".to_string()); + } + const PREFIX : Option<&str> = Some("upload_"); let public_upload_path = MediaFileBucketPath::generate_new(PREFIX, extension.as_deref()); From 5933387f5c650c6a16161efcbaf9363e05b88050 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 13 Jan 2024 22:49:14 -0500 Subject: [PATCH 35/54] tts_models migration: weights db column, list tts models --- .../down.sql | 5 + .../up.sql | 7 + .../model_weights_extension_tts_details.sql | 2 +- crates/lib/migration/src/lib.rs | 1 + .../list_tts_models_for_migration.rs | 155 ++++++++++++++++++ .../lib/migration/src/text_to_speech/mod.rs | 1 + .../list_model_weights_for_text_to_speech.rs | 153 +++++++++++++++++ .../src/queries/model_weights/list/mod.rs | 3 +- .../queries/tts/tts_models/list_tts_models.rs | 2 +- .../endpoints/tts/list_tts_models.rs | 60 ++++--- .../service/web/storyteller_web/src/main.rs | 1 + .../web/storyteller_web/src/server_state.rs | 5 + 12 files changed, 361 insertions(+), 34 deletions(-) create mode 100644 _sql/migrations/2024-01-14-030220_alter_model_weights_extension_tts_add_pipeline/down.sql create mode 100644 _sql/migrations/2024-01-14-030220_alter_model_weights_extension_tts_add_pipeline/up.sql create mode 100644 crates/lib/migration/src/text_to_speech/list_tts_models_for_migration.rs create mode 100644 crates/lib/migration/src/text_to_speech/mod.rs create mode 100644 crates/lib/mysql_queries/src/queries/model_weights/list/list_model_weights_for_text_to_speech.rs diff --git a/_sql/migrations/2024-01-14-030220_alter_model_weights_extension_tts_add_pipeline/down.sql b/_sql/migrations/2024-01-14-030220_alter_model_weights_extension_tts_add_pipeline/down.sql new file mode 100644 index 0000000000..b9a1398228 --- /dev/null +++ b/_sql/migrations/2024-01-14-030220_alter_model_weights_extension_tts_add_pipeline/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE model_weights_extension_tts_details DROP COLUMN text_pipeline_type; diff --git a/_sql/migrations/2024-01-14-030220_alter_model_weights_extension_tts_add_pipeline/up.sql b/_sql/migrations/2024-01-14-030220_alter_model_weights_extension_tts_add_pipeline/up.sql new file mode 100644 index 0000000000..30aefdf219 --- /dev/null +++ b/_sql/migrations/2024-01-14-030220_alter_model_weights_extension_tts_add_pipeline/up.sql @@ -0,0 +1,7 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE model_weights_extension_tts_details + ADD COLUMN text_pipeline_type VARCHAR(64) DEFAULT NULL + AFTER ietf_primary_language_subtag; diff --git a/_sql/migrations_squashed/model_weights/model_weights_extension_tts_details.sql b/_sql/migrations_squashed/model_weights/model_weights_extension_tts_details.sql index 01478a453f..edb2534d2b 100644 --- a/_sql/migrations_squashed/model_weights/model_weights_extension_tts_details.sql +++ b/_sql/migrations_squashed/model_weights/model_weights_extension_tts_details.sql @@ -28,7 +28,7 @@ CREATE TABLE model_weights_extension_tts_details ( -- live in the `storyteller-ml` repository. -- Older models and brand new models will have this field set to null, and we'll infer which value to -- lazily backfill based on the creation date and language setting. - -- tts_text_pipeline_type VARCHAR(64) DEFAULT NULL, + text_pipeline_type VARCHAR(64) DEFAULT NULL, -- If true, multiply the mel outputs before being vocoded by a globally default constant. use_default_mel_multiply_factor BOOLEAN NOT NULL DEFAULT FALSE, diff --git a/crates/lib/migration/src/lib.rs b/crates/lib/migration/src/lib.rs index 89f29d3e6b..ff860f3c6f 100644 --- a/crates/lib/migration/src/lib.rs +++ b/crates/lib/migration/src/lib.rs @@ -18,4 +18,5 @@ #![allow(dead_code)] #![allow(non_snake_case)] +pub mod text_to_speech; pub mod voice_conversion; diff --git a/crates/lib/migration/src/text_to_speech/list_tts_models_for_migration.rs b/crates/lib/migration/src/text_to_speech/list_tts_models_for_migration.rs new file mode 100644 index 0000000000..d2a7439949 --- /dev/null +++ b/crates/lib/migration/src/text_to_speech/list_tts_models_for_migration.rs @@ -0,0 +1,155 @@ +use chrono::{DateTime, Utc}; +use sqlx::MySql; +use sqlx::pool::PoolConnection; + +use enums::common::visibility::Visibility; +use errors::AnyhowResult; +use mysql_queries::queries::model_weights::list::list_model_weights_for_text_to_speech::{list_model_weights_for_text_to_speech, ModelWeightForTts}; +use mysql_queries::queries::tts::tts_models::list_tts_models::{list_tts_models_with_connection, TtsModelRecordForList}; + +/// List TTS models +/// This is for the tts model list page. +/// Since we're listing, we have to use a flag to determine which query to perform. +pub async fn list_tts_models_for_migration( + mysql_connection: &mut PoolConnection, + use_weights_table: bool, +) -> AnyhowResult> { + // NB: This is temporary migration code as we switch from the `tts_models` table to the `model_weights` table. + if use_weights_table { + let models = list_model_weights_for_text_to_speech( + mysql_connection).await?; + + Ok(models.into_iter() + .map(|model| TtsModelForList::ModelWeight(model)) + .collect()) + + } else { + let models = list_tts_models_with_connection( + mysql_connection, None, false).await?; + + Ok(models.into_iter() + .map(|model| TtsModelForList::LegacyTts(model)) + .collect()) + } +} + +/// Union over the legacy table and the new table to support an easier migration. +/// This enum can hold a record of either type and present a unified accessor interface. +#[derive(Clone)] +pub enum TtsModelForList { + /// Old type from the `tts_models` table, on the way out + LegacyTts(TtsModelRecordForList), + /// New type, replacing the `tts_models` table. + ModelWeight(ModelWeightForTts), +} + +impl TtsModelForList { + pub fn token(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => model.model_token.as_str(), + TtsModelForList::ModelWeight(ref model) => model.token.as_str(), + } + } + + pub fn title(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => &model.title, + TtsModelForList::ModelWeight(ref model) => &model.title, + } + } + + pub fn ietf_language_tag(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => &model.ietf_language_tag, + TtsModelForList::ModelWeight(ref model) => &model.ietf_language_tag, + } + } + + pub fn ietf_primary_language_subtag(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => &model.ietf_primary_language_subtag, + TtsModelForList::ModelWeight(ref model) => &model.ietf_primary_language_subtag, + } + } + + pub fn creator_user_token(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => &model.creator_user_token, + TtsModelForList::ModelWeight(ref model) => model.creator_user_token.as_str(), + } + } + + pub fn creator_username(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => &model.creator_username, + TtsModelForList::ModelWeight(ref model) => &model.creator_username, + } + } + + pub fn creator_display_name(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => &model.creator_display_name, + TtsModelForList::ModelWeight(ref model) => &model.creator_display_name, + } + } + + pub fn creator_gravatar_hash(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => &model.creator_gravatar_hash, + TtsModelForList::ModelWeight(ref model) => &model.creator_gravatar_hash, + } + } + + pub fn creator_set_visibility(&self) -> Visibility{ + match self { + TtsModelForList::LegacyTts(ref model) => model.creator_set_visibility, + TtsModelForList::ModelWeight(ref model) => model.creator_set_visibility, + } + } + + pub fn created_at(&self) -> &DateTime { + match self { + TtsModelForList::LegacyTts(ref model) => &model.created_at, + TtsModelForList::ModelWeight(ref model) => &model.created_at, + } + } + + pub fn updated_at(&self) -> &DateTime { + match self { + TtsModelForList::LegacyTts(ref model) => &model.updated_at, + TtsModelForList::ModelWeight(ref model) => &model.updated_at, + } + } + + // Not supported in the new table + pub fn is_front_page_featured(&self) -> bool { + match self { + TtsModelForList::LegacyTts(ref model) => model.is_front_page_featured, + TtsModelForList::ModelWeight(ref _model) => false, + } + } + + // Not supported in the new table + pub fn user_ratings_positive_count(&self) -> u32 { + match self { + TtsModelForList::LegacyTts(ref model) => model.user_ratings_positive_count, + TtsModelForList::ModelWeight(ref _model) => 0, + } + } + + // Not supported in the new table + pub fn user_ratings_negative_count(&self) -> u32 { + match self { + TtsModelForList::LegacyTts(ref model) => model.user_ratings_negative_count, + TtsModelForList::ModelWeight(ref _model) => 0, + } + } + + // Not supported in the new table + pub fn user_ratings_total_count(&self) -> u32 { + match self { + TtsModelForList::LegacyTts(ref model) => model.user_ratings_total_count, + TtsModelForList::ModelWeight(ref _model) => 0, + } + } +} diff --git a/crates/lib/migration/src/text_to_speech/mod.rs b/crates/lib/migration/src/text_to_speech/mod.rs new file mode 100644 index 0000000000..2e243cfce7 --- /dev/null +++ b/crates/lib/migration/src/text_to_speech/mod.rs @@ -0,0 +1 @@ +pub mod list_tts_models_for_migration; \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/queries/model_weights/list/list_model_weights_for_text_to_speech.rs b/crates/lib/mysql_queries/src/queries/model_weights/list/list_model_weights_for_text_to_speech.rs new file mode 100644 index 0000000000..014c2c4ba0 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/model_weights/list/list_model_weights_for_text_to_speech.rs @@ -0,0 +1,153 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use sqlx::MySql; +use sqlx::pool::PoolConnection; + +use enums::by_table::model_weights::weights_types::WeightsType; +use enums::common::visibility::Visibility; +use errors::AnyhowResult; +use tokens::tokens::model_weights::ModelWeightToken; +use tokens::tokens::users::UserToken; + +#[derive(Serialize, Clone)] +pub struct ModelWeightForTts { + pub token: ModelWeightToken, + pub weight_type: WeightsType, + + pub title: String, + + // We only supported tacotron2 in the legacy case + //pub tts_model_type: String, + + pub ietf_language_tag: String, + pub ietf_primary_language_subtag: String, + + pub creator_user_token: UserToken, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + // NB: We won't port these... + //pub is_locked_from_use: bool, + //pub is_front_page_featured: bool, + //pub is_twitch_featured: bool, + //pub maybe_suggested_unique_bot_command: Option, + //pub user_ratings_positive_count: u32, + //pub user_ratings_negative_count: u32, + //pub user_ratings_total_count: u32, // NB: Does not include "neutral" ratings. + + pub creator_set_visibility: Visibility, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +// TODO(bt,2024-01-13): This is written to support the migration to the `model_weights` +// table from the `tts_models` table. We should write a more cross-cutting +// query to treat this not as a special cased thing. All weights should be listable. + +/// This is to support the tts model list page. +/// Later this will be migrated or replaced with a more generic query. +pub async fn list_model_weights_for_text_to_speech( + mysql_connection: &mut PoolConnection, +) -> AnyhowResult> { + + let models = + list_tts_models_for_all_creators(mysql_connection).await?; + + Ok(models.into_iter() + .map(|model| { + ModelWeightForTts { + token: model.token, + weight_type: model.weight_type, + creator_user_token: model.creator_user_token, + creator_username: model.creator_username, + creator_display_name: model.creator_display_name, + creator_gravatar_hash: model.creator_gravatar_hash, + title: model.title, + ietf_language_tag: model.ietf_language_tag.unwrap_or("en".to_string()), + ietf_primary_language_subtag: model.ietf_primary_language_subtag.unwrap_or("en".to_string()), + creator_set_visibility: model.creator_set_visibility, + created_at: model.created_at, + updated_at: model.updated_at, + } + }) + .collect::>()) +} + +async fn list_tts_models_for_all_creators( + mysql_connection: &mut PoolConnection, +) -> AnyhowResult> { + // NB: Scoped to only tt2 weights + let maybe_results = sqlx::query_as!( + RawModelWeightForTts, + r#" +SELECT + w.token as `token: tokens::tokens::model_weights::ModelWeightToken`, + w.weights_type as `weight_type: enums::by_table::model_weights::weights_types::WeightsType`, + w.title, + w.creator_user_token as `creator_user_token: tokens::tokens::users::UserToken`, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + w_extension.ietf_language_tag, + w_extension.ietf_primary_language_subtag, + w.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`, + w.created_at, + w.updated_at +FROM model_weights as w +LEFT OUTER JOIN model_weights_extension_tts_details as w_extension + ON w_extension.model_weights_token = w.token +JOIN users + ON users.token = w.creator_user_token +WHERE + w.weights_type IN ("tt2") + AND w.user_deleted_at IS NULL + AND w.mod_deleted_at IS NULL + "# + ) + .fetch_all(&mut **mysql_connection) + .await; + + match maybe_results { + Ok(results) => Ok(results), + Err(err) => match err { + sqlx::Error::RowNotFound => Ok(Vec::new()), + _ => { + Err(anyhow!("error querying : {:?}", err)) + } + } + } +} + +struct RawModelWeightForTts { + pub token: ModelWeightToken, + pub weight_type: WeightsType, + + pub title: String, + + // We only supported tacotron2 in the legacy case + //pub tts_model_type: String, + + pub ietf_language_tag: Option, + pub ietf_primary_language_subtag: Option, + + pub creator_user_token: UserToken, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + // NB: We won't port these... + //pub is_locked_from_use: bool, + //pub is_front_page_featured: bool, + //pub is_twitch_featured: bool, + //pub maybe_suggested_unique_bot_command: Option, + //pub user_ratings_positive_count: u32, + //pub user_ratings_negative_count: u32, + //pub user_ratings_total_count: u32, // NB: Does not include "neutral" ratings. + + pub creator_set_visibility: Visibility, + + pub created_at: DateTime, + pub updated_at: DateTime, +} diff --git a/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs b/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs index 3b5b9b35b2..f56b2896be 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs @@ -1,5 +1,6 @@ +pub mod list_model_weights_for_elastic_search_backfill_using_cursor; +pub mod list_model_weights_for_text_to_speech; pub mod list_model_weights_for_voice_conversion; pub mod list_weights_by_tokens; pub mod list_weights_by_user; -pub mod list_model_weights_for_elastic_search_backfill_using_cursor; pub mod list_weights_query_builder; \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/queries/tts/tts_models/list_tts_models.rs b/crates/lib/mysql_queries/src/queries/tts/tts_models/list_tts_models.rs index 3c0c393b82..c9ee839d9a 100644 --- a/crates/lib/mysql_queries/src/queries/tts/tts_models/list_tts_models.rs +++ b/crates/lib/mysql_queries/src/queries/tts/tts_models/list_tts_models.rs @@ -12,7 +12,7 @@ use crate::helpers::boolean_converters::i8_to_bool; // FIXME: This is the old style of query scoping and shouldn't be copied. // TODO/FIXME : This struct is returned publicly in some endpoints! -#[derive(Serialize)] +#[derive(Serialize, Clone)] pub struct TtsModelRecordForList { pub model_token: String, pub tts_model_type: String, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs index 94a3854335..f3fbfdbe8b 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs @@ -16,10 +16,11 @@ use log::{error, info, warn}; use sqlx::MySql; use sqlx::pool::PoolConnection; +use enums::by_table::tts_models::tts_model_type::TtsModelType; use enums::common::visibility::Visibility; use errors::AnyhowResult; +use migration::text_to_speech::list_tts_models_for_migration::list_tts_models_for_migration; use mysql_queries::queries::tts::tts_category_assignments::fetch_and_build_tts_model_category_map::fetch_and_build_tts_model_category_map_with_connection; -use mysql_queries::queries::tts::tts_models::list_tts_models::list_tts_models_with_connection; use mysql_queries::queries::users::user_sessions::get_user_session_by_token::SessionUserRecord; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; @@ -133,7 +134,7 @@ pub async fn list_tts_models_handler( })?; // TODO: Fail open in case the DB is down. Pull from expired cache if query fails. - let models = get_all_models(&mut mysql_connection) + let models = get_all_models(&mut mysql_connection, server_state.flags.switch_tts_to_model_weights) .await .map_err(|e| { error!("Error querying database: {:?}", e); @@ -212,47 +213,44 @@ pub fn render_response_payload(response: ListTtsModelsSuccessResponse) -> Result Ok(body) } -async fn get_all_models(mysql_connection: &mut PoolConnection) -> AnyhowResult> { - let mut models = list_tts_models_with_connection( - mysql_connection, - None, - false - ).await?; +async fn get_all_models(mysql_connection: &mut PoolConnection, use_weights_table: bool) -> AnyhowResult> { + let mut models = list_tts_models_for_migration(mysql_connection, use_weights_table).await?; + + // Make the list nice for human readers. + models.sort_by(|a, b| + natural_lexical_cmp(&a.title(), &b.title())); let model_categories_map = fetch_and_build_tts_model_category_map_with_connection(mysql_connection).await?; - // Make the list nice for human readers. - models.sort_by(|a, b| - natural_lexical_cmp(&a.title, &b.title)); - let models_for_response = models.into_iter() .map(|model| { - let model_token = model.model_token.clone(); + // NB: All the cloning is just for the migration of tts_models --> model_weights + let model_token = model.token().to_string(); TtsModelRecordForResponse { - model_token: model.model_token, - tts_model_type: model.tts_model_type, - creator_user_token: model.creator_user_token, - creator_username: model.creator_username, - creator_display_name: model.creator_display_name, - creator_gravatar_hash: model.creator_gravatar_hash, - title: model.title, - ietf_language_tag: model.ietf_language_tag, - ietf_primary_language_subtag: model.ietf_primary_language_subtag, - is_front_page_featured: model.is_front_page_featured, - is_twitch_featured: model.is_twitch_featured, - maybe_suggested_unique_bot_command: model.maybe_suggested_unique_bot_command, - creator_set_visibility: model.creator_set_visibility, + model_token: model.token().to_string(), + tts_model_type: TtsModelType::Tacotron2.to_string(), // NB(bt): We only ever supported the TT2 value + creator_user_token: model.creator_user_token().to_string(), + creator_username: model.creator_username().to_string(), + creator_display_name: model.creator_display_name().to_string(), + creator_gravatar_hash: model.creator_gravatar_hash().to_string(), + title: model.title().to_string(), + ietf_language_tag: model.ietf_language_tag().to_string(), + ietf_primary_language_subtag: model.ietf_primary_language_subtag().to_string(), + is_front_page_featured: model.is_front_page_featured(), + is_twitch_featured: false, // NB: This is no longer used. + maybe_suggested_unique_bot_command: None, // NB: This is no longer used. + creator_set_visibility: model.creator_set_visibility(), user_ratings: UserRatingsStats { - positive_count: model.user_ratings_positive_count, - negative_count: model.user_ratings_negative_count, - total_count: model.user_ratings_total_count, + positive_count: model.user_ratings_positive_count(), + negative_count: model.user_ratings_negative_count(), + total_count: model.user_ratings_total_count(), }, category_tokens: model_categories_map.model_to_category_tokens.get(&model_token) .map(|hash| hash.clone()) .unwrap_or(HashSet::new()), - created_at: model.created_at, - updated_at: model.updated_at, + created_at: model.created_at().clone(), + updated_at: model.updated_at().clone(), } }) .collect::>(); diff --git a/crates/service/web/storyteller_web/src/main.rs b/crates/service/web/storyteller_web/src/main.rs index 457b4dc148..67bfaaac05 100644 --- a/crates/service/web/storyteller_web/src/main.rs +++ b/crates/service/web/storyteller_web/src/main.rs @@ -359,6 +359,7 @@ async fn main() -> AnyhowResult<()> { // Temporary flags enable_enqueue_generic_tts_job: easyenv::get_env_bool_or_default("FF_ENABLE_ENQUEUE_GENERIC_TTS_JOB", false), switch_voice_conversion_to_model_weights: easyenv::get_env_bool_or_default("FF_SWITCH_VOICE_CONVERSION_TO_MODEL_WEIGHTS", false), + switch_tts_to_model_weights: easyenv::get_env_bool_or_default("FF_SWITCH_TTS_TO_MODEL_WEIGHTS", false), }; let third_party_url_redirector = ThirdPartyUrlRedirector::new(server_environment); diff --git a/crates/service/web/storyteller_web/src/server_state.rs b/crates/service/web/storyteller_web/src/server_state.rs index 38264e89c2..2d9e96693b 100644 --- a/crates/service/web/storyteller_web/src/server_state.rs +++ b/crates/service/web/storyteller_web/src/server_state.rs @@ -263,6 +263,11 @@ pub struct StaticFeatureFlags { /// TEMPORARY: Move voice control model listing over to `model_weights` from `voice_conversion_models` /// This will control all downstream enqueuing, jobs, etc. pub switch_voice_conversion_to_model_weights: bool, + + // TODO(2024-01-13): Remove temporary flag when done. + /// TEMPORARY: Move voice control model listing over to `model_weights` from `tts_models` + /// This will control all downstream enqueuing, jobs, etc. + pub switch_tts_to_model_weights: bool, } /// Instead of top level service denial, these are bans against entities that instead return From 89cd104551bbd90e25092f4043877b8b736b1ff9 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 13 Jan 2024 22:58:39 -0500 Subject: [PATCH 36/54] codegen --- ...509100e22c614367b67ab026b5edd70927ea5.json | 146 ++++++++++++++++++ 1 file changed, 146 insertions(+) create mode 100644 .sqlx/query-4d45f5ce213627465e991eaec9b509100e22c614367b67ab026b5edd70927ea5.json diff --git a/.sqlx/query-4d45f5ce213627465e991eaec9b509100e22c614367b67ab026b5edd70927ea5.json b/.sqlx/query-4d45f5ce213627465e991eaec9b509100e22c614367b67ab026b5edd70927ea5.json new file mode 100644 index 0000000000..0f795c35d7 --- /dev/null +++ b/.sqlx/query-4d45f5ce213627465e991eaec9b509100e22c614367b67ab026b5edd70927ea5.json @@ -0,0 +1,146 @@ +{ + "db_name": "MySQL", + "query": "\nSELECT\n w.token as `token: tokens::tokens::model_weights::ModelWeightToken`,\n w.weights_type as `weight_type: enums::by_table::model_weights::weights_types::WeightsType`,\n w.title,\n w.creator_user_token as `creator_user_token: tokens::tokens::users::UserToken`,\n users.username as creator_username,\n users.display_name as creator_display_name,\n users.email_gravatar_hash as creator_gravatar_hash,\n w_extension.ietf_language_tag,\n w_extension.ietf_primary_language_subtag,\n w.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`,\n w.created_at,\n w.updated_at\nFROM model_weights as w\nLEFT OUTER JOIN model_weights_extension_tts_details as w_extension\n ON w_extension.model_weights_token = w.token\nJOIN users\n ON users.token = w.creator_user_token\nWHERE\n w.weights_type IN (\"tt2\")\n AND w.user_deleted_at IS NULL\n AND w.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "token: tokens::tokens::model_weights::ModelWeightToken", + "type_info": { + "type": "VarString", + "flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "weight_type: enums::by_table::model_weights::weights_types::WeightsType", + "type_info": { + "type": "VarString", + "flags": "NOT_NULL | MULTIPLE_KEY | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "title", + "type_info": { + "type": "VarString", + "flags": "NOT_NULL | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 3, + "name": "creator_user_token: tokens::tokens::users::UserToken", + "type_info": { + "type": "VarString", + "flags": "NOT_NULL | MULTIPLE_KEY | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": "NOT_NULL | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 6, + "name": "creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": "NOT_NULL | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 7, + "name": "ietf_language_tag", + "type_info": { + "type": "VarString", + "flags": "BINARY", + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 8, + "name": "ietf_primary_language_subtag", + "type_info": { + "type": "VarString", + "flags": "BINARY", + "char_set": 224, + "max_size": 48 + } + }, + { + "ordinal": 9, + "name": "creator_set_visibility: enums::common::visibility::Visibility", + "type_info": { + "type": "String", + "flags": "NOT_NULL | MULTIPLE_KEY | BINARY | ENUM", + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 10, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": "NOT_NULL | BINARY | TIMESTAMP", + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 11, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": "NOT_NULL | BINARY | TIMESTAMP | ON_UPDATE_NOW", + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + true, + true, + false, + false, + false + ] + }, + "hash": "4d45f5ce213627465e991eaec9b509100e22c614367b67ab026b5edd70927ea5" +} From eebedbce0705d8793f63489d1e51bda6e115cbd5 Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Sun, 14 Jan 2024 18:39:33 -0500 Subject: [PATCH 37/54] Runs inference but the checkpoint is corrupt? --- .../config/inference-job.development.env | 2 +- .../job_loop/determine_dependency_status.rs | 2 +- .../image_generation/sd/process_job.rs | 5 +++ .../sd/sd_inference_command.rs | 33 ++++++++++--------- 4 files changed, 25 insertions(+), 17 deletions(-) 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 bc9a8165ff..8e542d0ef6 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/crates/service/job/inference_job/config/inference-job.development.env @@ -110,7 +110,7 @@ RERENDER_INFERENCE_EXECUTABLE_OR_COMMAND="venv/bin/python rerender.py" # # # ================================= # STABLE_DIFFUSION_INFERENCE_ROOT_DIRECTORY="/home/tensor/code/storyteller/storyteller-ml/image_generation/Storyteller-Studio-Filter/" -STABLE_DIFFUSION_INFERENCE_COMMAND="venv/bin/python main.py" +STABLE_DIFFUSION_INFERENCE_COMMAND=".venv/bin/python main.py" # ================================= # # # diff --git a/crates/service/job/inference_job/src/job/job_loop/determine_dependency_status.rs b/crates/service/job/inference_job/src/job/job_loop/determine_dependency_status.rs index 9a79c77746..57e99ae5b5 100644 --- a/crates/service/job/inference_job/src/job/job_loop/determine_dependency_status.rs +++ b/crates/service/job/inference_job/src/job/job_loop/determine_dependency_status.rs @@ -188,7 +188,7 @@ async fn get_model_record_from_cacheable_query(job_dependencies: &JobDependencie } } _ => { - warn!("Job does not have a category or a model token: {:?} - category: {:?}, token: {:?}", + warn!("Only applies to older job workflows. Job does not have a category or a model token: {:?} - category: {:?}, token: {:?}", job.id, job.inference_category, job.maybe_model_token); MaybeInferenceModel::None 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 5b698036e9..71c57cdcd9 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,5 +1,7 @@ use std::ops::Deref; use std::path::PathBuf; +use std::thread; +use std::thread::Thread; use std::time::Duration; use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails; use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; @@ -290,6 +292,9 @@ 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. Ok(JobSuccessResult { 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 e76ce47cab..265b35f8f7 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 @@ -196,26 +196,25 @@ impl StableDiffusionInferenceCommand { } // ===== Begin Python Args ===== - command.push_str("--prompt "); - command.push_str(&path_to_string(args.prompt)); + command.push_str(format!("--prompt \"{}\"",&path_to_string(args.prompt)).as_ref()); - command.push_str(" --result_dir "); + command.push_str(" --output-dir "); command.push_str(&path_to_string(args.work_dir)); - command.push_str(" --result_file "); + command.push_str(" --output-name "); command.push_str(&path_to_string(args.output_file)); - command.push_str(" --stderr_output_file "); - command.push_str(&path_to_string(args.stderr_output_file.clone())); + // command.push_str(" --stderr-output-file "); + // command.push_str(&path_to_string(args.stderr_output_file.clone())); - command.push_str(" --negative_prompt "); - command.push_str(&path_to_string(args.negative_prompt)); + command.push_str(format!(" --negative-prompt \"{}\"",&path_to_string(args.negative_prompt)).as_ref()); + //command.push_str(&path_to_string(args.negative_prompt)); - command.push_str(" --number_of_samples "); + command.push_str(" --number-of-samples "); command.push_str(args.number_of_samples.to_string().as_str()); - command.push_str(" --samplers "); - command.push_str(&path_to_string(args.samplers)); + command.push_str(format!(" --samplers \"{}\"",&path_to_string(args.samplers)).as_ref()); + //command.push_str(&path_to_string(args.samplers)); command.push_str(" --width "); command.push_str(args.width.to_string().as_str()); @@ -223,23 +222,27 @@ impl StableDiffusionInferenceCommand { command.push_str(" --height "); command.push_str(args.height.to_string().as_str()); - command.push_str(" --cfg_scale "); + command.push_str(" --cfg-scale "); command.push_str(args.cfg_scale.to_string().as_str()); command.push_str(" --seed "); command.push_str(args.seed.to_string().as_str()); - command.push_str(" --lora_path "); + command.push_str(" --loRA-path "); command.push_str(&path_to_string(args.lora_path)); - command.push_str(" --check_point "); + command.push_str(" --check-point "); command.push_str(&path_to_string(args.checkpoint_path)); command.push_str(" --vae "); command.push_str(&path_to_string(args.vae)); - command.push_str(" --batch_count "); + command.push_str(" --batch-count "); command.push_str(args.batch_count.to_string().as_str()); + + command.push_str(" --batch-size "); + command.push_str(1.to_string().as_str()); + info!("Command: {:?}", command); From 5e98049c501eafae0caf91ed1d9f6540c6c037c7 Mon Sep 17 00:00:00 2001 From: Kasisnu <5044120+kasisnu@users.noreply.github.com> Date: Mon, 15 Jan 2024 17:44:29 +0530 Subject: [PATCH 38/54] switch from vallex -> styletts2 for zero shot --- .../media_files/media_file_origin_model_type.rs | 11 ++++++++++- .../create/insert_media_file_from_zero_shot_tts.rs | 6 +++--- .../src/job/job_types/tts/process_single_tts_job.rs | 11 +++++++++-- .../tts/styletts2/process_job_inference_voice.rs | 4 +++- .../tts/styletts2/styletts2_inference_command.rs | 6 ++++++ .../tts/vall_e_x/process_job_inference_voice.rs | 2 ++ .../voice_designer/inference/enqueue_tts_request.rs | 2 +- .../endpoints/voice_designer/voices/create_voice.rs | 2 +- 8 files changed, 35 insertions(+), 9 deletions(-) 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 d9fbd7cdd5..1a3a9e58f8 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 @@ -38,6 +38,9 @@ pub enum MediaFileOriginModelType { #[serde(rename = "mocap_net")] MocapNet, + + #[serde(rename = "styletts2")] + StyleTTS2, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -56,6 +59,7 @@ impl MediaFileOriginModelType { Self::VallEX => "vall_e_x", Self::Rerender => "rerender", Self::MocapNet => "mocap_net", + Self::StyleTTS2 => "styletts2", } } @@ -68,6 +72,7 @@ impl MediaFileOriginModelType { "vall_e_x" => Ok(Self::VallEX), "rerender" => Ok(Self::Rerender), "mocap_net" => Ok(Self::MocapNet), + "styletts2" => Ok(Self::StyleTTS2), _ => Err(format!("invalid value: {:?}", value)), } } @@ -83,6 +88,7 @@ impl MediaFileOriginModelType { Self::VallEX, Self::Rerender, Self::MocapNet, + Self::StyleTTS2, ]) } } @@ -109,6 +115,7 @@ mod tests { assert_eq!(MediaFileOriginModelType::VallEX.to_str(), "vall_e_x"); assert_eq!(MediaFileOriginModelType::Rerender.to_str(), "rerender"); assert_eq!(MediaFileOriginModelType::MocapNet.to_str(), "mocap_net"); + assert_eq!(MediaFileOriginModelType::StyleTTS2.to_str(), "styletts2"); } #[test] @@ -120,13 +127,14 @@ mod tests { assert_eq!(MediaFileOriginModelType::from_str("vall_e_x").unwrap(), MediaFileOriginModelType::VallEX); assert_eq!(MediaFileOriginModelType::from_str("rerender").unwrap(), MediaFileOriginModelType::Rerender); assert_eq!(MediaFileOriginModelType::from_str("mocap_net").unwrap(), MediaFileOriginModelType::MocapNet); + assert_eq!(MediaFileOriginModelType::from_str("styletts2").unwrap(), MediaFileOriginModelType::StyleTTS2); assert!(MediaFileOriginModelType::from_str("foo").is_err()); } #[test] fn all_variants() { let mut variants = MediaFileOriginModelType::all_variants(); - assert_eq!(variants.len(), 7); + assert_eq!(variants.len(), 8); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::RvcV2)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::SadTalker)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::SoVitsSvc)); @@ -134,6 +142,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::VallEX)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::Rerender)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::MocapNet)); + assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::StyleTTS2)); assert_eq!(variants.pop_first(), None); } } diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs index e250919752..ff73ed35d5 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs @@ -18,6 +18,7 @@ pub struct InsertArgs<'a> { pub pool: &'a MySqlPool, pub job: &'a AvailableInferenceJob, + pub origin_model_type: MediaFileOriginModelType, pub maybe_mime_type: Option<&'a str>, pub file_size_bytes: u64, pub sha256_checksum: &'a str, @@ -69,7 +70,6 @@ pub async fn insert_media_file_from_zero_shot( // 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 record_id = { @@ -111,8 +111,8 @@ pub async fn insert_media_file_from_zero_shot( ORIGIN_CATEGORY.to_str(), ORIGIN_PRODUCT_CATEGORY.to_str(), - ORIGIN_MODEL_TYPE.to_str(), - + args.origin_model_type.to_str(), + MEDIA_TYPE.to_str(), args.maybe_mime_type, args.file_size_bytes, diff --git a/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs index d9a6cf1f53..495901d1dc 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs @@ -8,7 +8,8 @@ use mysql_queries::queries::tts::tts_models::get_tts_model_for_inference_improve use crate::job::job_loop::job_success_result::JobSuccessResult; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; -use crate::job::job_types::tts::{tacotron2_v2_early_fakeyou, vits}; +use crate::job::job_types::tts::{styletts2, tacotron2_v2_early_fakeyou, vits}; +use crate::job::job_types::tts::styletts2::process_job::StyleTTS2ProcessJobArgs; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::process_job::ProcessJobArgs; use crate::job::job_types::tts::vall_e_x::process_job::VALLEXProcessJobArgs; use crate::job::job_types::tts::vits::process_job::VitsProcessJobArgs; @@ -24,7 +25,7 @@ pub async fn process_single_tts_job( // TODO: Move common checks for slurs, etc. here. match job.maybe_model_type { - Some(InferenceModelType::VallEX) => { + Some(InferenceModelType::VallEX) | Some(InferenceModelType::StyleTTS2) => { // Zero-shot TTS does not need a fine-tuned model token. dispatch_zero_shot_model( job_dependencies, @@ -64,6 +65,12 @@ async fn dispatch_zero_shot_model( job, }).await } + Some(InferenceModelType::StyleTTS2) => { + styletts2::process_job::process_job(StyleTTS2ProcessJobArgs { + job_dependencies, + job, + }).await + } _ => { Err(ProcessSingleJobError::Other(anyhow!("not a zero-shot model"))) } diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs index b7fcd141de..b602d04fbf 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs @@ -11,6 +11,7 @@ use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; use cloud_storage::bucket_client::BucketClient; use cloud_storage::bucket_path_unifier::BucketPathUnifier; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType; use filesys::file_size::file_size; use hashing::sha256::sha256_hash_file::sha256_hash_file; use mysql_queries::queries::media_files::create::insert_media_file_from_zero_shot_tts::insert_media_file_from_zero_shot; @@ -135,7 +136,7 @@ pub async fn process_inference_voice( let command_exit_status = model_dependencies.inference_command.execute_inference( InferenceArgs { - input_embedding_file_path: &workdir, + input_embedding_file_path: &downloaded_weights_path, input_text_file_path: &text_input_fs_path, // text output_file_name: &output_file_name, // output file name in the output folder stderr_output_file: &stderr_output_file, @@ -208,6 +209,7 @@ pub async fn process_inference_voice( let (media_file_token, id) = insert_media_file_from_zero_shot(InsertArgs { pool: &args.job_dependencies.db.mysql_pool, job: &job, + origin_model_type: MediaFileOriginModelType::StyleTTS2, maybe_mime_type: Some(&MIME_TYPE), file_size_bytes, sha256_checksum: &file_checksum, diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs index bda3b67f5d..814cf97939 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs @@ -423,6 +423,12 @@ impl StyleTTS2CreateEmbeddingCommand { command.push_str(" --voice "); command.push_str(&args.audio_file); + command.push_str(" --vcsteps "); + command.push_str(format!("{}", 20).as_str()); + + command.push_str(" --input "); + command.push_str("/dev/null"); + // ===== End Python Args ===== if let Some(docker_options) = self.maybe_docker_options.as_ref() { diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs index 72272f4fab..7646fd7652 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs @@ -11,6 +11,7 @@ use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; use cloud_storage::bucket_client::BucketClient; use cloud_storage::bucket_path_unifier::BucketPathUnifier; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType; use filesys::file_size::file_size; use hashing::sha256::sha256_hash_file::sha256_hash_file; use mysql_queries::queries::media_files::create::insert_media_file_from_zero_shot_tts::insert_media_file_from_zero_shot; @@ -224,6 +225,7 @@ pub async fn process_inference_voice( let (media_file_token, id) = insert_media_file_from_zero_shot(InsertArgs { pool: &args.job_dependencies.db.mysql_pool, job: &job, + origin_model_type: MediaFileOriginModelType::VallEX, maybe_mime_type: Some(&MIME_TYPE), file_size_bytes, sha256_checksum: &file_checksum, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs index 7e09f40782..5a8e7bb3f0 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs @@ -189,7 +189,7 @@ pub async fn enqueue_tts_request( let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, inference_category: InferenceCategory::TextToSpeech, - maybe_model_type: Some(InferenceModelType::VallEX), // NB: Model is static during inference + maybe_model_type: Some(InferenceModelType::StyleTTS2), // NB: Model is static during inference maybe_model_token: None, // NB: Model is static during inference maybe_input_source_token: None, maybe_input_source_token_type: None, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs index 6654843007..f8dd1274b9 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs @@ -177,7 +177,7 @@ pub async fn create_voice_handler( let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, inference_category: InferenceCategory::TextToSpeech, - maybe_model_type: Some(InferenceModelType::VallEX), // NB: Model is static during inference + maybe_model_type: Some(InferenceModelType::StyleTTS2), // NB: Model is static during inference maybe_model_token: None, // NB: Model is static during inference maybe_input_source_token: None, maybe_input_source_token_type: None, From cf4c335d41f9864aa54cbe46def08c20bf15f5b9 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 15 Jan 2024 21:27:44 -0500 Subject: [PATCH 39/54] migration tool for tts models --- .../public/weight_files/bucket_file_path.rs | 10 + crates/cli/migration-tool/src/main.rs | 6 +- .../cli/migration-tool/src/migrations/mod.rs | 1 + .../tts_models_to_weights/copy_cloud_files.rs | 47 ++++ .../tts_models_to_weights/migrate.rs | 93 +++++++ .../migrations/tts_models_to_weights/mod.rs | 2 + .../queries/model_weights/migration/mod.rs | 1 + .../upsert_model_weight_from_tts_model.rs | 249 ++++++++++++++++++ 8 files changed, 406 insertions(+), 3 deletions(-) create mode 100644 crates/cli/migration-tool/src/migrations/tts_models_to_weights/copy_cloud_files.rs create mode 100644 crates/cli/migration-tool/src/migrations/tts_models_to_weights/migrate.rs create mode 100644 crates/cli/migration-tool/src/migrations/tts_models_to_weights/mod.rs create mode 100644 crates/lib/mysql_queries/src/queries/model_weights/migration/upsert_model_weight_from_tts_model.rs diff --git a/crates/_types/buckets/src/public/weight_files/bucket_file_path.rs b/crates/_types/buckets/src/public/weight_files/bucket_file_path.rs index e0e959ae0c..e1d7babedf 100644 --- a/crates/_types/buckets/src/public/weight_files/bucket_file_path.rs +++ b/crates/_types/buckets/src/public/weight_files/bucket_file_path.rs @@ -25,6 +25,11 @@ impl PublicPath for WeightFileBucketPath {} impl WeightFileBucketPath { + // TODO(bt,2024-01-14): This is temporary standardization. Clean this up. + pub fn generate_for_tt2_model() -> Self { + Self::generate_new(Some("tt2_"), Some(".pt")) + } + // TODO(bt,2023-12-19): This is temporary standardization. Clean this up. pub fn generate_for_svc_model() -> Self { Self::generate_new(Some("svc_"), Some(".pt")) @@ -35,6 +40,11 @@ impl WeightFileBucketPath { Self::generate_new(Some("rvc_"), Some(".pt")) } + // TODO(bt,2024-01-15): This is temporary standardization. Clean this up. + pub fn tt2_model_file_from_object_hash(hash: &str) -> Self { + Self::from_object_hash(hash, Some("tt2_"), Some(".pt")) + } + // TODO(bt,2023-12-19): This is temporary standardization. Clean this up. pub fn svc_model_file_from_object_hash(hash: &str) -> Self { Self::from_object_hash(hash, Some("svc_"), Some(".pt")) diff --git a/crates/cli/migration-tool/src/main.rs b/crates/cli/migration-tool/src/main.rs index 9fdf2a3452..79ac3a11a7 100644 --- a/crates/cli/migration-tool/src/main.rs +++ b/crates/cli/migration-tool/src/main.rs @@ -5,16 +5,16 @@ use std::time::Duration; + use log::info; use sqlx::{MySql, Pool}; use sqlx::mysql::MySqlPoolOptions; -use cloud_storage::bucket_client::BucketClient; +use cloud_storage::bucket_client::BucketClient; use config::shared_constants::DEFAULT_RUST_LOG; use errors::AnyhowResult; use crate::deps::Deps; -use crate::migrations::voice_conversion_to_weights::migrate::migrate_voice_conversion_to_weights; pub mod deps; pub mod migrations; @@ -40,7 +40,7 @@ pub async fn main() -> AnyhowResult<()> { bucket_production_private: get_bucket_client("PRODUCTION_PRIVATE")?, }; - migrate_voice_conversion_to_weights(&deps).await?; + //migrate_voice_conversion_to_weights(&deps).await?; Ok(()) } diff --git a/crates/cli/migration-tool/src/migrations/mod.rs b/crates/cli/migration-tool/src/migrations/mod.rs index c0183ac4c9..63f3ec6607 100644 --- a/crates/cli/migration-tool/src/migrations/mod.rs +++ b/crates/cli/migration-tool/src/migrations/mod.rs @@ -1 +1,2 @@ +pub mod tts_models_to_weights; pub mod voice_conversion_to_weights; \ No newline at end of file diff --git a/crates/cli/migration-tool/src/migrations/tts_models_to_weights/copy_cloud_files.rs b/crates/cli/migration-tool/src/migrations/tts_models_to_weights/copy_cloud_files.rs new file mode 100644 index 0000000000..41880f1e24 --- /dev/null +++ b/crates/cli/migration-tool/src/migrations/tts_models_to_weights/copy_cloud_files.rs @@ -0,0 +1,47 @@ +use tempdir::TempDir; + +use buckets::public::weight_files::bucket_file_path::WeightFileBucketPath; +use cloud_storage::bucket_path_unifier::BucketPathUnifier; +use errors::AnyhowResult; +use filesys::safe_delete_temp_directory::safe_delete_temp_directory; +use filesys::safe_delete_temp_file::safe_delete_temp_file; +use hashing::sha256::sha256_hash_file::sha256_hash_file; +use mysql_queries::queries::model_weights::migration::upsert_model_weight_from_tts_model::CopiedTtsFileData; +use mysql_queries::queries::tts::tts_models::migration::list_whole_tts_models_using_cursor::WholeTtsModelRecord; + +use crate::deps::Deps; + +pub async fn copy_cloud_files(model: &WholeTtsModelRecord, deps: &Deps) -> AnyhowResult { + let copied_file_data = copy_model(model, deps).await?; + Ok(copied_file_data) +} + +async fn copy_model(model: &WholeTtsModelRecord, deps: &Deps) -> AnyhowResult { + let bucket_path_unifier = BucketPathUnifier::default_paths(); + + let old_model_bucket_path = bucket_path_unifier.tts_synthesizer_path(&model.private_bucket_hash); + + // TODO(bt,2023-12-19): Probably faster to stream between buckets, but whatever. + let temp_dir = TempDir::new("model_transfer")?; + let model_temp_fs_path = temp_dir.path().join("model.bin"); + + deps.bucket_production_private.download_file_to_disk(&old_model_bucket_path, &model_temp_fs_path).await?; + + let file_checksum = sha256_hash_file(&model_temp_fs_path)?; + + let new_model_bucket_path = WeightFileBucketPath::generate_for_tt2_model(); + + deps.bucket_production_public.upload_filename_with_content_type( + &new_model_bucket_path.get_full_object_path_str(), + &model_temp_fs_path, + "application/octet-stream").await?; + + safe_delete_temp_file(&model_temp_fs_path); + safe_delete_temp_directory(&temp_dir); + + Ok(CopiedTtsFileData { + bucket_path: new_model_bucket_path, + file_sha_hash: file_checksum, + }) +} + diff --git a/crates/cli/migration-tool/src/migrations/tts_models_to_weights/migrate.rs b/crates/cli/migration-tool/src/migrations/tts_models_to_weights/migrate.rs new file mode 100644 index 0000000000..50f10e04c9 --- /dev/null +++ b/crates/cli/migration-tool/src/migrations/tts_models_to_weights/migrate.rs @@ -0,0 +1,93 @@ +use std::thread; +use std::time::Duration; + +use log::{info, warn}; + +use buckets::public::weight_files::bucket_file_path::WeightFileBucketPath; +use enums::by_table::model_weights::weights_types::WeightsType; +use errors::AnyhowResult; +use mysql_queries::queries::model_weights::get_weight::get_weight_by_token; +use mysql_queries::queries::model_weights::migration::upsert_model_weight_from_tts_model::{CopiedTtsFileData, upsert_model_weight_from_tts_model}; +use mysql_queries::queries::tts::tts_models::migration::list_whole_tts_models_using_cursor::{list_whole_tts_models_using_cursor, WholeTtsModelRecord}; +use tokens::tokens::model_weights::ModelWeightToken; + +use crate::deps::Deps; +use crate::migrations::tts_models_to_weights::copy_cloud_files::copy_cloud_files; + +const PAGE_SIZE: u64 = 10; + +// TODO(bt,2024-01-15): This hasn't been fully vetted / tested. Do not run until it's combed over. + +pub async fn migrate_tts_to_weights(deps: &Deps) -> AnyhowResult<()> { + + let mut cursor = 0; + + loop { + info!("Querying {PAGE_SIZE} models at cursor = {cursor}"); + + let results + = list_whole_tts_models_using_cursor(&deps.mysql_production, PAGE_SIZE, cursor).await?; + + if results.is_empty() { + warn!("No more results found; exiting."); + break; + } + + for result in results.iter() { + println!("\n\nmigrating over result {:?} : {:?} ...", result.token, result); + copy_model_record_and_bucket_files(result, &deps).await?; + } + + if let Some(last_id) = results.last().map(|result| result.id) { + cursor = last_id as u64; + } + + thread::sleep(Duration::from_secs(2)); + } + + Ok(()) +} + +async fn copy_model_record_and_bucket_files(record: &WholeTtsModelRecord, deps: &Deps) -> AnyhowResult<()> { + let mut maybe_copied_data = None; + + match &record.maybe_migration_new_model_weights_token { + None => { + info!("Copying bucket files for record {:?} ...", record.token); + let result = copy_cloud_files(record, &deps).await?; + maybe_copied_data = Some(result); + } + Some(token) => { + info!("Existing migrated record; skipping cloud bucket copy..."); + maybe_copied_data = lookup_bucket_details_from_record(token, &deps).await?; + } + } + + if let Some(copied_data) = &maybe_copied_data { + info!("Upserting records..."); + upsert_model_weight_from_tts_model(record, &deps.mysql_production, copied_data).await?; + } + + Ok(()) +} + +/// Look up the bucket storage details from a previously-migrated record +async fn lookup_bucket_details_from_record(model_weight_token: &ModelWeightToken, deps: &Deps) -> AnyhowResult> { + let maybe_model_weight_record = + get_weight_by_token(model_weight_token, true, &deps.mysql_production).await?; + + let model_weight_record = match maybe_model_weight_record { + None => return Ok(None), + Some(record) => record, + }; + + let bucket_path = match model_weight_record.weights_type { + WeightsType::Tacotron2 => WeightFileBucketPath::tt2_model_file_from_object_hash(&model_weight_record.public_bucket_hash), + _ => return Ok(None), // NB: SVC, RVC, Stable Diffusion, etc. aren't valid for TTS migration + }; + + Ok(Some(CopiedTtsFileData { + bucket_path, + file_sha_hash: model_weight_record.file_checksum_sha2, + })) +} \ No newline at end of file diff --git a/crates/cli/migration-tool/src/migrations/tts_models_to_weights/mod.rs b/crates/cli/migration-tool/src/migrations/tts_models_to_weights/mod.rs new file mode 100644 index 0000000000..a2f14892d9 --- /dev/null +++ b/crates/cli/migration-tool/src/migrations/tts_models_to_weights/mod.rs @@ -0,0 +1,2 @@ +pub mod copy_cloud_files; +pub mod migrate; \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/queries/model_weights/migration/mod.rs b/crates/lib/mysql_queries/src/queries/model_weights/migration/mod.rs index 2789fecc14..939e70a942 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/migration/mod.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/migration/mod.rs @@ -1 +1,2 @@ +pub mod upsert_model_weight_from_tts_model; pub mod upsert_model_weight_from_voice_conversion_model; \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/queries/model_weights/migration/upsert_model_weight_from_tts_model.rs b/crates/lib/mysql_queries/src/queries/model_weights/migration/upsert_model_weight_from_tts_model.rs new file mode 100644 index 0000000000..ec9f0fd24e --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/model_weights/migration/upsert_model_weight_from_tts_model.rs @@ -0,0 +1,249 @@ +use sqlx::{MySql, MySqlPool, Transaction}; + +use buckets::public::weight_files::bucket_file_path::WeightFileBucketPath; +use enums::by_table::model_weights::weights_category::WeightsCategory; +use enums::by_table::model_weights::weights_types::WeightsType; +use errors::AnyhowResult; +use tokens::tokens::model_weights::ModelWeightToken; + +use crate::queries::tts::tts_models::migration::list_whole_tts_models_using_cursor::WholeTtsModelRecord; + +pub struct CopiedTtsFileData { + pub bucket_path: WeightFileBucketPath, + pub file_sha_hash: String, +} + +/// Migrate `tts_models` records to `model_weights` + `model_weights_extension_tts_details` records. +/// This is designed to be idempotent and re-runnable. Any time we re-run the query, we should get the same result. +/// This will enable us to perfect the query and get the write flows online and switched over. +/// +pub async fn upsert_model_weight_from_tts_model( + record: &WholeTtsModelRecord, + mysql_pool: &MySqlPool, + copied_data: &CopiedTtsFileData, +) -> AnyhowResult<()> { + + let mut transaction = mysql_pool.begin().await?; + + let model_weight_token = create_or_generate_token(record); + + upsert_model_weights_record(record, &model_weight_token, copied_data, &mut transaction).await?; + upsert_model_weights_extension_record(record, &model_weight_token, &mut transaction).await?; + + update_original_record(record, &model_weight_token, &mut transaction).await?; + + transaction.commit().await?; + + Ok(()) +} + +pub fn create_or_generate_token(record: &WholeTtsModelRecord) -> ModelWeightToken { + match record.maybe_migration_new_model_weights_token { + Some(ref token) => token.clone(), + None => ModelWeightToken::generate(), + } +} + +pub async fn upsert_model_weights_record( + record: &WholeTtsModelRecord, + model_weight_token: &ModelWeightToken, + copied_data: &CopiedTtsFileData, + mut transaction: &mut Transaction<'_, MySql>, +) -> AnyhowResult<()> { + + // NB: We never supported other TTS models in the tts_models table. + const WEIGHTS_TYPE : WeightsType = WeightsType::Tacotron2; + const WEIGHTS_CATEGORY : WeightsCategory = WeightsCategory::TextToSpeech; + + // NB: Not setting a few fields (for now) + // maybe_last_update_user_token - seems like bad design + // TODO(bt): file checksum + // TODO(bt): rename maybe_public_bucket_extension to maybe_public_bucket_suffix (!!!) + // TODO(bt): do we need model_weights.ip_address_last_update without audit logs? + // TODO(bt): rename creator_ip_address to ip_address_creation (and add ip_address_last_update) + // TODO(bt): Check model_weights column integer types - signed vs unsigned + let query = sqlx::query!( + r#" +INSERT INTO model_weights +SET + token = ?, + weights_type = ?, + weights_category = ?, + title = ?, + maybe_cover_image_media_file_token = NULL, + maybe_description_markdown = ?, + maybe_description_rendered_html = ?, + creator_user_token = ?, + creator_ip_address = ?, + + creator_set_visibility = ?, + maybe_last_update_user_token = NULL, + original_download_url = ?, + original_filename = ?, + file_size_bytes = ?, + file_checksum_sha2 = ?, + + public_bucket_hash = ?, + maybe_public_bucket_prefix = ?, + maybe_public_bucket_extension = ?, + + maybe_migration_old_model_token = ?, + version = ?, + created_at = ?, + updated_at = ?, + user_deleted_at = ?, + mod_deleted_at = ? + +ON DUPLICATE KEY UPDATE + weights_type = ?, + weights_category = ?, + title = ?, + maybe_cover_image_media_file_token = NULL, + maybe_description_markdown = ?, + maybe_description_rendered_html = ?, + creator_user_token = ?, + creator_ip_address = ?, + creator_set_visibility = ?, + maybe_last_update_user_token = NULL, + original_download_url = ?, + original_filename = ?, + file_size_bytes = ?, + file_checksum_sha2 = ?, + public_bucket_hash = ?, + maybe_public_bucket_prefix = ?, + maybe_public_bucket_extension = ?, + maybe_migration_old_model_token = ?, + version = ?, + created_at = ?, + updated_at = ?, + user_deleted_at = ?, + mod_deleted_at = ? + "#, + // Insert + model_weight_token, + WEIGHTS_TYPE, + WEIGHTS_CATEGORY, + record.title, + record.description_markdown, + record.description_rendered_html, + record.creator_user_token, + record.creator_ip_address_creation, + record.creator_set_visibility.to_str(), + record.original_download_url, + record.original_filename, + record.file_size_bytes, + + copied_data.file_sha_hash, + copied_data.bucket_path.get_object_hash(), + copied_data.bucket_path.get_optional_prefix(), + copied_data.bucket_path.get_optional_extension(), + + record.token.as_str(), + record.version, + record.created_at, + record.updated_at, + record.user_deleted_at, + record.mod_deleted_at, + + // Update + WEIGHTS_TYPE, + WEIGHTS_CATEGORY, + record.title, + record.description_markdown, + record.description_rendered_html, + record.creator_user_token, + record.creator_ip_address_creation, + record.creator_set_visibility.to_str(), + record.original_download_url, + record.original_filename, + record.file_size_bytes, + + copied_data.file_sha_hash, + copied_data.bucket_path.get_object_hash(), + copied_data.bucket_path.get_optional_prefix(), + copied_data.bucket_path.get_optional_extension(), + + record.token.as_str(), + record.version, + record.created_at, + record.updated_at, + record.user_deleted_at, + record.mod_deleted_at, + ); + + let _r = query.execute(&mut **transaction).await?; + + Ok(()) +} + +pub async fn upsert_model_weights_extension_record( + record: &WholeTtsModelRecord, + model_weight_token: &ModelWeightToken, + transaction: &mut Transaction<'_, MySql> +) -> AnyhowResult<()> { + let query = sqlx::query!( + r#" +INSERT INTO model_weights_extension_tts_details +SET + model_weights_token = ?, + ietf_language_tag = ?, + ietf_primary_language_subtag = ?, + text_pipeline_type = ?, + use_default_mel_multiply_factor = ?, + maybe_custom_mel_multiply_factor = ?, + maybe_default_pretrained_vocoder = ?, + maybe_custom_vocoder_token = ? +ON DUPLICATE KEY UPDATE + ietf_language_tag = ?, + ietf_primary_language_subtag = ?, + text_pipeline_type = ?, + use_default_mel_multiply_factor = ?, + maybe_custom_mel_multiply_factor = ?, + maybe_default_pretrained_vocoder = ?, + maybe_custom_vocoder_token = ? + "#, + // Insert + &model_weight_token, + record.ietf_language_tag, + record.ietf_primary_language_subtag, + record.text_pipeline_type, + record.use_default_mel_multiply_factor, + record.maybe_custom_mel_multiply_factor, + record.maybe_default_pretrained_vocoder, + record.maybe_custom_vocoder_token, + + // Update + record.ietf_language_tag, + record.ietf_primary_language_subtag, + record.text_pipeline_type, + record.use_default_mel_multiply_factor, + record.maybe_custom_mel_multiply_factor, + record.maybe_default_pretrained_vocoder, + record.maybe_custom_vocoder_token + ); + + let _r = query.execute(&mut **transaction).await?; + + Ok(()) +} + +pub async fn update_original_record( + record: &WholeTtsModelRecord, + model_weight_token: &ModelWeightToken, + transaction: &mut Transaction<'_, MySql> +) -> AnyhowResult<()> { + let query = sqlx::query!( + r#" +UPDATE tts_models +SET + maybe_migration_new_model_weights_token = ? +WHERE token = ? + "#, + model_weight_token, + record.token, + ); + + let _r = query.execute(&mut **transaction).await?; + + Ok(()) +} From 2507cf03329659993e488ac78114a6b9e0b362d3 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 15 Jan 2024 21:48:45 -0500 Subject: [PATCH 40/54] analysis of AI streamer bug reports --- _sql/analytics/tts/tts_inference_jobs.sql | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/_sql/analytics/tts/tts_inference_jobs.sql b/_sql/analytics/tts/tts_inference_jobs.sql index 17af4d8bde..171de2b1a3 100644 --- a/_sql/analytics/tts/tts_inference_jobs.sql +++ b/_sql/analytics/tts/tts_inference_jobs.sql @@ -2,6 +2,20 @@ -- noinspection SqlNoDataSourceInspectionForFile -- noinspection SqlResolveForFile + +-- Find jobs that are broken for a given user +select id, last_assigned_worker, internal_debugging_failure_reason,created_at +from tts_inference_jobs +where maybe_creator_user_token IN ( + select token + from users + where username IN ('knifecat ', 'rewritten_code', 'johnloberger') +) + and attempt_count > 1 + and status IN ('complete_failure', 'dead', 'attempt_failed') +and created_at > ( CURDATE() - INTERVAL 5 DAY ); + + -- Total pending jobs select count(*) from tts_inference_jobs From d516058802435d4315f305153d711521247a40a4 Mon Sep 17 00:00:00 2001 From: Kasisnu <5044120+kasisnu@users.noreply.github.com> Date: Tue, 16 Jan 2024 11:36:28 +0530 Subject: [PATCH 41/54] codegen --- ...6841acb708e165e1c7855b441b9623a2151197578929.json | 12 ++++++++++++ ...64f34173f99db1e515e370c47065f932aaec56f89112.json | 12 ++++++++++++ 2 files changed, 24 insertions(+) create mode 100644 .sqlx/query-1ce9219b8f05fe54adaa6841acb708e165e1c7855b441b9623a2151197578929.json create mode 100644 .sqlx/query-aa7ec1c1e32f773212b264f34173f99db1e515e370c47065f932aaec56f89112.json diff --git a/.sqlx/query-1ce9219b8f05fe54adaa6841acb708e165e1c7855b441b9623a2151197578929.json b/.sqlx/query-1ce9219b8f05fe54adaa6841acb708e165e1c7855b441b9623a2151197578929.json new file mode 100644 index 0000000000..8fc92594f2 --- /dev/null +++ b/.sqlx/query-1ce9219b8f05fe54adaa6841acb708e165e1c7855b441b9623a2151197578929.json @@ -0,0 +1,12 @@ +{ + "db_name": "MySQL", + "query": "\nINSERT INTO model_weights_extension_tts_details\nSET\n model_weights_token = ?,\n ietf_language_tag = ?,\n ietf_primary_language_subtag = ?,\n text_pipeline_type = ?,\n use_default_mel_multiply_factor = ?,\n maybe_custom_mel_multiply_factor = ?,\n maybe_default_pretrained_vocoder = ?,\n maybe_custom_vocoder_token = ?\nON DUPLICATE KEY UPDATE\n ietf_language_tag = ?,\n ietf_primary_language_subtag = ?,\n text_pipeline_type = ?,\n use_default_mel_multiply_factor = ?,\n maybe_custom_mel_multiply_factor = ?,\n maybe_default_pretrained_vocoder = ?,\n maybe_custom_vocoder_token = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 15 + }, + "nullable": [] + }, + "hash": "1ce9219b8f05fe54adaa6841acb708e165e1c7855b441b9623a2151197578929" +} diff --git a/.sqlx/query-aa7ec1c1e32f773212b264f34173f99db1e515e370c47065f932aaec56f89112.json b/.sqlx/query-aa7ec1c1e32f773212b264f34173f99db1e515e370c47065f932aaec56f89112.json new file mode 100644 index 0000000000..34731cd396 --- /dev/null +++ b/.sqlx/query-aa7ec1c1e32f773212b264f34173f99db1e515e370c47065f932aaec56f89112.json @@ -0,0 +1,12 @@ +{ + "db_name": "MySQL", + "query": "\nUPDATE tts_models\nSET\n maybe_migration_new_model_weights_token = ?\nWHERE token = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + }, + "hash": "aa7ec1c1e32f773212b264f34173f99db1e515e370c47065f932aaec56f89112" +} From cec0cfc907aa5ca75671847a7a8836c406041061 Mon Sep 17 00:00:00 2001 From: Kasisnu <5044120+kasisnu@users.noreply.github.com> Date: Tue, 16 Jan 2024 11:39:46 +0530 Subject: [PATCH 42/54] remove misleading downloaders from styletts2 --- .../tts/styletts2/model_downloaders.rs | 71 +------------------ 1 file changed, 1 insertion(+), 70 deletions(-) diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs index c50a6dcfff..242811b95d 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs @@ -4,51 +4,6 @@ pub struct StyleTTS2Downloaders { // pub styletts2: StyleTTS2ModelMapping, } -crate::impl_model_downloader!( - StyleTTS2ModelMapping, - "Vall-e-x Model Mapping", - "VALL_E_X_MAPPING_BUCKET_PATH", - "/dependencies/zero_shot_tts/vall-e-x_1.0/StyleTTS2-checkpoint.pt", - "VALL_E_X_MAPPING_FILESYSTEM_PATH", - "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/StyleTTS2-checkpoint.pt" -); - -crate::impl_model_downloader!( - WhisperLargeModelMapping, - "whisper Large Model Mapping", - "WHISPER_LARGE_MAPPING_BUCKET_PATH", - "/dependencies/zero_shot_tts/vall-e-x_1.0/whisper-large/model.safetensors", - "WHISPER_LARGE_MAPPING_FILESYSTEM_PATH", - "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper-large/model.safetensors" -); - -crate::impl_model_downloader!( - WhisperMediumModelMapping, - "whisper Medium Model Mapping", - "WHISPER_MEDIUM_TENSOR_BUCKET_PATH", - "/dependencies/zero_shot_tts/vall-e-x_1.0/whisper/medium.pt", - "WHISPER_MEDIUM_TENSOR_FILESYSTEM_PATH", - "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/whisper/medium.pt" -); - -crate::impl_model_downloader!( - VocosEncodecModelMapping, - "Vocos Encodec Model Mapping", - "VOCOS_ENCODEC_TENSOR_BUCKET_PATH", - "/dependencies/zero_shot_tts/vall-e-x_1.0/vocos-encodec/encodec_pytorch_model.bin", - "VOCOS_ENCODEC_TENSOR_FILESYSTEM_PATH", - "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/encodec_pytorch_model.bin" -); - -crate::impl_model_downloader!( - VocosEncodecConfigMapping, - "Vocos Config Mapping", - "VOCOS_ENCODEC_CONFIG_BUCKET_PATH", - "/dependencies/zero_shot_tts/vall-e-x_1.0/vocos-encodec/config.yaml", - "VOCOS_ENCODEC_CONFIG_FILESYSTEM_PATH", - "/tmp/downloads/zero_shot_tts/vall-e-x_1.0/vocos-encodec/config.yaml" -); - impl StyleTTS2Downloaders { // TODO(KS): All checkpoints are in the container, so we don't need to download them here. // Maybe they should be here instead @@ -60,28 +15,4 @@ impl StyleTTS2Downloaders { pub fn all_downloaders(&self) -> Vec<&dyn ModelDownloader> { vec![] } -} - -mod test { - // use std::path::Path; - - // use crate::job::job_types::lipsync::sad_talker::model_downloaders::{SadTalkerModelFirstMapping, SadTalkerModelSecondMapping}; - // use crate::util::model_downloader::ModelDownloader; - - // #[test] - // fn test_downloader_model_first_mapping() { - // let model = SadTalkerModelFirstMapping::default(); - // assert_eq!(model.get_model_name(), "sad talker model first mapping"); - // assert_eq!(model.get_cloud_bucket_path(), "/animation_sadtalker/model/mapping_00109-model.pth.tar"); - // assert_eq!(model.get_filesystem_path(), Path::new("/tmp/downloads/sadtalker/mapping_00109-model.pth.tar")); - // } - - // #[test] - // fn test_downloader_model_second_mapping() { - // let model = SadTalkerModelSecondMapping::default(); - // assert_eq!(model.get_model_name(), "sad talker model second mapping"); - // assert_eq!(model.get_cloud_bucket_path(), "/animation_sadtalker/model/mapping_00229-model.pth.tar"); - // assert_eq!(model.get_filesystem_path(), Path::new("/tmp/downloads/sadtalker/mapping_00229-model.pth.tar")); - // } - -} +} \ No newline at end of file From 18d46a0f5762a432dbfedfeec3c6c85a991eaeed Mon Sep 17 00:00:00 2001 From: Kasisnu <5044120+kasisnu@users.noreply.github.com> Date: Tue, 16 Jan 2024 13:06:29 +0530 Subject: [PATCH 43/54] bugfix: ignore more than one samples for zs --- .../src/job/job_types/tts/styletts2/process_job_create_voice.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs index 3c890b3cea..f6ef449880 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs @@ -160,7 +160,7 @@ pub async fn process_create_voice( // TODO(KS): This is a hack to get only the first file to download // Ideally we could have improved logic to pick the "best" media file from a dataset - if (index > 0) { + if (index >= 0) { break; } } From 7cc310764e22bdf7a1172bbc58a4a0a84bd0ce88 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 16 Jan 2024 19:10:23 -0500 Subject: [PATCH 44/54] new tables: idempotency_tokens, batch_generations; new column --- .../down.sql | 5 +++ .../up.sql | 19 +++++++++ .../down.sql | 5 +++ .../up.sql | 42 +++++++++++++++++++ .../down.sql | 5 +++ .../up.sql | 7 ++++ _sql/migrations_squashed/media_files.sql | 5 +++ 7 files changed, 88 insertions(+) create mode 100644 _sql/migrations/2024-01-16-231833_idempotency_tokens/down.sql create mode 100644 _sql/migrations/2024-01-16-231833_idempotency_tokens/up.sql create mode 100644 _sql/migrations/2024-01-16-231841_batch_generations/down.sql create mode 100644 _sql/migrations/2024-01-16-231841_batch_generations/up.sql create mode 100644 _sql/migrations/2024-01-16-231852_alter_media_files_add_batch_token/down.sql create mode 100644 _sql/migrations/2024-01-16-231852_alter_media_files_add_batch_token/up.sql diff --git a/_sql/migrations/2024-01-16-231833_idempotency_tokens/down.sql b/_sql/migrations/2024-01-16-231833_idempotency_tokens/down.sql new file mode 100644 index 0000000000..ab73e06165 --- /dev/null +++ b/_sql/migrations/2024-01-16-231833_idempotency_tokens/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +drop table idempotency_tokens; \ No newline at end of file diff --git a/_sql/migrations/2024-01-16-231833_idempotency_tokens/up.sql b/_sql/migrations/2024-01-16-231833_idempotency_tokens/up.sql new file mode 100644 index 0000000000..5667756d56 --- /dev/null +++ b/_sql/migrations/2024-01-16-231833_idempotency_tokens/up.sql @@ -0,0 +1,19 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- A form of optimistic locking +CREATE TABLE idempotency_tokens ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- This can be a UUID or any entropic value. + idempotency_value VARCHAR(36) NOT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (idempotency_value) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_sql/migrations/2024-01-16-231841_batch_generations/down.sql b/_sql/migrations/2024-01-16-231841_batch_generations/down.sql new file mode 100644 index 0000000000..419c98f550 --- /dev/null +++ b/_sql/migrations/2024-01-16-231841_batch_generations/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +drop table batch_generations; diff --git a/_sql/migrations/2024-01-16-231841_batch_generations/up.sql b/_sql/migrations/2024-01-16-231841_batch_generations/up.sql new file mode 100644 index 0000000000..d738abb9f5 --- /dev/null +++ b/_sql/migrations/2024-01-16-231841_batch_generations/up.sql @@ -0,0 +1,42 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- Batch generations +CREATE TABLE batch_generations ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- The batch token + -- This is **NOT** a unique token. + token VARCHAR(32) NOT NULL, + + -- ========== ENTITY LINKAGE ========== + + -- The type of record this is attached to + -- Currently supported values: + -- * 'model_weight' + -- * 'media_file' + -- * 'comment' + entity_type VARCHAR(32) NOT NULL, + + -- The token for the record. + entity_token VARCHAR(32) NOT NULL, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token, entity_type, entity_token), + KEY index_token (token), + KEY index_entity_token (entity_token) + + -- NB: Realistically, we probably don't need this index until we have a lot of entity types, and the entity_token + -- entropy is probably sufficient enough as-is. + -- KEY index_entity_type_and_token (entity_type, entity_token), + + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_sql/migrations/2024-01-16-231852_alter_media_files_add_batch_token/down.sql b/_sql/migrations/2024-01-16-231852_alter_media_files_add_batch_token/down.sql new file mode 100644 index 0000000000..924a36d70c --- /dev/null +++ b/_sql/migrations/2024-01-16-231852_alter_media_files_add_batch_token/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE media_files DROP COLUMN is_batch_generated; diff --git a/_sql/migrations/2024-01-16-231852_alter_media_files_add_batch_token/up.sql b/_sql/migrations/2024-01-16-231852_alter_media_files_add_batch_token/up.sql new file mode 100644 index 0000000000..1cbc612de5 --- /dev/null +++ b/_sql/migrations/2024-01-16-231852_alter_media_files_add_batch_token/up.sql @@ -0,0 +1,7 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE media_files + ADD COLUMN is_batch_generated BOOLEAN NOT NULL DEFAULT FALSE + AFTER maybe_origin_filename; diff --git a/_sql/migrations_squashed/media_files.sql b/_sql/migrations_squashed/media_files.sql index e04cb2de30..35d7334ca9 100644 --- a/_sql/migrations_squashed/media_files.sql +++ b/_sql/migrations_squashed/media_files.sql @@ -58,6 +58,11 @@ CREATE TABLE media_files ( -- The original filename of the media (if uploaded by a user) maybe_origin_filename VARCHAR(255) DEFAULT NULL, + -- Whether this media file was generated as part of a batch + -- If so, we can look up the batch in a separate query/call to the `batch_generations` table. + -- We won't hold the batch token in this table since it'll be very sparsely populated. + is_batch_generated BOOLEAN NOT NULL DEFAULT FALSE, + -- ========== METADATA ========== -- We can store text transcripts, etc. in other tables and refer back via foreign key. From 2d4d2c667f794d6de694280004a58086f34d4ff1 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 16 Jan 2024 19:49:08 -0500 Subject: [PATCH 45/54] reference schemas --- .../migrations_squashed/batch_generations.sql | 42 +++++++++++++++++++ .../idempotency_tokens.sql | 19 +++++++++ 2 files changed, 61 insertions(+) create mode 100644 _sql/migrations_squashed/batch_generations.sql create mode 100644 _sql/migrations_squashed/idempotency_tokens.sql diff --git a/_sql/migrations_squashed/batch_generations.sql b/_sql/migrations_squashed/batch_generations.sql new file mode 100644 index 0000000000..d738abb9f5 --- /dev/null +++ b/_sql/migrations_squashed/batch_generations.sql @@ -0,0 +1,42 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- Batch generations +CREATE TABLE batch_generations ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- The batch token + -- This is **NOT** a unique token. + token VARCHAR(32) NOT NULL, + + -- ========== ENTITY LINKAGE ========== + + -- The type of record this is attached to + -- Currently supported values: + -- * 'model_weight' + -- * 'media_file' + -- * 'comment' + entity_type VARCHAR(32) NOT NULL, + + -- The token for the record. + entity_token VARCHAR(32) NOT NULL, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token, entity_type, entity_token), + KEY index_token (token), + KEY index_entity_token (entity_token) + + -- NB: Realistically, we probably don't need this index until we have a lot of entity types, and the entity_token + -- entropy is probably sufficient enough as-is. + -- KEY index_entity_type_and_token (entity_type, entity_token), + + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_sql/migrations_squashed/idempotency_tokens.sql b/_sql/migrations_squashed/idempotency_tokens.sql new file mode 100644 index 0000000000..5667756d56 --- /dev/null +++ b/_sql/migrations_squashed/idempotency_tokens.sql @@ -0,0 +1,19 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- A form of optimistic locking +CREATE TABLE idempotency_tokens ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- This can be a UUID or any entropic value. + idempotency_value VARCHAR(36) NOT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (idempotency_value) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; From 1771e1dffe73da67ee142e6614fd709471be6e72 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 16 Jan 2024 20:52:11 -0500 Subject: [PATCH 46/54] fix lint/test failures --- .../inference_model_type.rs | 2 +- crates/cli/dev-database-seed/src/main.rs | 3 +- crates/cli/migration-tool/src/cli_args.rs | 47 +++++++++++++++++++ crates/cli/migration-tool/src/main.rs | 16 ++++++- .../tts/styletts2/model_downloaders.rs | 7 +-- .../tts/styletts2/process_job_create_voice.rs | 2 +- .../styletts2/process_job_inference_voice.rs | 2 +- 7 files changed, 70 insertions(+), 9 deletions(-) create mode 100644 crates/cli/migration-tool/src/cli_args.rs diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs index 48d7de9dce..a779157b7e 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs @@ -139,7 +139,7 @@ mod tests { fn all_variants() { // Static check let mut variants = InferenceModelType::all_variants(); - assert_eq!(variants.len(), 8); + assert_eq!(variants.len(), 9); assert_eq!(variants.pop_first(), Some(InferenceModelType::RvcV2)); assert_eq!(variants.pop_first(), Some(InferenceModelType::SadTalker)); assert_eq!(variants.pop_first(), Some(InferenceModelType::SoVitsSvc)); diff --git a/crates/cli/dev-database-seed/src/main.rs b/crates/cli/dev-database-seed/src/main.rs index 12d16ddd0b..6450f5c052 100644 --- a/crates/cli/dev-database-seed/src/main.rs +++ b/crates/cli/dev-database-seed/src/main.rs @@ -8,12 +8,11 @@ use errors::AnyhowResult; use crate::bucket_clients::get_bucket_clients; use crate::cli_args::parse_cli_args; 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::seed_model_weights::seed_weights_files; // V2 v3 will come later on when everything is more conslidated and clear. use crate::seeding::tts_tacotron2::seed_tts_tacotron2; use crate::seeding::users::seed_user_accounts; use crate::seeding::voice_conversion::seed_voice_conversion; -use crate::seeding::zero_shot_tts::seed_zero_shot_tts; pub mod bucket_clients; diff --git a/crates/cli/migration-tool/src/cli_args.rs b/crates/cli/migration-tool/src/cli_args.rs new file mode 100644 index 0000000000..0d39c15587 --- /dev/null +++ b/crates/cli/migration-tool/src/cli_args.rs @@ -0,0 +1,47 @@ +use std::str::FromStr; + +use clap::Parser; +use strum::{EnumCount, EnumString, IntoEnumIterator}; +use strum::EnumIter; + +use errors::{anyhow, AnyhowResult}; + +pub struct ParsedArgs { + pub action: Action, +} + +#[derive(Clone, Copy, Debug, EnumIter, EnumCount, EnumString, strum::Display)] +#[strum(serialize_all = "snake_case")] +pub enum Action { + /// Migrate VC + MigrateVoiceConversion, + /// Migrate TTS + MigrateTts, +} + +#[derive(Parser, Debug)] +#[command(name="migration-tool")] +pub struct Args { + #[arg(name="action", long="action", help="action to take", required=true)] + action: String, +} + +pub fn parse_cli_args() -> AnyhowResult { + let args = Args::parse(); + + Ok(ParsedArgs { + action: action_from_str(&args.action)?, + }) +} + + +fn action_from_str(value: &str) -> AnyhowResult { + let action = Action::from_str(value) + .map_err(|err| { + let choices = Action::iter() + .map(|e| e.to_string()) + .collect::>(); + anyhow!("parse error: {:?}, provided: \"{}\" choices: {:?}", err, value, choices) + })?; + Ok(action) +} diff --git a/crates/cli/migration-tool/src/main.rs b/crates/cli/migration-tool/src/main.rs index 79ac3a11a7..71a4270e68 100644 --- a/crates/cli/migration-tool/src/main.rs +++ b/crates/cli/migration-tool/src/main.rs @@ -14,8 +14,12 @@ use cloud_storage::bucket_client::BucketClient; use config::shared_constants::DEFAULT_RUST_LOG; use errors::AnyhowResult; +use crate::cli_args::{Action, parse_cli_args}; use crate::deps::Deps; +use crate::migrations::tts_models_to_weights::migrate::migrate_tts_to_weights; +use crate::migrations::voice_conversion_to_weights::migrate::migrate_voice_conversion_to_weights; +pub mod cli_args; pub mod deps; pub mod migrations; @@ -40,10 +44,20 @@ pub async fn main() -> AnyhowResult<()> { bucket_production_private: get_bucket_client("PRODUCTION_PRIVATE")?, }; - //migrate_voice_conversion_to_weights(&deps).await?; + let args = parse_cli_args()?; + + match args.action { + Action::MigrateVoiceConversion => { + migrate_voice_conversion_to_weights(&deps).await?; + } + Action::MigrateTts => { + migrate_tts_to_weights(&deps).await?; + } + } Ok(()) } + async fn get_mysql(env_var_name: &str) -> AnyhowResult> { info!("Connecting to MySQL {env_var_name}..."); diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs index 242811b95d..fd32b2b2f8 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs @@ -1,8 +1,9 @@ use crate::util::model_downloader::ModelDownloader; -pub struct StyleTTS2Downloaders { - // pub styletts2: StyleTTS2ModelMapping, -} +pub struct StyleTTS2Downloaders; +//{ +// pub styletts2: StyleTTS2ModelMapping, +//} impl StyleTTS2Downloaders { // TODO(KS): All checkpoints are in the container, so we don't need to download them here. diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs index f6ef449880..92b36939f8 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs @@ -160,7 +160,7 @@ pub async fn process_create_voice( // TODO(KS): This is a hack to get only the first file to download // Ideally we could have improved logic to pick the "best" media file from a dataset - if (index >= 0) { + if downloaded_dataset.len() >= 1 { break; } } diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs index b602d04fbf..199181d6f1 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs @@ -102,7 +102,7 @@ pub async fn process_inference_voice( .new_tempdir(&work_temp_dir) .map_err(|e| ProcessSingleJobError::from_io_error(e))?; - let workdir = work_temp_dir.path().to_path_buf(); + let _workdir = work_temp_dir.path().to_path_buf(); let filename = "style.npz".to_string(); let mut downloaded_weights_path = work_temp_dir.path().to_path_buf(); From 2487a8d64687263bef9e02b40a06a3472ece846f Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 16 Jan 2024 21:03:56 -0500 Subject: [PATCH 47/54] refactor query modules --- crates/lib/mysql_queries/src/queries/mod.rs | 1 - .../mysql_queries/src/queries/users/mod.rs | 1 + .../batch_get_user_bookmarks.rs | 0 .../user_bookmarks/delete_user_bookmark.rs | 0 .../get_total_bookmark_count_for_entity.rs | 4 +- .../user_bookmarks/get_user_bookmark.rs | 0 ...get_user_bookmark_transactional_locking.rs | 2 +- .../user_bookmarks/list_user_bookmarks.rs | 4 +- .../list_user_bookmarks_for_entity.rs | 2 +- .../list_user_bookmarks_result.rs | 74 +++++++++---------- .../queries/{ => users}/user_bookmarks/mod.rs | 0 .../user_bookmarks/upsert_user_bookmark.rs | 2 +- .../user_bookmark_entity_token.rs | 0 .../batch_get_user_bookmarks_handler.rs | 2 +- .../create_user_bookmark_handler.rs | 8 +- .../delete_user_bookmark_handler.rs | 4 +- .../list_user_bookmarks_for_entity_handler.rs | 4 +- ...list_user_bookmarks_for_session_handler.rs | 2 +- .../list_user_bookmarks_for_user_handler.rs | 2 +- 19 files changed, 56 insertions(+), 56 deletions(-) rename crates/lib/mysql_queries/src/queries/{ => users}/user_bookmarks/batch_get_user_bookmarks.rs (100%) rename crates/lib/mysql_queries/src/queries/{ => users}/user_bookmarks/delete_user_bookmark.rs (100%) rename crates/lib/mysql_queries/src/queries/{ => users}/user_bookmarks/get_total_bookmark_count_for_entity.rs (91%) rename crates/lib/mysql_queries/src/queries/{ => users}/user_bookmarks/get_user_bookmark.rs (100%) rename crates/lib/mysql_queries/src/queries/{ => users}/user_bookmarks/get_user_bookmark_transactional_locking.rs (98%) rename crates/lib/mysql_queries/src/queries/{ => users}/user_bookmarks/list_user_bookmarks.rs (98%) rename crates/lib/mysql_queries/src/queries/{ => users}/user_bookmarks/list_user_bookmarks_for_entity.rs (96%) rename crates/lib/mysql_queries/src/queries/{ => users}/user_bookmarks/list_user_bookmarks_result.rs (82%) rename crates/lib/mysql_queries/src/queries/{ => users}/user_bookmarks/mod.rs (100%) rename crates/lib/mysql_queries/src/queries/{ => users}/user_bookmarks/upsert_user_bookmark.rs (94%) rename crates/lib/mysql_queries/src/queries/{ => users}/user_bookmarks/user_bookmark_entity_token.rs (100%) diff --git a/crates/lib/mysql_queries/src/queries/mod.rs b/crates/lib/mysql_queries/src/queries/mod.rs index a02c0a3a01..55f1039245 100644 --- a/crates/lib/mysql_queries/src/queries/mod.rs +++ b/crates/lib/mysql_queries/src/queries/mod.rs @@ -18,7 +18,6 @@ pub mod stats; pub mod trending_model_analytics; pub mod tts; pub mod twitch; -pub mod user_bookmarks; pub mod users; pub mod vocoder; pub mod voice_clone_requests; diff --git a/crates/lib/mysql_queries/src/queries/users/mod.rs b/crates/lib/mysql_queries/src/queries/users/mod.rs index 1b5e04bc70..e64f9df9c5 100644 --- a/crates/lib/mysql_queries/src/queries/users/mod.rs +++ b/crates/lib/mysql_queries/src/queries/users/mod.rs @@ -1,5 +1,6 @@ pub mod user; pub mod user_badges; +pub mod user_bookmarks; pub mod user_password_resets; pub mod user_profiles; pub mod user_ratings; diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/batch_get_user_bookmarks.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/batch_get_user_bookmarks.rs similarity index 100% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/batch_get_user_bookmarks.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/batch_get_user_bookmarks.rs diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/delete_user_bookmark.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/delete_user_bookmark.rs similarity index 100% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/delete_user_bookmark.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/delete_user_bookmark.rs diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/get_total_bookmark_count_for_entity.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_total_bookmark_count_for_entity.rs similarity index 91% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/get_total_bookmark_count_for_entity.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_total_bookmark_count_for_entity.rs index c366eab630..6e153a5219 100644 --- a/crates/lib/mysql_queries/src/queries/user_bookmarks/get_total_bookmark_count_for_entity.rs +++ b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_total_bookmark_count_for_entity.rs @@ -1,11 +1,11 @@ use anyhow::anyhow; use log::warn; -use sqlx::{MySql}; +use sqlx::MySql; use sqlx::pool::PoolConnection; use errors::AnyhowResult; -use crate::queries::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; +use crate::queries::users::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; pub struct BookmarkCount { pub total_count: usize, diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/get_user_bookmark.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_user_bookmark.rs similarity index 100% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/get_user_bookmark.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_user_bookmark.rs diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/get_user_bookmark_transactional_locking.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_user_bookmark_transactional_locking.rs similarity index 98% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/get_user_bookmark_transactional_locking.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_user_bookmark_transactional_locking.rs index b849cc655d..bbc1490f27 100644 --- a/crates/lib/mysql_queries/src/queries/user_bookmarks/get_user_bookmark_transactional_locking.rs +++ b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_user_bookmark_transactional_locking.rs @@ -7,7 +7,7 @@ use errors::AnyhowResult; use tokens::tokens::user_bookmarks::UserBookmarkToken; use tokens::tokens::users::UserToken; -use crate::queries::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; +use crate::queries::users::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; pub struct UserBookmark { pub token: UserBookmarkToken, diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks.rs similarity index 98% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks.rs index 8dbe73c1dc..4d8a1b71ef 100644 --- a/crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks.rs +++ b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks.rs @@ -6,8 +6,8 @@ use enums::by_table::model_weights::weights_types::WeightsType; use enums::by_table::user_bookmarks::user_bookmark_entity_type::UserBookmarkEntityType; use errors::AnyhowResult; -use crate::queries::user_bookmarks::list_user_bookmarks_result::RawUserBookmarkRecord; -use crate::queries::user_bookmarks::list_user_bookmarks_result::UserBookmark; +use crate::queries::users::user_bookmarks::list_user_bookmarks_result::RawUserBookmarkRecord; +use crate::queries::users::user_bookmarks::list_user_bookmarks_result::UserBookmark; pub struct UserBookmarkListPage { pub results: Vec, diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks_for_entity.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks_for_entity.rs similarity index 96% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks_for_entity.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks_for_entity.rs index 3df0f0812c..785f1ac802 100644 --- a/crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks_for_entity.rs +++ b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks_for_entity.rs @@ -7,7 +7,7 @@ use errors::AnyhowResult; use tokens::tokens::user_bookmarks::UserBookmarkToken; use tokens::tokens::users::UserToken; -use crate::queries::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; +use crate::queries::users::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; pub struct UserBookmark { pub token: UserBookmarkToken, diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks_result.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks_result.rs similarity index 82% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks_result.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks_result.rs index 88c6954d45..2184350686 100644 --- a/crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks_result.rs +++ b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks_result.rs @@ -70,56 +70,56 @@ pub struct UserBookmark { } pub struct RawUserBookmarkRecord { - pub (crate) token: UserBookmarkToken, + pub(crate) token: UserBookmarkToken, - pub (crate) entity_type: UserBookmarkEntityType, - pub (crate) entity_token: String, + pub(crate) entity_type: UserBookmarkEntityType, + pub(crate) entity_token: String, // TODO(bt,2023-12-30): I don't think these user fields are necessary. // We already know the user we're querying on behalf of, so presumably we // already have their info (ie. profile page use). - pub (crate) user_token: UserToken, - pub (crate) username: String, - pub (crate) user_display_name: String, - pub (crate) user_gravatar_hash: String, + pub(crate) user_token: UserToken, + pub(crate) username: String, + pub(crate) user_display_name: String, + pub(crate) user_gravatar_hash: String, - pub (crate) created_at: DateTime, - pub (crate) updated_at: DateTime, - pub (crate) deleted_at: Option>, + pub(crate) created_at: DateTime, + pub(crate) updated_at: DateTime, + pub(crate) deleted_at: Option>, - pub (crate) maybe_ratings_positive_count: Option, - pub (crate) maybe_ratings_negative_count: Option, - pub (crate) maybe_bookmark_count: Option, + pub(crate) maybe_ratings_positive_count: Option, + pub(crate) maybe_ratings_negative_count: Option, + pub(crate) maybe_bookmark_count: Option, - pub (crate) maybe_media_file_type: Option, - pub (crate) maybe_media_file_origin_category: Option, - pub (crate) maybe_media_file_public_bucket_hash: Option, - pub (crate) maybe_media_file_public_bucket_prefix: Option, - pub (crate) maybe_media_file_public_bucket_extension: Option, + pub(crate) maybe_media_file_type: Option, + pub(crate) maybe_media_file_origin_category: Option, + pub(crate) maybe_media_file_public_bucket_hash: Option, + pub(crate) maybe_media_file_public_bucket_prefix: Option, + pub(crate) maybe_media_file_public_bucket_extension: Option, - pub (crate) maybe_media_file_creator_user_token: Option, - pub (crate) maybe_media_file_creator_username: Option, - pub (crate) maybe_media_file_creator_display_name: Option, - pub (crate) maybe_media_file_creator_gravatar_hash: Option, + pub(crate) maybe_media_file_creator_user_token: Option, + pub(crate) maybe_media_file_creator_username: Option, + pub(crate) maybe_media_file_creator_display_name: Option, + pub(crate) maybe_media_file_creator_gravatar_hash: Option, - pub (crate) maybe_model_weight_type: Option, - pub (crate) maybe_model_weight_category: Option, + pub(crate) maybe_model_weight_type: Option, + pub(crate) maybe_model_weight_category: Option, - pub (crate) maybe_model_weight_cover_image_public_bucket_hash: Option, - pub (crate) maybe_model_weight_cover_image_public_bucket_prefix: Option, - pub (crate) maybe_model_weight_cover_image_public_bucket_extension: Option, + pub(crate) maybe_model_weight_cover_image_public_bucket_hash: Option, + pub(crate) maybe_model_weight_cover_image_public_bucket_prefix: Option, + pub(crate) maybe_model_weight_cover_image_public_bucket_extension: Option, - pub (crate) maybe_model_weight_creator_user_token: Option, - pub (crate) maybe_model_weight_creator_username: Option, - pub (crate) maybe_model_weight_creator_display_name: Option, - pub (crate) maybe_model_weight_creator_gravatar_hash: Option, + pub(crate) maybe_model_weight_creator_user_token: Option, + pub(crate) maybe_model_weight_creator_username: Option, + pub(crate) maybe_model_weight_creator_display_name: Option, + pub(crate) maybe_model_weight_creator_gravatar_hash: Option, - pub (crate) maybe_descriptive_text_model_weight_title: Option, - pub (crate) maybe_descriptive_text_tts_model_title: Option, - pub (crate) maybe_descriptive_text_tts_result_inference_text: Option, - pub (crate) maybe_descriptive_text_user_display_name: Option, - pub (crate) maybe_descriptive_text_voice_conversion_model_title: Option, - pub (crate) maybe_descriptive_text_zs_voice_title: Option, + pub(crate) maybe_descriptive_text_model_weight_title: Option, + pub(crate) maybe_descriptive_text_tts_model_title: Option, + pub(crate) maybe_descriptive_text_tts_result_inference_text: Option, + pub(crate) maybe_descriptive_text_user_display_name: Option, + pub(crate) maybe_descriptive_text_voice_conversion_model_title: Option, + pub(crate) maybe_descriptive_text_zs_voice_title: Option, } impl RawUserBookmarkRecord { diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/mod.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/mod.rs similarity index 100% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/mod.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/mod.rs diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/upsert_user_bookmark.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/upsert_user_bookmark.rs similarity index 94% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/upsert_user_bookmark.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/upsert_user_bookmark.rs index ef361360a6..34062a5594 100644 --- a/crates/lib/mysql_queries/src/queries/user_bookmarks/upsert_user_bookmark.rs +++ b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/upsert_user_bookmark.rs @@ -7,7 +7,7 @@ use errors::AnyhowResult; use tokens::tokens::user_bookmarks::UserBookmarkToken; use tokens::tokens::users::UserToken; -use crate::queries::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; +use crate::queries::users::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; pub struct CreateUserBookmarkArgs<'e, 'c, E> where E: 'e + Executor<'c, Database = MySql> diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/user_bookmark_entity_token.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/user_bookmark_entity_token.rs similarity index 100% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/user_bookmark_entity_token.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/user_bookmark_entity_token.rs diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/batch_get_user_bookmarks_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/batch_get_user_bookmarks_handler.rs index 212d5b4e36..7f1ba945e6 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/batch_get_user_bookmarks_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/batch_get_user_bookmarks_handler.rs @@ -9,7 +9,7 @@ use utoipa::{IntoParams, ToSchema}; use enums::by_table::user_bookmarks::user_bookmark_entity_type::UserBookmarkEntityType; use http_server_common::response::serialize_as_json_error::serialize_as_json_error; -use mysql_queries::queries::user_bookmarks::batch_get_user_bookmarks::{batch_get_user_bookmarks, BatchUserBookmark}; +use mysql_queries::queries::users::user_bookmarks::batch_get_user_bookmarks::{batch_get_user_bookmarks, BatchUserBookmark}; use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::tts_models::TtsModelToken; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/create_user_bookmark_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/create_user_bookmark_handler.rs index 638ec1500b..49e9e391da 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/create_user_bookmark_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/create_user_bookmark_handler.rs @@ -16,10 +16,10 @@ use utoipa::ToSchema; use enums::by_table::user_bookmarks::user_bookmark_entity_type::UserBookmarkEntityType; use mysql_queries::queries::entity_stats::stats_entity_token::StatsEntityToken; use mysql_queries::queries::entity_stats::upsert_entity_stats_on_bookmark_event::{BookmarkAction, upsert_entity_stats_on_bookmark_event, UpsertEntityStatsArgs}; -use mysql_queries::queries::user_bookmarks::upsert_user_bookmark::{upsert_user_bookmark, CreateUserBookmarkArgs}; -use mysql_queries::queries::user_bookmarks::get_total_bookmark_count_for_entity::get_total_bookmark_count_for_entity; -use mysql_queries::queries::user_bookmarks::get_user_bookmark_transactional_locking::{BookmarkIdentifier, get_user_bookmark_transactional_locking}; -use mysql_queries::queries::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; +use mysql_queries::queries::users::user_bookmarks::upsert_user_bookmark::{upsert_user_bookmark, CreateUserBookmarkArgs}; +use mysql_queries::queries::users::user_bookmarks::get_total_bookmark_count_for_entity::get_total_bookmark_count_for_entity; +use mysql_queries::queries::users::user_bookmarks::get_user_bookmark_transactional_locking::{BookmarkIdentifier, get_user_bookmark_transactional_locking}; +use mysql_queries::queries::users::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::tts_models::TtsModelToken; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/delete_user_bookmark_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/delete_user_bookmark_handler.rs index 5b3fd42eb1..dbadd9ddf8 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/delete_user_bookmark_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/delete_user_bookmark_handler.rs @@ -16,8 +16,8 @@ use utoipa::ToSchema; use mysql_queries::queries::entity_stats::stats_entity_token::StatsEntityToken; use mysql_queries::queries::entity_stats::upsert_entity_stats_on_bookmark_event::{BookmarkAction, upsert_entity_stats_on_bookmark_event, UpsertEntityStatsArgs}; -use mysql_queries::queries::user_bookmarks::delete_user_bookmark::delete_user_bookmark; -use mysql_queries::queries::user_bookmarks::get_user_bookmark_transactional_locking::{BookmarkIdentifier, get_user_bookmark_transactional_locking}; +use mysql_queries::queries::users::user_bookmarks::delete_user_bookmark::delete_user_bookmark; +use mysql_queries::queries::users::user_bookmarks::get_user_bookmark_transactional_locking::{BookmarkIdentifier, get_user_bookmark_transactional_locking}; use tokens::tokens::user_bookmarks::UserBookmarkToken; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_entity_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_entity_handler.rs index abeb7fa0ca..b69dbf6a1e 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_entity_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_entity_handler.rs @@ -15,8 +15,8 @@ use log::warn; use utoipa::ToSchema; use enums::by_table::user_bookmarks::user_bookmark_entity_type::UserBookmarkEntityType; -use mysql_queries::queries::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; -use mysql_queries::queries::user_bookmarks::list_user_bookmarks_for_entity::list_user_bookmarks_for_entity; +use mysql_queries::queries::users::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; +use mysql_queries::queries::users::user_bookmarks::list_user_bookmarks_for_entity::list_user_bookmarks_for_entity; use tokens::tokens::user_bookmarks::UserBookmarkToken; use crate::http_server::common_responses::user_details_lite::{DefaultAvatarInfo, UserDetailsLight}; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_session_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_session_handler.rs index 5c4a30fe01..3b10b10df9 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_session_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_session_handler.rs @@ -10,7 +10,7 @@ use chrono::{DateTime, Utc}; use log::warn; use enums::by_table::user_bookmarks::user_bookmark_entity_type::UserBookmarkEntityType; -use mysql_queries::queries::user_bookmarks::list_user_bookmarks::{ list_user_bookmarks_by_maybe_entity_type}; +use mysql_queries::queries::users::user_bookmarks::list_user_bookmarks::{ list_user_bookmarks_by_maybe_entity_type}; use tokens::tokens::user_bookmarks::UserBookmarkToken; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_user_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_user_handler.rs index a1771118fe..0bc22eb325 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_user_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_user_handler.rs @@ -19,7 +19,7 @@ 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 enums::by_table::user_bookmarks::user_bookmark_entity_type::UserBookmarkEntityType; -use mysql_queries::queries::user_bookmarks::list_user_bookmarks::{list_user_bookmarks_by_maybe_entity_type, ListUserBookmarksForUserArgs}; +use mysql_queries::queries::users::user_bookmarks::list_user_bookmarks::{list_user_bookmarks_by_maybe_entity_type, ListUserBookmarksForUserArgs}; use tokens::tokens::user_bookmarks::UserBookmarkToken; use crate::http_server::common_responses::pagination_page::PaginationPage; From 33613cb5d8b6e6679533cd04c0c737d66f3b0779 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 16 Jan 2024 23:00:17 -0500 Subject: [PATCH 48/54] insert batch generation records --- Cargo.lock | 21 ++++ Cargo.toml | 1 + .../_types/composite_identifiers/Cargo.toml | 40 +++++++ .../batch_generation_entity.rs | 28 +++++ .../src/by_table/batch_generations/mod.rs | 1 + .../composite_identifiers/src/by_table/mod.rs | 1 + .../_types/composite_identifiers/src/lib.rs | 22 ++++ .../batch_generation_entity_type.rs | 109 ++++++++++++++++++ .../src/by_table/batch_generations/mod.rs | 1 + crates/_types/enums/src/by_table/mod.rs | 2 + crates/_types/tokens/src/prefixes.rs | 10 +- .../tokens/src/tokens/batch_generations.rs | 15 +++ crates/_types/tokens/src/tokens/mod.rs | 3 +- crates/lib/mysql_queries/Cargo.toml | 1 + .../insert_batch_generation_records.rs | 54 +++++++++ .../src/queries/batch_generations/mod.rs | 1 + crates/lib/mysql_queries/src/queries/mod.rs | 3 +- crates/service/job/inference_job/Cargo.toml | 1 + crates/service/web/storyteller_web/Cargo.toml | 1 + 19 files changed, 311 insertions(+), 4 deletions(-) create mode 100644 crates/_types/composite_identifiers/Cargo.toml create mode 100644 crates/_types/composite_identifiers/src/by_table/batch_generations/batch_generation_entity.rs create mode 100644 crates/_types/composite_identifiers/src/by_table/batch_generations/mod.rs create mode 100644 crates/_types/composite_identifiers/src/by_table/mod.rs create mode 100644 crates/_types/composite_identifiers/src/lib.rs create mode 100644 crates/_types/enums/src/by_table/batch_generations/batch_generation_entity_type.rs create mode 100644 crates/_types/enums/src/by_table/batch_generations/mod.rs create mode 100644 crates/_types/tokens/src/tokens/batch_generations.rs create mode 100644 crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs create mode 100644 crates/lib/mysql_queries/src/queries/batch_generations/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 5e29620c49..e0b6036e71 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1259,6 +1259,24 @@ dependencies = [ "my-workspace-hack", ] +[[package]] +name = "composite_identifiers" +version = "0.0.1" +dependencies = [ + "enums", + "my-workspace-hack", + "serde", + "serde_derive", + "serde_json", + "sqlx", + "sqlx-core", + "sqlx-mysql", + "strum", + "tokens", + "toml 0.5.11", + "utoipa", +] + [[package]] name = "concurrency" version = "0.0.1" @@ -3052,6 +3070,7 @@ dependencies = [ "chrono", "cloud_storage", "collections", + "composite_identifiers", "concurrency", "config", "dotenv", @@ -3857,6 +3876,7 @@ dependencies = [ "anyhow", "buckets", "chrono", + "composite_identifiers", "config", "container_common", "dotenv", @@ -6006,6 +6026,7 @@ dependencies = [ "buckets", "chrono", "cloud_storage", + "composite_identifiers", "config", "container_common", "cookies", diff --git a/Cargo.toml b/Cargo.toml index 0bd2f6e824..4b04d5e36c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ members = [ # Types: common strongly-typed tokens, ids, payloads, etc. # Also common cloud buckets and other "strongly typed" configs. "crates/_types/buckets", + "crates/_types/composite_identifiers", "crates/_types/enums", "crates/_types/reusable_types", # TODO(bt, 2022-12-21): Deprecate "crates/_types/tokens", diff --git a/crates/_types/composite_identifiers/Cargo.toml b/crates/_types/composite_identifiers/Cargo.toml new file mode 100644 index 0000000000..f9aa2156b2 --- /dev/null +++ b/crates/_types/composite_identifiers/Cargo.toml @@ -0,0 +1,40 @@ +[package] +name = "composite_identifiers" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[lib] +name = "composite_identifiers" +path = "src/lib.rs" + +[dependencies] +# Internal +enums = { path = "../enums" } +tokens = { path = "../tokens" } + +# Workspace hack (faster compile times) +my-workspace-hack = { version = "0.1", path = "../../my-workspace-hack" } + +# Serialization +serde = "1.0.125" +serde_derive = "1.0.125" + +sqlx = { version = "0.7.2", features = [ "mysql", "sqlite", "runtime-tokio-native-tls", "chrono" ] } +sqlx-core = { version = "0.7.2" } +sqlx-mysql= { version = "0.7.2" } + +utoipa = { version = "4", features = ["chrono","actix_extras"] } + +[dev-dependencies] + +# Enum utilities +strum = { version = "0.25", features = ["derive"] } + +# Test serialization +serde_json = "1.0.108" +toml = "0.5.8" diff --git a/crates/_types/composite_identifiers/src/by_table/batch_generations/batch_generation_entity.rs b/crates/_types/composite_identifiers/src/by_table/batch_generations/batch_generation_entity.rs new file mode 100644 index 0000000000..0c9a261314 --- /dev/null +++ b/crates/_types/composite_identifiers/src/by_table/batch_generations/batch_generation_entity.rs @@ -0,0 +1,28 @@ +use enums::by_table::batch_generations::batch_generation_entity_type::BatchGenerationEntityType; +use tokens::tokens::media_files::MediaFileToken; + +// TODO(bt, 2024-01-16): A ref type with <'a> lifetime of inner data instead of ownership? +pub enum BatchGenerationEntity { + MediaFile(MediaFileToken), +} + +// TODO(bt, 2024-01-16): Codegen? Make traits for these? Maybe overkill. +impl BatchGenerationEntity { + pub fn get_entity_type(&self) -> BatchGenerationEntityType { + match self { + Self::MediaFile(_) => BatchGenerationEntityType::MediaFile, + } + } + + pub fn get_entity_token_str(&self) -> &str { + match self { + Self::MediaFile(token) => token.as_str(), + } + } + + pub fn get_composite_keys(&self) -> (BatchGenerationEntityType, &str) { + match self { + Self::MediaFile(token) => (BatchGenerationEntityType::MediaFile, token.as_str()), + } + } +} diff --git a/crates/_types/composite_identifiers/src/by_table/batch_generations/mod.rs b/crates/_types/composite_identifiers/src/by_table/batch_generations/mod.rs new file mode 100644 index 0000000000..0f4d8b5201 --- /dev/null +++ b/crates/_types/composite_identifiers/src/by_table/batch_generations/mod.rs @@ -0,0 +1 @@ +pub mod batch_generation_entity; \ No newline at end of file diff --git a/crates/_types/composite_identifiers/src/by_table/mod.rs b/crates/_types/composite_identifiers/src/by_table/mod.rs new file mode 100644 index 0000000000..4bbf977cdb --- /dev/null +++ b/crates/_types/composite_identifiers/src/by_table/mod.rs @@ -0,0 +1 @@ +pub mod batch_generations; diff --git a/crates/_types/composite_identifiers/src/lib.rs b/crates/_types/composite_identifiers/src/lib.rs new file mode 100644 index 0000000000..08d90188da --- /dev/null +++ b/crates/_types/composite_identifiers/src/lib.rs @@ -0,0 +1,22 @@ +//! composite_identifiers +//! +//! The purpose of this library is to provide easy to construct, **typesafe**, composite identifiers. +//! This is mostly a combination of `tokens` and `enums`. +//! + +// Never allow these +#![forbid(private_bounds)] +#![forbid(private_interfaces)] +#![forbid(unused_must_use)] // NB: It's unsafe to not close/check some things + +// Okay to toggle +#![forbid(unreachable_patterns)] +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +// Always allow +#![allow(dead_code)] +#![allow(non_snake_case)] + +pub mod by_table; diff --git a/crates/_types/enums/src/by_table/batch_generations/batch_generation_entity_type.rs b/crates/_types/enums/src/by_table/batch_generations/batch_generation_entity_type.rs new file mode 100644 index 0000000000..41633225f9 --- /dev/null +++ b/crates/_types/enums/src/by_table/batch_generations/batch_generation_entity_type.rs @@ -0,0 +1,109 @@ +use std::collections::BTreeSet; + +use serde::Deserialize; +use serde::Serialize; +#[cfg(test)] +use strum::EnumCount; +#[cfg(test)] +use strum::EnumIter; +use utoipa::ToSchema; + +/// Used in the `batch_generations` table in a `VARCHAR(32)` field named `entity_type`. +#[cfg_attr(test, derive(EnumIter, EnumCount))] +#[derive(Clone, Copy, Eq, PartialEq, Hash, Deserialize, Serialize, Ord, PartialOrd, ToSchema)] +pub enum BatchGenerationEntityType { + /// Media files + /// This will probably be the only type supported, but we'll leave the possibility of future types. + #[serde(rename = "media_file")] + MediaFile, +} + +// TODO(bt, 2023-01-17): This desperately needs MySQL integration tests! +impl_enum_display_and_debug_using_to_str!(BatchGenerationEntityType); +impl_mysql_enum_coders!(BatchGenerationEntityType); +impl_mysql_from_row!(BatchGenerationEntityType); + +/// NB: Legacy API for older code. +impl BatchGenerationEntityType { + pub fn to_str(&self) -> &'static str { + match self { + Self::MediaFile => "media_file", + } + } + + pub fn from_str(value: &str) -> Result { + match value { + "media_file" => Ok(Self::MediaFile), + _ => Err(format!("invalid value: {:?}", value)), + } + } + + pub fn all_variants() -> BTreeSet { + // NB: BTreeSet is sorted + // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. + BTreeSet::from([ + Self::MediaFile, + ]) + } +} + +#[cfg(test)] +mod tests { + use crate::by_table::batch_generations::batch_generation_entity_type::BatchGenerationEntityType; + use crate::test_helpers::assert_serialization; + + mod serde { + use super::*; + + #[test] + fn test_serialization() { + assert_serialization(BatchGenerationEntityType::MediaFile, "media_file"); + } + } + + mod impl_methods { + use super::*; + + #[test] + fn to_str() { + assert_eq!(BatchGenerationEntityType::MediaFile.to_str(), "media_file"); + } + + #[test] + fn from_str() { + assert_eq!(BatchGenerationEntityType::from_str("media_file").unwrap(), BatchGenerationEntityType::MediaFile); + assert!(BatchGenerationEntityType::from_str("foo").is_err()); + } + } + + mod manual_variant_checks { + use super::*; + + #[test] + fn all_variants() { + let mut variants = BatchGenerationEntityType::all_variants(); + assert_eq!(variants.len(), 1); + assert_eq!(variants.pop_first(), Some(BatchGenerationEntityType::MediaFile)); + assert_eq!(variants.pop_first(), None); + } + } + + mod mechanical_checks { + use super::*; + + #[test] + fn variant_length() { + use strum::IntoEnumIterator; + assert_eq!(BatchGenerationEntityType::all_variants().len(), BatchGenerationEntityType::iter().len()); + } + + #[test] + fn round_trip() { + for variant in BatchGenerationEntityType::all_variants() { + assert_eq!(variant, BatchGenerationEntityType::from_str(variant.to_str()).unwrap()); + assert_eq!(variant, BatchGenerationEntityType::from_str(&format!("{}", variant)).unwrap()); + assert_eq!(variant, BatchGenerationEntityType::from_str(&format!("{:?}", variant)).unwrap()); + } + } + } +} diff --git a/crates/_types/enums/src/by_table/batch_generations/mod.rs b/crates/_types/enums/src/by_table/batch_generations/mod.rs new file mode 100644 index 0000000000..2babf4fabc --- /dev/null +++ b/crates/_types/enums/src/by_table/batch_generations/mod.rs @@ -0,0 +1 @@ +pub mod batch_generation_entity_type; \ No newline at end of file diff --git a/crates/_types/enums/src/by_table/mod.rs b/crates/_types/enums/src/by_table/mod.rs index b3134ad59d..0c0169b2ac 100644 --- a/crates/_types/enums/src/by_table/mod.rs +++ b/crates/_types/enums/src/by_table/mod.rs @@ -1,5 +1,6 @@ // ===== MySql ===== pub mod audit_logs; +pub mod batch_generations; pub mod comments; pub mod email_sender_jobs; pub mod entity_stats; @@ -17,6 +18,7 @@ pub mod user_ratings; pub mod voice_conversion_models; pub mod voice_conversion_results; pub mod zs_voices; + // ===== Sqlite ===== pub mod tts_render_tasks; pub mod web_rendition_targets; diff --git a/crates/_types/tokens/src/prefixes.rs b/crates/_types/tokens/src/prefixes.rs index 382492bae4..e4263e21bb 100644 --- a/crates/_types/tokens/src/prefixes.rs +++ b/crates/_types/tokens/src/prefixes.rs @@ -24,6 +24,7 @@ use strum::EnumIter; pub(crate) enum TokenPrefix { AnonymousVisitorTracking, // AVTs are not stored as primary keys in any table, but an index in many tables. AuditLog, + BatchGeneration, Comment, DownloadJob, EmailSenderJob, @@ -90,6 +91,7 @@ impl PrefixGenerator for TokenPrefix { match self { Self::AnonymousVisitorTracking => "avt_", Self::AuditLog => "audit_", + Self::BatchGeneration => "batch_g_", Self::Comment => "comment_", Self::DownloadJob => "jdown_", // NB: Previously "JGUP:" Self::EmailSenderJob => "email_job_", @@ -297,7 +299,13 @@ mod tests { #[test] fn test_all_prefixes_end_with_separator_length_one() { for prefix in get_all_prefixes().iter().map(|s| *s) { - if prefix == "news_story_" || prefix == "tts_task_" || prefix == "INT_API:" || prefix == "pw_reset_" || prefix == "email_job_" { + if prefix == "news_story_" + || prefix == "tts_task_" + || prefix == "INT_API:" + || prefix == "pw_reset_" + || prefix == "email_job_" + || prefix == "batch_g_" + { // TODO/FIXME: I'm too tired at 5AM to replacen from the left. Make this test valid. // These tokens are from the AIChatBot sidecar, so asserting their validity is less important. continue; diff --git a/crates/_types/tokens/src/tokens/batch_generations.rs b/crates/_types/tokens/src/tokens/batch_generations.rs new file mode 100644 index 0000000000..052681f21a --- /dev/null +++ b/crates/_types/tokens/src/tokens/batch_generations.rs @@ -0,0 +1,15 @@ +use std::fmt::Debug; + +use serde::Deserialize; +use serde::Serialize; +use utoipa::ToSchema; + +use crate::prefixes::TokenPrefix; + +/// The primary key for Audit Logs. +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, sqlx::Type, Debug, Serialize, Deserialize, ToSchema)] +#[sqlx(transparent)] +pub struct BatchGenerationToken(pub String); + +impl_string_token!(BatchGenerationToken); +impl_crockford_generator!(BatchGenerationToken, 32usize, TokenPrefix::BatchGeneration, CrockfordLower); diff --git a/crates/_types/tokens/src/tokens/mod.rs b/crates/_types/tokens/src/tokens/mod.rs index b7c4eed592..447cbedf43 100644 --- a/crates/_types/tokens/src/tokens/mod.rs +++ b/crates/_types/tokens/src/tokens/mod.rs @@ -2,6 +2,7 @@ pub mod anonymous_visitor_tracking; pub mod api_tokens_external; pub mod api_tokens_internal; pub mod audit_logs; +pub mod batch_generations; pub mod comments; pub mod email_sender_job_token; pub mod firehose_entries; @@ -20,6 +21,7 @@ pub mod tts_results; pub mod twitch_event_rule; pub mod twitch_oauth_tokens_grouping; pub mod twitch_oauth_tokens_internal; +pub mod user_bookmarks; pub mod user_sessions; pub mod user_subscriptions; pub mod users; @@ -34,4 +36,3 @@ pub mod w2l_templates; pub mod zs_voice_dataset_samples; pub mod zs_voice_datasets; pub mod zs_voices; -pub mod user_bookmarks; diff --git a/crates/lib/mysql_queries/Cargo.toml b/crates/lib/mysql_queries/Cargo.toml index 228c472322..2dcf665e29 100644 --- a/crates/lib/mysql_queries/Cargo.toml +++ b/crates/lib/mysql_queries/Cargo.toml @@ -16,6 +16,7 @@ path = "src/lib.rs" # Internal buckets = { path = "../../_types/buckets" } +composite_identifiers = { path = "../../_types/composite_identifiers" } config = { path = "../deprecated/config" } container_common = { path = "../deprecated/container_common" } enums = { path = "../../_types/enums" } diff --git a/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs b/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs new file mode 100644 index 0000000000..a572c42257 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs @@ -0,0 +1,54 @@ +use anyhow::anyhow; +use log::error; +use sqlx::{Executor, MySql, QueryBuilder, Transaction}; + +use composite_identifiers::by_table::batch_generations::batch_generation_entity::BatchGenerationEntity; +use errors::AnyhowResult; +use tokens::tokens::batch_generations::BatchGenerationToken; + +pub struct BatchEntry { + pub batch_token: BatchGenerationToken, + pub entity: BatchGenerationEntity, +} + +pub struct InsertBatchArgs<'a> { + pub entries: Vec, + pub transaction: &'a mut Transaction<'a, MySql>, +} + +/// NB: Caller is responsible for rolling back the transaction if this fails. +pub async fn insert_batch_generation_records(args: InsertBatchArgs<'_>) -> AnyhowResult<()> { + let mut query_builder = QueryBuilder::new(r#" + INSERT INTO batch_generations (token, entity_type, entity_token) VALUES + "#); + + for (i, entry) in args.entries.iter().enumerate() { + let (entity_type, entity_token) = entry.entity.get_composite_keys(); + + query_builder.push("("); + query_builder.push_bind(&entry.batch_token); + query_builder.push(","); + + query_builder.push_bind(entity_type); + query_builder.push(","); + + query_builder.push_bind(entity_token); + query_builder.push(")"); + + if i < args.entries.len() - 1 { + query_builder.push(","); + } + } + + let query = query_builder.build(); + + let query_result = query.execute(&mut **args.transaction).await; + + match query_result { + Ok(_) => Ok(()), + Err(err) => { + error!("Error with batch generation query: {:?}", &err); + Err(anyhow!("model category creation error: {:?}", &err)) + }, + } +} diff --git a/crates/lib/mysql_queries/src/queries/batch_generations/mod.rs b/crates/lib/mysql_queries/src/queries/batch_generations/mod.rs new file mode 100644 index 0000000000..5f34676f55 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/batch_generations/mod.rs @@ -0,0 +1 @@ +pub mod insert_batch_generation_records; diff --git a/crates/lib/mysql_queries/src/queries/mod.rs b/crates/lib/mysql_queries/src/queries/mod.rs index 55f1039245..86079b7494 100644 --- a/crates/lib/mysql_queries/src/queries/mod.rs +++ b/crates/lib/mysql_queries/src/queries/mod.rs @@ -1,5 +1,6 @@ pub mod api_tokens; pub mod audit_logs; +pub mod batch_generations; pub mod billing; pub mod comments; pub mod email_sender_jobs; @@ -24,5 +25,3 @@ pub mod voice_clone_requests; pub mod voice_conversion; pub mod voice_designer; pub mod w2l; - - diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml index 19e17e422b..5644d6e950 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/crates/service/job/inference_job/Cargo.toml @@ -20,6 +20,7 @@ bootstrap = { path = "../../../lib/bootstrap" } buckets = { path = "../../../_types/buckets" } cloud_storage = { path = "../../../lib/api_clients/cloud_storage" } collections = { path = "../../../lib/collections" } +composite_identifiers = { path = "../../../_types/composite_identifiers" } concurrency = { path = "../../../lib/concurrency" } config = { path = "../../../lib/deprecated/config" } easyenv = { path = "../../../lib/easyenv" } diff --git a/crates/service/web/storyteller_web/Cargo.toml b/crates/service/web/storyteller_web/Cargo.toml index d8224d0269..5c959b2966 100644 --- a/crates/service/web/storyteller_web/Cargo.toml +++ b/crates/service/web/storyteller_web/Cargo.toml @@ -24,6 +24,7 @@ actix_helpers = { path = "../../../lib/actix_helpers" } bootstrap = { path = "../../../lib/bootstrap" } buckets = { path = "../../../_types/buckets" } cloud_storage = { path = "../../../lib/api_clients/cloud_storage" } +composite_identifiers = { path = "../../../_types/composite_identifiers" } config = { path = "../../../lib/deprecated/config" } container_common = { path = "../../../lib/deprecated/container_common" } cookies = { path = "../../../lib/cookies" } From 2d91504eb5bc6cb7c939a6e314b37e9e0b182b7b Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 17 Jan 2024 02:21:23 -0500 Subject: [PATCH 49/54] test to show usage --- .../insert_batch_generation_records.rs | 60 ++++++++++++++++--- 1 file changed, 52 insertions(+), 8 deletions(-) diff --git a/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs b/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs index a572c42257..3cb66c11be 100644 --- a/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs +++ b/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs @@ -11,22 +11,27 @@ pub struct BatchEntry { pub entity: BatchGenerationEntity, } -pub struct InsertBatchArgs<'a> { - pub entries: Vec, - pub transaction: &'a mut Transaction<'a, MySql>, +pub struct InsertBatchArgs<'a, 'b> { + pub entries: Vec, + pub transaction: &'a mut Transaction<'b, MySql>, } -/// NB: Caller is responsible for rolling back the transaction if this fails. -pub async fn insert_batch_generation_records(args: InsertBatchArgs<'_>) -> AnyhowResult<()> { +/// Insert a list of entities into a "batch" together for grouping; returns the new batch token identifier to return +/// to the HTTP caller. +/// +/// NB: Calling code is responsible for rolling back the transaction if this fails. +pub async fn insert_batch_generation_records<'a, 'b>(args: InsertBatchArgs<'a, 'b>) -> AnyhowResult { + let batch_token = BatchGenerationToken::generate(); + let mut query_builder = QueryBuilder::new(r#" INSERT INTO batch_generations (token, entity_type, entity_token) VALUES "#); for (i, entry) in args.entries.iter().enumerate() { - let (entity_type, entity_token) = entry.entity.get_composite_keys(); + let (entity_type, entity_token) = entry.get_composite_keys(); query_builder.push("("); - query_builder.push_bind(&entry.batch_token); + query_builder.push_bind(&batch_token); query_builder.push(","); query_builder.push_bind(entity_type); @@ -45,10 +50,49 @@ pub async fn insert_batch_generation_records(args: InsertBatchArgs<'_>) -> Anyho let query_result = query.execute(&mut **args.transaction).await; match query_result { - Ok(_) => Ok(()), + Ok(_) => Ok(batch_token), Err(err) => { error!("Error with batch generation query: {:?}", &err); Err(anyhow!("model category creation error: {:?}", &err)) }, } } + +#[ignore] +#[cfg(test)] +mod tests { + + use sqlx::mysql::MySqlPoolOptions; + use composite_identifiers::by_table::batch_generations::batch_generation_entity::BatchGenerationEntity; + use config::shared_constants::DEFAULT_MYSQL_CONNECTION_STRING; + use tokens::tokens::media_files::MediaFileToken; + use crate::queries::batch_generations::insert_batch_generation_records::{insert_batch_generation_records, InsertBatchArgs}; + + async fn setup() -> sqlx::Pool { + MySqlPoolOptions::new() + .max_connections(3) + .connect(&DEFAULT_MYSQL_CONNECTION_STRING).await + .unwrap() + } + + #[ignore] + #[tokio::test] + async fn test() { + let pool = setup().await; + + let mut transaction = pool.begin().await.unwrap(); + + let entries = vec![ + BatchGenerationEntity::MediaFile(MediaFileToken::new_from_str("media_foo")), + BatchGenerationEntity::MediaFile(MediaFileToken::new_from_str("media_bar")), + BatchGenerationEntity::MediaFile(MediaFileToken::new_from_str("media_baz")), + ]; + + let r = insert_batch_generation_records(InsertBatchArgs { + entries, + transaction: &mut transaction, + }).await; + + transaction.commit().await.unwrap(); + } +} \ No newline at end of file From 24818f8ecf1ffaca80ce6209b73d0d3bdb4dd79e Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Wed, 17 Jan 2024 03:06:09 -0500 Subject: [PATCH 50/54] Ok this builds inference job. Fixes still needed --- .../remote_file_manager/media_descriptor.rs | 21 ++ .../remote_file_manager/weights_descriptor.rs | 2 +- .../src/google_drive_download_command.rs | 54 +++++ .../image_generation_payload.rs | 7 + .../insert_media_file_from_zero_shot_tts.rs | 2 +- .../create/insert_media_file_generic.rs | 172 +++++++++++++++ .../image_generation/sd/process_job.rs | 197 +++++++++++++++--- 7 files changed, 429 insertions(+), 26 deletions(-) create mode 100644 crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs 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 e69de29bb2..3135618c58 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 @@ -0,0 +1,21 @@ +use crate::remote_file_manager::file_descriptor::FileDescriptor; +// TODO ASK BRANDON WHAT path should this be? +const MEDIA_FILE_DIRECTORY: &str = "/media"; + +pub struct MediaImagePngDescriptor {} +impl FileDescriptor for MediaImagePngDescriptor { + fn remote_directory_path(&self) -> &str { + MEDIA_FILE_DIRECTORY + } + fn get_suffix(&self)->String { + "png".to_string() + } + + fn get_prefix(&self)->String { + "image".to_string() + } + + fn is_public(&self) -> bool { + true + } +} diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs index 6c59c0f30b..d302941864 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs @@ -3,7 +3,6 @@ use super::file_descriptor::FileDescriptor; pub struct WeightsLoRADescriptor; const REMOTE_FILE_DIRECTORY: &str = "/weights"; - impl FileDescriptor for WeightsLoRADescriptor { fn remote_directory_path(&self) -> &str { REMOTE_FILE_DIRECTORY @@ -26,6 +25,7 @@ impl FileDescriptor for WeightsLoRADescriptor { } } + pub struct WeightsSD15Descriptor; impl FileDescriptor for WeightsSD15Descriptor { diff --git a/crates/lib/api_clients/google_drive_common/src/google_drive_download_command.rs b/crates/lib/api_clients/google_drive_common/src/google_drive_download_command.rs index 292de11bc0..a78e9f9236 100644 --- a/crates/lib/api_clients/google_drive_common/src/google_drive_download_command.rs +++ b/crates/lib/api_clients/google_drive_common/src/google_drive_download_command.rs @@ -57,6 +57,60 @@ impl GoogleDriveDownloadCommand { } } + pub async fn download_file_with_file_name(&self, + download_url: &str, + temp_dir: &TempDir, + file_name:&str) -> AnyhowResult + { + let temp_dir_path = temp_dir.path() + .to_str() + .unwrap_or("/tmp") + .to_string(); + + let temp_filename = random_crockford_token(10); + let temp_filename = format!("{}/{}", temp_dir_path, file_name); + + info!("Downloading {} to: {}", download_url, temp_filename); + + let mut command = format!("{} --url \"{}\" --output_filename {}", + &self.download_script, + download_url, + &temp_filename); + + if let Some(venv_activation_script) = self.maybe_venv_activation_script.as_deref() { + // NB: "." is source for non-bash shells + command = format!(". {} && {}", + venv_activation_script, + &command); + } + + if let Some(change_to_directory) = self.maybe_change_to_directory.as_deref() { + command = format!("cd {} && {}", + change_to_directory, + &command); + } + + if let Some(docker_options) = self.maybe_docker_options.as_ref() { + command = docker_options.to_command_string(&command); + } + + info!("Running command: {}", command); + + let result = Command::new("sh") + .arg("-c") + .arg(command) + .output()?; + + info!("Downloader Result: {:?}", result); + + if !result.status.success() { + let reason = String::from_utf8(result.stderr).unwrap_or("UNKNOWN".to_string()); + return Err(anyhow!("Failure to execute command: {:?}", reason)) + } + + Ok(temp_filename) + } + /// Download file from Google Drive into the `TempDir`. /// Return the local filesystem filename. pub async fn download_file(&self, 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 233f802533..e14d49ec33 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 @@ -74,6 +74,13 @@ pub struct StableDiffusionArgs { #[serde(rename = "t")] pub type_of_inference: String, + + // loRA / checkpoint description + #[serde(rename = "de")] + pub description: Option, + // loRA / checkpoint name and description + #[serde(rename = "na")] + pub name: Option } diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs index e250919752..478afd3b7e 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs @@ -65,7 +65,7 @@ pub async fn insert_media_file_from_zero_shot( maybe_creator_category_synthetic_id = Some(next_lipsync_id); } - // TODO: Ask if it make sense to use ONE function to do the insert + // 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; 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 new file mode 100644 index 0000000000..ec51b21e76 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs @@ -0,0 +1,172 @@ +use anyhow::anyhow; +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_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::media_files::MediaFileToken; +use tokens::tokens::weight_model_weight; +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; +// 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 public_bucket_directory_hash: &'a str, + pub maybe_public_bucket_prefix: Option<&'a str>, + pub maybe_public_bucket_extension: Option<&'a str>, + + pub is_on_prem: bool, + + pub worker_hostname: &'a str, + pub worker_cluster: &'a str, + + 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 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, + &mut transaction + ).await?; + + let next_lipsync_id = transactional_increment_generic_synthetic_id( + &user_token, + IdCategory::ZeroShotTtsResult, + &mut transaction + ).await?; + + maybe_creator_file_synthetic_id = Some(next_media_file_id); + maybe_creator_category_synthetic_id = Some(next_lipsync_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 record_id = { + let query_result = sqlx::query!( + r#" + INSERT INTO media_files + SET + token = ?, + + origin_category = ?, + origin_product_category = ?, + maybe_origin_model_type = ?, + + media_type = ?, + maybe_mime_type = ?, + file_size_bytes = ?, + + checksum_sha2 = ?, + + public_bucket_directory_hash = ?, + maybe_public_bucket_prefix = ?, + maybe_public_bucket_extension = ?, + + maybe_creator_user_token = ?, + maybe_creator_anonymous_visitor_token = ?, + creator_ip_address = ?, + + creator_set_visibility = ?, + + maybe_creator_file_synthetic_id = ?, + maybe_creator_category_synthetic_id = ?, + + is_generated_on_prem = ?, + generated_by_worker = ?, + generated_by_cluster = ? + + "#, + result_token.as_str(), + + ORIGIN_CATEGORY.to_str(), + ORIGIN_PRODUCT_CATEGORY.to_str(), + ORIGIN_MODEL_TYPE.to_str(), + + MEDIA_TYPE.to_str(), + args.maybe_mime_type, + args.file_size_bytes, + + args.sha256_checksum, + + args.public_bucket_directory_hash, + args.maybe_public_bucket_prefix, + args.maybe_public_bucket_extension, + + args.job.maybe_creator_user_token, + args.job.maybe_creator_anonymous_visitor_token, + args.job.creator_ip_address, + + args.job.creator_set_visibility.to_str(), + + 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 + }; + + transaction.commit().await?; + + Ok((result_token, record_id)) +} \ No newline at end of file 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 71c57cdcd9..2056c21327 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,8 +1,9 @@ +use std::env::temp_dir; use std::ops::Deref; use std::path::PathBuf; use std::thread; use std::thread::Thread; -use std::time::Duration; +use std::time::{Duration, Instant}; use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails; use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; @@ -11,16 +12,36 @@ use crate::job::job_loop::job_success_result::{ JobSuccessResult, ResultEntity } use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::job_dependencies::JobDependencies; use anyhow::anyhow; +use log::info; +use sqlx::MySqlPool; +use cloud_storage::remote_file_manager::media_descriptor::MediaImagePngDescriptor; use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; +use cloud_storage::remote_file_manager::weights_descriptor::WeightsLoRADescriptor; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType::UploadModel; +use enums::by_table::generic_synthetic_ids::id_category::IdCategory::MediaFile; +use enums::by_table::media_files::media_file_type::MediaFileType; +use enums::by_table::media_uploads::media_upload_source::MediaUploadSource; +use enums::by_table::media_uploads::media_upload_type::MediaUploadType; +use enums::by_table::model_weights::weights_category::WeightsCategory; +use enums::by_table::model_weights::weights_types::WeightsType; +use enums::common::visibility::Visibility; use filesys::path_to_string::path_to_string; +use filesys::safe_delete_temp_directory::safe_delete_temp_directory; +use mysql_queries::payloads::media_upload_modification_details::MediaUploadModificationDetails; +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_uploads::insert_media_upload::{Args, insert_media_upload}; +use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken; use tokens::tokens::model_weights::ModelWeightToken; +use tokens::tokens::users::UserToken; use crate::job::job_types::image_generation::sd::sd_inference_command::InferenceArgs; use crate::job::job_types::image_generation::sd::stable_diffusion_dependencies::StableDiffusionDependencies; use crate::job::job_types::image_generation::sd::validate_inputs::validate_inputs; +use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand; +use mysql_queries::queries::model_weights::create::create_weight::{create_weight, CreateModelWeightsArgs}; + pub struct StableDiffusionProcessArgs<'a> { pub job_dependencies: &'a JobDependencies, pub job: &'a AvailableInferenceJob, @@ -30,8 +51,6 @@ pub struct StableDiffusionProcessArgs<'a> { // run inference // insert record into the db with the inference job token complete. -pub async fn download_from_google_link() {} - pub async fn sd_args_from_job( args: &StableDiffusionProcessArgs<'_> ) -> Result { @@ -87,7 +106,6 @@ pub async fn process_job_checkpoint(args: &StableDiffusionProcessArgs<'_>) -> Re // entity_type: InferenceResultType::UploadModel, // entity_token: // }; - Ok(JobSuccessResult { maybe_result_entity: None, inference_duration: Duration::from_secs(0), @@ -105,6 +123,8 @@ pub async fn process_job_inference( ) -> 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 { @@ -116,6 +136,17 @@ pub async fn process_job_inference( } }; + 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 mut job_progress_reporter = args.job_dependencies.clients.job_progress_reporter .new_generic_inference(job.inference_job_token.as_str()) .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; @@ -126,28 +157,66 @@ pub async fn process_job_inference( // 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 + let mut 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"); - let lora_path = work_temp_dir.path().join("lora"); - let vae_path = work_temp_dir.path().join("vae"); + let sd_checkpoint_path = work_temp_dir.path().join("sd_checkpoint.safetensors"); + let 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"); - println!("Paths to download to:"); - println!("sd_checkpoint_path:{}",sd_checkpoint_path.display()); - println!("lora_path:{}",lora_path.display()); - println!("vae_path:{}",vae_path.display()); - println!("output_path:{}",output_path.display()); + 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()); + + // 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()); + } // // 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 retrieved_sd_record = match weight_token { Some(token) => { let retrieved_sd_record = get_weight_by_token( @@ -272,11 +341,11 @@ pub async fn process_job_inference( 20 } }; + let inference_start_time = Instant::now(); - // #TODO allow for batch and saving later. output files sd_deps.inference_command.execute_inference(InferenceArgs { work_dir: work_temp_dir.path().to_path_buf(), - output_file: output_path, + output_file: output_path.clone(), stderr_output_file: stderr_output_file, prompt: prompt, negative_prompt:sd_args.maybe_n_prompt.unwrap_or_default(), @@ -292,15 +361,95 @@ 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 + + + 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_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?; + } + + // 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_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, + cached_user_ratings_total_count: 0, + cached_user_ratings_positive_count: 0, + cached_user_ratings_negative_count: 0, + maybe_cached_user_ratings_ratio: None, + version: 0, + mysql_pool, + }).await?; + + Ok(JobSuccessResult {// TODO: batch token needs to go here + maybe_result_entity: None, + inference_duration: Duration::from_secs(0), + }) + } + + - Ok(JobSuccessResult { - maybe_result_entity: None, - inference_duration: Duration::from_secs(0), - }) } #[cfg(test)] From 7b90b11ce6121bb6211da18077ea90064b86c2ed Mon Sep 17 00:00:00 2001 From: ArEnSc Date: Wed, 17 Jan 2024 17:33:55 -0500 Subject: [PATCH 51/54] Fixed build issues ready to merge to master after review. --- .../generic_inference_jobs/inference_model_type.rs | 2 +- .../generic_inference_args/generic_inference_args.rs | 2 +- .../image_generation_payload.rs | 2 ++ .../job/job_types/image_generation/sd/process_job.rs | 4 ---- .../inference_job/src/job_specific_dependencies.rs | 2 +- .../endpoints/image_gen/enqueue_image_generation.rs | 11 ++++++++--- 6 files changed, 13 insertions(+), 10 deletions(-) diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs index a50c3f9c54..a7945f166f 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs @@ -56,7 +56,7 @@ impl InferenceModelType { Self::Vits => "vits", Self::VallEX => "vall_e_x", Self::RerenderAVideo => "rerender_a_video", - Self::StableDiffusion => "stable_diffusion" + Self::StableDiffusion => "stable_diffusion", Self::MocapNet => "mocap_net", Self::StyleTTS2 => "styletts2", } 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 1f1397a0fd..46aa3caad9 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 @@ -138,7 +138,7 @@ impl InferenceCategoryAbbreviated { Self::TextToSpeech => InferenceCategory::TextToSpeech, Self::VoiceConversion => InferenceCategory::VoiceConversion, Self::VideoFilter => InferenceCategory::VideoFilter, - Self::ImageGeneration =>InferenceCategory::ImageGeneration + Self::ImageGeneration =>InferenceCategory::ImageGeneration, Self::Mocap => InferenceCategory::Mocap, } } 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 e14d49ec33..3fbdb1034e 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 @@ -77,9 +77,11 @@ pub struct StableDiffusionArgs { // loRA / checkpoint description #[serde(rename = "de")] + #[serde(skip_serializing_if = "Option::is_none")] pub description: Option, // loRA / checkpoint name and description #[serde(rename = "na")] + #[serde(skip_serializing_if = "Option::is_none")] pub name: Option } 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 2056c21327..062c09eb05 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 @@ -434,10 +434,6 @@ pub async fn process_job_inference( public_bucket_hash: "".to_string(), maybe_public_bucket_prefix: None, maybe_public_bucket_extension: None, - cached_user_ratings_total_count: 0, - cached_user_ratings_positive_count: 0, - cached_user_ratings_negative_count: 0, - maybe_cached_user_ratings_ratio: None, version: 0, mysql_pool, }).await?; diff --git a/crates/service/job/inference_job/src/job_specific_dependencies.rs b/crates/service/job/inference_job/src/job_specific_dependencies.rs index 2ec984dc97..c0b1969bc8 100644 --- a/crates/service/job/inference_job/src/job_specific_dependencies.rs +++ b/crates/service/job/inference_job/src/job_specific_dependencies.rs @@ -100,7 +100,7 @@ impl JobSpecificDependencies { maybe_vall_e_x_dependencies, maybe_vits_dependencies, maybe_rerender_dependencies, - maybe_stable_diffusion_dependencies + maybe_stable_diffusion_dependencies, maybe_mocapnet_dependencies, maybe_styletts2_dependencies, }) 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 9f09505293..9660a17c0e 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,4 +1,4 @@ -#![forbid(unused_imports)] +// #![forbid(unused_imports)] #![forbid(unused_mut)] #![forbid(unused_variables)] @@ -119,6 +119,8 @@ pub struct EnqueueImageGenRequest { maybe_cfg_scale: Option, maybe_number_of_samples: Option, maybe_batch_count: Option, + name: String, + description: String, } #[derive(Serialize,ToSchema)] @@ -332,7 +334,8 @@ pub async fn enqueue_image_generation_request( } 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 inference_args = StableDiffusionArgs { maybe_sd_model_token: Some(sd_weight_token), maybe_lora_model_token: Some(lora_token), @@ -347,7 +350,9 @@ pub async fn enqueue_image_generation_request( maybe_batch_count: Some(batch_count), maybe_width: Some(width), maybe_height: Some(height), - maybe_sampler: Some(sampler) + maybe_sampler: Some(sampler), + description: Some(description), + name: Some(name), }; // create the inference args here From cecd3dd6df872528ad70d9bb168d3a49ad15c8f5 Mon Sep 17 00:00:00 2001 From: ArEnSc <6252325+ArEnSc@users.noreply.github.com> Date: Wed, 17 Jan 2024 18:03:20 -0500 Subject: [PATCH 52/54] Update .env --- .env | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.env b/.env index c41492d472..a39f1b1d8d 100644 --- a/.env +++ b/.env @@ -76,4 +76,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. From 8dc7c81a786c99ff603d37e11a9baad18fe2cb52 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 17 Jan 2024 19:48:03 -0500 Subject: [PATCH 53/54] fix tests --- .../src/by_table/generic_inference_jobs/inference_category.rs | 3 ++- .../by_table/generic_inference_jobs/inference_model_type.rs | 3 +-- .../by_table/generic_inference_jobs/inference_result_type.rs | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs index 7506775638..3a1a44f553 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs @@ -31,6 +31,7 @@ pub enum InferenceCategory { #[serde(rename = "image_generation")] ImageGeneration, + #[serde(rename = "mocap")] Mocap, } @@ -117,7 +118,7 @@ mod tests { fn all_variants() { // Static check let mut variants = InferenceCategory::all_variants(); - assert_eq!(variants.len(), 5); + assert_eq!(variants.len(), 6); assert_eq!(variants.pop_first(), Some(InferenceCategory::LipsyncAnimation)); assert_eq!(variants.pop_first(), Some(InferenceCategory::TextToSpeech)); assert_eq!(variants.pop_first(), Some(InferenceCategory::VoiceConversion)); diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs index a7945f166f..32a8757f30 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs @@ -147,8 +147,7 @@ mod tests { fn all_variants() { // Static check let mut variants = InferenceModelType::all_variants(); - assert_eq!(variants.len(), 8); - assert_eq!(variants.len(), 9); + assert_eq!(variants.len(), 10); assert_eq!(variants.pop_first(), Some(InferenceModelType::RvcV2)); assert_eq!(variants.pop_first(), Some(InferenceModelType::SadTalker)); assert_eq!(variants.pop_first(), Some(InferenceModelType::SoVitsSvc)); diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs index 0b2b7ffaaa..e411079f40 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs @@ -110,7 +110,7 @@ mod tests { fn all_variants() { // Static check let mut variants = InferenceResultType::all_variants(); - assert_eq!(variants.len(), 4); + assert_eq!(variants.len(), 5); assert_eq!(variants.pop_first(), Some(InferenceResultType::MediaFile)); assert_eq!(variants.pop_first(), Some(InferenceResultType::TextToSpeech)); assert_eq!(variants.pop_first(), Some(InferenceResultType::VoiceConversion)); From f6da4f32478bf69a7b1020bb2998656b723bbe9a Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 17 Jan 2024 20:34:13 -0500 Subject: [PATCH 54/54] fix lints --- Cargo.lock | 1 + crates/cli/dev-database-seed/src/main.rs | 4 +- .../src/seeding/model_weights.rs | 20 +--- .../remote_file_manager/media_descriptor.rs | 3 +- .../remote_cloud_file_manager.rs | 4 +- .../src/google_drive_download_command.rs | 1 - .../generic_inference_args.rs | 15 +-- .../image_generation_payload.rs | 2 +- .../insert_batch_generation_records.rs | 5 +- .../image_generation/process_single_ig_job.rs | 6 +- .../image_generation/sd/process_job.rs | 104 ++++++++---------- .../src/job_specific_dependencies.rs | 6 +- crates/service/web/storyteller_web/Cargo.toml | 2 + .../image_gen/enqueue_generation_type.rs | 2 +- .../image_gen/enqueue_image_generation.rs | 10 +- 15 files changed, 74 insertions(+), 111 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b1002167dd..fcc378de3e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6090,6 +6090,7 @@ dependencies = [ "sha2 0.9.9", "speculoos", "sqlx", + "strum", "symphonia", "time 0.3.30", "tokens", diff --git a/crates/cli/dev-database-seed/src/main.rs b/crates/cli/dev-database-seed/src/main.rs index f8ba5d5ef5..e123dc3c24 100644 --- a/crates/cli/dev-database-seed/src/main.rs +++ b/crates/cli/dev-database-seed/src/main.rs @@ -7,10 +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; // V1 +use crate::seeding::model_weights::seed_weights; use crate::seeding::seed_media_files::seed_media_files; -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::model_weights::test_seed_weights_files; use crate::seeding::tts_tacotron2::seed_tts_tacotron2; use crate::seeding::users::seed_user_accounts; use crate::seeding::voice_conversion::seed_voice_conversion; diff --git a/crates/cli/dev-database-seed/src/seeding/model_weights.rs b/crates/cli/dev-database-seed/src/seeding/model_weights.rs index 6992fc6e69..cc67f89f96 100644 --- a/crates/cli/dev-database-seed/src/seeding/model_weights.rs +++ b/crates/cli/dev-database-seed/src/seeding/model_weights.rs @@ -3,13 +3,10 @@ use std::vec; use log::info; use sqlx::{MySql, Pool}; -use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; + use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails; - -use cloud_storage::remote_file_manager::weights_descriptor; +use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; use cloud_storage::remote_file_manager::weights_descriptor::{WeightsLoRADescriptor, WeightsSD15Descriptor}; -use easyenv::from_filename; - use enums::by_table::model_weights::{ weights_category::WeightsCategory, weights_types::WeightsType, @@ -17,7 +14,6 @@ use enums::by_table::model_weights::{ use enums::common::visibility::Visibility; use errors::{anyhow, AnyhowResult}; use mysql_queries::queries::model_weights::create::create_weight::{create_weight, CreateModelWeightsArgs}; -use mysql_queries::queries::users::user::create_account::create_account; 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::model_weights::ModelWeightToken; @@ -631,10 +627,6 @@ Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix public_bucket_hash: metadata1.bucket_details.clone().unwrap().object_hash, maybe_public_bucket_prefix: Some(metadata1.bucket_details.clone().unwrap().prefix), maybe_public_bucket_extension: Some(metadata1.bucket_details.clone().unwrap().suffix), - cached_user_ratings_total_count: 10, - cached_user_ratings_positive_count: 9, - cached_user_ratings_negative_count: 1, - maybe_cached_user_ratings_ratio: Some(0.9), version: 1, mysql_pool: &mysql_pool, // replace with actual MySqlPool }; @@ -657,10 +649,6 @@ Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix public_bucket_hash: metadata2.bucket_details.clone().unwrap().object_hash.clone(), maybe_public_bucket_prefix: Some(metadata2.bucket_details.clone().unwrap().prefix), maybe_public_bucket_extension: Some(metadata2.bucket_details.clone().unwrap().suffix), - cached_user_ratings_total_count: 20, - cached_user_ratings_positive_count: 9, - cached_user_ratings_negative_count: 2, - maybe_cached_user_ratings_ratio: Some(0.9), version: 2, mysql_pool: &mysql_pool, // replace with actual MySqlPool }; @@ -683,10 +671,6 @@ Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix public_bucket_hash: metadata3.bucket_details.clone().unwrap().object_hash, maybe_public_bucket_prefix: Some(metadata3.bucket_details.clone().unwrap().prefix), maybe_public_bucket_extension: Some(metadata3.bucket_details.clone().unwrap().suffix), - cached_user_ratings_total_count: 10, - cached_user_ratings_positive_count: 9, - cached_user_ratings_negative_count: 1, - maybe_cached_user_ratings_ratio: Some(0.9), version: 1, mysql_pool: &mysql_pool, // replace with actual MySqlPool }; 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 3135618c58..d969193271 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 @@ -2,7 +2,8 @@ use crate::remote_file_manager::file_descriptor::FileDescriptor; // TODO ASK BRANDON WHAT path should this be? const MEDIA_FILE_DIRECTORY: &str = "/media"; -pub struct MediaImagePngDescriptor {} +pub struct MediaImagePngDescriptor; + impl FileDescriptor for MediaImagePngDescriptor { fn remote_directory_path(&self) -> &str { MEDIA_FILE_DIRECTORY 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 daf6da1ee6..61e1f14e19 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 @@ -68,8 +68,8 @@ impl RemoteCloudFileClient { result.bucket_details = Some(RemoteCloudBucketDetails { object_hash: directory.get_file_object_hash().to_string(), - suffix:suffix, - prefix:prefix, + suffix, + prefix, }); println!("Uploading media file to bucket path: {:?}",directory.get_full_remote_cloud_file_path()); diff --git a/crates/lib/api_clients/google_drive_common/src/google_drive_download_command.rs b/crates/lib/api_clients/google_drive_common/src/google_drive_download_command.rs index a78e9f9236..56ab8d4931 100644 --- a/crates/lib/api_clients/google_drive_common/src/google_drive_download_command.rs +++ b/crates/lib/api_clients/google_drive_common/src/google_drive_download_command.rs @@ -67,7 +67,6 @@ impl GoogleDriveDownloadCommand { .unwrap_or("/tmp") .to_string(); - let temp_filename = random_crockford_token(10); let temp_filename = format!("{}/{}", temp_dir_path, file_name); info!("Downloading {} to: {}", download_url, temp_filename); 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 46aa3caad9..eae5b1de77 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 @@ -1,11 +1,11 @@ use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; use errors::AnyhowResult; +use crate::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; use crate::payloads::generic_inference_args::lipsync_payload::LipsyncArgs; use crate::payloads::generic_inference_args::mocap_payload::MocapArgs; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; use crate::payloads::generic_inference_args::videofilter_payload::RerenderArgs; -use crate::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; /// Used to encode extra state for the `generic_inference_jobs` table. /// This should act somewhat like a serialized protobuf stored inside a record. @@ -146,11 +146,12 @@ impl InferenceCategoryAbbreviated { #[cfg(test)] mod tests { + use tokens::tokens::{media_files::MediaFileToken, model_weights::ModelWeightToken}; + use crate::payloads::generic_inference_args::generic_inference_args::{FundamentalFrequencyMethodForJob, GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; + use crate::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; use crate::payloads::generic_inference_args::lipsync_payload::{LipsyncAnimationAudioSource, LipsyncAnimationImageSource, LipsyncArgs}; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; - use crate::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; - use tokens::tokens::{model_weights::ModelWeightToken, media_files::MediaFileToken}; #[test] fn typical_lipsync_animation_args_serialize() { @@ -218,7 +219,7 @@ mod tests { args: Some(PolymorphicInferenceArgs::Ig(StableDiffusionArgs { maybe_sd_model_token: Some(sd_model_token), maybe_lora_model_token: Some(lora_model_token), - maybe_sampler: Some("sampler".as_ref()), + maybe_sampler: Some("sampler".to_string()), maybe_height: Some(512), maybe_width: Some(512), maybe_cfg_scale: Some(7), @@ -229,9 +230,9 @@ mod tests { maybe_seed: Some(seed), maybe_upload_path: Some(upload_path), maybe_lora_upload_path: Some(lora_upload_path), - type_of_inference: type_of_inference, - description: Some("Option".as_ref()), - title: Some("Model Name".as_ref()) + type_of_inference, + description: Some("Option".to_string()), + name: Some("Model Name".to_string()) })), }; 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 3fbdb1034e..cdae13bdd2 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 @@ -1,4 +1,4 @@ -use tokens::tokens::{model_weights::ModelWeightToken, media_files::MediaFileToken}; +use tokens::tokens::model_weights::ModelWeightToken; /// Video sources can be one of several: /// - F: media_files (todo) diff --git a/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs b/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs index 3cb66c11be..4b9c6f4d56 100644 --- a/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs +++ b/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs @@ -1,6 +1,6 @@ use anyhow::anyhow; use log::error; -use sqlx::{Executor, MySql, QueryBuilder, Transaction}; +use sqlx::{MySql, QueryBuilder, Transaction}; use composite_identifiers::by_table::batch_generations::batch_generation_entity::BatchGenerationEntity; use errors::AnyhowResult; @@ -61,11 +61,12 @@ pub async fn insert_batch_generation_records<'a, 'b>(args: InsertBatchArgs<'a, ' #[ignore] #[cfg(test)] mod tests { - use sqlx::mysql::MySqlPoolOptions; + use composite_identifiers::by_table::batch_generations::batch_generation_entity::BatchGenerationEntity; use config::shared_constants::DEFAULT_MYSQL_CONNECTION_STRING; use tokens::tokens::media_files::MediaFileToken; + use crate::queries::batch_generations::insert_batch_generation_records::{insert_batch_generation_records, InsertBatchArgs}; async fn setup() -> sqlx::Pool { diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs index a5f09ef9f9..bb26c58933 100644 --- a/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs @@ -1,14 +1,14 @@ use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; +use errors::anyhow; use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; + use crate::job::job_loop::job_success_result::JobSuccessResult; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::job::job_types::image_generation::sd::process_job::process_job_selection; -use crate::job_dependencies::JobDependencies; use crate::job::job_types::image_generation::sd::process_job::StableDiffusionProcessArgs; -use errors::anyhow; +use crate::job_dependencies::JobDependencies; pub async fn process_single_ig_job(job_dependencies: &JobDependencies, job: &AvailableInferenceJob) -> Result { - let job = job; match job.maybe_model_type { Some(InferenceModelType::StableDiffusion) => { dispatch_sd_job( 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 062c09eb05..b7ea33cf16 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,30 @@ -use std::env::temp_dir; -use std::ops::Deref; -use std::path::PathBuf; use std::thread; -use std::thread::Thread; use std::time::{Duration, Instant}; -use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails; -use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; -use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; -use mysql_queries::queries::model_weights::get_weight::get_weight_by_token; -use crate::job::job_loop::job_success_result::{ JobSuccessResult, ResultEntity }; -use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; -use crate::job_dependencies::JobDependencies; + use anyhow::anyhow; use log::info; -use sqlx::MySqlPool; -use cloud_storage::remote_file_manager::media_descriptor::MediaImagePngDescriptor; -use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; +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::weights_descriptor::WeightsLoRADescriptor; -use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType::UploadModel; -use enums::by_table::generic_synthetic_ids::id_category::IdCategory::MediaFile; use enums::by_table::media_files::media_file_type::MediaFileType; -use enums::by_table::media_uploads::media_upload_source::MediaUploadSource; -use enums::by_table::media_uploads::media_upload_type::MediaUploadType; use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; -use enums::common::visibility::Visibility; use filesys::path_to_string::path_to_string; -use filesys::safe_delete_temp_directory::safe_delete_temp_directory; -use mysql_queries::payloads::media_upload_modification_details::MediaUploadModificationDetails; +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::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::media_uploads::insert_media_upload::{Args, insert_media_upload}; -use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken; +use mysql_queries::queries::model_weights::create::create_weight::{create_weight, CreateModelWeightsArgs}; +use mysql_queries::queries::model_weights::get_weight::get_weight_by_token; use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::users::UserToken; -use crate::job::job_types::image_generation::sd::sd_inference_command::InferenceArgs; -use crate::job::job_types::image_generation::sd::stable_diffusion_dependencies::StableDiffusionDependencies; -use crate::job::job_types::image_generation::sd::validate_inputs::validate_inputs; -use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand; -use mysql_queries::queries::model_weights::create::create_weight::{create_weight, CreateModelWeightsArgs}; +use crate::job::job_loop::job_success_result::JobSuccessResult; +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; pub struct StableDiffusionProcessArgs<'a> { pub job_dependencies: &'a JobDependencies, @@ -82,7 +66,6 @@ pub async fn sd_args_from_job( pub async fn process_job_selection (args: StableDiffusionProcessArgs<'_> ) -> Result { - let job = args.job; let sd_args = sd_args_from_job(&args).await?; if sd_args.type_of_inference == "inference" { process_job_inference(&args).await @@ -98,7 +81,7 @@ pub async fn process_job_selection (args: StableDiffusionProcessArgs<'_> } } -pub async fn process_job_checkpoint(args: &StableDiffusionProcessArgs<'_>) -> Result { +pub async fn process_job_checkpoint(_args: &StableDiffusionProcessArgs<'_>) -> Result { // insert model record in create model record @@ -111,7 +94,7 @@ pub async fn process_job_checkpoint(args: &StableDiffusionProcessArgs<'_>) -> Re inference_duration: Duration::from_secs(0), }) } -pub async fn process_job_lora(args: &StableDiffusionProcessArgs<'_>) -> Result { +pub async fn process_job_lora(_args: &StableDiffusionProcessArgs<'_>) -> Result { Ok(JobSuccessResult { maybe_result_entity: None, inference_duration: Duration::from_secs(0), @@ -147,9 +130,10 @@ pub async fn process_job_inference( } - let mut job_progress_reporter = args.job_dependencies.clients.job_progress_reporter + let _job_progress_reporter = args.job_dependencies.clients.job_progress_reporter .new_generic_inference(job.inference_job_token.as_str()) .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + //==================== TEMP DIR ==================== // let work_temp_dir = format!("temp_stable_diffusion_inference_{}", job.id.0); @@ -192,9 +176,10 @@ pub async fn process_job_inference( } } - if download_url.len() <= 0 { + 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", @@ -202,9 +187,10 @@ pub async fn process_job_inference( 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 { + + 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) => { + Err(_e) => { return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Download"))); } }; @@ -300,9 +286,7 @@ pub async fn process_job_inference( ).await?; let vae_weight_record = match vae_weight_record { - Some(val) => { - val - }, + Some(val) => val, None => { return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))) } @@ -325,31 +309,26 @@ pub async fn process_job_inference( ).await?; let prompt = match sd_args.maybe_prompt { - Some(val) => { - val - } + Some(val) => val, None => return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("No Prompt provided!"))) }; let stderr_output_file = work_temp_dir.path().join("sd_err.txt"); let number_of_samples = match sd_args.maybe_number_of_samples { - Some(val) => { - val - }, - None => { - 20 - } + Some(val) => val, + None => 20, }; + let inference_start_time = Instant::now(); sd_deps.inference_command.execute_inference(InferenceArgs { work_dir: work_temp_dir.path().to_path_buf(), output_file: output_path.clone(), - stderr_output_file: stderr_output_file, - prompt: prompt, + stderr_output_file, + prompt, negative_prompt:sd_args.maybe_n_prompt.unwrap_or_default(), - number_of_samples: number_of_samples, + number_of_samples, samplers: 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), @@ -387,7 +366,7 @@ pub async fn process_job_inference( pool: mysql_pool, maybe_creator_user_token: Some(&creator_user_token), maybe_creator_anonymous_visitor_token: None, - creator_ip_address: creator_ip_address, + creator_ip_address, creator_set_visibility: Default::default(), upload_type: UploadType::Filesystem, media_file_type: MediaFileType::Image, @@ -424,7 +403,7 @@ pub async fn process_job_inference( maybe_description_markdown: None, maybe_description_rendered_html: None, creator_user_token: Some(&creator_user_token), - creator_ip_address: creator_ip_address, + creator_ip_address, creator_set_visibility: Default::default(), maybe_last_update_user_token: None, original_download_url: None, @@ -450,14 +429,17 @@ pub async fn process_job_inference( #[cfg(test)] mod tests { -use anyhow::anyhow; -use std::path::PathBuf; -use cloud_storage::remote_file_manager::{ -remote_cloud_file_manager::{RemoteCloudFileClient, self}, -remote_cloud_bucket_details::RemoteCloudBucketDetails, -}; -use errors::{AnyhowError, AnyhowResult}; -#[tokio::test] + use std::path::PathBuf; + + use anyhow::anyhow; + + use cloud_storage::remote_file_manager::{ + remote_cloud_bucket_details::RemoteCloudBucketDetails, + remote_cloud_file_manager::RemoteCloudFileClient, + }; + use errors::AnyhowResult; + + #[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; diff --git a/crates/service/job/inference_job/src/job_specific_dependencies.rs b/crates/service/job/inference_job/src/job_specific_dependencies.rs index c0b1969bc8..17be1ab77d 100644 --- a/crates/service/job/inference_job/src/job_specific_dependencies.rs +++ b/crates/service/job/inference_job/src/job_specific_dependencies.rs @@ -1,8 +1,9 @@ use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; use errors::AnyhowResult; -use sqlx::pool::maybe; +use crate::job::job_types::image_generation::sd::stable_diffusion_dependencies::StableDiffusionDependencies; use crate::job::job_types::lipsync::sad_talker::sad_talker_dependencies::SadTalkerDependencies; +use crate::job::job_types::mocap::mocap_net::mocapnet_dependencies::MocapNetDependencies; use crate::job::job_types::tts::styletts2::styletts2_dependencies::StyleTTS2Dependencies; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::tacotron2_dependencies::Tacotron2Dependencies; use crate::job::job_types::tts::vall_e_x::vall_e_x_dependencies::VallExDependencies; @@ -10,9 +11,6 @@ use crate::job::job_types::tts::vits::vits_dependencies::VitsDependencies; use crate::job::job_types::vc::rvc_v2::rvc_v2_dependencies::RvcV2Dependencies; use crate::job::job_types::vc::so_vits_svc::svc_dependencies::SvcDependencies; use crate::job::job_types::videofilter::rerender_a_video::rerender_dependencies::RerenderDependencies; -use crate::job::job_types::image_generation::sd::stable_diffusion_dependencies::StableDiffusionDependencies; - -use crate::job::job_types::mocap::mocap_net::mocapnet_dependencies::MocapNetDependencies; use crate::util::scoped_execution::ScopedExecution; pub struct JobSpecificDependencies { diff --git a/crates/service/web/storyteller_web/Cargo.toml b/crates/service/web/storyteller_web/Cargo.toml index 5c959b2966..0926381123 100644 --- a/crates/service/web/storyteller_web/Cargo.toml +++ b/crates/service/web/storyteller_web/Cargo.toml @@ -177,3 +177,5 @@ diesel = { version = "2.1.3", features = ["mysql"] } [dev-dependencies] speculoos = "0.9.0" # Fluent assertions +strum = { version = "0.25", features = ["derive"] } + diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_generation_type.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_generation_type.rs index da85580539..752c69fc86 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_generation_type.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_generation_type.rs @@ -8,7 +8,7 @@ use utoipa::ToSchema; #[cfg_attr(test, derive(EnumIter, EnumCount))] -#[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Deserialize, Serialize, ToSchema)] +#[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Debug, Deserialize, Serialize, ToSchema)] pub enum EnqueueImageGenType { #[serde(rename = "upload_lora")] UploadLoRA, 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 9660a17c0e..2fe137a43f 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 @@ -2,6 +2,7 @@ #![forbid(unused_mut)] #![forbid(unused_variables)] +use std::fmt::{Display, Formatter}; use std::fmt::Debug; use std::sync::Arc; @@ -11,8 +12,6 @@ use actix_web::http::StatusCode; use log::warn; use serde::Deserialize; use serde::Serialize; - -use tokens::tokens::model_weights::ModelWeightToken; use utoipa::ToSchema; use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; @@ -24,22 +23,19 @@ use mysql_queries::payloads::generic_inference_args::generic_inference_args::{ InferenceCategoryAbbreviated, PolymorphicInferenceArgs, }; - use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{ insert_generic_inference_job, InsertGenericInferenceArgs, }; - - use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::users::UserToken; use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::server_state::ServerState; -use std::fmt::{Display, Formatter}; /// Debug requests can get routed to special "debug-only" workers, which can /// be used to trial new code, run debugging, etc. const DEBUG_HEADER_NAME: &str = "enable-debug-mode"; @@ -344,7 +340,7 @@ pub async fn enqueue_image_generation_request( maybe_seed: Some(seed), maybe_upload_path: Some(upload_path), maybe_lora_upload_path: Some(lora_upload_path), - type_of_inference: type_of_inference, + type_of_inference, maybe_cfg_scale: Some(cfg_scale), maybe_number_of_samples: Some(number_of_samples), maybe_batch_count: Some(batch_count),