Merge pull request #268 from storytold/jan-04-mithoo

adds styletts2 as voice designer model
This commit is contained in:
Kasisnu authored and GitHub committed 2024-01-16 10:52:06 +05:30
commit b77810cd29
19 files changed
+1450 -10

No files matched your search

@@ -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",
}
}
}
@@ -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
@@ -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);
}
}
@@ -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<Self, String> {
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);
}
}
@@ -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,
@@ -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;
@@ -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")))
}
@@ -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;
@@ -0,0 +1,87 @@
use crate::util::model_downloader::ModelDownloader;
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
pub fn build_all_from_env() -> Self {
Self {
}
}
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"));
// }
}
@@ -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<JobSuccessResult, ProcessSingleJobError> {
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?")))
}
}
}
}
@@ -0,0 +1,349 @@
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<JobSuccessResult, ProcessSingleJobError> {
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<DatasetSampleRecordForList> = match result {
Ok(val) => val,
Err(e) => return Err(ProcessSingleJobError::from_anyhow_error(e)),
};
info!("Dataset length info: {}", dataset.len());
let mut downloaded_dataset: Vec<PathBuf> = 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());
// 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: {}", 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
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();
// 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: &output_embedding_path,
audio_file: 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);
}
//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!("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("style.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<AudioFile, ProcessSingleJobError> {
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<PathBuf>) -> String {
paths
.into_iter()
.map(|p| format!("\"{}\"", p.display()))
.collect::<Vec<String>>()
.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);
}
}
@@ -0,0 +1,304 @@
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::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;
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;
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<JobSuccessResult, ProcessSingleJobError> {
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) => {
clean_symbols(value.as_str())
}
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);
}
}
// 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 = "style.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 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 = 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_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,
}
);
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);
}
// 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,
origin_model_type: MediaFileOriginModelType::StyleTTS2,
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<VoiceFile, ProcessSingleJobError> {
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<VoiceFile, ProcessSingleJobError> {
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)
}
@@ -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<Self> {
Ok(Self {
create_embedding_command: StyleTTS2CreateEmbeddingCommand::from_env()?,
downloaders: StyleTTS2Downloaders::build_all_from_env(),
inference_command: StyleTTS2InferenceCommand::from_env()?,
})
}
}
@@ -0,0 +1,497 @@
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<HashSet<String>> = 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::<HashSet<String>>()
});
// The inference command assumes everything has been setup to run inference
// so we have all the files downloaded for the inputs.
pub struct InferenceArgs<P: AsRef<Path>> {
/// --driven_audio: path to the input audio
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<P: AsRef<Path>> {
pub output_embedding_path: P,
pub audio_file: 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<String>,
/// Optional default config file to use
maybe_default_config_path: Option<PathBuf>,
/// If this is run under Docker (eg. in development), these are the options.
maybe_docker_options: Option<DockerOptions>,
/// If the execution should be ended after a certain point.
maybe_execution_timeout: Option<Duration>,
/// Inference arg.
/// --checkpoint_dir: optional location for checkpoints directory
pub alternate_checkpoint_dir: Option<PathBuf>,
}
impl StyleTTS2InferenceCommand {
pub fn new<P: AsRef<Path>>(
root_code_directory: P,
executable_or_command: ExecutableOrCommand,
maybe_virtual_env_activation_command: Option<&str>,
maybe_default_config_path: Option<P>,
maybe_docker_options: Option<DockerOptions>,
maybe_execution_timeout: Option<Duration>,
alternate_checkpoint_dir: Option<PathBuf>
) -> 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<Self> {
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<P: AsRef<Path>>(&self, args: InferenceArgs<P>) -> CommandExitStatus {
match self.do_execute_inference(args) {
Ok(exit_status) => exit_status,
Err(error) =>
CommandExitStatus::FailureWithReason { reason: format!("error: {:?}", error) },
}
}
fn do_execute_inference<P: AsRef<Path>>(
&self,
args: InferenceArgs<P>
) -> AnyhowResult<CommandExitStatus> {
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(" --input ");
command.push_str(&path_to_string(args.input_text_file_path));
command.push_str(" --input-style-npz ");
command.push_str(&path_to_string(args.input_embedding_file_path));
command.push_str(" --output ");
command.push_str(&path_to_string(args.output_file_name));
//TODO(KS): Make this a parameter
command.push_str(" --vcsteps ");
command.push_str(format!("{}", 20).as_str());
// ===== 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<String>,
/// Optional default config file to use
maybe_default_config_path: Option<PathBuf>,
/// If this is run under Docker (eg. in development), these are the options.
maybe_docker_options: Option<DockerOptions>,
/// If the execution should be ended after a certain point.
maybe_execution_timeout: Option<Duration>,
/// Inference arg.
/// --checkpoint_dir: optional location for checkpoints directory
pub alternate_checkpoint_dir: Option<PathBuf>,
}
impl StyleTTS2CreateEmbeddingCommand {
pub fn new<P: AsRef<Path>>(
root_code_directory: P,
executable_or_command: ExecutableOrCommand,
maybe_virtual_env_activation_command: Option<&str>,
maybe_default_config_path: Option<P>,
maybe_docker_options: Option<DockerOptions>,
maybe_execution_timeout: Option<Duration>,
alternate_checkpoint_dir: Option<PathBuf>
) -> 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<Self> {
// 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<P: AsRef<Path>>(&self, args: CreateVoiceInferenceArgs<P>) -> CommandExitStatus {
match self.do_execute_inference(args) {
Ok(exit_status) => exit_status,
Err(error) =>
CommandExitStatus::FailureWithReason { reason: format!("error: {:?}", error) },
}
}
fn do_execute_inference<P: AsRef<Path>>(
&self,
args: CreateVoiceInferenceArgs<P>
) -> AnyhowResult<CommandExitStatus> {
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(" --output-style-npz ");
command.push_str(&path_to_string(&args.output_embedding_path));
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() {
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))
}
}
}
}
}
}
@@ -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<String>, // infers
pub voice_dataset_token: Option<String>, // creates
pub job_type: JobType
}
pub fn validate_job(job: &AvailableInferenceJob) -> Result<JobArgs, ProcessSingleJobError> {
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, ProcessSingleJobError>(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,ProcessSingleJobError>(JobArgs {
voice_token: None,
voice_dataset_token: Some(dataset_token.clone()),
job_type: JobType::Create
});
}
Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Missing Voice Token!")))
}
@@ -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,
@@ -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<VitsDependencies>,
pub maybe_rerender_dependencies: Option<RerenderDependencies>,
pub maybe_mocapnet_dependencies: Option<MocapNetDependencies>,
pub maybe_styletts2_dependencies: Option<StyleTTS2Dependencies>,
}
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,
})
}
}
@@ -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,
@@ -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,