diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs index deb3f4fe29..3192780832 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::model_weights::ModelWeightToken; +use enums::common::visibility::Visibility; #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] #[serde(untagged)] @@ -37,8 +38,25 @@ pub struct WorkflowArgs { #[serde(skip_serializing_if = "Option::is_none")] pub maybe_output_path: Option, + // Upload information // google drive link for uploads #[serde(rename = "gd")] #[serde(skip_serializing_if = "Option::is_none")] - pub maybe_google_drive_link: Option + pub maybe_google_drive_link: Option, + + #[serde(rename = "ti")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_title: Option, + + #[serde(rename = "de")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_description: Option, + + #[serde(rename = "ch")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_commit_hash: Option, + + #[serde(rename = "cv")] + #[serde(skip_serializing_if = "Option::is_none")] + pub creator_visibility: Option, } diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/upload_prompt.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/upload_prompt.rs index 21ad4e990a..66a0c4c96d 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/upload_prompt.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/upload_prompt.rs @@ -1,25 +1,29 @@ -use std::time::Duration; -use anyhow::{anyhow, Result}; -use log::info; - -use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; -use cloud_storage::remote_file_manager::weights_descriptor::WeightsWorkflowDescriptor; -use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use std::time::{Duration}; +use cloud_storage::remote_file_manager::weights_descriptor::{WeightsSD15Descriptor, WeightsWorkflowDescriptor}; use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; use enums::common::visibility::Visibility; use filesys::file_exists::file_exists; +use anyhow::{anyhow, Result}; use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand; +use log::{debug, error, info, warn}; +use mysql_queries::queries::model_weights::create::create_model_weight_from_voice_conversion_download_job::CreateModelWeightArgs; use mysql_queries::queries::model_weights::create::create_weight; -use mysql_queries::queries::model_weights::create::create_weight::CreateModelWeightsArgs; use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::users::UserToken; +use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; +use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use mysql_queries::queries::media_files::create::insert_media_file_from_comfy_ui::{insert_media_file_from_comfy_ui, InsertArgs}; +use mysql_queries::queries::model_weights::create::create_weight::CreateModelWeightsArgs; + + 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::workflow::comfy_ui::process_job::ComfyProcessJobArgs; use crate::job::job_types::workflow::process_single_wf_job::get_workflow_args_from_job; +use crate::job_dependencies::JobDependencies; +use crate::job::job_types::workflow::comfy_ui::process_job::ComfyProcessJobArgs; pub async fn upload_prompt(args: ComfyProcessJobArgs<'_>) -> Result{ let job = args.job; @@ -28,6 +32,35 @@ pub async fn upload_prompt(args: ComfyProcessJobArgs<'_>) -> Result { + val + }, + None => { "".to_string() } + }; + + let description = match wf_args.maybe_description { + Some(val) => { + val + }, + None => { "".to_string() } + }; + + let commit_hash = match wf_args.maybe_commit_hash { + Some(val) => { + val + }, + None => { "".to_string() } + }; + + let visibility = match wf_args.creator_visibility { + Some(val) => { + val + }, + None => { Visibility::Public } + }; + + let file_name = "workflow.json"; let download_script = easyenv::get_env_string_or_default( @@ -56,9 +89,8 @@ pub async fn upload_prompt(args: ComfyProcessJobArgs<'_>) -> Result) -> Result) -> Result, + pub creator_set_visibility:Option } #[derive(Serialize, ToSchema)] @@ -118,8 +122,22 @@ pub async fn enqueue_workflow_upload_request( return Err(EnqueueWorkFlowRequestError::BadInput("Missing Google Drive Link".to_string())); } + if request.title.is_empty() { + return Err(EnqueueWorkFlowRequestError::BadInput("Missing Title".to_string())); + } + + if request.description.is_empty() { + return Err(EnqueueWorkFlowRequestError::BadInput("Missing Description".to_string())); + } + + if request.commit_hash.is_empty() { + return Err(EnqueueWorkFlowRequestError::BadInput("Missing commit hash".to_string())); + } + let mut maybe_user_token: Option = None; - let visibility = Visibility::Private; + + let visibility = request.creator_set_visibility + .unwrap_or(Visibility::Private); let mut mysql_connection = server_state.mysql_pool.acquire().await.map_err(|err| { warn!("MySql pool error: {:?}", err); @@ -201,7 +219,10 @@ pub async fn enqueue_workflow_upload_request( // ==================== INFERENCE ARGS ==================== // let google_drive_link = request.google_drive_link.clone(); - + let title = request.title.clone(); + let description = request.description.clone(); + let commit_hash = request.commit_hash.clone(); + let inference_args = WorkflowArgs { maybe_sd_model: None, maybe_lora_model: None, @@ -209,7 +230,11 @@ pub async fn enqueue_workflow_upload_request( maybe_workflow_config: None, maybe_input_file: None, maybe_output_path: None, - maybe_google_drive_link: Some(google_drive_link) + maybe_google_drive_link: Some(google_drive_link), + maybe_title: Some(title), + maybe_description: Some(description), + maybe_commit_hash: Some(commit_hash), + creator_visibility: Some(visibility), }; // create the inference args here