From 8d8e366f49119bfa4a6be4e7e1eb535c73fe1927 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 15 Mar 2024 01:00:23 -0400 Subject: [PATCH] add input media file to enqueue VST --- .../inference_input_source_token_type.rs | 11 ++++++++++- .../endpoints/workflows/enqueue_comfy_ui_handler.rs | 5 +++-- .../enqueue_video_style_transfer_handler.rs | 13 +++++++++---- 3 files changed, 22 insertions(+), 7 deletions(-) diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_input_source_token_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_input_source_token_type.rs index bbddcce2f7..76d98627fb 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_input_source_token_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_input_source_token_type.rs @@ -11,6 +11,8 @@ use strum::EnumIter; #[cfg_attr(test, derive(EnumIter, EnumCount))] #[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Deserialize, Serialize)] pub enum InferenceInputSourceTokenType { + #[serde(rename = "media_file")] + MediaFile, #[serde(rename = "media_upload")] MediaUpload, } @@ -23,12 +25,14 @@ impl_mysql_enum_coders!(InferenceInputSourceTokenType); impl InferenceInputSourceTokenType { pub fn to_str(&self) -> &'static str { match self { + Self::MediaFile => "media_file", Self::MediaUpload => "media_upload", } } pub fn from_str(value: &str) -> Result { match value { + "media_file" => Ok(Self::MediaFile), "media_upload" => Ok(Self::MediaUpload), _ => Err(format!("invalid value: {:?}", value)), } @@ -38,6 +42,7 @@ impl InferenceInputSourceTokenType { // NB: BTreeSet is sorted // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. BTreeSet::from([ + InferenceInputSourceTokenType::MediaFile, InferenceInputSourceTokenType::MediaUpload, ]) } @@ -50,16 +55,19 @@ mod tests { #[test] fn test_serialization() { + assert_serialization(InferenceInputSourceTokenType::MediaFile, "media_file"); assert_serialization(InferenceInputSourceTokenType::MediaUpload, "media_upload"); } #[test] fn to_str() { + assert_eq!(InferenceInputSourceTokenType::MediaFile.to_str(), "media_file"); assert_eq!(InferenceInputSourceTokenType::MediaUpload.to_str(), "media_upload"); } #[test] fn from_str() { + assert_eq!(InferenceInputSourceTokenType::from_str("media_file").unwrap(), InferenceInputSourceTokenType::MediaFile); assert_eq!(InferenceInputSourceTokenType::from_str("media_upload").unwrap(), InferenceInputSourceTokenType::MediaUpload); } @@ -67,7 +75,8 @@ mod tests { fn all_variants() { // Static check let mut variants = InferenceInputSourceTokenType::all_variants(); - assert_eq!(variants.len(), 1); + assert_eq!(variants.len(), 2); + assert_eq!(variants.pop_first(), Some(InferenceInputSourceTokenType::MediaFile)); assert_eq!(variants.pop_first(), Some(InferenceInputSourceTokenType::MediaUpload)); assert_eq!(variants.pop_first(), None); diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui_handler.rs index 8ce3ee2546..bad951ca69 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui_handler.rs @@ -12,6 +12,7 @@ use actix_web::http::StatusCode; use log::{error, info, warn}; use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; +use enums::by_table::generic_inference_jobs::inference_input_source_token_type::InferenceInputSourceTokenType; use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType; use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; use enums::common::visibility::Visibility; @@ -283,8 +284,8 @@ pub async fn enqueue_comfy_ui_handler( inference_category: InferenceCategory::Workflow, maybe_model_type: Some(InferenceModelType::ComfyUi), // 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_input_source_token: request.maybe_input_file.as_ref().map(|f| f.as_str()), + maybe_input_source_token_type: Some(InferenceInputSourceTokenType::MediaFile), maybe_download_url: None, maybe_cover_image_media_file_token: None, maybe_raw_inference_text: None, // No text diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_video_style_transfer_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_video_style_transfer_handler.rs index e5e84af6e9..ce8c4ed4c8 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_video_style_transfer_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_video_style_transfer_handler.rs @@ -12,6 +12,7 @@ use log::{error, info, warn}; use utoipa::ToSchema; use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; +use enums::by_table::generic_inference_jobs::inference_input_source_token_type::InferenceInputSourceTokenType; use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType; use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; use enums::common::visibility::Visibility; @@ -23,6 +24,7 @@ use mysql_queries::payloads::generic_inference_args::workflow_payload::WorkflowA use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{insert_generic_inference_job, InsertGenericInferenceArgs}; use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::media_files::MediaFileToken; use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session; use crate::configs::plans::plan_category::PlanCategory; @@ -44,9 +46,12 @@ pub struct EnqueueVideoStyleTransferRequest { /// Entropy for request de-duplication (required) uuid_idempotency_token: String, + /// The video to use as input (required) + input_media_file_token: MediaFileToken, + /// The name of the style to invoke (required) style: StyleTransferName, - + /// The positive prompt (optional) prompt: Option, @@ -244,6 +249,7 @@ pub async fn enqueue_video_style_transfer_handler( } let inference_args = WorkflowArgs { + maybe_input_file: Some(request.input_media_file_token.clone()), style_name: Some(request.style), creator_visibility: Some(set_visibility), trim_start_milliseconds: Some(trim_start_millis), @@ -255,7 +261,6 @@ pub async fn enqueue_video_style_transfer_handler( maybe_lora_model: None, maybe_json_modifications: None, maybe_workflow_config: None, - maybe_input_file: None, maybe_output_path: None, maybe_google_drive_link: None, maybe_title: None, @@ -274,8 +279,8 @@ pub async fn enqueue_video_style_transfer_handler( inference_category: InferenceCategory::Workflow, maybe_model_type: Some(InferenceModelType::ComfyUi), // 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_input_source_token: Some(request.input_media_file_token.as_str()), + maybe_input_source_token_type: Some(InferenceInputSourceTokenType::MediaFile), maybe_download_url: None, maybe_cover_image_media_file_token: None, maybe_raw_inference_text: None, // No text