daniel543x/artcraft
1
0 Fork 0
mirror of https://github.com/storytold/artcraft.git synced 2026-11-05 17:58:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #205 from storytold/feature/rerender

Browse Source 
rerender a video
This commit is contained in:
saltacc authored and GitHub committed 2023-12-13 17:34:52 -08:00
2 parents 0ed04839d3 b30c352d79
commit 7c96d3b75b
34 files changed
+1501 -14
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

Cargo.lock
Generated
+4
View File
@@ -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",
crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs
+11 -1
View File
@@ -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
crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs
+12 -3
View File
@@ -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
crates/_types/enums/src/by_table/generic_synthetic_ids/id_category.rs
+12 -1
View File
@@ -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));
crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs
+12 -3
View File
@@ -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);
}
}
crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs
+12 -1
View File
@@ -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);
}
}
crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs
+10
View File
@@ -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<i32>,
},
/// 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,
}
}
}
crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs
+2 -1
View File
@@ -1,3 +1,4 @@
pub mod generic_inference_args;
pub mod lipsync_payload;
pub mod tts_payload;
pub mod tts_payload;
pub mod videofilter_payload;
crates/lib/mysql_queries/src/payloads/generic_inference_args/videofilter_payload.rs
+36
View File
@@ -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<VideofilterVideoSource>,
#[serde(rename = "sd")]
#[serde(skip_serializing_if = "Option::is_none")]
pub maybe_sd_model: Option<String>,
}
crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs
+2
View File
@@ -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
crates/lib/mysql_queries/src/queries/media_files/insert_media_file_from_rerender.rs
+158
View File
@@ -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<u64> = None;
let mut maybe_creator_category_synthetic_id : Option<u64> = 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))
}
crates/lib/mysql_queries/src/queries/media_files/mod.rs
+1
View File
@@ -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;
crates/service/job/inference_job/Cargo.toml
+4 -3
View File
@@ -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 = [] }
crates/service/job/inference_job/config/inference-job.development.env
+9
View File
@@ -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 #
crates/service/job/inference_job/src/job/job_loop/process_single_job.rs
+4
View File
@@ -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
crates/service/job/inference_job/src/job/job_types/mod.rs
+2 -1
View File
@@ -1,3 +1,4 @@
pub mod lipsync;
pub mod tts;
pub mod vc;
pub mod vc;
pub(crate) mod videofilter;
crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs
+2
View File
@@ -239,6 +239,7 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result<JobSuccessResu
PolymorphicInferenceArgs::Tts { .. } => 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<JobSuccessResu
PolymorphicInferenceArgs::Tts { .. } => None,
PolymorphicInferenceArgs::La(_) => None,
PolymorphicInferenceArgs::Vc { transpose, .. } => *transpose,
PolymorphicInferenceArgs::Rr(_) => None,
})
.flatten();
crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs
+2
View File
@@ -168,6 +168,7 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result<JobSuccess
PolymorphicInferenceArgs::Tts { .. } => 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<JobSuccess
PolymorphicInferenceArgs::Tts { .. } => None,
PolymorphicInferenceArgs::La(_) => None,
PolymorphicInferenceArgs::Vc { transpose, .. } => *transpose,
PolymorphicInferenceArgs::Rr(_) => None,
})
.flatten();
crates/service/job/inference_job/src/job/job_types/videofilter/mod.rs
+2
View File
@@ -0,0 +1,2 @@
pub(crate) mod rerender_a_video;
pub(crate) mod process_single_vf_job;
crates/service/job/inference_job/src/job/job_types/videofilter/process_single_vf_job.rs
+78
View File
@@ -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<JobSuccessResult, ProcessSingleJobError> {
// 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)
}
crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/categorize_error.rs
+5
View File
@@ -0,0 +1,5 @@
use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
pub fn categorize_error(stderr_contents: &str) -> Option<ProcessSingleJobError> {
None
}
crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs
+108
View File
@@ -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<dyn JobProgressReporter>,
job: &AvailableInferenceJob,
temp_dir_creator: &ScopedTempDirCreator,
work_temp_dir: &TempDir,
mysql_pool: &MySqlPool,
) -> Result<VideoFile, ProcessSingleJobError> {
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,
})
}
crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/mod.rs
+6
View File
@@ -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;
crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs
+409
View File
@@ -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<JobSuccessResult, ProcessSingleJobError> {
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::<i64>());
let seed = rand::random::<i32>();
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<String>,
#[serde(rename = "output")]
pub output: Option<String>,
#[serde(rename = "work_dir")]
pub work_dir: Option<String>,
#[serde(rename = "key_subdir")]
pub key_subdir: Option<String>,
#[serde(rename = "sd_model")]
pub sd_model: Option<String>,
#[serde(rename = "lora_path")]
pub lora_path: Option<String>,
#[serde(rename = "interval")]
pub interval: Option<i32>,
#[serde(rename = "prompt")]
pub prompt: Option<String>,
#[serde(rename = "a_prompt")]
pub a_prompt: Option<String>,
#[serde(rename = "n_prompt")]
pub n_prompt: Option<String>,
#[serde(rename = "x0_strength")]
pub x0_strength: Option<f32>,
#[serde(rename = "control_type")]
pub control_type: Option<String>,
#[serde(rename = "canny_low")]
pub canny_low: Option<i32>,
#[serde(rename = "canny_high")]
pub canny_high: Option<i32>,
#[serde(rename = "control_strength")]
pub control_strength: Option<f32>,
#[serde(rename = "seed")]
pub seed: Option<i64>,
#[serde(rename = "warp_period")]
pub warp_period: Option<Vec<f32>>,
#[serde(rename = "ada_period")]
pub ada_period: Option<Vec<f32>>,
}
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(", ")))
}
}
}
crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/rerender_dependencies.rs
+15
View File
@@ -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<Self> {
Ok(Self {
inference_command: RerenderInferenceCommand::from_env()?,
})
}
}
crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/rerender_inference_command.rs
+229
View File
@@ -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<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>>()
});
#[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<PathBuf>,
/// eg. `source python/bin/activate`
maybe_virtual_env_activation_command: Option<String>,
/// 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>,
}
#[derive(Clone)]
pub enum ExecutableOrCommand {
/// Eg. `inference.py`
Executable(PathBuf),
/// Eg. `python3 inference.py`
Command(String),
}
pub struct InferenceArgs<'s, P: AsRef<Path>> {
pub config_file: P,
pub stderr_output_file: &'s Path,
}
impl RerenderInferenceCommand {
pub fn from_env() -> AnyhowResult<Self> {
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<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.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))
}
}
}
}
}
}
crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/validate_job.rs
+85
View File
@@ -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<JobArgs, ProcessSingleJobError> {
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,
})
}
crates/service/job/inference_job/src/job_specific_dependencies.rs
+9
View File
@@ -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<Tacotron2Dependencies>,
pub maybe_vall_e_x_dependencies: Option<VallExDependencies>,
pub maybe_vits_dependencies: Option<VitsDependencies>,
pub maybe_rerender_dependencies: Option<RerenderDependencies>,
}
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,
})
}
}
crates/service/web/storyteller_web/Cargo.toml
+1
View File
@@ -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
crates/service/web/storyteller_web/src/configs/plans/plan.rs
+19
View File
@@ -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
crates/service/web/storyteller_web/src/http_server/endpoints/animation/enqueue_rerender_animation.rs
+229
View File
@@ -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<MediaFileToken>,
sd_model: String,
creator_set_visibility: Option<Visibility>,
}
/// 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<EnqueueRerenderAnimationRequest>,
server_state: web::Data<Arc<ServerState>>) -> Result<HttpResponse, EnqueueRerenderAnimationError>
{
let mut maybe_user_token : Option<UserToken> = 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<Visibility> = 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))
}
crates/service/web/storyteller_web/src/http_server/endpoints/animation/mod.rs
+1
View File
@@ -1 +1,2 @@
pub mod enqueue_rerender_animation;
pub mod enqueue_face_animation;
crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs
+8
View File
@@ -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" => {
crates/service/web/storyteller_web/src/routes.rs
+2
View File
@@ -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<T, B> (app: App<T>, 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 ====================