From a80ef65d7023cdff40416743f20ccc63b73c0f1f Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 28 May 2024 20:53:54 -0400 Subject: [PATCH] organize / cleanup comfy code --- Cargo.lock | 1 + crates/lib/files/filesys/Cargo.toml | 1 + crates/lib/files/filesys/src/lib.rs | 5 +- .../src/safe_recursively_delete_files.rs | 21 ++ .../comfy_ui/video_style_transfer/mod.rs | 1 + .../video_style_transfer/process_job.rs | 253 ++------------ .../validate_and_save_results.rs | 314 ++++++++++++++++++ 7 files changed, 369 insertions(+), 227 deletions(-) create mode 100644 crates/lib/files/filesys/src/safe_recursively_delete_files.rs create mode 100644 crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/validate_and_save_results.rs diff --git a/Cargo.lock b/Cargo.lock index aaf397f67b..f1c21866b8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2514,6 +2514,7 @@ dependencies = [ "errors", "log", "my-workspace-hack", + "walkdir", ] [[package]] diff --git a/crates/lib/files/filesys/Cargo.toml b/crates/lib/files/filesys/Cargo.toml index b46f886e53..574bf6df78 100644 --- a/crates/lib/files/filesys/Cargo.toml +++ b/crates/lib/files/filesys/Cargo.toml @@ -26,6 +26,7 @@ my-workspace-hack = { version = "0.1", path = "../../../my-workspace-hack" } # External # KEEP THIS TO A MINIMUM OR NIL. log = "0.4.20" +walkdir = "2.4.0" [dev-dependencies] diff --git a/crates/lib/files/filesys/src/lib.rs b/crates/lib/files/filesys/src/lib.rs index 86362f52b4..433ecfc3e0 100644 --- a/crates/lib/files/filesys/src/lib.rs +++ b/crates/lib/files/filesys/src/lib.rs @@ -19,6 +19,7 @@ #![allow(dead_code)] #![allow(non_snake_case)] +pub mod check_directory_exists; pub mod check_file_exists; pub mod create_dir_all_if_missing; pub mod directory_exists; @@ -29,7 +30,7 @@ pub mod filename_concat; pub mod is_filesystem_full_error; pub mod path_to_string; pub mod rename_across_devices; -pub mod check_directory_exists; -pub mod safe_delete_temp_file; pub mod safe_delete_possible_temp_file; pub mod safe_delete_temp_directory; +pub mod safe_delete_temp_file; +pub mod safe_recursively_delete_files; diff --git a/crates/lib/files/filesys/src/safe_recursively_delete_files.rs b/crates/lib/files/filesys/src/safe_recursively_delete_files.rs new file mode 100644 index 0000000000..656ff2723c --- /dev/null +++ b/crates/lib/files/filesys/src/safe_recursively_delete_files.rs @@ -0,0 +1,21 @@ +use std::path::Path; + +use walkdir::WalkDir; + +use crate::safe_delete_temp_file::safe_delete_temp_file; + +pub fn safe_recursively_delete_files(path: &Path) { + if !path.exists() || !path.is_dir() { + return; + } + + let walk = WalkDir::new(path); + + for entry in walk { + if let Ok(entry) = entry { + if entry.path().is_file() { + safe_delete_temp_file(entry.path()); + } + } + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/mod.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/mod.rs index f0f134b55a..5dad8269cc 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/mod.rs @@ -3,3 +3,4 @@ pub mod download_input_video; pub mod job_outputs; pub mod process_job; pub mod validate_job; +pub mod validate_and_save_results; diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/process_job.rs index 115ed12205..4d8bef296e 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/video_style_transfer/process_job.rs @@ -23,6 +23,7 @@ use filesys::file_size::file_size; use filesys::path_to_string::path_to_string; use filesys::safe_delete_temp_directory::safe_delete_temp_directory; use filesys::safe_delete_temp_file::safe_delete_temp_file; +use filesys::safe_recursively_delete_files::safe_recursively_delete_files; 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::Cu; @@ -45,34 +46,12 @@ use crate::job::job_types::workflow::comfy_ui::comfy_process_job_args::ComfyProc use crate::job::job_types::workflow::comfy_ui::video_style_transfer::comfy_ui_inference_command::{InferenceArgs, InferenceDetails}; use crate::job::job_types::workflow::comfy_ui::video_style_transfer::download_input_video::{download_input_video, DownloadInputVideoArgs}; use crate::job::job_types::workflow::comfy_ui::video_style_transfer::job_outputs::JobOutputs; +use crate::job::job_types::workflow::comfy_ui::video_style_transfer::validate_and_save_results::{SaveResultsArgs, validate_and_save_results}; use crate::job::job_types::workflow::comfy_ui::video_style_transfer::validate_job::validate_job; use crate::job_dependencies::JobDependencies; use crate::util::common_commands::ffmpeg_audio_replace_args::FfmpegAudioReplaceArgs; use crate::util::common_commands::ffmpeg_logo_watermark_command::WatermarkArgs; -fn get_file_extension(mimetype: &str) -> Result<&'static str> { - let ext = match mimetype { - "video/mp4" => "mp4", - "image/png" => "png", - "image/jpeg" => "jpg", - "image/gif" => "gif", - _ => return Err(anyhow!("Mimetype not supported: {}", mimetype)), - }; - Ok(ext) -} - - -fn recursively_delete_files_in(path: &Path) -> std::io::Result<()> { - for entry in WalkDir::new(path) { - let entry = entry?; - if entry.path().is_file() { - safe_delete_temp_file(entry.path()); - } - } - Ok(()) -} - - pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result { let job = args.job; let deps = args.job_dependencies; @@ -534,7 +513,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result) -> Result MediaFileType::Video, - // "image/png" => MediaFileType::Image, - // "image/jpeg" => MediaFileType::Image, - // _ => return Err(ProcessSingleJobError::Other(anyhow!("Mimetype not supported: {}", mimetype))), - //}; - - info!("Calculating sha256..."); - - let file_checksum = sha256_hash_file(final_video) - .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(PREFIX), - Some(&ext)); - - let result_bucket_object_pathbuf = result_bucket_location.to_full_object_pathbuf(); - - info!("Output file 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, - &final_video, - &mimetype) // TODO: We should check the mimetype to make sure bad payloads can't get uploaded - .await - .map_err(|e| ProcessSingleJobError::Other(e))?; - - // generate thumbnail using thumbnail service - let thumbnail_types = match mimetype.as_str() { - "video/mp4" => vec!["image/gif", "image/jpeg"], - _ => vec![], - }; - - info!("Generating thumbnail tasks..."); - - for output_type in thumbnail_types { - let thumbnail_task_result = ThumbnailTaskBuilder::new() - .with_bucket(&*args.job_dependencies.buckets.public_bucket_client.bucket_name()) - .with_path(&*path_to_string(result_bucket_object_pathbuf.clone())) - .with_source_mimetype(mimetype.as_str()) - .with_output_mimetype(output_type) - .with_output_suffix("thumb") - .with_output_extension(get_file_extension(output_type).unwrap()) - .with_event_id(&job.id.0.to_string()) - .send() - .await; - - match thumbnail_task_result { - Ok(thumbnail_task) => { - debug!("Thumbnail task created: {:?}", thumbnail_task); - }, - Err(e) => { - error!("Failed to create thumbnail task: {:?}", e); - } - } - } - - // Also upload the non-watermarked copy - info!("Uploading non-watermarked copy..."); - - let result = args.job_dependencies.buckets.public_bucket_client.upload_filename_with_content_type( - &result_bucket_object_pathbuf.with_extension("no_watermark.mp4"), - &videos.get_non_watermarked_video_to_upload(), - &mimetype) // TODO: We should check the mimetype to make sure bad payloads can't get uploaded - .await; - - if let Err(err) = result { - error!("Failed to upload non-watermarked copy: {:?}", err); - } + // ==================== (OPTIONAL) DEBUG SLEEP ==================== // if let Some(sleep_millis) = comfy_args.sleep_millis { info!("Sleeping for millis: {sleep_millis}"); @@ -804,115 +697,25 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result anyhow::Result<&'static str> { + let ext = match mimetype { + "video/mp4" => "mp4", + "image/png" => "png", + "image/jpeg" => "jpg", + "image/gif" => "gif", + _ => return Err(anyhow!("Mimetype not supported: {}", mimetype)), + }; + + Ok(ext) +} + +pub struct SaveResultsArgs<'a> { + pub job: &'a AvailableInferenceJob, + pub deps: &'a JobDependencies, + pub job_args: &'a JobArgs<'a>, + pub comfy_args: &'a WorkflowArgs, + pub videos: &'a JobOutputs, + pub job_progress_reporter: &'a mut Box, + pub work_temp_dir: &'a TempDir, + pub workflow_path: &'a str, + pub root_comfy_path: &'a Path, + pub maybe_lora_path: Option<&'a Path>, + pub download_video: VideoDownloadDetails, + pub should_insert_prompt_record: bool, + pub maybe_style_name: Option, + pub inference_duration: Duration, + pub maybe_positive_prompt: Option<&'a str>, + pub maybe_negative_prompt: Option<&'a str>, +} + +pub async fn validate_and_save_results(args: SaveResultsArgs<'_>) -> Result { + + // ==================== GET METADATA ==================== // + + info!("Interrogating result file size ..."); + + let final_video = args.videos.get_final_video_to_upload(); + + let file_size_bytes = file_size(final_video) + .map_err(|err| ProcessSingleJobError::Other(err))?; + + info!("Interrogating result mimetype ..."); + + let mimetype = get_mimetype_for_file(final_video) + .map_err(|err| ProcessSingleJobError::from_io_error(err))? + .map(|mime| mime.to_string()) + .ok_or(ProcessSingleJobError::Other(anyhow!("Mimetype could not be determined")))?; + + // create ext from mimetype + let ext = get_file_extension(mimetype.as_str()) + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // Extension is really a "suffix" and should have the leading period to act as an extension. + let ext = if ext.starts_with(".") { + ext.to_string() + } else { + format!(".{ext}") + }; + + const PREFIX: &str = "storyteller_"; + + //// determine media type from mime type + //let media_type = match mimetype.as_str() { + // "video/mp4" => MediaFileType::Video, + // "image/png" => MediaFileType::Image, + // "image/jpeg" => MediaFileType::Image, + // _ => return Err(ProcessSingleJobError::Other(anyhow!("Mimetype not supported: {}", mimetype))), + //}; + + info!("Calculating sha256..."); + + let file_checksum = sha256_hash_file(final_video) + .map_err(|err| { + ProcessSingleJobError::Other(anyhow!("Error hashing file: {:?}", err)) + })?; + + // ==================== UPLOAD VIDEO TO BUCKET ==================== // + + args.job_progress_reporter.log_status("uploading result") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + let result_bucket_location = MediaFileBucketPath::generate_new( + Some(PREFIX), + Some(&ext)); + + let result_bucket_object_pathbuf = result_bucket_location.to_full_object_pathbuf(); + + info!("Output file destination bucket path: {:?}", &result_bucket_object_pathbuf); + + info!("Uploading media ..."); + + args.deps.buckets.public_bucket_client.upload_filename_with_content_type( + &result_bucket_object_pathbuf, + &final_video, + &mimetype) // TODO: We should check the mimetype to make sure bad payloads can't get uploaded + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // generate thumbnail using thumbnail service + let thumbnail_types = match mimetype.as_str() { + "video/mp4" => vec!["image/gif", "image/jpeg"], + _ => vec![], + }; + + info!("Generating thumbnail tasks..."); + + for output_type in thumbnail_types { + let thumbnail_task_result = ThumbnailTaskBuilder::new() + .with_bucket(&*args.deps.buckets.public_bucket_client.bucket_name()) + .with_path(&*path_to_string(result_bucket_object_pathbuf.clone())) + .with_source_mimetype(mimetype.as_str()) + .with_output_mimetype(output_type) + .with_output_suffix("thumb") + .with_output_extension(&ext) + .with_event_id(&args.job.id.0.to_string()) + .send() + .await; + + match thumbnail_task_result { + Ok(thumbnail_task) => { + debug!("Thumbnail task created: {:?}", thumbnail_task); + }, + Err(e) => { + error!("Failed to create thumbnail task: {:?}", e); + } + } + } + + // Also upload the non-watermarked copy + info!("Uploading non-watermarked copy..."); + + let result = args.deps.buckets.public_bucket_client.upload_filename_with_content_type( + &result_bucket_object_pathbuf.with_extension("no_watermark.mp4"), + &args.videos.get_non_watermarked_video_to_upload(), + &mimetype) // TODO: We should check the mimetype to make sure bad payloads can't get uploaded + .await; + + if let Err(err) = result { + error!("Failed to upload non-watermarked copy: {:?}", err); + } + + if let Some(sleep_millis) = args.comfy_args.sleep_millis { + info!("Sleeping for millis: {sleep_millis}"); + thread::sleep(Duration::from_millis(sleep_millis)); + } + + // ==================== CLEANUP/ DELETE TEMP FILES ==================== // + + info!("Cleaning up temporary files..."); + + safe_delete_temp_file(&args.videos.original_video_path); + safe_delete_temp_file(&args.videos.trimmed_resampled_video_path); + safe_delete_temp_file(&args.videos.comfy_output_video_path); + safe_delete_temp_file(args.videos.video_to_watermark()); + safe_delete_temp_file(args.videos.get_final_video_to_upload()); + safe_delete_temp_file(args.videos.get_non_watermarked_video_to_upload()); + + // TODO(bt,2024-03-01): Do we really want to delete the workflow, models, etc.? + + safe_delete_temp_file(&args.workflow_path); + + // TODO(bt,2024-04-21): Not sure we want to delete the LoRA? + if let Some(lora_path) = args.maybe_lora_path { + safe_delete_temp_file(lora_path); + } + + let output_dir = args.root_comfy_path.join("output"); + safe_recursively_delete_files(&output_dir); + + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + safe_delete_temp_directory(&args.work_temp_dir); + + // ==================== SAVE RECORDS ==================== // + + // create a json detailing the args used to create the media + let args_json = serde_json::to_string(&args.job_args) + .map_err(|e| ProcessSingleJobError::Other(e.into()))?; + + info!("Saving ComfyUI result (media_files table record) ..."); + + // NB: We do this to avoid deep-frying the video. + // This also lets us hide the engine renders from users. + // This shouldn't ever become a deeply nested tree of children, but rather a single root + // with potentially many direct children. + let style_transfer_source_media_file_token = args.download_video + .input_video_media_file + .maybe_style_transfer_source_media_file_token + .as_ref() + .unwrap_or_else(|| &args.download_video.input_video_media_file.token); + + let prompt_token = PromptToken::generate(); + + let (media_file_token, id) = insert_media_file_from_comfy_ui(InsertArgs { + pool: &args.deps.db.mysql_pool, + job: &args.job, + maybe_mime_type: Some(&mimetype), + maybe_title: args.download_video.input_video_media_file.maybe_title.as_deref(), + maybe_style_transfer_source_media_file_token: Some(&style_transfer_source_media_file_token), + maybe_scene_source_media_file_token: args.download_video.input_video_media_file.maybe_scene_source_media_file_token.as_ref(), + file_size_bytes, + sha256_checksum: &file_checksum, + maybe_prompt_token: Some(&prompt_token), + public_bucket_directory_hash: result_bucket_location.get_object_hash(), + maybe_public_bucket_prefix: Some(PREFIX), + maybe_public_bucket_extension: Some(&ext), + is_on_prem: args.deps.job.info.container.is_on_prem, + worker_hostname: &args.deps.job.info.container.hostname, + worker_cluster: &args.deps.job.info.container.cluster_name, + extra_file_modification_info: Some(&args_json), + }) + .await + .map_err(|e| { + error!("Error saving media file record: {:?}", e); + ProcessSingleJobError::Other(e) + })?; + + if args.should_insert_prompt_record { + info!("Saving prompt record"); + + let mut other_args_builder = PromptInnerPayloadBuilder::new(); + + if let Some(style_name) = args.maybe_style_name { + info!("building PromptInnerPayload with style_name = {:?}", style_name); + other_args_builder.set_style_name(style_name); + } + + if args.comfy_args.use_face_detailer.unwrap_or(false) { + other_args_builder.set_used_face_detailer(true); + } + + if args.comfy_args.use_upscaler.unwrap_or(false) { + other_args_builder.set_used_upscaler(true); + } + + other_args_builder.set_strength(args.comfy_args.strength); + + if let Ok(duration) = chrono::Duration::from_std(args.inference_duration) { + // NB: Fail open. + other_args_builder.set_inference_duration(Some(duration)); + } + + let maybe_other_args = other_args_builder.build(); + + info!("maybe other prompt args: {:?}", maybe_other_args); + + // NB: Don't fail the job if the query fails. + let prompt_result = insert_prompt(InsertPromptArgs { + maybe_apriori_prompt_token: Some(&prompt_token), + prompt_type: PromptType::ComfyUi, + maybe_creator_user_token: args.job.maybe_creator_user_token_typed.as_ref(), + maybe_positive_prompt: args.maybe_positive_prompt.as_deref(), + maybe_negative_prompt: args.maybe_negative_prompt.as_deref(), + maybe_other_args: maybe_other_args.as_ref(), + creator_ip_address: &args.job.creator_ip_address, + mysql_executor: &args.deps.db.mysql_pool, + phantom: Default::default(), + }).await; + + if let Err(err) = prompt_result { + error!("No prompt result token? something has failed: {:?} (we'll ignore this error)", err); + } + } + + info!("ComfyUI Done."); + + args.job_progress_reporter.log_status("done") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + info!("Result video media token: {:?}", &media_file_token); + + info!("Job {:?} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}, Result Token: {}", + args.job.id, id, &media_file_token); + + Ok(media_file_token) +}