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 eb193ca55f..542385bf17 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 @@ -47,6 +47,9 @@ pub enum MediaFileOriginModelType { #[serde(rename = "stable_diffusion_1_5")] StableDiffusion15, + + #[serde(rename = "gpt_sovits")] + GptSovits, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -67,7 +70,8 @@ impl MediaFileOriginModelType { Self::MocapNet => "mocap_net", Self::ComfyUi => "comfy_ui", Self::StyleTTS2 => "styletts2", - Self::StableDiffusion15 => "stable_diffusion_1_5" + Self::StableDiffusion15 => "stable_diffusion_1_5", + Self::GptSovits => "gpt_sovits", } } @@ -83,6 +87,7 @@ impl MediaFileOriginModelType { "comfy_ui" => Ok(Self::ComfyUi), "styletts2" => Ok(Self::StyleTTS2), "stable_diffusion_1_5" => Ok(Self::StableDiffusion15), + "gpt_sovits" => Ok(Self::GptSovits), _ => Err(format!("invalid value: {:?}", value)), } } @@ -100,7 +105,8 @@ impl MediaFileOriginModelType { Self::MocapNet, Self::ComfyUi, Self::StyleTTS2, - Self::StableDiffusion15 + Self::StableDiffusion15, + Self::GptSovits, ]) } } @@ -130,6 +136,7 @@ mod tests { assert_eq!(MediaFileOriginModelType::ComfyUi.to_str(), "comfy_ui"); assert_eq!(MediaFileOriginModelType::StyleTTS2.to_str(), "styletts2"); assert_eq!(MediaFileOriginModelType::StableDiffusion15.to_str(), "stable_diffusion_1_5"); + assert_eq!(MediaFileOriginModelType::GptSovits.to_str(), "gpt_sovits"); } #[test] @@ -144,13 +151,14 @@ mod tests { assert_eq!(MediaFileOriginModelType::from_str("comfy_ui").unwrap(), MediaFileOriginModelType::ComfyUi); assert_eq!(MediaFileOriginModelType::from_str("styletts2").unwrap(), MediaFileOriginModelType::StyleTTS2); assert_eq!(MediaFileOriginModelType::from_str("stable_diffusion_1_5").unwrap(), MediaFileOriginModelType::StableDiffusion15); + assert_eq!(MediaFileOriginModelType::from_str("gpt_sovits").unwrap(), MediaFileOriginModelType::GptSovits); assert!(MediaFileOriginModelType::from_str("foo").is_err()); } #[test] fn all_variants() { let mut variants = MediaFileOriginModelType::all_variants(); - assert_eq!(variants.len(), 10); + assert_eq!(variants.len(), 11); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::RvcV2)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::SadTalker)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::SoVitsSvc)); @@ -161,6 +169,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::ComfyUi)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::StyleTTS2)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::StableDiffusion15)); + assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::GptSovits)); assert_eq!(variants.pop_first(), None); } } diff --git a/crates/lib/jobs_common/src/semi_persistent_cache_dir.rs b/crates/lib/jobs_common/src/semi_persistent_cache_dir.rs index 0646bf640b..437ba4757b 100644 --- a/crates/lib/jobs_common/src/semi_persistent_cache_dir.rs +++ b/crates/lib/jobs_common/src/semi_persistent_cache_dir.rs @@ -10,6 +10,7 @@ pub struct SemiPersistentCacheDir { tts_synthesizer_model_root: PathBuf, tts_pretrained_vocoder_model_root: PathBuf, // Pretrained, non-user uploaded vocoders vocoder_model_root: PathBuf, // User-uploaded vocoders + gpt_sovits_model_root: PathBuf, // GPT-SoViT models voice_conversion_model_root: PathBuf, @@ -37,6 +38,7 @@ impl SemiPersistentCacheDir { tts_synthesizer_model_root: cache_root.join("tts/synthesizer_models/"), tts_pretrained_vocoder_model_root: cache_root.join("tts/vocoder_models_pretrained/"), vocoder_model_root: cache_root.join("tts/user_uploaded_vocoder_models/"), + gpt_sovits_model_root: cache_root.join("tts/gpt_sovits_models/"), voice_conversion_model_root: cache_root.join("voice_conversion/models/"), @@ -191,6 +193,20 @@ impl SemiPersistentCacheDir { Ok(()) } + // ==================== GPT SOVITS MODELS ==================== + pub fn gpt_sovits_model_path(&self, model_filename: &str) -> PathBuf { + self.gpt_sovits_model_root.join(model_filename) + } + + pub fn gpt_sovits_model_directory(&self) -> &Path { + &self.gpt_sovits_model_root + } + + pub fn create_gpt_sovits_model_path(&self) -> AnyhowResult<()> { + let _ = fs::create_dir_all(self.gpt_sovits_model_directory())?; + Ok(()) + } + // ==================== W2L OUTPUT RESULTS ==================== // We cache W2L faces here. diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/gptsovits_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/gptsovits_payload.rs index d70e26feb5..928fcc705e 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/gptsovits_payload.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/gptsovits_payload.rs @@ -1,10 +1,4 @@ -use std::collections::HashMap; - use enums::common::visibility::Visibility; -use enums::no_table::style_transfer::style_transfer_name::StyleTransferName; -use tokens::tokens::media_files::MediaFileToken; -use tokens::tokens::model_weights::ModelWeightToken; - #[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)] pub struct GptSovitsPayload { @@ -20,3 +14,4 @@ pub struct GptSovitsPayload { #[serde(skip_serializing_if = "Option::is_none")] pub creator_visibility: Option, } + diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/gpt_sovits_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/gpt_sovits_dependencies.rs new file mode 100644 index 0000000000..304bd26866 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/gpt_sovits_dependencies.rs @@ -0,0 +1,16 @@ +use errors::AnyhowResult; + +use crate::job::job_types::gpt_sovits::gpt_sovits_inference_command::GptSovitsInferenceCommand; + +pub struct GptSovitsDependencies { + pub inference_command: GptSovitsInferenceCommand, +} + + +impl GptSovitsDependencies { + pub fn setup() -> AnyhowResult { + Ok(Self { + inference_command: GptSovitsInferenceCommand::from_env()?, + }) + } +} \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/gpt_sovits_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/gpt_sovits_inference_command.rs new file mode 100644 index 0000000000..46d9858936 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/gpt_sovits_inference_command.rs @@ -0,0 +1,252 @@ +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 filesys::path_to_string::path_to_string; +use subprocess_common::command_exit_status::CommandExitStatus; +use subprocess_common::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions}; + +// 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 GptSovitsInferenceCommand { + /// Where the code lives + pub(crate) gpt_sovits_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, + + /// 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, + +} +#[derive(Clone)] +pub enum ExecutableOrCommand { + /// Eg. `inference.py` + Executable(PathBuf), + + /// Eg. `python3 inference.py` + Command(String), +} + +#[derive(Debug)] +pub struct InferenceArgs<'s> { + pub stderr_output_file: &'s Path, + pub stdout_output_file: &'s Path, + + pub input_text_file: &'s Path, + pub gpt_model_path: &'s Path, + pub sovits_model_path: &'s Path, + pub reference_audio_path: &'s Path, + pub output_audio_directory: &'s Path, + + pub maybe_reference_free: Option, + pub maybe_temperature: Option, + pub maybe_target_language: Option, +} + +impl GptSovitsInferenceCommand { + pub fn new( + gpt_sovits_code_directory: PathBuf, + executable_or_command: ExecutableOrCommand, + maybe_virtual_env_activation_command: Option, + maybe_docker_options: Option, + maybe_execution_timeout: Option, + ) -> Self { + Self { + gpt_sovits_code_directory, + executable_or_command, + maybe_virtual_env_activation_command, + maybe_docker_options, + maybe_execution_timeout, + } + } + + pub fn from_env() -> anyhow::Result { + let gpt_sovits_code_directory = easyenv::get_env_pathbuf_required("GPT_SOVITS_CODE_DIRECTORY")?; + + let maybe_inference_command = easyenv::get_env_string_optional("GPT_SOVITS_INFERENCE_COMMAND"); + + let maybe_inference_executable = easyenv::get_env_pathbuf_optional("GPT_SOVITS_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("GPT_SOVITS_VENV_ACTIVATION_COMMAND"); + let maybe_docker_options = easyenv::get_env_string_optional("GPT_SOVITS_DOCKER_OPTIONS") + .map(|image_name| { + DockerOptions { + image_name, + maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()), + maybe_environment_variables: None, + maybe_gpu: Some(DockerGpu::All), + } + }); + let maybe_execution_timeout = easyenv::get_env_duration_seconds_optional("GPT_SOVITS_EXECUTION_TIMEOUT"); + + Ok(Self { + gpt_sovits_code_directory, + executable_or_command, + maybe_virtual_env_activation_command, + maybe_docker_options, + maybe_execution_timeout, + }) + } + + 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) }, + } + } + + pub fn do_execute_inference( + &self, + args: InferenceArgs, + ) -> anyhow::Result { + let mut command = String::new(); + + + command.push_str(&format!("cd {}", path_to_string(&self.gpt_sovits_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(" infer "); + } + ExecutableOrCommand::Command(ref cmd) => { + command.push_str(cmd); + command.push_str(" "); + } + } + + command.push_str(&format!(" --target_text {}", path_to_string(args.input_text_file))); + command.push_str(&format!(" --gpt_model {}", path_to_string(args.gpt_model_path))); + command.push_str(&format!(" --sovits_model {}", path_to_string(args.sovits_model_path))); + command.push_str(&format!(" --ref_audio {}", path_to_string(args.reference_audio_path))); + command.push_str(&format!(" --output_path {}", path_to_string(args.output_audio_directory))); + + if let Some(reference_free) = args.maybe_reference_free { + command.push_str(&format!(" --ref_free {}", reference_free)); + } + + if let Some(temperature) = args.maybe_temperature { + command.push_str(&format!(" --temperature {}", temperature)); + } + + if let Some(target_language) = args.maybe_target_language { + command.push_str(&format!(" --target_language {}", target_language)); + } + + if let Some(maybe_docker_options) = self.maybe_docker_options.as_ref() { + command = maybe_docker_options.to_command_string(&command); + } + + info!("Running command: {}", 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_os_str()); + + 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)?; + + 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)) + } + } + } + } + } +} \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/mod.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/mod.rs index 88a5b47531..769aeee3d2 100644 --- a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/mod.rs @@ -1,2 +1,7 @@ pub mod process_single_gpt_sovits_job; -mod upload_model; \ No newline at end of file +mod upload_model; +mod gpt_sovits_inference_command; +pub mod gpt_sovits_dependencies; +mod tts_inference; + +mod model_package; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/model_package/download_package.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/model_package/download_package.rs new file mode 100644 index 0000000000..577b5647ef --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/model_package/download_package.rs @@ -0,0 +1,79 @@ +use std::collections::HashSet; +use std::path::{Path, PathBuf}; + +use anyhow::anyhow; +use log::{debug, error, info}; +use once_cell::sync::Lazy; + +use buckets::public::weight_files::bucket_directory::WeightFileBucketDirectory; +use buckets::public::weight_files::bucket_file_path::WeightFileBucketPath; +use cloud_storage::bucket_client::BucketClient; +use filesys::rename_across_devices::{rename_across_devices, RenameError}; +use filesys::safe_delete_temp_directory::safe_delete_temp_directory; +use tokens::tokens::model_weights::ModelWeightToken; + +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::gpt_sovits::model_package::model_package::GptSovitsPackageFileType; +use crate::util::filesystem::scoped_temp_dir_creator::ScopedTempDirCreator; + +pub async fn download_package( + model_weight_token: ModelWeightToken, + weights_file_bucket_directory: &WeightFileBucketDirectory, + download_directory: &Path, + bucket_client: &BucketClient, + scoped_tempdir_creator: &ScopedTempDirCreator, + force_download: bool, +) -> Result<(), ProcessSingleJobError> { + + if !download_directory.exists() { + std::fs::create_dir_all(&download_directory) + .map_err(|err| ProcessSingleJobError::IoError(err))?; + } + + let temp_dir_name = format!("model_download_{}", model_weight_token.entropy_suffix()); + let temp_dir = scoped_tempdir_creator.new_tempdir(&temp_dir_name) + .map_err(|e| anyhow!("problem creating tempdir: {:?}", e))?; + + for entry in GptSovitsPackageFileType::all_variants() { + let suffix_with_package_identifier = entry.get_expected_package_suffix(); + let temp_path = temp_dir.path().join(format!("{}{}", model_weight_token.to_string(), suffix_with_package_identifier)); + let final_download_path = download_directory.join(format!("{}{}", model_weight_token.to_string(), suffix_with_package_identifier)); + + if final_download_path.exists() && !force_download { + info!("Skipping download of {} because it already exists", final_download_path.display().to_string()); + continue; + } + + if final_download_path.exists() { + debug!("Deleting existing file at {:?}", final_download_path.display().to_string()); + std::fs::remove_file(&final_download_path) + .map_err(|err| ProcessSingleJobError::IoError(err))?; + } + + let bucket_public_download_path = WeightFileBucketPath::from_object_hash( + weights_file_bucket_directory.get_object_hash(), + Some("weight_"), + Some(&suffix_with_package_identifier), + ); + + debug!("Downloading to {} from bucket path: {:?}", temp_path.display().to_string(), &bucket_public_download_path.get_full_object_path_str()); + bucket_client.download_file_to_disk(bucket_public_download_path.get_full_object_path_str(), &temp_path) + .await + .map_err(|e| { + anyhow!("couldn't download {} cloud object to disk: {:?}", &bucket_public_download_path.get_full_object_path_str(), e) + })?; + + debug!("Renaming file from {:?} to {:?}!", &temp_path, final_download_path.display().to_string()); + + rename_across_devices(&temp_path, final_download_path) + .map_err(|err| { + safe_delete_temp_directory(&temp_dir); + match err { + RenameError::StorageFull => ProcessSingleJobError::FilesystemFull, + RenameError::IoError(err) => ProcessSingleJobError::from_io_error(err), + }})?; + } + + safe_delete_temp_directory(&temp_dir); + Ok(()) +} diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/model_package/mod.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/model_package/mod.rs new file mode 100644 index 0000000000..70d297ff4a --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/model_package/mod.rs @@ -0,0 +1,2 @@ +pub mod download_package; +pub mod model_package; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/model_package/model_package.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/model_package/model_package.rs new file mode 100644 index 0000000000..b88c6718ec --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/model_package/model_package.rs @@ -0,0 +1,102 @@ +use std::collections::{BTreeSet, HashSet}; + +use once_cell::sync::Lazy; + +use buckets::public::weight_files::bucket_file_path::WeightFileBucketPath; + +pub const SUFFIX: &str = ".bin"; + +static ALLOWED_TYPES_GPT : Lazy> = Lazy::new(|| { + HashSet::from([ + "ckpt", + ]) +}); + +static ALLOWED_TYPES_SOVITS : Lazy> = Lazy::new(|| { + HashSet::from([ + "pth", + ]) +}); + +static ALLOWED_TYPES_REF_AUDIO : Lazy> = Lazy::new(|| { + HashSet::from([ + "wav", + ]) +}); + +#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq)] +pub enum GptSovitsPackageFileType { + GptModel, + SovitsCheckpoint, + ReferenceAudio, +} + +impl GptSovitsPackageFileType { + pub fn for_extension(extension: &str) -> Option { + if ALLOWED_TYPES_GPT.contains(extension) { + Some(GptSovitsPackageFileType::GptModel) + } else if ALLOWED_TYPES_SOVITS.contains(extension) { + Some(GptSovitsPackageFileType::SovitsCheckpoint) + } else if ALLOWED_TYPES_REF_AUDIO.contains(extension) { + Some(GptSovitsPackageFileType::ReferenceAudio) + } else { + None + } + } + + pub fn extension_is_allowed(&self, extension: &str) -> bool { + match self { + GptSovitsPackageFileType::GptModel => ALLOWED_TYPES_GPT.contains(extension), + GptSovitsPackageFileType::SovitsCheckpoint => ALLOWED_TYPES_SOVITS.contains(extension), + GptSovitsPackageFileType::ReferenceAudio => ALLOWED_TYPES_REF_AUDIO.contains(extension), + } + } + + + pub fn all_variants() -> BTreeSet { + // NB: BTreeSet is sorted + // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. + BTreeSet::from([ + Self::GptModel, + Self::SovitsCheckpoint, + Self::ReferenceAudio, + ]) + } + + pub fn package_identifier(&self) -> &str { + match self { + GptSovitsPackageFileType::GptModel => "gpt_model", + GptSovitsPackageFileType::SovitsCheckpoint => "sovits_checkpoint", + GptSovitsPackageFileType::ReferenceAudio => "reference_audio", + } + } + + pub fn get_expected_package_suffix(&self) -> String { + format!(".{}{}", self.package_identifier(), SUFFIX) + } +} + +#[derive(Debug)] +pub enum GptSovitsPackageError { + InvalidArchive, + InvalidGPTModel(String), + InvalidSovitsCheckpoint(String), + InvalidReferenceAudio(String), + UploadError, + TooManyFiles, + ExtractionError, + FileError, +} + +pub struct GptSovitsPackageFile { + pub public_upload_path: WeightFileBucketPath, + pub sha256_checksum: String, + pub file_size_bytes: u64, +} + +pub struct GptSovitsPackageDetails { + pub gpt_model: GptSovitsPackageFile, + pub sovits_checkpoint: GptSovitsPackageFile, + pub reference_audio: GptSovitsPackageFile, +} + diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/process_single_gpt_sovits_job.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/process_single_gpt_sovits_job.rs index 7fcec1ca59..fd694e292e 100644 --- a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/process_single_gpt_sovits_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/process_single_gpt_sovits_job.rs @@ -5,39 +5,29 @@ use mysql_queries::queries::generic_inference::job::list_available_generic_infer use crate::job::job_loop::job_success_result::JobSuccessResult; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::gpt_sovits::tts_inference::process_single_gpt_sovits_tts_job::process_single_gpt_sovits_tts_job; use crate::job::job_types::gpt_sovits::upload_model::process_gpt_sovits_upload_job::process_gpt_sovits_upload_job; -use crate::job::job_types::workflow::upload_workflow::process_upload_workflow_job::process_upload_workflow_job; -use crate::job::job_types::workflow::video_style_transfer::extract_vst_workflow_payload_from_job::extract_vst_workflow_payload_from_job; -use crate::job::job_types::workflow::video_style_transfer::process_video_style_transfer_job::process_video_style_transfer_job; use crate::state::job_dependencies::JobDependencies; pub async fn process_single_gpt_sovits_job( job_dependencies: &JobDependencies, job: &AvailableInferenceJob ) -> Result { - - // First try to dispatch with the newer "job type". - match job.job_type { + return match job.job_type { InferenceJobType::GptSovits => { - let job_success_result = process_gpt_sovits_upload_job(job_dependencies, job).await?; - return Ok(job_success_result); + match job.maybe_download_url { + Some(_) => { + let job_success_result = process_gpt_sovits_upload_job(job_dependencies, job).await?; + Ok(job_success_result) + } + None => { + let job_success_result = process_single_gpt_sovits_tts_job(job_dependencies, job).await?; + Ok(job_success_result) + } + } } _ => { - return Err(ProcessSingleJobError::Other(anyhow!("tts model type not set"))) + Err(ProcessSingleJobError::Other(anyhow!("tts model type not set"))) } } - - // If we couldn't dispatch with "job type", fall back to older heuristics. - let workflow_args = extract_vst_workflow_payload_from_job(&job)?; - - let job_success_result = match workflow_args.maybe_google_drive_link { - Some(_link) => { - process_upload_workflow_job(job_dependencies, job).await? - } - None => { - process_video_style_transfer_job(job_dependencies, job).await? - } - }; - - Ok(job_success_result) } \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/tts_inference/check_and_validate_job.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/tts_inference/check_and_validate_job.rs new file mode 100644 index 0000000000..727fc426ac --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/tts_inference/check_and_validate_job.rs @@ -0,0 +1,75 @@ +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 tokens::tokens::model_weights::ModelWeightToken; + +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; + +#[derive(Serialize, Debug)] +pub struct JobArgs { + pub input_text: String, + pub gpt_sovits_model: ModelWeightToken, + pub reference_free: Option, + pub temperature: Option, + pub target_language: Option, +} + +pub fn check_and_validate_job(job: &AvailableInferenceJob) -> Result { + let inference_args = job.maybe_inference_args + .as_ref() + .map(|args| args.args.as_ref()) + .flatten(); + + let inference_category = job.maybe_inference_args + .as_ref() + .map(|args| args.inference_category) + .flatten(); + + match inference_category { + Some(InferenceCategoryAbbreviated::GptSovits) => {}, // Valid + 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 = match inference_args { + Some(args) => args, + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no inference args for job!"))); + } + }; + + let inference_args = match inference_args { + PolymorphicInferenceArgs::Gs(inference_args) => inference_args, + _ => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!"))); + } + }; + + let input_text = match &job.maybe_raw_inference_text { + Some(text) => text.clone(), + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no raw inference text for job!"))); + } + }; + + let gpt_sovits_model = match &job.maybe_model_token { + Some(token) => ModelWeightToken::new_from_str(token), + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no model token for job!"))); + } + }; + + + Ok(JobArgs { + input_text: input_text, + gpt_sovits_model: gpt_sovits_model, + reference_free: None, + temperature: None, + target_language: None, + }) +} \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/tts_inference/mod.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/tts_inference/mod.rs new file mode 100644 index 0000000000..2a0acf625b --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/tts_inference/mod.rs @@ -0,0 +1,2 @@ +pub mod process_single_gpt_sovits_tts_job; +mod check_and_validate_job; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/tts_inference/process_single_gpt_sovits_tts_job.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/tts_inference/process_single_gpt_sovits_tts_job.rs new file mode 100644 index 0000000000..cf5e4877d1 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/tts_inference/process_single_gpt_sovits_tts_job.rs @@ -0,0 +1,191 @@ +use std::time::Instant; + +use anyhow::anyhow; +use log::info; + +use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; +use buckets::public::weight_files::bucket_directory::WeightFileBucketDirectory; +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::check_file_exists::check_file_exists; +use filesys::file_size::file_size; +use hashing::sha256::sha256_hash_file::sha256_hash_file; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use mysql_queries::queries::media_files::create::insert_media_file_from_zero_shot_tts::{insert_media_file_from_zero_shot, InsertArgs}; +use mysql_queries::queries::model_weights::get::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::job_types::gpt_sovits::gpt_sovits_inference_command::InferenceArgs; +use crate::job::job_types::gpt_sovits::model_package::download_package::download_package; +use crate::job::job_types::gpt_sovits::model_package::model_package::GptSovitsPackageFileType; +use crate::job::job_types::gpt_sovits::tts_inference::check_and_validate_job::check_and_validate_job; +use crate::state::job_dependencies::JobDependencies; + +const BUCKET_FILE_PREFIX: &str = "fakeyou_"; +const BUCKET_FILE_EXTENSION: &str = ".wav"; +const MIME_TYPE: &str = "audio/wav"; + +pub async fn process_single_gpt_sovits_tts_job( + job_dependencies: &JobDependencies, + job: &AvailableInferenceJob +) -> Result { + + let mut job_progress_reporter = job_dependencies + .clients + .job_progress_reporter + .new_generic_inference(job.inference_job_token.as_str()) + .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + + + let gpt_sovits_deps = job_dependencies + .job + .job_specific_dependencies + .maybe_gpt_sovits_dependencies + .as_ref() + .ok_or_else(|| ProcessSingleJobError::JobSystemMisconfiguration(Some("Missing GPT-Sovits dependencies".to_string())))?; + + let job_args = check_and_validate_job(job)?; + + let work_temp_dir = format!("temp_gpt_sovits_tts_{}", job.id.0); + + // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem. + let work_temp_dir = job_dependencies + .fs + .scoped_temp_dir_creator_for_work + .new_tempdir(&work_temp_dir) + .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + + let output_dir = work_temp_dir.path().join("output"); + let output_file_path = output_dir.join("GPTSoVITS_1.wav"); + + if !output_dir.exists() { + std::fs::create_dir_all(&output_dir) + .map_err(|err| ProcessSingleJobError::IoError(err))?; + } + + let input_text = job_args.input_text; + let gpt_sovits_model = job_args.gpt_sovits_model; + + let model_record = get_weight_by_token(&gpt_sovits_model, false, &job_dependencies.db.mysql_pool) + .await?.ok_or_else(|| ProcessSingleJobError::ModelDeleted)?; + + let model_token = model_record.token.clone(); + let weights_file_bucket_directory = WeightFileBucketDirectory::from_object_hash(&model_record.public_bucket_hash); + + let weights_directory = &job_dependencies.fs.semi_persistent_cache.gpt_sovits_model_path(&model_token.as_str()); + + // TODO(KS): Expose this as configuration + let force_download = job.attempt_count > 2; + + download_package( + gpt_sovits_model, + &weights_file_bucket_directory, + &weights_directory, + &job_dependencies.buckets.public_bucket_client, + &job_dependencies.fs.scoped_temp_dir_creator_for_short_lived_downloads, + force_download, + ).await?; + + let stderr_output_file = work_temp_dir.path().join("stderr.txt"); + let stdout_output_file = work_temp_dir.path().join("stdout.txt"); + let text_input_fs_path = work_temp_dir.path().join("inference_input.txt"); + std::fs::write(&text_input_fs_path, &input_text) + .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + + let gpt_model_path = weights_directory.join(format!("{}{}", model_token.as_str(), GptSovitsPackageFileType::GptModel.get_expected_package_suffix())); + let sovits_model_path = weights_directory.join(format!("{}{}", model_token.as_str(), GptSovitsPackageFileType::SovitsCheckpoint.get_expected_package_suffix())); + let reference_audio_path = weights_directory.join(format!("{}{}", model_token.as_str(), GptSovitsPackageFileType::ReferenceAudio.get_expected_package_suffix())); + + let inference_start_time = Instant::now(); + let command_exit_status = gpt_sovits_deps + .inference_command + .execute_inference(InferenceArgs{ + stderr_output_file: &stderr_output_file, + stdout_output_file: &stdout_output_file, + input_text_file: &text_input_fs_path, + gpt_model_path: &gpt_model_path, + sovits_model_path: &sovits_model_path, + reference_audio_path: &reference_audio_path, + output_audio_directory: &output_dir, + maybe_reference_free: None, + maybe_temperature: None, + maybe_target_language: Some("english".to_string()), + }); + let inference_duration = Instant::now().duration_since(inference_start_time); + + info!("Inference command exited with status: {:?}", command_exit_status); + info!("Inference took duration to complete: {:?}", &inference_duration); + + // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== // + + info!("Checking that output files exist..."); + + check_file_exists(&output_file_path).map_err(|e| ProcessSingleJobError::Other(e))?; + + // 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 + info!("Upload Bucket Path: {:?}", result_bucket_object_pathbuf); + info!("Upload File Path: {:?}", &output_file_path); + + job_dependencies.buckets.public_bucket_client + .upload_filename_with_content_type( + &result_bucket_object_pathbuf, + &output_file_path, + &MIME_TYPE + ) + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // ==================== UPLOAD AUDIO TO BUCKET ==================== + + info!("Calculating sha256..."); + + let file_checksum = sha256_hash_file(&output_file_path).map_err(|err| { + ProcessSingleJobError::Other(anyhow!("Error hashing file: {:?}", err)) + })?; + + let file_size_bytes = file_size(&output_file_path).map_err(|err| + ProcessSingleJobError::Other(err) + )?; + + job_progress_reporter.log_status("done").map_err(|e| ProcessSingleJobError::Other(e))?; + let (media_file_token, id) = insert_media_file_from_zero_shot(InsertArgs { + pool: &job_dependencies.db.mysql_pool, + job: &job, + origin_model_type: MediaFileOriginModelType::GptSovits, + 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: job_dependencies.job.info.container.is_on_prem, + worker_hostname: &job_dependencies.job.info.container.hostname, + worker_cluster: &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, + }) +} \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/extract_and_verify_gpt_sovits_package.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/extract_and_upload_gpt_sovits_package.rs similarity index 68% rename from crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/extract_and_verify_gpt_sovits_package.rs rename to crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/extract_and_upload_gpt_sovits_package.rs index 97f1895aba..6514c18571 100644 --- a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/extract_and_verify_gpt_sovits_package.rs +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/extract_and_upload_gpt_sovits_package.rs @@ -1,9 +1,7 @@ -use std::collections::{BTreeSet, HashSet}; use std::io::{BufReader, Cursor, Read}; -use std::path::{Path, PathBuf}; +use std::path::PathBuf; use log::{error, info, warn}; -use once_cell::sync::Lazy; use zip::ZipArchive; use buckets::public::weight_files::bucket_directory::WeightFileBucketDirectory; @@ -12,107 +10,16 @@ use cloud_storage::bucket_client::BucketClient; use filesys::path_to_string::path_to_string; use hashing::sha256::sha256_hash_bytes::sha256_hash_bytes; -static ALLOWED_TYPES_GPT : Lazy> = Lazy::new(|| { - HashSet::from([ - "pth", - ]) -}); +use crate::job::job_types::gpt_sovits::model_package::model_package::{GptSovitsPackageDetails, GptSovitsPackageError, GptSovitsPackageFile, GptSovitsPackageFileType}; -static ALLOWED_TYPES_SOVITS : Lazy> = Lazy::new(|| { - HashSet::from([ - "ckpt", - ]) -}); - -static ALLOWED_TYPES_REF_AUDIO : Lazy> = Lazy::new(|| { - HashSet::from([ - "wav", - ]) -}); - -#[derive(Debug, Clone, Ord, PartialOrd, Eq, PartialEq)] -pub enum GptSovitsPackageFileType { - GptModel, - SovitsCheckpoint, - ReferenceAudio, -} - -impl GptSovitsPackageFileType { - pub fn for_extension(extension: &str) -> Option { - if ALLOWED_TYPES_GPT.contains(extension) { - Some(GptSovitsPackageFileType::GptModel) - } else if ALLOWED_TYPES_SOVITS.contains(extension) { - Some(GptSovitsPackageFileType::SovitsCheckpoint) - } else if ALLOWED_TYPES_REF_AUDIO.contains(extension) { - Some(GptSovitsPackageFileType::ReferenceAudio) - } else { - None - } - } - - pub fn extension_is_allowed(&self, extension: &str) -> bool { - match self { - GptSovitsPackageFileType::GptModel => ALLOWED_TYPES_GPT.contains(extension), - GptSovitsPackageFileType::SovitsCheckpoint => ALLOWED_TYPES_SOVITS.contains(extension), - GptSovitsPackageFileType::ReferenceAudio => ALLOWED_TYPES_REF_AUDIO.contains(extension), - } - } - - - pub fn all_variants() -> BTreeSet { - // NB: BTreeSet is sorted - // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. - BTreeSet::from([ - Self::GptModel, - Self::SovitsCheckpoint, - Self::ReferenceAudio, - ]) - } - - pub fn package_identifier(&self) -> &str { - match self { - GptSovitsPackageFileType::GptModel => "gpt_model", - GptSovitsPackageFileType::SovitsCheckpoint => "sovits_checkpoint", - GptSovitsPackageFileType::ReferenceAudio => "reference_audio", - } - } -} - -#[derive(Debug)] -pub enum GptSovitsPackageError { - InvalidArchive, - InvalidGPTModel(String), - InvalidSovitsCheckpoint(String), - InvalidReferenceAudio(String), - UploadError, - TooManyFiles, - ExtractionError, - FileError, -} - -pub struct GptSovitsPackageFile { - pub public_upload_path: WeightFileBucketPath, - pub sha256_checksum: String, - pub file_size_bytes: u64, -} - -pub struct GptSovitsPackageDetails { - pub gpt_model: GptSovitsPackageFile, - pub sovits_checkpoint: GptSovitsPackageFile, - pub reference_audio: GptSovitsPackageFile, -} - - -pub async fn extract_and_verify_gpt_sovits_package( +pub async fn extract_and_upload_gpt_sovits_package_files( zip_container_file_bytes: &[u8], bucket_client: &BucketClient, weights_file_bucket_directory: &WeightFileBucketDirectory, - prefix: Option<&str>, - suffix: Option<&str> ) -> Result { - let mut cursor = Cursor::new(zip_container_file_bytes); - let reader = std::io::BufReader::new(cursor); - let mut archive = zip::ZipArchive::new(reader) + let cursor = Cursor::new(zip_container_file_bytes); + let reader = BufReader::new(cursor); + let mut archive = ZipArchive::new(reader) .map_err(|err| { error!("Error reading zip archive: {:?}", err); GptSovitsPackageError::InvalidArchive @@ -128,7 +35,6 @@ pub async fn extract_and_verify_gpt_sovits_package( let entries = get_relevant_zip_entries(&mut archive)?; - for entry in entries.iter() { info!("Entry: {:?}", entry); @@ -150,17 +56,15 @@ pub async fn extract_and_verify_gpt_sovits_package( match entry.package_type { GptSovitsPackageFileType::GptModel => { - let suffix_with_package_identifier = format!(".{}{}", entry.package_type.package_identifier(), suffix.unwrap_or("")); let bucket_public_upload_path = WeightFileBucketPath::from_object_hash( weights_file_bucket_directory.get_object_hash(), Some("weight_"), - Some(&suffix_with_package_identifier), + Some(&entry.package_type.get_expected_package_suffix()), ); info!("Uploading GPT model to: {}", bucket_public_upload_path.get_full_object_path_str()); - bucket_client.upload_file_with_content_type( + bucket_client.upload_file( bucket_public_upload_path.get_full_object_path_str(), - zip_item_bytes.as_ref(), - "application/octet-stream") + zip_item_bytes.as_ref()) .await .map_err(|e| { warn!("Upload gpt package to bucket error: {:?}", e); @@ -182,17 +86,15 @@ pub async fn extract_and_verify_gpt_sovits_package( }); }, GptSovitsPackageFileType::SovitsCheckpoint => { - let suffix_with_package_identifier = format!(".{}{}", entry.package_type.package_identifier(), suffix.unwrap_or("")); let bucket_public_upload_path = WeightFileBucketPath::from_object_hash( weights_file_bucket_directory.get_object_hash(), Some("weight_"), - Some(&suffix_with_package_identifier), + Some(&entry.package_type.get_expected_package_suffix()), ); info!("Uploading sovits checkpoint to {:?}", bucket_public_upload_path.get_full_object_path_str()); - bucket_client.upload_file_with_content_type( + bucket_client.upload_file( bucket_public_upload_path.get_full_object_path_str(), - zip_item_bytes.as_ref(), - "application/octet-stream") + zip_item_bytes.as_ref()) .await .map_err(|e| { warn!("Upload sovits package to bucket error: {:?}", e); @@ -214,17 +116,15 @@ pub async fn extract_and_verify_gpt_sovits_package( }); }, GptSovitsPackageFileType::ReferenceAudio => { - let suffix_with_package_identifier = format!(".{}{}", entry.package_type.package_identifier(), suffix.unwrap_or("")); let bucket_public_upload_path = WeightFileBucketPath::from_object_hash( weights_file_bucket_directory.get_object_hash(), Some("weight_"), - Some(&suffix_with_package_identifier), + Some(&entry.package_type.get_expected_package_suffix()), ); info!("Uploading reference audio package to {:?}", bucket_public_upload_path.get_full_object_path_str()); - bucket_client.upload_file_with_content_type( + bucket_client.upload_file( bucket_public_upload_path.get_full_object_path_str(), - zip_item_bytes.as_ref(), - "application/octet-stream") + zip_item_bytes.as_ref()) .await .map_err(|e| { warn!("Upload reference audio package to bucket error: {:?}", e); @@ -267,10 +167,10 @@ struct PackageZipEntryDetails { fn get_relevant_zip_entries(archive: &mut ZipArchive>>) -> Result, GptSovitsPackageError> { let mut entries: Vec = Vec::new(); - for i in 0..(archive.len()) { + for i in 0..archive.len() { info!("Reading file {}...", i); - let mut file = archive.by_index(i) + let file = archive.by_index(i) .map_err(|err| { error!("Problem reading file from archive: {:?}", err); GptSovitsPackageError::InvalidArchive diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/mod.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/mod.rs index 3e4e096b60..f935589f3f 100644 --- a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/mod.rs @@ -1,3 +1,3 @@ pub mod process_gpt_sovits_upload_job; -mod extract_gpt_sovits_payload_from_job; -mod extract_and_verify_gpt_sovits_package; \ No newline at end of file +pub mod extract_gpt_sovits_payload_from_job; +mod extract_and_upload_gpt_sovits_package; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/process_gpt_sovits_upload_job.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/process_gpt_sovits_upload_job.rs index e138c2408d..19059bd99e 100644 --- a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/process_gpt_sovits_upload_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/process_gpt_sovits_upload_job.rs @@ -5,8 +5,6 @@ use log::{info, warn}; use buckets::public::weight_files::bucket_directory::WeightFileBucketDirectory; use buckets::public::weight_files::bucket_file_path::WeightFileBucketPath; -use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; -use cloud_storage::remote_file_manager::weights_descriptor::WeightsWorkflowDescriptor; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; @@ -24,9 +22,9 @@ use tokens::tokens::users::UserToken; use crate::job::job_loop::job_success_result::{JobSuccessResult, ResultEntity}; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; -use crate::job::job_types::gpt_sovits::upload_model::extract_and_verify_gpt_sovits_package::{extract_and_verify_gpt_sovits_package, GptSovitsPackageError}; +use crate::job::job_types::gpt_sovits::model_package::model_package::SUFFIX; +use crate::job::job_types::gpt_sovits::upload_model::extract_and_upload_gpt_sovits_package::extract_and_upload_gpt_sovits_package_files; use crate::job::job_types::gpt_sovits::upload_model::extract_gpt_sovits_payload_from_job::extract_gpt_sovits_payload_from_job; -use crate::job::job_types::workflow::video_style_transfer::extract_vst_workflow_payload_from_job::extract_vst_workflow_payload_from_job; use crate::state::job_dependencies::JobDependencies; pub async fn process_gpt_sovits_upload_job(deps: &JobDependencies, job: &AvailableInferenceJob) -> Result{ @@ -112,15 +110,16 @@ pub async fn process_gpt_sovits_upload_job(deps: &JobDependencies, job: &Availab info!("File Retrieved at {}", download_file_path.display()); const PREFIX: Option<&str> = Some("weight_"); - const SUFFIX: Option<&str> = Some(".bin"); let mut file_bytes = Vec::new(); file_bytes = file_read_bytes(&download_file_path) .map_err(|e| ProcessSingleJobError::from_anyhow_error(anyhow!("Processing archive failed")))?; - let weight_file_bucket_directory = WeightFileBucketDirectory::generate_new(); + let model_weight_token: &ModelWeightToken = &ModelWeightToken::generate(); + let weight_file_bucket_directory = WeightFileBucketDirectory::from_object_hash(model_weight_token.entropy_suffix()); - let gpt_sovits_package_details = extract_and_verify_gpt_sovits_package(&file_bytes, bucket_client, &weight_file_bucket_directory, PREFIX, SUFFIX) + // TODO(KS, 22-07-2023): We should decide if these details are worth persisting anywhere + let _ = extract_and_upload_gpt_sovits_package_files(&file_bytes, bucket_client, &weight_file_bucket_directory) .await.map_err(|easyenv| { warn!("Failed to extract and verify GPT-SoViTS package: {:?}", easyenv); ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to extract and verify GPT-SoViTS package")) @@ -142,12 +141,10 @@ pub async fn process_gpt_sovits_upload_job(deps: &JobDependencies, job: &Availab let file_checksum = sha256_hash_file(&download_file_path) .map_err(|e| ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to process archive checksum")))?; - let model_weight_token: &ModelWeightToken = &ModelWeightToken::new_from_str(weight_file_bucket_directory.get_object_hash()); - let model_weight_token_result = create_weight::create_weight(CreateModelWeightsArgs { token: &model_weight_token, - weights_type: WeightsType::ComfyUi, - weights_category: WeightsCategory::WorkflowConfig, + weights_type: WeightsType::GptSoVits, + weights_category: WeightsCategory::TextToSpeech, title, maybe_cover_image_media_file_token: job.maybe_cover_image_media_file_token.clone(), maybe_description_markdown: Some(description), @@ -163,7 +160,7 @@ pub async fn process_gpt_sovits_upload_job(deps: &JobDependencies, job: &Availab file_checksum_sha2: file_checksum, public_bucket_hash: bucket_public_upload_path.get_object_hash().to_string(), maybe_public_bucket_prefix: Some(PREFIX.unwrap().to_string()), - maybe_public_bucket_extension: Some(SUFFIX.unwrap().to_string()), + maybe_public_bucket_extension: Some(".zip".to_string()), version: 0, mysql_pool, }).await?; diff --git a/crates/service/job/inference_job/src/main.rs b/crates/service/job/inference_job/src/main.rs index 5c39c2e6b8..6453c4716a 100644 --- a/crates/service/job/inference_job/src/main.rs +++ b/crates/service/job/inference_job/src/main.rs @@ -51,12 +51,12 @@ use crate::http_server::run_http_server::launch_http_server; use crate::job::job_loop::main_loop::main_loop; use crate::state::job_dependencies::{BucketDependencies, ClientDependencies, DatabaseDependencies, FileSystemDetails, JobCaches, JobDependencies, JobInstanceInfo, JobSystemControls, JobSystemDependencies}; use crate::state::job_specific_dependencies::JobSpecificDependencies; +use crate::state::scoped_job_type_execution::ScopedJobTypeExecution; +use crate::state::scoped_model_type_execution::ScopedModelTypeExecution; use crate::util::filesystem::scoped_temp_dir_creator::ScopedTempDirCreator; use crate::util::instrumentation::{init_otel_metrics_pipeline, JobInstrumentLabels}; use crate::util::instrumentation::JobInstruments; use crate::util::model_weights_cache::model_weights_cache_directory::ModelWeightsCacheDirectory; -use crate::state::scoped_job_type_execution::ScopedJobTypeExecution; -use crate::state::scoped_model_type_execution::ScopedModelTypeExecution; pub mod http_server; pub mod job; @@ -367,6 +367,7 @@ fn set_up_directories(job_dependencies: &JobDependencies) -> AnyhowResult<()> { fs.semi_persistent_cache.create_tts_pretrained_vocoder_model_path()?; fs.semi_persistent_cache.create_tts_synthesizer_model_path()?; fs.semi_persistent_cache.create_voice_conversion_model_path()?; + fs.semi_persistent_cache.create_gpt_sovits_model_path()?; Ok(()) } diff --git a/crates/service/job/inference_job/src/state/job_specific_dependencies.rs b/crates/service/job/inference_job/src/state/job_specific_dependencies.rs index 10ff62fe4f..364aabbb52 100644 --- a/crates/service/job/inference_job/src/state/job_specific_dependencies.rs +++ b/crates/service/job/inference_job/src/state/job_specific_dependencies.rs @@ -3,6 +3,7 @@ use enums::by_table::generic_inference_jobs::inference_model_type::InferenceMode use errors::AnyhowResult; use crate::job::job_types::format_conversion::fbx_to_gltf::dependencies::FbxToGltfDependencies; +use crate::job::job_types::gpt_sovits::gpt_sovits_dependencies::GptSovitsDependencies; 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; @@ -32,6 +33,7 @@ pub struct JobSpecificDependencies { pub maybe_comfy_ui_dependencies: Option, pub maybe_convert_fbx_to_gltf_dependencies: Option, pub maybe_convert_bvh_to_workflow_dependencies: Option, + pub maybe_gpt_sovits_dependencies: Option, } impl JobSpecificDependencies { @@ -53,6 +55,7 @@ impl JobSpecificDependencies { let mut maybe_comfy_ui_dependencies = None; let mut maybe_convert_fbx_to_gltf_dependencies = None; let mut maybe_convert_bvh_to_workflow_dependencies = None; + let mut maybe_gpt_sovits_dependencies = None; if scoped_model_type_execution.can_run_job(InferenceModelType::ComfyUi) || scoped_job_type_execution.can_run_job(InferenceJobType::LivePortrait) @@ -63,6 +66,11 @@ impl JobSpecificDependencies { maybe_comfy_ui_dependencies = Some(ComfyDependencies::setup().await?); } + if scoped_job_type_execution.can_run_job(InferenceJobType::GptSovits) { + print_with_space("Setting GPT-SoViTS dependencies..."); + maybe_gpt_sovits_dependencies = Some(GptSovitsDependencies::setup()?); + } + if scoped_model_type_execution.can_run_job(InferenceModelType::RvcV2) { print_with_space("Setting RVCv2 dependencies..."); maybe_rvc_v2_dependencies = Some(RvcV2Dependencies::setup()?); @@ -137,6 +145,7 @@ impl JobSpecificDependencies { maybe_comfy_ui_dependencies, maybe_convert_fbx_to_gltf_dependencies, maybe_convert_bvh_to_workflow_dependencies, + maybe_gpt_sovits_dependencies, }) } } diff --git a/crates/service/job/tts_inference_job/src/main.rs b/crates/service/job/tts_inference_job/src/main.rs index 79df62e046..ddf83246f4 100644 --- a/crates/service/job/tts_inference_job/src/main.rs +++ b/crates/service/job/tts_inference_job/src/main.rs @@ -186,6 +186,7 @@ async fn main() -> AnyhowResult<()> { semi_persistent_cache.create_tts_synthesizer_model_path()?; semi_persistent_cache.create_tts_pretrained_vocoder_model_path()?; semi_persistent_cache.create_custom_vocoder_model_path()?; + semi_persistent_cache.create_gpt_sovits_model_path()?; let waveglow_vocoder_model_filename = easyenv::get_env_string_or_default( "TTS_WAVEGLOW_VOCODER_MODEL_FILENAME", "waveglow.pth");