diff --git a/Cargo.lock b/Cargo.lock index 25191d742f..9f21b922eb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3016,11 +3016,13 @@ dependencies = [ "log", "media", "memory_caching", + "mime", "mimetypes", "my-workspace-hack", "mysql_queries", "once_cell", "r2d2_redis", + "rand 0.8.5", "redis_common", "serde", "serde_derive", @@ -3029,6 +3031,7 @@ dependencies = [ "subprocess", "subprocess_common", "tempdir", + "time 0.3.30", "tokens", "tokio", "tts_common", @@ -5902,6 +5905,7 @@ dependencies = [ "hex", "hmac 0.10.1", "hostname", + "http", "http_server_common", "hyper", "infer 0.5.0", diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs index b38f181460..d1a420544b 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs @@ -25,6 +25,9 @@ pub enum InferenceCategory { #[serde(rename = "voice_conversion")] VoiceConversion, + + #[serde(rename = "rerender_a_video")] + RerenderAVideo, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -38,6 +41,7 @@ impl InferenceCategory { Self::LipsyncAnimation => "lipsync_animation", Self::TextToSpeech => "text_to_speech", Self::VoiceConversion => "voice_conversion", + Self::RerenderAVideo => "rerender_a_video", } } @@ -46,6 +50,7 @@ impl InferenceCategory { "lipsync_animation" => Ok(Self::LipsyncAnimation), "text_to_speech" => Ok(Self::TextToSpeech), "voice_conversion" => Ok(Self::VoiceConversion), + "rerender_a_video" => Ok(Self::RerenderAVideo), _ => Err(format!("invalid value: {:?}", value)), } } @@ -57,6 +62,7 @@ impl InferenceCategory { Self::LipsyncAnimation, Self::TextToSpeech, Self::VoiceConversion, + Self::RerenderAVideo, ]) } } @@ -71,6 +77,7 @@ mod tests { assert_serialization(InferenceCategory::LipsyncAnimation, "lipsync_animation"); assert_serialization(InferenceCategory::TextToSpeech, "text_to_speech"); assert_serialization(InferenceCategory::VoiceConversion, "voice_conversion"); + assert_serialization(InferenceCategory::RerenderAVideo, "rerender_a_video"); } #[test] @@ -78,6 +85,7 @@ mod tests { assert_eq!(InferenceCategory::LipsyncAnimation.to_str(), "lipsync_animation"); assert_eq!(InferenceCategory::TextToSpeech.to_str(), "text_to_speech"); assert_eq!(InferenceCategory::VoiceConversion.to_str(), "voice_conversion"); + assert_eq!(InferenceCategory::RerenderAVideo.to_str(), "rerender_a_video"); } #[test] @@ -85,16 +93,18 @@ mod tests { assert_eq!(InferenceCategory::from_str("lipsync_animation").unwrap(), InferenceCategory::LipsyncAnimation); assert_eq!(InferenceCategory::from_str("text_to_speech").unwrap(), InferenceCategory::TextToSpeech); assert_eq!(InferenceCategory::from_str("voice_conversion").unwrap(), InferenceCategory::VoiceConversion); + assert_eq!(InferenceCategory::from_str("rerender_a_video").unwrap(), InferenceCategory::RerenderAVideo); } #[test] fn all_variants() { // Static check let mut variants = InferenceCategory::all_variants(); - assert_eq!(variants.len(), 3); + assert_eq!(variants.len(), 4); assert_eq!(variants.pop_first(), Some(InferenceCategory::LipsyncAnimation)); assert_eq!(variants.pop_first(), Some(InferenceCategory::TextToSpeech)); assert_eq!(variants.pop_first(), Some(InferenceCategory::VoiceConversion)); + assert_eq!(variants.pop_first(), Some(InferenceCategory::RerenderAVideo)); assert_eq!(variants.pop_first(), None); // Generated check diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs index 2a939038d6..a3ea0bcd29 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs @@ -31,6 +31,8 @@ pub enum InferenceModelType { Vits, #[serde(rename = "vall_e_x")] VallEX, + #[serde(rename = "rerender_a_video")] + RerenderAVideo, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -46,7 +48,8 @@ impl InferenceModelType { Self::SoVitsSvc => "so_vits_svc", Self::Tacotron2 => "tacotron2", Self::Vits => "vits", - Self::VallEX => "vall_e_x" + Self::VallEX => "vall_e_x", + Self::RerenderAVideo => "rerender_a_video", } } @@ -58,6 +61,7 @@ impl InferenceModelType { "tacotron2" => Ok(Self::Tacotron2), "vits" => Ok(Self::Vits), "vall_e_x" => Ok(Self::VallEX), + "rerender_a_video" => Ok(Self::RerenderAVideo), _ => Err(format!("invalid value: {:?}", value)), } } @@ -71,7 +75,8 @@ impl InferenceModelType { InferenceModelType::SoVitsSvc, InferenceModelType::Tacotron2, InferenceModelType::Vits, - InferenceModelType::VallEX + InferenceModelType::VallEX, + InferenceModelType::RerenderAVideo ]) } } @@ -89,6 +94,7 @@ mod tests { assert_serialization(InferenceModelType::Tacotron2, "tacotron2"); assert_serialization(InferenceModelType::Vits, "vits"); assert_serialization(InferenceModelType::VallEX, "vall_e_x"); + assert_serialization(InferenceModelType::RerenderAVideo, "rerender_a_video"); } #[test] @@ -99,6 +105,7 @@ mod tests { assert_eq!(InferenceModelType::Tacotron2.to_str(), "tacotron2"); assert_eq!(InferenceModelType::Vits.to_str(), "vits"); assert_eq!(InferenceModelType::VallEX.to_str(), "vall_e_x"); + assert_eq!(InferenceModelType::RerenderAVideo.to_str(), "rerender_a_video"); } #[test] @@ -109,19 +116,21 @@ mod tests { assert_eq!(InferenceModelType::from_str("tacotron2").unwrap(), InferenceModelType::Tacotron2); assert_eq!(InferenceModelType::from_str("vits").unwrap(), InferenceModelType::Vits); assert_eq!(InferenceModelType::from_str("vall_e_x").unwrap(), InferenceModelType::VallEX); + assert_eq!(InferenceModelType::from_str("rerender_a_video").unwrap(), InferenceModelType::RerenderAVideo); } #[test] fn all_variants() { // Static check let mut variants = InferenceModelType::all_variants(); - assert_eq!(variants.len(), 6); + assert_eq!(variants.len(), 7); assert_eq!(variants.pop_first(), Some(InferenceModelType::RvcV2)); assert_eq!(variants.pop_first(), Some(InferenceModelType::SadTalker)); assert_eq!(variants.pop_first(), Some(InferenceModelType::SoVitsSvc)); assert_eq!(variants.pop_first(), Some(InferenceModelType::Tacotron2)); assert_eq!(variants.pop_first(), Some(InferenceModelType::Vits)); assert_eq!(variants.pop_first(), Some(InferenceModelType::VallEX)); + assert_eq!(variants.pop_first(), Some(InferenceModelType::RerenderAVideo)); assert_eq!(variants.pop_first(), None); // Generated check diff --git a/crates/_types/enums/src/by_table/generic_synthetic_ids/id_category.rs b/crates/_types/enums/src/by_table/generic_synthetic_ids/id_category.rs index eccff8cf4d..e5b2e90771 100644 --- a/crates/_types/enums/src/by_table/generic_synthetic_ids/id_category.rs +++ b/crates/_types/enums/src/by_table/generic_synthetic_ids/id_category.rs @@ -21,6 +21,10 @@ pub enum IdCategory { #[serde(rename = "lipsync_animation")] LipsyncAnimationResult, + // Results from rerender a video + #[serde(rename = "rerender")] + RerenderResult, + /// Results from tacotron2 /// Applies for RVC and SVC #[serde(rename = "tts_result")] @@ -62,6 +66,7 @@ impl IdCategory { Self::ZeroShotVoiceDataset => "zs_dataset", Self::ZeroShotVoiceEmbedding => "zs_voice", Self::ZeroShotTtsResult => "zs_tts_result", + Self::RerenderResult => "rerender", Self::ModelWeights => "model_weights", } } @@ -75,6 +80,7 @@ impl IdCategory { "zs_dataset" => Ok(Self::ZeroShotVoiceDataset), "zs_voice" => Ok(Self::ZeroShotVoiceEmbedding), "zs_tts_result" => Ok(Self::ZeroShotTtsResult), + "rerender" => Ok(Self::RerenderResult), "model_weights" => Ok(Self::ModelWeights), _ => Err(format!("invalid value: {:?}", value)), } @@ -86,6 +92,7 @@ impl IdCategory { BTreeSet::from([ Self::MediaFile, Self::LipsyncAnimationResult, + Self::RerenderResult, Self::TtsResult, Self::VoiceConversionResult, Self::ZeroShotTtsResult, @@ -108,6 +115,7 @@ mod tests { fn test_serialization() { assert_serialization(IdCategory::MediaFile, "media_file"); assert_serialization(IdCategory::LipsyncAnimationResult, "lipsync_animation"); + assert_serialization(IdCategory::RerenderResult, "rerender"); assert_serialization(IdCategory::TtsResult, "tts_result"); assert_serialization(IdCategory::VoiceConversionResult, "voice_conversion"); assert_serialization(IdCategory::ZeroShotVoiceDataset, "zs_dataset"); @@ -120,6 +128,7 @@ mod tests { fn to_str() { assert_eq!(IdCategory::MediaFile.to_str(), "media_file"); assert_eq!(IdCategory::LipsyncAnimationResult.to_str(), "lipsync_animation"); + assert_eq!(IdCategory::RerenderResult.to_str(), "rerender"); assert_eq!(IdCategory::TtsResult.to_str(), "tts_result"); assert_eq!(IdCategory::VoiceConversionResult.to_str(), "voice_conversion"); assert_eq!(IdCategory::ZeroShotVoiceDataset.to_str(), "zs_dataset"); @@ -132,6 +141,7 @@ mod tests { fn from_str() { assert_eq!(IdCategory::from_str("media_file").unwrap(), IdCategory::MediaFile); assert_eq!(IdCategory::from_str("lipsync_animation").unwrap(), IdCategory::LipsyncAnimationResult); + assert_eq!(IdCategory::from_str("rerender").unwrap(), IdCategory::RerenderResult); assert_eq!(IdCategory::from_str("tts_result").unwrap(), IdCategory::TtsResult); assert_eq!(IdCategory::from_str("voice_conversion").unwrap(), IdCategory::VoiceConversionResult); assert_eq!(IdCategory::from_str("zs_dataset").unwrap(), IdCategory::ZeroShotVoiceDataset); @@ -144,9 +154,10 @@ mod tests { fn all_variants() { // Static check let mut variants = IdCategory::all_variants(); - assert_eq!(variants.len(), 8); + assert_eq!(variants.len(), 9); assert_eq!(variants.pop_first(), Some(IdCategory::MediaFile)); assert_eq!(variants.pop_first(), Some(IdCategory::LipsyncAnimationResult)); + assert_eq!(variants.pop_first(), Some(IdCategory::RerenderResult)); assert_eq!(variants.pop_first(), Some(IdCategory::TtsResult)); assert_eq!(variants.pop_first(), Some(IdCategory::VoiceConversionResult)); assert_eq!(variants.pop_first(), Some(IdCategory::ZeroShotTtsResult)); 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 8cb9685bc7..fa75c7f791 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 @@ -30,6 +30,9 @@ pub enum MediaFileOriginModelType { #[serde(rename = "vall_e_x")] VallEX, + + #[serde(rename = "rerender")] + Rerender, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -45,7 +48,8 @@ impl MediaFileOriginModelType { Self::SadTalker => "sad_talker", Self::SoVitsSvc => "so_vits_svc", Self::Tacotron2 => "tacotron2", - Self::VallEX => "vall_e_x" + Self::VallEX => "vall_e_x", + Self::Rerender => "rerender", } } @@ -56,6 +60,7 @@ impl MediaFileOriginModelType { "so_vits_svc" => Ok(Self::SoVitsSvc), "tacotron2" => Ok(Self::Tacotron2), "vall_e_x" => Ok(Self::VallEX), + "rerender" => Ok(Self::Rerender), _ => Err(format!("invalid value: {:?}", value)), } } @@ -68,7 +73,8 @@ impl MediaFileOriginModelType { Self::SadTalker, Self::SoVitsSvc, Self::Tacotron2, - Self::VallEX + Self::VallEX, + Self::Rerender, ]) } } @@ -93,6 +99,7 @@ mod tests { assert_eq!(MediaFileOriginModelType::SoVitsSvc.to_str(), "so_vits_svc"); assert_eq!(MediaFileOriginModelType::Tacotron2.to_str(), "tacotron2"); assert_eq!(MediaFileOriginModelType::VallEX.to_str(), "vall_e_x"); + assert_eq!(MediaFileOriginModelType::Rerender.to_str(), "rerender"); } #[test] @@ -102,18 +109,20 @@ mod tests { assert_eq!(MediaFileOriginModelType::from_str("so_vits_svc").unwrap(), MediaFileOriginModelType::SoVitsSvc); assert_eq!(MediaFileOriginModelType::from_str("tacotron2").unwrap(), MediaFileOriginModelType::Tacotron2); assert_eq!(MediaFileOriginModelType::from_str("vall_e_x").unwrap(), MediaFileOriginModelType::VallEX); + assert_eq!(MediaFileOriginModelType::from_str("rerender").unwrap(), MediaFileOriginModelType::Rerender); assert!(MediaFileOriginModelType::from_str("foo").is_err()); } #[test] fn all_variants() { let mut variants = MediaFileOriginModelType::all_variants(); - assert_eq!(variants.len(), 5); + assert_eq!(variants.len(), 6); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::RvcV2)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::SadTalker)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::SoVitsSvc)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::Tacotron2)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::VallEX)); + assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::Rerender)); assert_eq!(variants.pop_first(), None); } } diff --git a/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs b/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs index 7b6f2a029a..103bab0042 100644 --- a/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs +++ b/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs @@ -38,6 +38,10 @@ pub enum MediaFileOriginProductCategory { /// Media files created by (or uploaded for) the Zero Shot voice product. #[serde(rename = "zs_voice")] ZeroShotVoice, + + // Media files for Rerender + #[serde(rename = "rerender")] + Rerender, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -54,6 +58,7 @@ impl MediaFileOriginProductCategory { Self::Unknown => "unknown", Self::VoiceConversion => "voice_conversion", Self::ZeroShotVoice => "zs_voice", + Self::Rerender => "rerender", } } @@ -64,6 +69,7 @@ impl MediaFileOriginProductCategory { "unknown" => Ok(Self::Unknown), "voice_conversion" => Ok(Self::VoiceConversion), "zs_voice" => Ok(Self::ZeroShotVoice), + "rerender" => Ok(Self::Rerender), _ => Err(format!("invalid value: {:?}", value)), } } @@ -74,6 +80,7 @@ impl MediaFileOriginProductCategory { BTreeSet::from([ Self::FaceAnimator, Self::TextToSpeech, + Self::Rerender, Self::Unknown, Self::VoiceConversion, Self::ZeroShotVoice, @@ -96,6 +103,7 @@ mod tests { assert_serialization(MediaFileOriginProductCategory::Unknown, "unknown"); assert_serialization(MediaFileOriginProductCategory::VoiceConversion, "voice_conversion"); assert_serialization(MediaFileOriginProductCategory::ZeroShotVoice, "zs_voice"); + assert_serialization(MediaFileOriginProductCategory::Rerender, "rerender") } #[test] @@ -105,6 +113,7 @@ mod tests { assert_eq!(MediaFileOriginProductCategory::Unknown.to_str(), "unknown"); assert_eq!(MediaFileOriginProductCategory::VoiceConversion.to_str(), "voice_conversion"); assert_eq!(MediaFileOriginProductCategory::ZeroShotVoice.to_str(), "zs_voice"); + assert_eq!(MediaFileOriginProductCategory::Rerender.to_str(), "rerender"); } #[test] @@ -114,17 +123,19 @@ mod tests { assert_eq!(MediaFileOriginProductCategory::from_str("unknown").unwrap(), MediaFileOriginProductCategory::Unknown); assert_eq!(MediaFileOriginProductCategory::from_str("voice_conversion").unwrap(), MediaFileOriginProductCategory::VoiceConversion); assert_eq!(MediaFileOriginProductCategory::from_str("zs_voice").unwrap(), MediaFileOriginProductCategory::ZeroShotVoice); + assert_eq!(MediaFileOriginProductCategory::from_str("rerender").unwrap(), MediaFileOriginProductCategory::Rerender); } #[test] fn all_variants() { let mut variants = MediaFileOriginProductCategory::all_variants(); - assert_eq!(variants.len(), 5); + assert_eq!(variants.len(), 6); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::FaceAnimator)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::TextToSpeech)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Unknown)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::VoiceConversion)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::ZeroShotVoice)); + assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Rerender)); assert_eq!(variants.pop_first(), None); } } diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs index 541de77aa4..f34a4469c3 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs @@ -3,6 +3,7 @@ use errors::AnyhowResult; use crate::payloads::generic_inference_args::lipsync_payload::LipsyncArgs; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; +use crate::payloads::generic_inference_args::videofilter_payload::{RerenderArgs}; /// Used to encode extra state for the `generic_inference_jobs` table. @@ -35,6 +36,10 @@ pub enum InferenceCategoryAbbreviated { #[serde(rename = "vc")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. #[serde(alias = "voice_conversion")] VoiceConversion, + + #[serde(rename = "rr")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + #[serde(alias = "rerender_a_video")] + RerenderAVideo, } #[derive(Serialize, Deserialize, Clone, Copy, Debug)] @@ -84,6 +89,9 @@ pub enum PolymorphicInferenceArgs { #[serde(skip_serializing_if = "Option::is_none")] transpose: Option, }, + + /// Rerender a video. (Short name to save space when serializing.) + Rr(RerenderArgs) } @@ -104,6 +112,7 @@ impl InferenceCategoryAbbreviated { InferenceCategory::LipsyncAnimation => Self::LipsyncAnimation, InferenceCategory::TextToSpeech => Self::TextToSpeech, InferenceCategory::VoiceConversion => Self::VoiceConversion, + InferenceCategory::RerenderAVideo => Self::RerenderAVideo, } } @@ -112,6 +121,7 @@ impl InferenceCategoryAbbreviated { Self::LipsyncAnimation => InferenceCategory::LipsyncAnimation, Self::TextToSpeech => InferenceCategory::TextToSpeech, Self::VoiceConversion => InferenceCategory::VoiceConversion, + Self::RerenderAVideo => InferenceCategory::RerenderAVideo, } } } diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs index 37980d38d1..e68fc5a71c 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs @@ -1,3 +1,4 @@ pub mod generic_inference_args; pub mod lipsync_payload; -pub mod tts_payload; \ No newline at end of file +pub mod tts_payload; +pub mod videofilter_payload; \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/videofilter_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/videofilter_payload.rs new file mode 100644 index 0000000000..12d316933f --- /dev/null +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/videofilter_payload.rs @@ -0,0 +1,36 @@ +/// Video sources can be one of several: +/// - F: media_files (todo) +/// - U: media_uploads (legacy) +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub enum VideofilterVideoSource { + // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + /// Media File Token (media_files table) + /// Serde cannot yet rename enum variants. + F(String), + + // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + /// Media Upload Token (media_uploads table) + /// Serde cannot yet rename enum variants. + U(String), +} + + +impl VideofilterVideoSource { + pub fn media_file_token(token: &str) -> Self { + VideofilterVideoSource::F(token.to_string()) + } + pub fn media_upload_token(token: &str) -> Self { + VideofilterVideoSource::U(token.to_string()) + } +} + +#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct RerenderArgs { + #[serde(rename = "vs")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_video_source: Option, + + #[serde(rename = "sd")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_sd_model: Option, +} diff --git a/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs b/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs index 237a9f7c59..489ddd8391 100644 --- a/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs +++ b/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs @@ -178,6 +178,7 @@ fn raw_record_to_public_result(record: RawGenericInferenceJobStatus) -> GenericI InferenceCategory::LipsyncAnimation => Some("lipsync animation"), InferenceCategory::TextToSpeech => record.maybe_tts_model_title.as_deref(), InferenceCategory::VoiceConversion => record.maybe_voice_conversion_model_title.as_deref(), + InferenceCategory::RerenderAVideo => Some("rerender animation"), }; // NB: A bit of a hack. We store TTS results with a full path. @@ -186,6 +187,7 @@ fn raw_record_to_public_result(record: RawGenericInferenceJobStatus) -> GenericI InferenceCategory::LipsyncAnimation => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), InferenceCategory::TextToSpeech => (false, record.maybe_tts_public_bucket_path.as_deref()), InferenceCategory::VoiceConversion => (true, record.maybe_voice_conversion_public_bucket_hash.as_deref()), + InferenceCategory::RerenderAVideo => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), }; // NB: We've moved voice conversion out of their own table and into media_files diff --git a/crates/lib/mysql_queries/src/queries/media_files/insert_media_file_from_rerender.rs b/crates/lib/mysql_queries/src/queries/media_files/insert_media_file_from_rerender.rs new file mode 100644 index 0000000000..c599a4fad0 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/media_files/insert_media_file_from_rerender.rs @@ -0,0 +1,158 @@ +use anyhow::anyhow; +use sqlx; +use sqlx::MySqlPool; + +use enums::by_table::generic_synthetic_ids::id_category::IdCategory; +use enums::by_table::media_files::media_file_origin_category::MediaFileOriginCategory; +use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType; +use enums::by_table::media_files::media_file_origin_product_category::MediaFileOriginProductCategory; +use enums::by_table::media_files::media_file_type::MediaFileType; +use errors::AnyhowResult; +use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::users::UserToken; + +use crate::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use crate::queries::generic_synthetic_ids::transactional_increment_generic_synthetic_id::transactional_increment_generic_synthetic_id; + +pub struct InsertArgs<'a> { + pub pool: &'a MySqlPool, + pub job: &'a AvailableInferenceJob, + + pub maybe_mime_type: Option<&'a str>, + pub file_size_bytes: u64, + pub sha256_checksum: &'a str, + // TODO: Media duration. + //pub duration_millis: u64, + + pub public_bucket_directory_hash: &'a str, + pub maybe_public_bucket_prefix: Option<&'a str>, + pub maybe_public_bucket_extension: Option<&'a str>, + + pub is_on_prem: bool, + pub worker_hostname: &'a str, + pub worker_cluster: &'a str, +} + +pub async fn insert_media_file_from_rerender( + args: InsertArgs<'_> +) -> AnyhowResult<(MediaFileToken, u64)> +{ + let result_token = MediaFileToken::generate(); + + let mut maybe_creator_file_synthetic_id : Option = None; + let mut maybe_creator_category_synthetic_id : Option = None; + + let mut transaction = args.pool.begin().await?; + + if let Some(creator_user_token) = args.job.maybe_creator_user_token.as_deref() { + let user_token = UserToken::new_from_str(creator_user_token); + + let next_media_file_id = transactional_increment_generic_synthetic_id( + &user_token, + IdCategory::MediaFile, + &mut transaction + ).await?; + + let next_rerender_id = transactional_increment_generic_synthetic_id( + &user_token, + IdCategory::RerenderResult, + &mut transaction + ).await?; + + maybe_creator_file_synthetic_id = Some(next_media_file_id); + maybe_creator_category_synthetic_id = Some(next_rerender_id); + } + + let creator_ip_address = args.job.creator_ip_address.as_str(); + let creator_set_visibility = args.job.creator_set_visibility.clone(); + + const ORIGIN_CATEGORY : MediaFileOriginCategory = MediaFileOriginCategory::Inference; + const ORIGIN_PRODUCT_CATEGORY : MediaFileOriginProductCategory = MediaFileOriginProductCategory::Rerender; + const ORIGIN_MODEL_TYPE : MediaFileOriginModelType = MediaFileOriginModelType::Rerender; + const MEDIA_TYPE : MediaFileType = MediaFileType::Video; + + let record_id = { + let query_result = sqlx::query!( + r#" +INSERT INTO media_files +SET + token = ?, + + origin_category = ?, + origin_product_category = ?, + maybe_origin_model_type = ?, + + media_type = ?, + maybe_mime_type = ?, + file_size_bytes = ?, + + checksum_sha2 = ?, + + public_bucket_directory_hash = ?, + maybe_public_bucket_prefix = ?, + maybe_public_bucket_extension = ?, + + maybe_creator_user_token = ?, + maybe_creator_anonymous_visitor_token = ?, + creator_ip_address = ?, + + creator_set_visibility = ?, + + maybe_creator_file_synthetic_id = ?, + maybe_creator_category_synthetic_id = ?, + + is_generated_on_prem = ?, + generated_by_worker = ?, + generated_by_cluster = ? + + "#, + result_token.as_str(), + + ORIGIN_CATEGORY.to_str(), + ORIGIN_PRODUCT_CATEGORY.to_str(), + ORIGIN_MODEL_TYPE.to_str(), + + MEDIA_TYPE.to_str(), + args.maybe_mime_type, + args.file_size_bytes, + + args.sha256_checksum, + + args.public_bucket_directory_hash, + args.maybe_public_bucket_prefix, + args.maybe_public_bucket_extension, + + args.job.maybe_creator_user_token, + args.job.maybe_creator_anonymous_visitor_token, + args.job.creator_ip_address, + + args.job.creator_set_visibility.to_str(), + + maybe_creator_file_synthetic_id, + maybe_creator_category_synthetic_id, + + args.is_on_prem, + args.worker_hostname, + args.worker_cluster + ) + .execute(&mut *transaction) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + // TODO: handle better + //transaction.rollback().await?; + return Err(anyhow!("Mysql error: {:?}", err)); + } + }; + + record_id + }; + + transaction.commit().await?; + + Ok((result_token, record_id)) +} diff --git a/crates/lib/mysql_queries/src/queries/media_files/mod.rs b/crates/lib/mysql_queries/src/queries/media_files/mod.rs index 8930a3ae7c..ff431d1cb4 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/mod.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/mod.rs @@ -8,3 +8,4 @@ pub mod insert_media_file_from_zero_shot_tts; pub mod list_media_files; pub mod list_media_files_for_user; pub mod update_media_file; +pub mod insert_media_file_from_rerender; diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml index 199d70f522..cdc9aaf01a 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/crates/service/job/inference_job/Cargo.toml @@ -38,9 +38,6 @@ subprocess_common = { path = "../../../lib/subprocess_common" } tokens = { path = "../../../_types/tokens" } tts_common = { path = "../../../lib/tts_common" } -# Workspace hack (faster compile times) -my-workspace-hack = { version = "0.1", path = "../../../my-workspace-hack" } - # External anyhow = "1.0.75" chrono = { version = "0.4.22", features = ["serde"] } @@ -78,3 +75,7 @@ serde_json = "1.0.108" # The initial round of stabilizations for the async/await language feature in Rust 1.39 did not # include support for async fn in traits. async-trait = "0.1.58" +my-workspace-hack = { version = "0.1", path = "../../../my-workspace-hack" } +time = "0.3.30" +mime = "0.3.17" +rand = { version = "0.8.5", features = [] } diff --git a/crates/service/job/inference_job/config/inference-job.development.env b/crates/service/job/inference_job/config/inference-job.development.env index a8427bea58..b911d80f6e 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/crates/service/job/inference_job/config/inference-job.development.env @@ -95,6 +95,15 @@ VITS_INFERENCE_SCRIPT="infer_ts_job.py" VITS_INFERENCE_MAYBE_PYTHON_INTERPRETER="python3" VITS_INFERENCE_MAYBE_DOCKER_IMAGE_SHA="03017896367b" + +# ================================= # +# # +# Rerender # +# # +# ================================= # +RERENDER_INFERENCE_ROOT_DIRECTORY="/home/salt/clone/storyteller-ml/animation/rerender-lora/" +RERENDER_INFERENCE_EXECUTABLE_OR_COMMAND="venv/bin/python rerender.py" + # ================================= # # # # CACHE SETTINGS # diff --git a/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs b/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs index 2d8f04efa9..c7edcdbb7f 100644 --- a/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs +++ b/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs @@ -16,6 +16,7 @@ use crate::job::job_loop::process_single_job_success_case::ProcessSingleJobSucce use crate::job::job_types::lipsync::process_single_lipsync_job::process_single_lipsync_job; use crate::job::job_types::tts::process_single_tts_job::process_single_tts_job; use crate::job::job_types::vc::process_single_vc_job::process_single_vc_job; +use crate::job::job_types::videofilter::process_single_vf_job::process_single_rr_job; use crate::job_dependencies::JobDependencies; pub async fn process_single_job( @@ -182,6 +183,9 @@ async fn do_process_single_job( InferenceCategory::VoiceConversion => { process_single_vc_job(job_dependencies, job).await? } + InferenceCategory::RerenderAVideo => { + process_single_rr_job(job_dependencies, job).await? + } }; let maybe_entity_type = job_success_result.maybe_result_entity diff --git a/crates/service/job/inference_job/src/job/job_types/mod.rs b/crates/service/job/inference_job/src/job/job_types/mod.rs index f21a93a492..a688e281ba 100644 --- a/crates/service/job/inference_job/src/job/job_types/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/mod.rs @@ -1,3 +1,4 @@ pub mod lipsync; pub mod tts; -pub mod vc; \ No newline at end of file +pub mod vc; +pub(crate) mod videofilter; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs b/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs index 158f304a2b..fb444a2a07 100644 --- a/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs @@ -239,6 +239,7 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::La(_) => None, PolymorphicInferenceArgs::Vc { override_f0_method, .. } => *override_f0_method, + PolymorphicInferenceArgs::Rr(_) => None, }) .flatten(); @@ -247,6 +248,7 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::La(_) => None, PolymorphicInferenceArgs::Vc { transpose, .. } => *transpose, + PolymorphicInferenceArgs::Rr(_) => None, }) .flatten(); diff --git a/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs b/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs index e02716e323..65b6fa34a6 100644 --- a/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs @@ -168,6 +168,7 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::La(_) => None, PolymorphicInferenceArgs::Vc { auto_predict_f0, .. } => *auto_predict_f0, + PolymorphicInferenceArgs::Rr(_) => None, }) .flatten() .unwrap_or(false); @@ -177,6 +178,7 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::La(_) => None, PolymorphicInferenceArgs::Vc { transpose, .. } => *transpose, + PolymorphicInferenceArgs::Rr(_) => None, }) .flatten(); diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/mod.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/mod.rs new file mode 100644 index 0000000000..305ceb8a28 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/mod.rs @@ -0,0 +1,2 @@ +pub(crate) mod rerender_a_video; +pub(crate) mod process_single_vf_job; diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/process_single_vf_job.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/process_single_vf_job.rs new file mode 100644 index 0000000000..9d7cc0163c --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/process_single_vf_job.rs @@ -0,0 +1,78 @@ +use anyhow::anyhow; +use log::info; + +use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs::Rr; +use mysql_queries::payloads::generic_inference_args::videofilter_payload::VideofilterVideoSource; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use mysql_queries::queries::media_files::get_media_file::get_media_file; +use tokens::tokens::media_files::MediaFileToken; + +use crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::videofilter::rerender_a_video; +use crate::job::job_types::videofilter::rerender_a_video::process_job::RerenderProcessJobArgs; +use crate::job_dependencies::JobDependencies; + +pub async fn process_single_rr_job(job_dependencies: &JobDependencies, job: &AvailableInferenceJob) -> Result { + // let maybe_inference_args = job.maybe_inference_args.as_ref().ok_or(ProcessSingleJobError::Other(anyhow!("Inference args not found")))?; + // + // let maybe_rerender_args = maybe_inference_args.args.as_ref().map(|args| match args { + // Rr(args) => Some(args), + // _ => None, + // }).flatten(); + // + // let rerender_args = match maybe_rerender_args { + // None => return Err(ProcessSingleJobError::Other(anyhow!("Rerender args not found"))), + // Some(args) => args, + // }; + // + // let videofilter_source = rerender_args.maybe_video_source.as_ref().ok_or(ProcessSingleJobError::Other(anyhow!("Video source not found")))?; + // + // let media_file_token = match videofilter_source { + // VideofilterVideoSource::F(token) => token, + // _ => return Err(ProcessSingleJobError::Other(anyhow!("Video source not found"))), + // }; + // + // let media_file_token = MediaFileToken::new_from_str(media_file_token); + // + // info!("media_file_token: {:?}", media_file_token); + // + // + // let media_file = match media_file_result { + // Ok(Some(result)) => result, + // Ok(None) => { + // return Err(ProcessSingleJobError::Other(anyhow!("Media file not found"))); + // } + // Err(e) => { + // return Err(ProcessSingleJobError::Other(anyhow!("Media file not found"))); + // } + // }; + + // info!("Source media upload file size (bytes): {}", &media_upload.original_file_size_bytes); + // info!("Source media upload duration (millis): {}", &media_upload.original_duration_millis); + // info!("Source media upload duration (seconds): {}", (media_upload.original_duration_millis as f32 / 1000.0)); + + + + // let job_success_result = match vc_model.model_type { + // VoiceConversionModelType::RvcV2 => { + // rvc_v2::process_job::process_job(RvcV2ProcessJobArgs { + // job_dependencies, + // job, + // vc_model: &vc_model, + // media_upload_token: &media_upload_token, + // media_upload: &media_upload, + // }).await? + // } + // }; + + let job_success_result = rerender_a_video::process_job::process_job( + RerenderProcessJobArgs { + job_dependencies, + job, + // media_file + } + ).await?; + + Ok(job_success_result) +} diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/categorize_error.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/categorize_error.rs new file mode 100644 index 0000000000..a1effe90c7 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/categorize_error.rs @@ -0,0 +1,5 @@ +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; + +pub fn categorize_error(stderr_contents: &str) -> Option { + None +} diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs new file mode 100644 index 0000000000..6ec512d1ea --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs @@ -0,0 +1,108 @@ +use std::path::PathBuf; + +use anyhow::anyhow; +use log::{error, warn}; +use sqlx::MySqlPool; +use tempdir::TempDir; +use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; +use buckets::public::media_uploads::bucket_file_path::MediaUploadOriginalFilePath; +use cloud_storage::bucket_client::BucketClient; +use jobs_common::job_progress_reporter::job_progress_reporter::JobProgressReporter; +use mysql_queries::payloads::generic_inference_args::videofilter_payload::VideofilterVideoSource; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use mysql_queries::queries::media_files::get_media_file::get_media_file; +use mysql_queries::queries::media_uploads::get_media_upload_for_inference::get_media_upload_for_inference; +use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::media_uploads::MediaUploadToken; + +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::util::maybe_download_file_from_bucket::{maybe_download_file_from_bucket, MaybeDownloadArgs}; +use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; + +pub struct VideoFile { + pub filesystem_path: PathBuf, +} + +pub async fn download_video_file( + video_source: &VideofilterVideoSource, + public_bucket_client: &BucketClient, + job_progress_reporter: &mut Box, + job: &AvailableInferenceJob, + temp_dir_creator: &ScopedTempDirCreator, + work_temp_dir: &TempDir, + mysql_pool: &MySqlPool, +) -> Result { + + let bucket_object_path; + + match video_source { + VideofilterVideoSource::F(media_file_token) => { + let media_file_token = MediaFileToken::new_from_str(media_file_token); + let media_file_result = get_media_file( + &media_file_token, + false, + mysql_pool + ).await; + match media_file_result { + Ok(Some(result)) => { + let media_file_bucket_path = + MediaFileBucketPath::from_object_hash( + &result.public_bucket_directory_hash, + result.maybe_public_bucket_prefix.as_deref(), + result.maybe_public_bucket_extension.as_deref(), + ); + bucket_object_path = media_file_bucket_path.to_full_object_pathbuf(); + } + Ok(None) => { + return Err(ProcessSingleJobError::from_anyhow_error( + anyhow!("could not find media file: {:?}", media_file_token))) + } + Err(e) => { + error!("could not query media file: {:?}", e); + return Err(ProcessSingleJobError::from_anyhow_error(e)) + } + } + } + VideofilterVideoSource::U(media_upload_token) => { + let media_upload_token = MediaUploadToken::new_from_str(media_upload_token); + let media_upload_result = get_media_upload_for_inference( + &media_upload_token, + mysql_pool + ).await; + + let media_upload_result = match media_upload_result { + Ok(Some(result)) => result, + Ok(None) => { + warn!("could not find media upload: {:?}", media_upload_token); + return Err(ProcessSingleJobError::from_anyhow_error( + anyhow!("could not find media upload: {:?}", media_upload_token))) + } + Err(e) => { + error!("could not query media upload: {:?}", e); + return Err(ProcessSingleJobError::from_anyhow_error(e)) + } + }; + + let media_upload_bucket_path = MediaUploadOriginalFilePath::from_object_hash(&media_upload_result.public_bucket_directory_hash); + bucket_object_path = media_upload_bucket_path.to_full_object_pathbuf(); + } + } + + let downloaded_filesystem_path = work_temp_dir.path().join("video.mp4"); + + maybe_download_file_from_bucket(MaybeDownloadArgs { + name_or_description_of_file: "video file", + final_filesystem_file_path: &downloaded_filesystem_path, + bucket_object_path: &bucket_object_path, + bucket_client: public_bucket_client, + job_progress_reporter, + job_progress_update_description: "downloading", + job_id: job.id.0, + scoped_tempdir_creator: &temp_dir_creator, + maybe_existing_file_minimum_size_required: None, + }).await?; + + Ok(VideoFile { + filesystem_path: downloaded_filesystem_path, + }) +} diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/mod.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/mod.rs new file mode 100644 index 0000000000..28c9dea48a --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/mod.rs @@ -0,0 +1,6 @@ +pub(crate) mod process_job; +mod validate_job; +mod download_video_file; +mod rerender_inference_command; +pub(crate) mod rerender_dependencies; +mod categorize_error; diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs new file mode 100644 index 0000000000..705671b1c1 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs @@ -0,0 +1,409 @@ +use std::fs::read_to_string; +use std::time::Instant; + +use anyhow::{anyhow, Error}; +use log::{error, info, warn}; +use serde::{Deserialize, Serialize}; +use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; + +use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use filesys::check_file_exists::check_file_exists; +use filesys::file_size::file_size; +use filesys::safe_delete_temp_directory::safe_delete_temp_directory; +use filesys::safe_delete_temp_file::safe_delete_temp_file; +use hashing::sha256::sha256_hash_file::sha256_hash_file; +use mimetypes::mimetype_for_file::get_mimetype_for_file; +use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs::Rr; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use mysql_queries::queries::media_files::insert_media_file_from_rerender::{insert_media_file_from_rerender, InsertArgs}; + +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::videofilter::rerender_a_video::download_video_file::download_video_file; +use crate::job::job_types::videofilter::rerender_a_video::rerender_inference_command::InferenceArgs; +use crate::job::job_types::videofilter::rerender_a_video::validate_job::validate_job; +use crate::job_dependencies::JobDependencies; + +/// The maximum that either width or height can be +const MAX_DIMENSION : u32 = 1500; + +const BUCKET_FILE_PREFIX: &str = "fakeyou_"; +const BUCKET_FILE_EXTENSION: &str = ".mp4"; + +pub struct RerenderProcessJobArgs<'a> { + pub job_dependencies: &'a JobDependencies, + pub job: &'a AvailableInferenceJob, + // pub media_file: MediaFile, +} + +pub async fn process_job(args: RerenderProcessJobArgs<'_>) -> Result { + let job = args.job; + let deps = args.job_dependencies; + + let mut job_progress_reporter = args.job_dependencies + .clients + .job_progress_reporter + .new_generic_inference(job.inference_job_token.as_str()) + .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + + let model_dependencies = args + .job_dependencies + .job + .job_specific_dependencies + .maybe_rerender_dependencies + .as_ref() + .ok_or_else(|| ProcessSingleJobError::JobSystemMisconfiguration(Some("missing Rerender dependencies".to_string())))?; + + // ==================== UNPACK + VALIDATE INFERENCE ARGS ==================== // + + let job_args = validate_job(job)?; + + // ==================== TEMP DIR ==================== // + + let work_temp_dir = format!("temp_rerender_inference_{}", job.id.0); + + // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem. + let work_temp_dir = args.job_dependencies + .fs + .scoped_temp_dir_creator_for_work + .new_tempdir(&work_temp_dir) + .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + + + // ==================== QUERY AND DOWNLOAD FILES ==================== // + + let video_path = download_video_file( + &job_args.video_source, + &args.job_dependencies.buckets.public_bucket_client, + &mut job_progress_reporter, + job, + &args.job_dependencies.fs.scoped_temp_dir_creator_for_work, + &work_temp_dir, + &deps.db.mysql_pool + ).await?; + + info!("Downloaded video file: {:?}", video_path.filesystem_path); + + // ==================== TRANSCODE MEDIA (IF NECESSARY) ==================== // + + let mut usable_video_path = video_path.filesystem_path.clone(); + + //TODO: re encode with ffmpeg + + info!("Used video file: {:?}", usable_video_path); + + // ==================== SETUP FOR INFERENCE ==================== // + + info!("Ready for Rerender A Video inference..."); + + job_progress_reporter.log_status("running inference") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + let output_video_fs_path = work_temp_dir.path().join("blend.mp4"); + + info!("Running Rerender A Video inference..."); + + info!("Expected output video filename: {:?}", &output_video_fs_path); + + // TODO: Limit output length for non-premium (???) + + let maybe_args = job.maybe_inference_args + .as_ref() + .map(|args| args.args.as_ref()) + .flatten(); + + let poly_args = match maybe_args { + None => return Err(ProcessSingleJobError::Other(anyhow!("Rerender args not found"))), + Some(args) => args, + }; + + let rr_args = match poly_args { + Rr(args) => args, + _ => return Err(ProcessSingleJobError::Other(anyhow!("Rerender args not found"))), + }; + + let sd_model = rr_args.maybe_sd_model.as_ref().ok_or(ProcessSingleJobError::Other(anyhow!("SD model not found")))?; + + // ============== CREATE AND WRITE CONFIG TO DISK ============== // + + // create seed if not exists + // let seed = rr_args.seed.unwrap_or_else(|| rand::random::()); + let seed = rand::random::(); + + let config = RerenderConfig { + input: Some(usable_video_path.to_str().unwrap().to_string()), + output: Some(output_video_fs_path.to_str().unwrap().to_string()), + work_dir: Some(work_temp_dir.path().to_str().unwrap().to_string()), + sd_model: Some(sd_model.to_string()), + prompt: Some("1girl, anime style, red hair".to_string()), + a_prompt: Some("1girl, 3d, realistic, read hair".to_string()), + n_prompt: Some("Bad quality, low quality, photo, artifacts, watermark, signature, faceless".to_string()), + seed: Some(seed as i64), + ..RerenderConfig::new() + }; + + config.validate().map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + + let config_path = work_temp_dir.path().join("config.json"); + let config_json = serde_json::to_string(&config).unwrap(); + std::fs::write(&config_path, config_json).unwrap(); + + // ==================== RUN INFERENCE SCRIPT ==================== // + let stderr_output_file = work_temp_dir.path().join("stderr.txt"); + let inference_start_time = Instant::now(); + + let command_exit_status = model_dependencies + .inference_command + .execute_inference(InferenceArgs { + config_file: &config_path, + stderr_output_file: &stderr_output_file, + }); + + let inference_duration = Instant::now().duration_since(inference_start_time); + + info!("Inference took duration to complete: {:?}", &inference_duration); + + if !command_exit_status.is_success() { + error!("Inference 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); + + // match categorize_error(&contents) { + // Some(ProcessSingleJobError::FaceDetectionFailure) => { + // warn!("Face not detected in source image"); + // error = ProcessSingleJobError::FaceDetectionFailure; + // } + // _ => {} + // } + } + + safe_delete_temp_file(&video_path.filesystem_path); + safe_delete_temp_file(&config_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); + } + + // ==================== CHECK NON-WATERMARKED RESULT ==================== // + + info!("Checking that output file exists: {:?} ...", output_video_fs_path); + + check_file_exists(&output_video_fs_path).map_err(|e| ProcessSingleJobError::Other(e))?; + + // ==================== OPTIONAL WATERMARK ==================== // + + let finished_file = output_video_fs_path.clone(); + + // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== // + + info!("Checking that output watermark file exists: {:?} ...", finished_file); + check_file_exists(&finished_file).map_err(|e| ProcessSingleJobError::Other(e))?; + + info!("Interrogating result file size ..."); + + let file_size_bytes = file_size(&finished_file) + .map_err(|err| ProcessSingleJobError::Other(err))?; + + info!("Interrogating result mimetype ..."); + + let mimetype = get_mimetype_for_file(&finished_file) + .map_err(|err| ProcessSingleJobError::from_io_error(err))? + .map(|mime| mime.to_string()) + .ok_or(ProcessSingleJobError::Other(anyhow!("Mimetype could not be determined")))?; + + info!("Calculating sha256..."); + + let file_checksum = sha256_hash_file(&finished_file) + .map_err(|err| { + ProcessSingleJobError::Other(anyhow!("Error hashing file: {:?}", err)) + })?; + + // ==================== UPLOAD VIDEO TO BUCKET ==================== // + + job_progress_reporter.log_status("uploading result") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + 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(); + + info!("Audio destination bucket path: {:?}", &result_bucket_object_pathbuf); + + info!("Uploading media ..."); + + args.job_dependencies.buckets.public_bucket_client.upload_filename_with_content_type( + &result_bucket_object_pathbuf, + &finished_file, + &mimetype) // TODO: We should check the mimetype to make sure bad payloads can't get uploaded + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // ==================== DELETE TEMP FILES ==================== // + + safe_delete_temp_file(&video_path.filesystem_path); + safe_delete_temp_file(&config_path); + safe_delete_temp_file(&output_video_fs_path); + safe_delete_temp_file(&stderr_output_file); + safe_delete_temp_directory(&work_temp_dir); + + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + safe_delete_temp_directory(&work_temp_dir); + + // ==================== SAVE RECORDS ==================== // + + info!("Saving Rerender result (media_files table record) ..."); + + let (media_file_token, id) = insert_media_file_from_rerender(InsertArgs { + pool: &args.job_dependencies.db.mysql_pool, + job: &job, + maybe_mime_type: Some(&mimetype), + 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!("Rerender Done."); + + job_progress_reporter.log_status("done") + .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, + }) +} + +#[derive(Serialize, Deserialize, Debug)] +pub struct RerenderConfig { + #[serde(rename = "input")] + pub input: Option, + + #[serde(rename = "output")] + pub output: Option, + + #[serde(rename = "work_dir")] + pub work_dir: Option, + + #[serde(rename = "key_subdir")] + pub key_subdir: Option, + + #[serde(rename = "sd_model")] + pub sd_model: Option, + + #[serde(rename = "lora_path")] + pub lora_path: Option, + + #[serde(rename = "interval")] + pub interval: Option, + + #[serde(rename = "prompt")] + pub prompt: Option, + + #[serde(rename = "a_prompt")] + pub a_prompt: Option, + + #[serde(rename = "n_prompt")] + pub n_prompt: Option, + + #[serde(rename = "x0_strength")] + pub x0_strength: Option, + + #[serde(rename = "control_type")] + pub control_type: Option, + + #[serde(rename = "canny_low")] + pub canny_low: Option, + + #[serde(rename = "canny_high")] + pub canny_high: Option, + + #[serde(rename = "control_strength")] + pub control_strength: Option, + + #[serde(rename = "seed")] + pub seed: Option, + + #[serde(rename = "warp_period")] + pub warp_period: Option>, + + #[serde(rename = "ada_period")] + pub ada_period: Option>, +} + +impl RerenderConfig { + pub fn new() -> Self { + Self { + input: None, + output: None, + work_dir: None, + key_subdir: Some("keys".to_string()), + sd_model: None, + lora_path: None, + interval: Some(8), + prompt: None, + a_prompt: None, + n_prompt: None, + x0_strength: Some(0.95), + control_type: Some("canny".to_string()), + canny_low: Some(50), + canny_high: Some(100), + control_strength: Some(0.7), + seed: None, + warp_period: Some(vec![0., 0.1]), + ada_period: Some(vec![0.8, 1.]), + } + } + + pub fn validate(&self) -> Result<(), Error> { + let fields = &[ + (self.input.is_some(), "input"), + (self.output.is_some(), "output"), + (self.work_dir.is_some(), "work_dir"), + (self.key_subdir.is_some(), "key_subdir"), + (self.sd_model.is_some(), "sd_model"), + (self.interval.is_some(), "interval"), + (self.prompt.is_some(), "prompt"), + (self.a_prompt.is_some(), "a_prompt"), + (self.n_prompt.is_some(), "n_prompt"), + (self.x0_strength.is_some(), "x0_strength"), + (self.control_type.is_some(), "control_type"), + (self.canny_low.is_some(), "canny_low"), + (self.canny_high.is_some(), "canny_high"), + (self.control_strength.is_some(), "control_strength"), + (self.seed.is_some(), "seed"), + (self.warp_period.is_some(), "warp_period"), + (self.ada_period.is_some(), "ada_period"), + ]; + + let missing_fields: Vec<&str> = fields + .iter() + .filter_map(|(is_some, field_name)| if *is_some { None } else { Some(*field_name) }) + .collect(); + + if missing_fields.is_empty() { + Ok(()) + } else { + Err(anyhow!("Missing fields: {}", missing_fields.join(", "))) + } + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/rerender_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/rerender_dependencies.rs new file mode 100644 index 0000000000..c50fe62ece --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/rerender_dependencies.rs @@ -0,0 +1,15 @@ +use errors::AnyhowResult; + +use crate::job::job_types::videofilter::rerender_a_video::rerender_inference_command::RerenderInferenceCommand; + +pub struct RerenderDependencies { + pub inference_command: RerenderInferenceCommand, +} + +impl RerenderDependencies { + pub fn setup() -> AnyhowResult { + Ok(Self { + inference_command: RerenderInferenceCommand::from_env()?, + }) + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/rerender_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/rerender_inference_command.rs new file mode 100644 index 0000000000..a3ba8953ae --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/rerender_inference_command.rs @@ -0,0 +1,229 @@ +use std::collections::HashSet; +use std::env; +use std::ffi::OsString; +use std::fs::File; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use anyhow::anyhow; +use log::info; +use once_cell::sync::Lazy; +use subprocess::{Popen, PopenConfig, Redirection}; + +use errors::AnyhowResult; +use filesys::path_to_string::path_to_string; +use subprocess_common::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions}; + +use crate::job::job_loop::command_exit_status::CommandExitStatus; + +// These environment vars are not copied over to the subprocess +// TODO/FIXME(bt, 2023-05-28): This is horrific security! +static IGNORED_ENVIRONMENT_VARS : Lazy> = Lazy::new(|| { + let env_var_names= [ + "MYSQL_URL", + "ACCESS_KEY", + "SECRET_KEY", + "NEWRELIC_API_KEY", + ]; + + env_var_names.iter() + .map(|value| value.to_string()) + .collect::>() +}); + +#[derive(Clone)] +pub struct RerenderInferenceCommand { + /// Where the code lives + rerender_root_code_directory: PathBuf, + + /// A single executable script or a much larger bash command. + executable_or_command: ExecutableOrCommand, + + /// Config file to use + config_path: Option, + + /// 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), +} + +pub struct InferenceArgs<'s, P: AsRef> { + pub config_file: P, + pub stderr_output_file: &'s Path, +} + +impl RerenderInferenceCommand { + pub fn from_env() -> AnyhowResult { + let rerender_root_code_directory = easyenv::get_env_pathbuf_required( + "RERENDER_INFERENCE_ROOT_DIRECTORY")?; + + let config_path = easyenv::get_env_pathbuf_optional( + "RERENDER_INFERENCE_CONFIG_PATH"); + + let executable_or_command = match easyenv::get_env_string_optional( + "RERENDER_INFERENCE_EXECUTABLE_OR_COMMAND") { + None => { + return Err(anyhow!("RERENDER_INFERENCE_EXECUTABLE_OR_COMMAND is required")); + } + Some(executable_or_command) => { + if executable_or_command.contains(" ") { + ExecutableOrCommand::Command(executable_or_command) + } else { + ExecutableOrCommand::Executable(PathBuf::from(executable_or_command)) + } + } + }; + + let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( + "RERENDER_INFERENCE_MAYBE_VENV_COMMAND"); + + let maybe_execution_timeout = + easyenv::get_env_duration_seconds_optional("RERENDER_TIMEOUT_SECONDS"); + + let maybe_docker_options = easyenv::get_env_string_optional( + "RERENDER_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), + } + }); + + Ok(Self { + rerender_root_code_directory, + executable_or_command, + config_path, + 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) }, + } + } + + fn do_execute_inference>( + &self, + args: InferenceArgs

, + ) -> AnyhowResult { + + let mut command = String::new(); + command.push_str(&format!("cd {}", path_to_string(&self.rerender_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(" --cfg "); + command.push_str(&path_to_string(args.config_file.as_ref())); + + // ===== 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_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.clone())?; + + match exit_status { + None => { + // NB: If the program didn't successfully terminate, kill it. + info!("Subprocess didn't end after timeout: {:?}; terminating...", &timeout); + let _r = p.terminate()?; + Ok(CommandExitStatus::Timeout) + } + Some(exit_status) => { + info!("Subprocess timed wait exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + } + } + } + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/validate_job.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/validate_job.rs new file mode 100644 index 0000000000..dd9bc75c2f --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/validate_job.rs @@ -0,0 +1,85 @@ +use anyhow::anyhow; + +use mysql_queries::payloads::generic_inference_args::generic_inference_args::{InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; +use mysql_queries::payloads::generic_inference_args::videofilter_payload::{VideofilterVideoSource}; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; + +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; + +pub struct JobArgs<'a> { + pub video_source: &'a VideofilterVideoSource, + pub sd_model: &'a str, +} + +pub fn 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::RerenderAVideo) => {}, // 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::Rr(inference_args) => inference_args, + _ => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!"))); + } + }; + + let video_source = match &inference_args.maybe_video_source { + Some(args) => args, + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no video source!"))); + } + }; + + let sd_model = match &inference_args.maybe_sd_model { + Some(args) => args, + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no sd model!"))); + } + }; + + + // let preprocess = match inference_args.maybe_preprocess { + // None => Some("full".to_string()), + // Some(Preprocess::F) => Some("full".to_string()), + // Some(Preprocess::EF) => Some("extfull".to_string()), + // Some(Preprocess::C) => Some("crop".to_string()), + // Some(Preprocess::EC) => Some("extcrop".to_string()), + // }; + // + // let enhancer = match inference_args.maybe_face_enhancer { + // None => None, + // Some(FaceEnhancer::G) => Some("gfpgan".to_string()), + // Some(FaceEnhancer::R) => Some("RestoreFormer".to_string()), + // }; + // + // let mut width = inference_args.maybe_resize_width; + // let mut height = inference_args.maybe_resize_height; + + Ok(JobArgs { + video_source, + sd_model, + }) +} diff --git a/crates/service/job/inference_job/src/job_specific_dependencies.rs b/crates/service/job/inference_job/src/job_specific_dependencies.rs index ddd313a46a..a863b3872c 100644 --- a/crates/service/job/inference_job/src/job_specific_dependencies.rs +++ b/crates/service/job/inference_job/src/job_specific_dependencies.rs @@ -7,6 +7,7 @@ use crate::job::job_types::tts::vall_e_x::vall_e_x_dependencies::VallExDependenc use crate::job::job_types::tts::vits::vits_dependencies::VitsDependencies; use crate::job::job_types::vc::rvc_v2::rvc_v2_dependencies::RvcV2Dependencies; use crate::job::job_types::vc::so_vits_svc::svc_dependencies::SvcDependencies; +use crate::job::job_types::videofilter::rerender_a_video::rerender_dependencies::RerenderDependencies; use crate::util::scoped_execution::ScopedExecution; pub struct JobSpecificDependencies { @@ -16,6 +17,7 @@ pub struct JobSpecificDependencies { pub maybe_tacotron2_dependencies: Option, pub maybe_vall_e_x_dependencies: Option, pub maybe_vits_dependencies: Option, + pub maybe_rerender_dependencies: Option, } impl JobSpecificDependencies { @@ -27,6 +29,7 @@ impl JobSpecificDependencies { let mut maybe_tacotron2_dependencies = None; let mut maybe_vall_e_x_dependencies = None; let mut maybe_vits_dependencies = None; + let mut maybe_rerender_dependencies = None; if scoped_execution.can_run_job(InferenceModelType::RvcV2) { print_with_space("Setting RVCv2 dependencies..."); @@ -58,6 +61,11 @@ impl JobSpecificDependencies { maybe_vits_dependencies = Some(VitsDependencies::setup()?); } + if scoped_execution.can_run_job(InferenceModelType::RerenderAVideo) { + print_with_space("Setting Rerender dependencies..."); + maybe_rerender_dependencies = Some(RerenderDependencies::setup()?); + } + Ok(JobSpecificDependencies { maybe_rvc_v2_dependencies, maybe_sad_talker_dependencies, @@ -65,6 +73,7 @@ impl JobSpecificDependencies { maybe_tacotron2_dependencies, maybe_vall_e_x_dependencies, maybe_vits_dependencies, + maybe_rerender_dependencies, }) } } diff --git a/crates/service/web/storyteller_web/Cargo.toml b/crates/service/web/storyteller_web/Cargo.toml index c8d81cef9a..bc0e4b61bc 100644 --- a/crates/service/web/storyteller_web/Cargo.toml +++ b/crates/service/web/storyteller_web/Cargo.toml @@ -165,6 +165,7 @@ async-trait = "0.1.58" utoipa = { version = "4", features = ["chrono","actix_extras"] } utoipa-swagger-ui = { version = "5", features = ["actix-web"] } +http = "0.2.11" [build-dependencies] # NB: Only using diesel for migrations support diff --git a/crates/service/web/storyteller_web/src/configs/plans/plan.rs b/crates/service/web/storyteller_web/src/configs/plans/plan.rs index 67dd178cc1..80d100a9fe 100644 --- a/crates/service/web/storyteller_web/src/configs/plans/plan.rs +++ b/crates/service/web/storyteller_web/src/configs/plans/plan.rs @@ -49,6 +49,9 @@ pub struct Plan { can_remove_visual_watermarks: bool, + // ========== Features for Videofilters ========== + rerender_requires_frontend_keepalive: bool, + // ========== Features for TTS ========== tts_base_priority_level: u8, @@ -91,6 +94,7 @@ impl Plan { is_development_plan: builder.is_development_plan, is_synthetic_plan: builder.is_synthetic_plan, lipsync_requires_frontend_keepalive: builder.lipsync_requires_frontend_keepalive, + rerender_requires_frontend_keepalive: builder.rerender_requires_frontend_keepalive, can_remove_visual_watermarks: builder.can_remove_visual_watermarks, tts_base_priority_level: builder.tts_base_priority_level, tts_max_duration: builder.tts_max_duration, @@ -146,6 +150,10 @@ impl Plan { self.lipsync_requires_frontend_keepalive } + pub fn rerender_requires_frontent_keepalive(&self) -> bool { + self.rerender_requires_frontend_keepalive + } + pub fn can_remove_visual_watermarks(&self) -> bool { self.can_remove_visual_watermarks } @@ -192,6 +200,9 @@ pub struct PlanBuilder { can_remove_visual_watermarks: bool, + // ========== Features for Videofilters ========== + rerender_requires_frontend_keepalive: bool, + // ========== Features for TTS ========== tts_base_priority_level: u8, @@ -239,6 +250,9 @@ impl PlanBuilder { lipsync_requires_frontend_keepalive: true, can_remove_visual_watermarks: false, + // Videofilters + rerender_requires_frontend_keepalive: true, + // TTS tts_base_priority_level : TTS_DEFAULT_PRIORITY_LEVEL, tts_max_duration: Duration::seconds(TTS_DEFAULT_DURATION_SECONDS), @@ -304,6 +318,11 @@ impl PlanBuilder { self } + pub fn rerender_requires_frontend_keepalive(mut self, value: bool) -> Self { + self.rerender_requires_frontend_keepalive = value; + self + } + pub fn can_remove_visual_watermarks(mut self, value: bool) -> Self { self.can_remove_visual_watermarks = value; self diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/animation/enqueue_rerender_animation.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/animation/enqueue_rerender_animation.rs new file mode 100644 index 0000000000..d90caa3b8b --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/animation/enqueue_rerender_animation.rs @@ -0,0 +1,229 @@ +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +use std::sync::Arc; + +use actix_web::{HttpRequest, HttpResponse, web}; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use log::{info, warn}; + +use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; +use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; +use enums::common::visibility::Visibility; +use http_server_common::request::get_request_header_optional::get_request_header_optional; +use http_server_common::request::get_request_ip::get_request_ip; +use mysql_queries::payloads::generic_inference_args::generic_inference_args::{GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; +use mysql_queries::payloads::generic_inference_args::videofilter_payload::{RerenderArgs, VideofilterVideoSource}; +use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{insert_generic_inference_job, InsertGenericInferenceArgs}; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::users::UserToken; + +use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; + +/// Debug requests can get routed to special "debug-only" workers, which can +/// be used to trial new code, run debugging, etc. +const DEBUG_HEADER_NAME : &str = "enable-debug-mode"; + +/// The routing tag header can send workloads to particular k8s hosts. +/// This is useful for catching the live logs or intercepting the job. +const ROUTING_TAG_HEADER_NAME : &str = "routing-tag"; + +#[derive(Deserialize)] +pub struct EnqueueRerenderAnimationRequest { + uuid_idempotency_token: String, + + video_source: Option, + sd_model: String, + creator_set_visibility: Option, +} + +/// Treated as an enum. Only one of these may be set. + +#[derive(Serialize)] +pub struct EnqueueRerenderAnimationSuccessResponse { + pub success: bool, + pub inference_job_token: InferenceJobToken, +} + +#[derive(Debug)] +pub enum EnqueueRerenderAnimationError { + BadInput(String), + NotAuthorized, + ServerError, + RateLimited, +} + +impl ResponseError for EnqueueRerenderAnimationError { + fn status_code(&self) -> StatusCode { + match *self { + EnqueueRerenderAnimationError::BadInput(_) => StatusCode::BAD_REQUEST, + EnqueueRerenderAnimationError::NotAuthorized => StatusCode::UNAUTHORIZED, + EnqueueRerenderAnimationError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + EnqueueRerenderAnimationError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + EnqueueRerenderAnimationError::BadInput(reason) => reason.to_string(), + EnqueueRerenderAnimationError::NotAuthorized => "unauthorized".to_string(), + EnqueueRerenderAnimationError::ServerError => "server error".to_string(), + EnqueueRerenderAnimationError::RateLimited => "rate limited".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl std::fmt::Display for EnqueueRerenderAnimationError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn enqueue_rerender_animation_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + let mut maybe_user_token : Option = None; + + let mut mysql_connection = server_state.mysql_pool + .acquire() + .await + .map_err(|err| { + warn!("MySql pool error: {:?}", err); + EnqueueRerenderAnimationError::ServerError + })?; + + let maybe_avt_token = server_state.avt_cookie_manager.get_avt_token_from_request(&http_request); + + // ==================== USER SESSION ==================== // + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session_extended_from_connection(&http_request, &mut mysql_connection) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EnqueueRerenderAnimationError::ServerError + })?; + + if let Some(user_session) = maybe_user_session.as_ref() { + maybe_user_token = Some(UserToken::new_from_str(&user_session.user_token)); + } + + // TODO: Plan should handle "first anonymous use" and "investor" cases. + let plan = get_correct_plan_for_session( + server_state.server_environment, + maybe_user_session.as_ref()); + + // TODO: Separate priority for animation. + let priority_level = plan.web_vc_base_priority_level(); + + // ==================== DEBUG MODE + ROUTING TAG ==================== // + + let is_debug_request = + get_request_header_optional(&http_request, DEBUG_HEADER_NAME) + .is_some(); + + let maybe_routing_tag= + get_request_header_optional(&http_request, ROUTING_TAG_HEADER_NAME) + .map(|routing_tag| routing_tag.trim().to_string()); + + // ==================== RATE LIMIT ==================== // + + let rate_limiter = match maybe_user_session { + None => &server_state.redis_rate_limiters.logged_out, + Some(ref user) => { + if user.role.is_banned { + return Err(EnqueueRerenderAnimationError::NotAuthorized); + } + &server_state.redis_rate_limiters.logged_in + }, + }; + + if let Err(_err) = rate_limiter.rate_limit_request(&http_request) { + return Err(EnqueueRerenderAnimationError::RateLimited); + } + + // ==================== LOOK UP MODEL INFO ==================== // + + // TODO(bt): CHECK DATABASE FOR TOKENS! + + let video_source: VideofilterVideoSource = { + if let Some(ref token) = request.video_source { + VideofilterVideoSource::media_file_token(token.as_str()) + } else { + return Err(EnqueueRerenderAnimationError::BadInput("video source not fully specified".to_string())); + } + }; + + let ip_address = get_request_ip(&http_request); + + let maybe_user_preferred_visibility : Option = maybe_user_session + .as_ref() + .map(|user_session: &users_component::utils::user_session_extended::UserSessionExtended| user_session.preferences.preferred_tts_result_visibility); // TODO: New setting for web-vc + + let set_visibility = request.creator_set_visibility + .or(maybe_user_preferred_visibility) + .unwrap_or(Visibility::Public); + + let inference_args = RerenderArgs { + maybe_video_source: Some(video_source), + maybe_sd_model: Some(request.sd_model.clone()), + }; + + info!("Creating rerender animation job record..."); + + let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { + uuid_idempotency_token: &request.uuid_idempotency_token, + inference_category: InferenceCategory::RerenderAVideo, + maybe_model_type: Some(InferenceModelType::RerenderAVideo), // NB: Model is static during inference + maybe_model_token: None, // NB: Model is static during inference + maybe_input_source_token: None, // TODO: Introduce a second foreign key ? + maybe_input_source_token_type: None, // TODO: Introduce a second foreign key ? + maybe_raw_inference_text: None, // No text + maybe_max_duration_seconds: None, + maybe_inference_args: Some(GenericInferenceArgs { + inference_category: Some(InferenceCategoryAbbreviated::RerenderAVideo), + args: Some(PolymorphicInferenceArgs::Rr(inference_args)), + }), + maybe_creator_user_token: maybe_user_token.as_ref(), + maybe_avt_token: maybe_avt_token.as_ref(), + creator_ip_address: &ip_address, + creator_set_visibility: set_visibility, + priority_level, + requires_keepalive: plan.rerender_requires_frontent_keepalive(), + is_debug_request, + maybe_routing_tag: maybe_routing_tag.as_deref(), + mysql_pool: &server_state.mysql_pool, + }).await; + + let job_token = match query_result { + Ok((job_token, _id)) => job_token, + Err(err) => { + warn!("New generic inference job creation DB error: {:?}", err); + return Err(EnqueueRerenderAnimationError::ServerError); + } + }; + + let response: EnqueueRerenderAnimationSuccessResponse = EnqueueRerenderAnimationSuccessResponse { + success: true, + inference_job_token: job_token, + }; + + let body = serde_json::to_string(&response) + .map_err(|_e| EnqueueRerenderAnimationError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/animation/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/animation/mod.rs index dd4f619b5e..1a000a595b 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/animation/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/animation/mod.rs @@ -1 +1,2 @@ +pub mod enqueue_rerender_animation; pub mod enqueue_face_animation; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs index ea3e6f8f15..5befd6c91e 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs @@ -256,6 +256,14 @@ fn record_to_payload( .get_full_object_path_str() .to_string() } + InferenceCategory::RerenderAVideo => { + MediaFileBucketPath::from_object_hash( + &result_details.public_bucket_location_or_hash, + result_details.maybe_media_file_public_bucket_prefix.as_deref(), + result_details.maybe_media_file_public_bucket_extension.as_deref()) + .get_full_object_path_str() + .to_string() + } InferenceCategory::TextToSpeech => { match result_details.entity_type.as_str() { "media_file" => { diff --git a/crates/service/web/storyteller_web/src/routes.rs b/crates/service/web/storyteller_web/src/routes.rs index 9dcdce3e35..26ffb503bc 100644 --- a/crates/service/web/storyteller_web/src/routes.rs +++ b/crates/service/web/storyteller_web/src/routes.rs @@ -12,6 +12,7 @@ use users_component::endpoints::edit_profile_handler::edit_profile_handler; use users_component::endpoints::get_profile_handler::get_profile_handler; use crate::http_server::endpoints::animation::enqueue_face_animation::enqueue_face_animation_handler; +use crate::http_server::endpoints::animation::enqueue_rerender_animation::enqueue_rerender_animation_handler; use crate::http_server::endpoints::api_tokens::create_api_token::create_api_token_handler; use crate::http_server::endpoints::api_tokens::delete_api_token::delete_api_token_handler; use crate::http_server::endpoints::api_tokens::edit_api_token::edit_api_token_handler; @@ -216,6 +217,7 @@ pub fn add_routes (app: App, server_environment: ServerEnvironment) -> let mut app = RouteBuilder::from_app(app) .add_post("/v1/animation/face_animation/create", enqueue_face_animation_handler) + .add_post("/v1/animation/rerender/create", enqueue_rerender_animation_handler) .into_app(); // ==================== "Generic" Inference ====================