From 57e2f70707a2ace6f5ead0407b0ca3c66bc1fb46 Mon Sep 17 00:00:00 2001 From: salt Date: Tue, 12 Dec 2023 23:07:27 -0800 Subject: [PATCH] WORKING! --- Cargo.lock | 1 + .../generic_synthetic_ids/id_category.rs | 6 + .../media_file_origin_model_type.rs | 7 +- .../media_file_origin_product_category.rs | 6 + .../insert_media_file_from_rerender.rs | 158 ++++++++++++ .../src/queries/media_files/mod.rs | 1 + crates/service/job/inference_job/Cargo.toml | 1 + .../config/inference-job.development.env | 2 +- .../videofilter/process_single_vf_job.rs | 76 +++--- .../rerender_a_video/download_video_file.rs | 32 ++- .../rerender_a_video/process_job.rs | 229 ++++++++++++------ 11 files changed, 415 insertions(+), 104 deletions(-) create mode 100644 crates/lib/mysql_queries/src/queries/media_files/insert_media_file_from_rerender.rs diff --git a/Cargo.lock b/Cargo.lock index 215225d64d..4193767e2c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3070,6 +3070,7 @@ dependencies = [ "newrelic-telemetry", "once_cell", "r2d2_redis", + "rand 0.8.5", "redis_common", "serde", "serde_derive", 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..01bddb94dd 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)), } 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..a79930d8c6 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)), } } 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..5965f9b3a9 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)), } } 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 65713e40b0..fb451aa8a1 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/crates/service/job/inference_job/Cargo.toml @@ -81,3 +81,4 @@ 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 359c8edf3e..15ab7f6497 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/crates/service/job/inference_job/config/inference-job.development.env @@ -101,7 +101,7 @@ VITS_INFERENCE_MAYBE_DOCKER_IMAGE_SHA="03017896367b" # # # ================================= # RERENDER_INFERENCE_ROOT_DIRECTORY="/home/salt/clone/storyteller-ml/animation/rerender-lora/" -RERENDER_INFERENCE_EXECUTABLE_OR_COMMAND="conda activate rerender && python rerender.py" +RERENDER_INFERENCE_EXECUTABLE_OR_COMMAND="venv/bin/python rerender.py" # ================================= # # # 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 index 6a35427d0e..9d7cc0163c 100644 --- 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 @@ -1,43 +1,56 @@ use anyhow::anyhow; -use log::{error, info}; +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_uploads::get_media_upload_for_inference::get_media_upload_for_inference; -use tokens::tokens::media_uploads::MediaUploadToken; +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_media_upload_token = job.maybe_input_source_token - .as_deref() - .map(|token| MediaUploadToken::new_from_str(token)); + // 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"))); + // } + // }; - let media_upload_token = match maybe_media_upload_token { - None => return Err(ProcessSingleJobError::Other(anyhow!("no associated media upload for vc job: {:?}", job.inference_job_token))), - Some(token) => token, - }; - - let maybe_media_upload_result = - get_media_upload_for_inference(&media_upload_token, &job_dependencies.db.mysql_pool).await; - - let media_upload = match maybe_media_upload_result { - Ok(Some(media_upload)) => media_upload, - Ok(None) => { - error!("no media upload record found for token: {:?}", media_upload_token); - return Err(ProcessSingleJobError::Other(anyhow!("no media upload record found for token: {:?}", media_upload_token))); - }, - Err(err) => { - error!("error fetching media upload record from db: {:?}", err); - return Err(ProcessSingleJobError::Other(err)); - }, - }; - - 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)); + // 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)); @@ -54,9 +67,10 @@ pub async fn process_single_rr_job(job_dependencies: &JobDependencies, job: &Ava // }; let job_success_result = rerender_a_video::process_job::process_job( - rerender_a_video::process_job::RerenderProcessJobArgs { + RerenderProcessJobArgs { job_dependencies, - job + job, + // media_file } ).await?; diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs index 344592bd08..e05d695841 100644 --- a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs @@ -4,13 +4,16 @@ use anyhow::anyhow; use log::{error, warn}; use sqlx::MySqlPool; use tempdir::TempDir; +use buckets::public::media_files::original_file::MediaFileBucketPath; use buckets::public::media_uploads::original_file::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; @@ -34,9 +37,32 @@ pub async fn download_video_file( let bucket_object_path; match video_source { - VideofilterVideoSource::F(_media_file_token) => { - // TODO(bt, 2023-09-08): Implement - return Err(ProcessSingleJobError::NotYetImplemented) + 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); diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs index affec1aee3..0df0d80325 100644 --- a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs @@ -1,8 +1,9 @@ use std::fs::read_to_string; use std::time::Instant; -use anyhow::anyhow; +use anyhow::{anyhow, Error}; use log::{error, info, warn}; +use serde::{Deserialize, Serialize}; use buckets::public::media_files::original_file::MediaFileBucketPath; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; @@ -12,19 +13,16 @@ 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_face_animation::{insert_media_file_from_face_animation, InsertArgs}; -use tokens::tokens::users::UserToken; +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::lipsync::sad_talker::categorize_error::categorize_error; -// use crate::job::job_types::lipsync::sad_talker::categorize_error::categorize_error; 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; -use crate::util::common_commands::ffmpeg_logo_watermark_command; /// The maximum that either width or height can be const MAX_DIMENSION : u32 = 1500; @@ -35,6 +33,7 @@ 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 { @@ -88,43 +87,8 @@ pub async fn process_job(args: RerenderProcessJobArgs<'_>) -> Result { - // error!("Image resize failure: {:?}", err); - // safe_delete_temp_file(&audio_path.filesystem_path); - // safe_delete_temp_file(&image_path.filesystem_path); - // safe_delete_temp_directory(&work_temp_dir); - // return Err(ProcessSingleJobError::Other(anyhow!("image resize failure: {:?}", err))); - // } - // Ok(resized_image_path) => { - // usable_image_path = resized_image_path; - // } - // } - // } + + //TODO: re encode with ffmpeg info!("Used video file: {:?}", usable_video_path); @@ -148,15 +112,43 @@ pub async fn process_job(args: RerenderProcessJobArgs<'_>) -> Result 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"); - - // TODO + let config_json = serde_json::to_string(&config).unwrap(); + std::fs::write(&config_path, config_json).unwrap(); // ==================== RUN INFERENCE SCRIPT ==================== // - - let workdir = work_temp_dir.path().to_path_buf(); let stderr_output_file = work_temp_dir.path().join("stderr.txt"); let inference_start_time = Instant::now(); @@ -174,7 +166,7 @@ pub async fn process_job(args: RerenderProcessJobArgs<'_>) -> Result) -> Result) -> Result) -> Result) -> Result) -> Result, + + #[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(", "))) + } + } +}