diff --git a/crates/service/job/inference_job/src/job/job_types/studio_gen2/process_single_studio_gen2_job.rs b/crates/service/job/inference_job/src/job/job_types/studio_gen2/process_single_studio_gen2_job.rs index cdc6adbac5..eb4a8a46b2 100644 --- a/crates/service/job/inference_job/src/job/job_types/studio_gen2/process_single_studio_gen2_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/studio_gen2/process_single_studio_gen2_job.rs @@ -3,7 +3,7 @@ use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::job::job_types::studio_gen2::animate_x::animate_x_dependencies::AnimateXDependencies; use crate::job::job_types::studio_gen2::animate_x::animate_x_inference_command::AnimateXInferenceArgs; use crate::job::job_types::studio_gen2::animate_x::animate_x_process_frames_command::{AnimateXProcessFramesCommand, ProcessFramesArgs}; -use crate::job::job_types::studio_gen2::download_file_for_studio::{download_file_for_studio, DownloadFileForStudioArgs}; +use crate::job::job_types::studio_gen2::download_file_for_studio::{download_file_for_studio, DownloadDetails, DownloadFileForStudioArgs}; use crate::job::job_types::studio_gen2::resize_image_for_studio::{resize_image_for_studio, ImageResizeType}; use crate::job::job_types::studio_gen2::stable_animator::stable_animator_command::InferenceArgs; use crate::job::job_types::studio_gen2::stable_animator::stable_animator_dependencies::StableAnimatorDependencies; @@ -29,6 +29,7 @@ use std::io::BufReader; use std::ops::Deref; use std::path::Path; use std::time::{Duration, Instant}; +use enums::by_table::media_files::media_file_type::MediaFileType; use subprocess_common::command_runner::command_args::CommandArgs; use subprocess_common::command_runner::command_runner_args::{RunAsSubprocessArgs, StreamRedirection}; use videos::ffprobe_get_dimensions::ffprobe_get_dimensions; @@ -118,6 +119,8 @@ pub async fn process_single_studio_gen2_job( info!("Downloaded image to {:?}", &unaltered_image_file.file_path); } } + + assert_file_is_image(&unaltered_image_file)?; // ==================== DOWNLOAD VIDEO ==================== // @@ -136,6 +139,8 @@ pub async fn process_single_studio_gen2_job( info!("Downloaded video to {:?}", &unaltered_video_file.file_path); } } + + assert_file_is_video(&unaltered_video_file)?; //if let Ok(Some(dimensions)) = ffprobe_get_dimensions(&videos.primary_video.original_download_path) { // info!("Download video dimensions: {}x{}", dimensions.width, dimensions.height); @@ -183,7 +188,7 @@ pub async fn process_single_studio_gen2_job( } Some(ffmpeg_args) => { info!("Resampling video file. Will use the following args to ffmpeg: {:?}", &ffmpeg_args.to_command_string()); - + let command_exit_status = gen2_deps.ffmpeg .run_with_subprocess(RunAsSubprocessArgs { args: Box::new(ffmpeg_args), @@ -204,13 +209,13 @@ pub async fn process_single_studio_gen2_job( // ========================= RESIZE IMAGE ======================== // let skip_image_resize = studio_args.skip_image_resize.unwrap_or(false); - + let mut inference_input_image_path = unaltered_image_file.file_path.clone(); let mut maybe_dimensions = None; - + if !skip_image_resize { info!("Resizing input image..."); - + let resized_image_path = work_paths.output_dir.path().join("resized_image.png"); let image_resize_type = match studio_job_type { @@ -401,7 +406,8 @@ pub async fn process_single_studio_gen2_job( Ok(token) => token, Err(err) => { error!("Error validating and saving results: {:?}", err); - + + print_work_dirs(&work_paths); maybe_debug_sleep(studio_args).await; // NB: Forcing generic type to `&Path` with turbofish @@ -416,6 +422,7 @@ pub async fn process_single_studio_gen2_job( // ==================== CLEANUP/ DELETE TEMP FILES ==================== // + print_work_dirs(&work_paths); maybe_debug_sleep(studio_args).await; info!("Cleaning up temporary files..."); @@ -449,7 +456,7 @@ pub async fn process_single_studio_gen2_job( async fn maybe_debug_sleep(args: &StudioGen2Payload) { if let Some(sleep_millis) = args.after_job_debug_sleep_millis { info!("Debug sleeping for millis: {sleep_millis}"); - // We wait for the specified timeout. If a SIGTERM is received during this time, + // We wait for the specified timeout. If a SIGTERM is received during this time, // we'll panic and exit the program. // https://tokio.rs/tokio/tutorial/select // https://tokio.rs/tokio/topics/shutdown @@ -461,3 +468,30 @@ async fn maybe_debug_sleep(args: &StudioGen2Payload) { } } } + +fn print_work_dirs(work_dirs: &StudioGen2Dirs) { + info!("Temp input directory: {:?}", work_dirs.input_dir.path()); + info!("Temp output directory: {:?}", work_dirs.input_dir.path()); +} + +fn assert_file_is_image(file: &DownloadDetails) -> Result<(), ProcessSingleJobError> { + match file.media_file.media_type { + MediaFileType::Image | + MediaFileType::Jpg | + MediaFileType::Png | + MediaFileType::Gif => Ok(()), + _ => Err(ProcessSingleJobError::Other(anyhow!( + "Wrong file type for image: {:?} and token: {:?}", + file.media_file.media_type, file.media_file.token))) + } +} + +fn assert_file_is_video(file: &DownloadDetails) -> Result<(), ProcessSingleJobError> { + match file.media_file.media_type { + MediaFileType::Video | + MediaFileType::Mp4 => Ok(()), + _ => Err(ProcessSingleJobError::Other(anyhow!( + "Wrong file type for video: {:?} and token: {:?}", + file.media_file.media_type, file.media_file.token))) + } +}