diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui_inference_command.rs index 80f0b0b501..6b3e2115c1 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui_inference_command.rs @@ -581,7 +581,7 @@ impl ComfyInferenceCommand { break; } - tokio::time::sleep(Duration::from_secs(2)).await; + tokio::time::sleep(Duration::from_secs(5)).await; } return Ok(status.unwrap()); diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/process_video_style_transfer_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/process_video_style_transfer_job.rs index 1502f024c9..29b91e0845 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/process_video_style_transfer_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/process_video_style_transfer_job.rs @@ -1,5 +1,5 @@ use std::collections::HashMap; -use std::fs::{File, read_to_string}; +use std::fs::{read_to_string, File}; use std::path::{Path, PathBuf}; use std::process::{Command, Stdio}; use std::thread; @@ -23,10 +23,10 @@ use enums::by_table::prompts::prompt_type::PromptType; use enums::common::job_status_plus::JobStatusPlus; use errors::AnyhowResult; use filesys::check_file_exists::check_file_exists; -use filesys::file_deletion::safe_delete_files_and_directories::safe_delete_files_and_directories; -use filesys::file_deletion::safe_delete_possible_files_and_directories::safe_delete_possible_files_and_directories; use filesys::file_deletion::safe_delete_directory::safe_delete_directory; use filesys::file_deletion::safe_delete_file::safe_delete_file; +use filesys::file_deletion::safe_delete_files_and_directories::safe_delete_files_and_directories; +use filesys::file_deletion::safe_delete_possible_files_and_directories::safe_delete_possible_files_and_directories; use filesys::file_deletion::safe_recursively_delete_files::safe_recursively_delete_files; use filesys::file_exists::file_exists; use filesys::file_read_bytes::file_read_bytes; @@ -60,11 +60,11 @@ use crate::job::job_types::workflow::video_style_transfer::steps::post_process_a use crate::job::job_types::workflow::video_style_transfer::steps::post_process_restore_audio::{post_process_restore_audio, PostProcessRestoreVideoArgs}; use crate::job::job_types::workflow::video_style_transfer::steps::preprocess_save_audio::{preprocess_save_audio, ProcessSaveAudioArgs}; use crate::job::job_types::workflow::video_style_transfer::steps::preprocess_trim_and_resample_videos::{preprocess_trim_and_resample_videos, ProcessTrimAndResampleVideoArgs}; -use crate::job::job_types::workflow::video_style_transfer::steps::validate_and_save_results::{SaveResultsArgs, validate_and_save_results}; +use crate::job::job_types::workflow::video_style_transfer::steps::validate_and_save_results::{validate_and_save_results, SaveResultsArgs}; use crate::job::job_types::workflow::video_style_transfer::util::comfy_dirs::ComfyDirs; use crate::job::job_types::workflow::video_style_transfer::util::process_preview_updates::PreviewProcessor; use crate::job::job_types::workflow::video_style_transfer::util::video_pathing::{PrimaryInputVideoAndPaths, SecondaryInputVideoAndPaths, VideoPathing}; -use crate::job::job_types::workflow::video_style_transfer::util::write_workflow_prompt::{WorkflowPromptArgs, write_workflow_prompt}; +use crate::job::job_types::workflow::video_style_transfer::util::write_workflow_prompt::{write_workflow_prompt, WorkflowPromptArgs}; use crate::state::job_dependencies::JobDependencies; pub async fn process_video_style_transfer_job(deps: &JobDependencies, job: &AvailableInferenceJob) -> Result { diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/util/process_preview_updates.rs b/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/util/process_preview_updates.rs index d4e091ffd3..efebacb1cf 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/util/process_preview_updates.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/util/process_preview_updates.rs @@ -125,6 +125,29 @@ impl StageDirectoryState { self.frame_states.insert(frame_number, PreviewFrameUploadResult::AttemptFailed); } + fn pending_frames(&self) -> Vec { + (0..self.expected_frame_count).filter(|frame_number| { + !self.frame_states.contains_key(frame_number) + }).collect() + } + + fn successfully_uploaded_frames(&self) -> Vec { + (0..self.expected_frame_count).filter(|frame_number| { + match self.frame_states.get(frame_number) { + Some(s) => *s == PreviewFrameUploadResult::UploadComplete, + None => false, + } + }).collect() + } + + fn is_frame_upload_pending(&self, frame_number: &u32) -> bool { + let state = self.frame_states.get(&frame_number); + match state { + Some(s) => *s != PreviewFrameUploadResult::UploadComplete, + None => true, + } + } + fn default_frame_extension(&self) -> &str { match self.stage { PreviewStage::FirstPass => "jpg", @@ -177,20 +200,13 @@ impl StageDirectoryState { } fn all_object_paths (&self) -> Vec { - (0..self.expected_frame_count).map(|frame_number| { + let frame_numbers = self.successfully_uploaded_frames(); + frame_numbers.into_iter().map(|frame_number| { self.expected_media_file_object_path(frame_number) }).collect() } - async fn upload_frame_from_disk(&mut self, bucket_client: &BucketClient, disk_path: PathBuf) -> PreviewFrameUpdate { - let frame_number = match disk_path.file_stem().and_then(|s| s.to_str()).and_then(|s| s.parse::().ok()) { - Some(n) => n, - None => { - log::error!("Invalid frame number in path: {:?}", disk_path); - return PreviewFrameUpdate { stage: self.stage, frame_number: 0, state: PreviewFrameUploadResult::InvalidInputName, disk_path, object_path: None }; - } - }; - + async fn upload_frame_from_disk(&mut self, bucket_client: &BucketClient, frame_number: u32, disk_path: PathBuf) -> PreviewFrameUpdate { let local_file_extension = match disk_path.extension().and_then(|s| s.to_str()) { Some(e) => { if ALLOWED_TYPES_FRAMES.contains(e) { @@ -286,6 +302,23 @@ impl PreviewProcessor { } } + fn infer_frame_number(&self, disk_path: &PathBuf) -> Option { + let relative_path = disk_path.strip_prefix(&self.base_directory); + match relative_path { + Ok(p) => { + let parts: Vec<&str> = p.to_str().unwrap().split('/').collect(); + if parts.len() != 2 { + return None; + } + let frame_number = parts[1].split('.').next().unwrap(); + match frame_number.parse::() { + Ok(n) => Some(n), + Err(_) => None, + } + } + Err(_) => None, + } + } async fn persist_redis_update(&mut self) { let base_url = easyenv::get_env_string_or_default( @@ -329,23 +362,37 @@ impl PreviewProcessor { return; } }; + let frame_number = match self.infer_frame_number(&disk_path) { + Some(n) => n.clone(), + None => { + log::error!("Invalid frame number in path: {:?}", disk_path); + return; + } + }; let stage_state = self.stages.get_mut(&stage).unwrap(); if stage_state.state == PreviewStageState::UploadComplete { log::debug!("Skipped re-uploading frame: {:?}", disk_path); return; } + if !stage_state.is_frame_upload_pending(&frame_number) { + log::debug!("Frame already uploaded: {:?}", disk_path); + return; + } if !stage_state.is_disk_path_valid(&disk_path) { log::warn!("Invalid disk path: {:?}", disk_path); return; } - let result = stage_state.upload_frame_from_disk(bucket_client, disk_path).await; + + let result = stage_state.upload_frame_from_disk(bucket_client, frame_number.clone(), disk_path).await; if result.state == PreviewFrameUploadResult::UploadComplete { // (KS): We can add frames one at a time but for now Comfy is generating all frames at once // so this check helps + log::info!("Frame upload complete: {:?} for stage {:?}", result.frame_number, stage); if stage_state.all_frames_uploaded() { stage_state.state = PreviewStageState::UploadComplete; } } + log::info!("Pending frames for stage {:?}: {:?}", stage, stage_state.pending_frames()); self.persist_redis_update().await; } }